Update context detection code

src/api/clients/gemini.ts

@@ -9,7 +9,6 @@ export class GeminiAPIClient implements APIClient {
     throw new Error('Method not implemented.');
   }
   fetchCompletions(
-    language: string,
     prefix: string,
     suffix: string,
   ): Promise<string | undefined> {

src/api/clients/openai-compatible.ts

@@ -61,17 +61,18 @@ export abstract class OpenAICompatibleAPIClient implements APIClient {
     }
   }
 
-  async fetchCompletions(language: string, prefix: string, suffix: string) {
+  async fetchCompletions(prefix: string, suffix: string) {
     const { settings } = this.plugin;
 
     if (this.openai === undefined) {
       return;
     }
 
     try {
-      const completions = await this.openai.completions.create({
-        prompt: `Continue the following code written in ${language} language:\n\n${prefix}`,
-        suffix,
+      // TODO:
+      // Get messages from the prompt generator.
+      const completions = await this.openai.chat.completions.create({
+        messages: [],
         model: settings.completions.model,
         max_tokens: settings.completions.maxTokens,
         temperature: settings.completions.temperature,
@@ -90,7 +91,7 @@ export abstract class OpenAICompatibleAPIClient implements APIClient {
         outputTokens,
       );
 
-      return completions.choices[0].text;
+      return completions.choices[0].message.content ?? undefined;
     } catch (error) {
       console.error(error);
       new Notice(

src/api/index.ts

@@ -2,9 +2,5 @@ import { ChatMessage } from '../types';
 
 export interface APIClient {
   fetchChat(messages: ChatMessage[]): AsyncGenerator<string | undefined>;
-  fetchCompletions(
-    language: string,
-    prefix: string,
-    suffix: string,
-  ): Promise<string | undefined>;
+  fetchCompletions(prefix: string, suffix: string): Promise<string | undefined>;
 }

src/api/prompts/block-quote/index.ts

@@ -0,0 +1,14 @@
+import { FewShotPrompt } from '..';
+import example1Assistant from './example1/assistant.txt';
+import example1User from './example1/user.md';
+import system from './system.txt';
+
+export const BlockQuotePrompt: FewShotPrompt = {
+  system,
+  examples: [
+    {
+      user: example1User,
+      assistant: example1Assistant,
+    },
+  ],
+};

src/api/prompts/block-quote/system.txt

@@ -0,0 +1,12 @@
+Complete the most suitable text at the location of the <MASK>.
+The <MASK> is located within a Markdown block quote. 
+Your answer must complete this quote in a way that fits the context of the surrounding text.
+Your answer must be written in the same language as the surrounding text.
+Your answer must not overlap with any text adjacent to the <MASK>.
+Your answer must have the following format:
+<LANGUAGE>
+Here, you write the language of your response e.g. English, Chinese, TypeScript, Python.
+<THOUGHT>
+Here, you reason about the answer, using the 80/20 rule for clarity and conciseness.
+<INSERT>
+Here, you write the text that should be inserted at the location of the <MASK>.

src/api/prompts/code-block/index.ts

@@ -1,8 +1,26 @@
-import { FewShowExample } from '../example';
-import assistant from './assistant.txt';
-import user from './user.md';
+import { FewShotPrompt } from '..';
+import example1Assistant from './example1/assistant.txt';
+import example1User from './example1/user.md';
+import example2Assistant from './example2/assistant.txt';
+import example2User from './example2/user.md';
+import example3Assistant from './example3/assistant.txt';
+import example3User from './example3/user.md';
+import system from './system.txt';
 
-export const CODE_BLOCK_EXAMPLE: FewShowExample = {
-  user,
-  assistant,
+export const CodeBlockPrompt: FewShotPrompt = {
+  system,
+  examples: [
+    {
+      user: example1User,
+      assistant: example1Assistant,
+    },
+    {
+      user: example2User,
+      assistant: example2Assistant,
+    },
+    {
+      user: example3User,
+      assistant: example3Assistant,
+    },
+  ],
 };

src/api/prompts/code-block/system.txt

@@ -0,0 +1,12 @@
+Complete the most suitable text at the location of the <MASK>.
+The <MASK> is located within a Markdown codeblock, written in the language {{LANGUAGE}}. 
+Your answer must complete this code block in the language {{LANGUAGE}}.
+Your answer must not complete any text outside this code block.
+Your answer must not overlap with any text adjacent to the <MASK>.
+Your answer must have the following format:
+<LANGUAGE>
+Here, you write the language of your response e.g. English, Chinese, TypeScript, Python.
+<THOUGHT>
+Here, you reason about the answer, using the 80/20 rule for clarity and conciseness.
+<INSERT>
+Here, you write the text that should be inserted at the location of the <MASK>.

src/api/prompts/context.ts

@@ -71,6 +71,9 @@ export function getLanguage(prefix: string, suffix: string): string {
 }
 
 function isCursorInBlock(text: string, regex: RegExp): boolean {
-  const blocks = text.match(regex) ?? [];
+  const blocks = text.match(regex) as string[] | null;
+  if (blocks === null) {
+    return false;
+  }
   return blocks.some((block) => block.includes(CURSOR_CHAR));
 }

src/api/prompts/generator.ts

@@ -0,0 +1,62 @@
+import Markpilot from 'src/main';
+import { ChatMessage } from 'src/types';
+import { FewShotPrompt } from '.';
+import { BlockQuotePrompt } from './block-quote';
+import { CodeBlockPrompt } from './code-block';
+import { Context, getContext, getLanguage } from './context';
+import { HeadingPrompt } from './heading';
+import { ListItemPrompt } from './list-item';
+import { MathBlockPrompt } from './math-block';
+import { ParagraphPrompt } from './paragraph';
+
+const PROMPTS: Record<Context, FewShotPrompt> = {
+  heading: HeadingPrompt,
+  paragraph: ParagraphPrompt,
+  'list-item': ListItemPrompt,
+  'block-quote': BlockQuotePrompt,
+  'math-block': MathBlockPrompt,
+  'code-block': CodeBlockPrompt,
+};
+
+export class PromptGenerator {
+  constructor(private plugin: Markpilot) {}
+
+  generate(prefix: string, suffix: string): ChatMessage[] {
+    const { settings } = this.plugin;
+
+    const windowSize = settings.completions.windowSize;
+    const truncatedPrefix = prefix.slice(
+      prefix.length - windowSize / 2,
+      prefix.length,
+    );
+    const truncatedSuffix = suffix.slice(0, windowSize / 2);
+
+    const context = getContext(prefix, suffix);
+    const prompt = PROMPTS[context];
+    if (context === 'code-block') {
+      const language = getLanguage(prefix, suffix);
+      prompt.system = prompt.system.replace('{{LANGUAGE}}', language);
+    }
+
+    return [
+      {
+        role: 'system',
+        content: prompt.system,
+      },
+      ...prompt.examples.flatMap((example) => [
+        {
+          role: 'user',
+          content: example.user,
+        },
+        {
+          role: 'assistant',
+          content: example.assistant,
+        },
+      ]),
+      {
+        role: 'user',
+        content: `${truncatedPrefix}<MASK>${truncatedSuffix}`,
+      },
+    ] as ChatMessage[];
+  }
+}

src/api/prompts/heading/index.ts

@@ -0,0 +1,20 @@
+import { FewShotPrompt } from '..';
+import example1Assistant from './example1/assistant.txt';
+import example1User from './example1/user.md';
+import example2Assistant from './example2/assistant.txt';
+import example2User from './example2/user.md';
+import system from './system.txt';
+
+export const HeadingPrompt: FewShotPrompt = {
+  system,
+  examples: [
+    {
+      user: example1User,
+      assistant: example1Assistant,
+    },
+    {
+      user: example2User,
+      assistant: example2Assistant,
+    },
+  ],
+};

src/api/prompts/heading/system.txt

@@ -0,0 +1,12 @@
+Complete the most suitable text at the location of the <MASK>.
+The <MASK> is located within a Markdown heading.
+Your answer must complete the title for this heading that fits the context of the surrounding text.
+Your answer must be written in the same language as the surrounding text.
+Your answer must not overlap with any text adjacent to the <MASK>.
+Your answer must have the following format:
+<LANGUAGE>
+Here, you write the language of your response e.g. English, Chinese, TypeScript, Python.
+<THOUGHT>
+Here, you reason about the answer, using the 80/20 rule for clarity and conciseness.
+<INSERT>
+Here, you write the text that should be inserted at the location of the <MASK>.

src/api/prompts/index.ts

@@ -1,21 +1,9 @@
-import systemPrompt from './system.txt';
+export interface FewShotPrompt {
+  system: string;
+  examples: FewShotExample[];
+}
 
-export interface FewShowExample {
+export interface FewShotExample {
   user: string;
   assistant: string;
 }
-
-export class PromptGenerator {
-  private systemPrompt = systemPrompt;
-
-  generate(prefix: string, suffix: string): string {
-    // TODO:
-    // 1. Determine the context from prefix and suffix.
-    // 2. Generate a prompt based on the context, with prefix and suffix trimmed according to window size.
-    const language = 'english';
-    if (language) {
-      return this.systemPrompt.replace('{{LANGUAGE}}', language);
-    }
-    return '';
-  }
-}

src/api/prompts/list-item/index.ts

@@ -0,0 +1,14 @@
+import { FewShotPrompt } from '..';
+import example1Assistant from './example1/assistant.txt';
+import example1User from './example1/user.md';
+import system from './system.txt';
+
+export const ListItemPrompt: FewShotPrompt = {
+  system,
+  examples: [
+    {
+      user: example1User,
+      assistant: example1Assistant,
+    },
+  ],
+};

src/api/prompts/list-item/system.txt

@@ -0,0 +1,13 @@
+Complete the most suitable text at the location of the <MASK>.
+The <MASK> is located within a Markdown list item.
+Your answer must complete one or multiple list items for this list that fits the context of the surrounding text.
+Your answer must not complete any text that is not part of this list.
+Your answer must be written in the same language as the surrounding text.
+Your answer must not overlap with any text adjacent to the <MASK>.
+Your answer must have the following format:
+<LANGUAGE>
+Here, you write the language of your response e.g. English, Chinese, TypeScript, Python.
+<THOUGHT>
+Here, you reason about the answer, using the 80/20 rule for clarity and conciseness.
+<INSERT>
+Here, you write the text that should be inserted at the location of the <MASK>.

src/api/prompts/math-block/index.ts

@@ -0,0 +1,20 @@
+import { FewShotPrompt } from '..';
+import example1Assistant from './example1/assistant.txt';
+import example1User from './example1/user.md';
+import example2Assistant from './example2/assistant.txt';
+import example2User from './example2/user.md';
+import system from './system.txt';
+
+export const MathBlockPrompt: FewShotPrompt = {
+  system,
+  examples: [
+    {
+      user: example1User,
+      assistant: example1Assistant,
+    },
+    {
+      user: example2User,
+      assistant: example2Assistant,
+    },
+  ],
+};

src/api/prompts/math-block/system.txt

@@ -0,0 +1,13 @@
+Complete the most suitable text at the location of the <MASK>.
+The <MASK> is located within a Markdown math block.
+Your answer must only contain LaTeX code that captures the math discussed in the surrounding text. 
+Your answer must not contain any text that is not part of the LaTeX code.
+Your answer must be written in the same language as the surrounding text.
+Your answer must not overlap with any text adjacent to the <MASK>.
+Your answer must have the following format:
+<LANGUAGE>
+Here, you write the language of your response e.g. English, Chinese, TypeScript, Python.
+<THOUGHT>
+Here, you reason about the answer, using the 80/20 rule for clarity and conciseness.
+<INSERT>
+Here, you write the text that should be inserted at the location of the <MASK>.

src/api/prompts/paragraph/index.ts

@@ -0,0 +1,14 @@
+import { FewShotPrompt } from '..';
+import example1Assistant from './example1/assistant.txt';
+import example1User from './example1/user.md';
+import system from './system.txt';
+
+export const ParagraphPrompt: FewShotPrompt = {
+  system,
+  examples: [
+    {
+      user: example1User,
+      assistant: example1Assistant,
+    },
+  ],
+};

src/api/prompts/paragraph/system.txt

@@ -0,0 +1,12 @@
+Complete the most suitable text at the location of the <MASK>.
+The <MASK> is located within a Markdown paragraph.
+Your answer must complete one or multiple sentences to this paragraph that fit the surrounding text.
+Your answer must be written in the same language as the surrounding text.
+Your answer must not overlap with any text adjacent to the <MASK>.
+Your answer must have the following format:
+<LANGUAGE>
+Here, you write the language of your response e.g. English, Chinese, TypeScript, Python.
+<THOUGHT>
+Here, you reason about the answer, using the 80/20 rule for clarity and conciseness.
+<INSERT>
+Here, you write the text that should be inserted at the location of the <MASK>.