add on completion method

xgyaiczy · Jul 28, 2023 · 7fd9799 · 7fd9799
1 parent 6e3ab97
commit 7fd9799
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 41 deletions.
diff --git a/features/chat/chat-api-data.ts b/features/chat/chat-api-data.ts
@@ -8,7 +8,6 @@ import {
   HumanMessagePromptTemplate,
   SystemMessagePromptTemplate,
 } from "langchain/prompts";
-import { ChainValues } from "langchain/schema";
 import {
   AzureCogDocument,
   AzureCogSearch,
@@ -32,12 +31,18 @@ export const PromptDataGPT = async (props: PromptGPTProps) => {
     streaming: true,
   });
 
-  const relevantDocuments = findRelevantDocuments(lastHumanMessage.content);
+  const relevantDocuments = await findRelevantDocuments(
+    lastHumanMessage.content
+  );
 
   const chain = loadQAMapReduceChain(chatModel, {
     combinePrompt: defineSystemPrompt(),
   });
-  const { stream, handlers } = LangChainStream();
+  const { stream, handlers } = LangChainStream({
+    onCompletion: async (completion: string) => {
+      await inertPromptAndResponse(id, lastHumanMessage.content, completion);
+    },
+  });
 
   const memory = buildMemory(chats);
 
@@ -47,24 +52,7 @@ export const PromptDataGPT = async (props: PromptGPTProps) => {
       question: lastHumanMessage.content,
       memory: memory,
     },
-    [
-      {
-        ...handlers,
-        handleChainEnd: async (
-          outputs: ChainValues,
-          runId: string,
-          parentRunId?: string,
-          tags?: string[]
-        ) => {
-          await handlers.handleChainEnd(outputs, runId);
-          await inertPromptAndResponse(
-            id,
-            lastHumanMessage.content,
-            outputs.text
-          );
-        },
-      },
-    ]
+    [handlers]
   );
 
   return new StreamingTextResponse(stream);

diff --git a/features/chat/chat-api.ts b/features/chat/chat-api.ts
@@ -1,14 +1,18 @@
 import { LangChainStream, StreamingTextResponse } from "ai";
 import { ChatOpenAI } from "langchain/chat_models/openai";
-import { HumanMessage, LLMResult, SystemMessage } from "langchain/schema";
+import { HumanMessage, SystemMessage } from "langchain/schema";
 import { PromptGPTProps, initAndGuardChatSession } from "./chat-api-helpers";
 import { mostRecentMemory } from "./chat-helpers";
 import { inertPromptAndResponse } from "./chat-service";
 
 export const PromptGPT = async (props: PromptGPTProps) => {
   const { lastHumanMessage, id, chats } = await initAndGuardChatSession(props);
 
-  const { stream, handlers } = LangChainStream();
+  const { stream, handlers } = LangChainStream({
+    onCompletion: async (completion: string) => {
+      await inertPromptAndResponse(id, lastHumanMessage.content, completion);
+    },
+  });
 
   const memory = mostRecentMemory(chats, 10);
 
@@ -26,24 +30,7 @@ export const PromptGPT = async (props: PromptGPTProps) => {
       new HumanMessage(lastHumanMessage.content),
     ],
     {},
-    [
-      {
-        ...handlers,
-        handleLLMEnd: async (
-          output: LLMResult,
-          runId: string,
-          parentRunId?: string | undefined,
-          tags?: string[] | undefined
-        ) => {
-          await handlers.handleLLMEnd(output, runId);
-          await inertPromptAndResponse(
-            id,
-            lastHumanMessage.content,
-            output.generations[0][0].text
-          );
-        },
-      },
-    ]
+    [handlers]
   );
 
   return new StreamingTextResponse(stream);