feat: improve copilot plugin (#3459)

apps/core/src/bootstrap/register-plugins.ts

@@ -76,7 +76,7 @@ await Promise.all(
         globalThis.__pluginPackageJson__.push(packageJson);
         logger.debug(`registering plugin ${pluginName}`);
         logger.debug(`package.json: ${packageJson}`);
-        if (!release) {
+        if (!release && process.env.NODE_ENV === 'production') {
           return Promise.resolve();
         }
         const pluginCompartment = new Compartment(createGlobalThis(), {});

apps/core/src/components/page-detail-editor.tsx

@@ -191,7 +191,12 @@ const LayoutPanel = memo(function LayoutPanel(
           </Suspense>
         </Panel>
         <PanelResizeHandle />
-        <Panel defaultSize={100 - node.splitPercentage}>
+        <Panel
+          defaultSize={100 - node.splitPercentage}
+          style={{
+            overflow: 'scroll',
+          }}
+        >
           <Suspense>
             <LayoutPanel node={node.second} editorProps={props.editorProps} />
           </Suspense>

plugins/copilot/src/UI/header-item.tsx

@@ -17,7 +17,7 @@ export const HeaderItem = (): ReactElement => {
                 return {
                   direction: 'horizontal',
                   first: 'editor',
-                  second: '@affine/copilot',
+                  second: '@affine/copilot-plugin',
                   splitPercentage: 70,
                 };
               } else {

plugins/copilot/src/core/chat.ts

@@ -11,22 +11,29 @@ import {
 
 import { IndexedDBChatMessageHistory } from './langchain/message-history';
 import { chatPrompt, followupQuestionPrompt } from './prompts';
+import { followupQuestionParser } from './prompts/output-parser';
 
-declare global {
+type ChatAI = {
-  interface WindowEventMap {
+  // Core chat AI
-    'llm-start': CustomEvent;
+  conversationChain: ConversationChain;
-    'llm-new-token': CustomEvent<{ token: string }>;
+  // Followup AI, used to generate followup questions
-  }
+  followupChain: LLMChain<string>;
-}
+  // Chat history, used to store messages
+  chatHistory: IndexedDBChatMessageHistory;
+};
+
+export type ChatAIConfig = {
+  events: {
+    llmStart: () => void;
+    llmNewToken: (token: string) => void;
+  };
+};
 
 export async function createChatAI(
   room: string,
-  openAIApiKey: string
+  openAIApiKey: string,
-): Promise<{
+  config: ChatAIConfig
-  conversationChain: ConversationChain;
+): Promise<ChatAI> {
-  followupChain: LLMChain<string>;
-  chatHistory: IndexedDBChatMessageHistory;
-}> {
   if (!openAIApiKey) {
     console.warn('OpenAI API key not set, chat will not work');
   }
@@ -44,25 +51,11 @@ export async function createChatAI(
     openAIApiKey: openAIApiKey,
     callbacks: [
       {
-        async handleLLMStart(llm, prompts, runId, parentRunId, extraParams) {
+        async handleLLMStart() {
-          console.log(
+          config.events.llmStart();
-            'handleLLMStart',
-            llm,
-            prompts,
-            runId,
-            parentRunId,
-            extraParams
-          );
-          window.dispatchEvent(new CustomEvent('llm-start'));
         },
-        async handleLLMNewToken(token, runId, parentRunId) {
+        async handleLLMNewToken(token) {
-          console.log('handleLLMNewToken', token, runId, parentRunId);
+          config.events.llmNewToken(token);
-          window.dispatchEvent(
-            new CustomEvent('llm-new-token', { detail: { token } })
-          );
-        },
-        async handleLLMEnd(output, runId, parentRunId) {
-          console.log('handleLLMEnd', output, runId, parentRunId);
         },
       },
     ],
@@ -77,6 +70,9 @@ export async function createChatAI(
   const followupPromptTemplate = new PromptTemplate({
     template: followupQuestionPrompt,
     inputVariables: ['human_conversation', 'ai_conversation'],
+    partialVariables: {
+      format_instructions: followupQuestionParser.getFormatInstructions(),
+    },
   });
 
   const followupChain = new LLMChain({
@@ -101,5 +97,5 @@ export async function createChatAI(
     conversationChain,
     followupChain,
     chatHistory,
-  } as const;
+  };
 }

plugins/copilot/src/core/hooks/index.ts

@@ -1,51 +1,56 @@
-import type { IndexedDBChatMessageHistory } from '@affine/copilot/core/langchain/message-history';
 import { atom, useAtomValue } from 'jotai';
 import { atomWithDefault, atomWithStorage } from 'jotai/utils';
 import type { WritableAtom } from 'jotai/vanilla';
+import type { PrimitiveAtom } from 'jotai/vanilla';
 import type { LLMChain } from 'langchain/chains';
 import { type ConversationChain } from 'langchain/chains';
 import { type BufferMemory } from 'langchain/memory';
 import type { BaseMessage } from 'langchain/schema';
 import { AIMessage } from 'langchain/schema';
 import { HumanMessage } from 'langchain/schema';
-import { z } from 'zod';
 
+import type { ChatAIConfig } from '../chat';
 import { createChatAI } from '../chat';
-
+import type { IndexedDBChatMessageHistory } from '../langchain/message-history';
-const followupResponseSchema = z.array(z.string());
+import { followupQuestionParser } from '../prompts/output-parser';
 
 export const openAIApiKeyAtom = atomWithStorage<string | null>(
   'com.affine.copilot.openai.token',
   null
 );
 
-export const chatAtom = atom(async get => {
+const conversationBaseWeakMap = new WeakMap<
-  const openAIApiKey = get(openAIApiKeyAtom);
+  ConversationChain,
-  if (!openAIApiKey) {
+  PrimitiveAtom<BaseMessage[]>
-    throw new Error('OpenAI API key not set, chat will not work');
+>();
-  }
-  return createChatAI('default-copilot', openAIApiKey);
-});
-
 const conversationWeakMap = new WeakMap<
   ConversationChain,
   WritableAtom<BaseMessage[], [string], Promise<void>>
 >();
 
-const getConversationAtom = (chat: ConversationChain) => {
+export const chatAtom = atom(async get => {
-  if (conversationWeakMap.has(chat)) {
+  const openAIApiKey = get(openAIApiKeyAtom);
-    return conversationWeakMap.get(chat) as WritableAtom<
+  if (!openAIApiKey) {
-      BaseMessage[],
+    throw new Error('OpenAI API key not set, chat will not work');
-      [string],
-      Promise<void>
-    >;
   }
-  const conversationBaseAtom = atom<BaseMessage[]>([]);
+  const events: ChatAIConfig['events'] = {
-  conversationBaseAtom.onMount = setAtom => {
+    llmStart: () => {
-    if (!chat) {
+      throw new Error('llmStart not set');
-      throw new Error();
+    },
-    }
+    llmNewToken: () => {
-    const memory = chat.memory as BufferMemory;
+      throw new Error('llmNewToken not set');
+    },
+  };
+  const chatAI = await createChatAI('default-copilot', openAIApiKey, {
+    events,
+  });
+  getOrCreateConversationAtom(chatAI.conversationChain);
+  const baseAtom = conversationBaseWeakMap.get(chatAI.conversationChain);
+  if (!baseAtom) {
+    throw new TypeError();
+  }
+  baseAtom.onMount = setAtom => {
+    const memory = chatAI.conversationChain.memory as BufferMemory;
     memory.chatHistory
       .getMessages()
       .then(messages => {
@@ -54,23 +59,27 @@ const getConversationAtom = (chat: ConversationChain) => {
       .catch(err => {
         console.error(err);
       });
-    const llmStart = (): void => {
+    events.llmStart = () => {
       setAtom(conversations => [...conversations, new AIMessage('')]);
     };
-    const llmNewToken = (event: CustomEvent<{ token: string }>): void => {
+    events.llmNewToken = token => {
       setAtom(conversations => {
         const last = conversations[conversations.length - 1] as AIMessage;
-        last.content += event.detail.token;
+        last.content += token;
         return [...conversations];
       });
     };
-    window.addEventListener('llm-start', llmStart);
-    window.addEventListener('llm-new-token', llmNewToken);
-    return () => {
-      window.removeEventListener('llm-start', llmStart);
-      window.removeEventListener('llm-new-token', llmNewToken);
-    };
   };
+  return chatAI;
+});
+
+const getOrCreateConversationAtom = (chat: ConversationChain) => {
+  if (conversationWeakMap.has(chat)) {
+    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+    return conversationWeakMap.get(chat)!;
+  }
+  const conversationBaseAtom = atom<BaseMessage[]>([]);
+  conversationBaseWeakMap.set(chat, conversationBaseAtom);
 
   const conversationAtom = atom<BaseMessage[], [string], Promise<void>>(
     get => get(conversationBaseAtom),
@@ -105,7 +114,9 @@ const getConversationAtom = (chat: ConversationChain) => {
 const followingUpWeakMap = new WeakMap<
   LLMChain<string>,
   {
-    questionsAtom: ReturnType<typeof atomWithDefault<Promise<string[]>>>;
+    questionsAtom: ReturnType<
+      typeof atomWithDefault<Promise<string[]> | string[]>
+    >;
     generateChatAtom: WritableAtom<null, [], void>;
   }
 >();
@@ -115,12 +126,10 @@ const getFollowingUpAtoms = (
   chatHistory: IndexedDBChatMessageHistory
 ) => {
   if (followingUpWeakMap.has(followupLLMChain)) {
-    return followingUpWeakMap.get(followupLLMChain) as {
+    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      questionsAtom: ReturnType<typeof atomWithDefault<Promise<string[]>>>;
+    return followingUpWeakMap.get(followupLLMChain)!;
-      generateChatAtom: WritableAtom<null, [], void>;
-    };
   }
-  const baseAtom = atomWithDefault<Promise<string[]>>(async () => {
+  const baseAtom = atomWithDefault<Promise<string[]> | string[]>(async () => {
     return chatHistory?.getFollowingUp() ?? [];
   });
   const setAtom = atom<null, [], void>(null, async (get, set) => {
@@ -137,10 +146,9 @@ const getFollowingUpAtoms = (
       ai_conversation: aiMessage,
       human_conversation: humanMessage,
     });
-    const followingUp = JSON.parse(response.text);
+    const followingUp = await followupQuestionParser.parse(response.text);
-    followupResponseSchema.parse(followingUp);
+    set(baseAtom, followingUp.followupQuestions);
-    set(baseAtom, followingUp);
+    chatHistory.saveFollowingUp(followingUp.followupQuestions).catch(() => {
-    chatHistory.saveFollowingUp(followingUp).catch(() => {
       console.error('failed to save followup');
     });
   });
@@ -155,11 +163,11 @@ const getFollowingUpAtoms = (
 };
 
 export function useChatAtoms(): {
-  conversationAtom: ReturnType<typeof getConversationAtom>;
+  conversationAtom: ReturnType<typeof getOrCreateConversationAtom>;
   followingUpAtoms: ReturnType<typeof getFollowingUpAtoms>;
 } {
   const chat = useAtomValue(chatAtom);
-  const conversationAtom = getConversationAtom(chat.conversationChain);
+  const conversationAtom = getOrCreateConversationAtom(chat.conversationChain);
   const followingUpAtoms = getFollowingUpAtoms(
     chat.followupChain,
     chat.chatHistory

plugins/copilot/src/core/prompts/index.ts

@@ -18,12 +18,10 @@ You can only give one reply for each conversation turn.
 `;
 
 export const followupQuestionPrompt = `Rules you must follow:
-- You only respond in JSON format
+Read the following conversation between AI and Human and generate at most 3 follow-up messages or questions the Human can ask
-- Read the following conversation between AI and Human and generate at most 3 follow-up messages or questions the Human can ask
+Each message in your response should be concise, no more than 15 words
-- Your response MUST be a valid JSON array of strings like this: ["some question", "another question"]
+You MUST reply in the same written language as the conversation
-- Each message in your response should be concise, no more than 15 words
+{format_instructions}
-- You MUST reply in the same written language as the conversation
-- Don't output anything other text
 The conversation is inside triple quotes:
 \`\`\`
 Human: {human_conversation}

plugins/copilot/src/core/prompts/output-parser.ts

@@ -0,0 +1,8 @@
+import { StructuredOutputParser } from 'langchain/output_parsers';
+import { z } from 'zod';
+
+export const followupQuestionParser = StructuredOutputParser.fromZodSchema(
+  z.object({
+    followupQuestions: z.array(z.string()),
+  })
+);