feat(app): migrate suggestion chain

src/app/api/suggestions/route.ts

@@ -1,7 +1,6 @@
 import generateSuggestions from '@/lib/agents/suggestions';
 import ModelRegistry from '@/lib/models/registry';
 import { ModelWithProvider } from '@/lib/models/types';
-import { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import { AIMessage, BaseMessage, HumanMessage } from '@langchain/core/messages';
 
 interface SuggestionsGenerationBody {
@@ -13,6 +12,13 @@ export const POST = async (req: Request) => {
   try {
     const body: SuggestionsGenerationBody = await req.json();
 
+    const registry = new ModelRegistry();
+
+    const llm = await registry.loadChatModel(
+      body.chatModel.providerId,
+      body.chatModel.key,
+    );
+
     const chatHistory = body.chatHistory
       .map((msg: any) => {
         if (msg.role === 'user') {
@@ -23,16 +29,9 @@ export const POST = async (req: Request) => {
       })
       .filter((msg) => msg !== undefined) as BaseMessage[];
 
-    const registry = new ModelRegistry();
-
-    const llm = await registry.loadChatModel(
-      body.chatModel.providerId,
-      body.chatModel.key,
-    );
-
     const suggestions = await generateSuggestions(
       {
-        chat_history: chatHistory,
+        chatHistory,
       },
       llm,
     );

src/lib/agents/suggestions/index.ts

@@ -1,55 +1,32 @@
-import { RunnableSequence, RunnableMap } from '@langchain/core/runnables';
 import ListLineOutputParser from '@/lib/outputParsers/listLineOutputParser';
-import { PromptTemplate } from '@langchain/core/prompts';
+import { ChatPromptTemplate, PromptTemplate } from '@langchain/core/prompts';
 import formatChatHistoryAsString from '@/lib/utils/formatHistory';
-import { BaseMessage } from '@langchain/core/messages';
+import { BaseMessage, HumanMessage, SystemMessage } from '@langchain/core/messages';
 import { BaseChatModel } from '@langchain/core/language_models/chat_models';
-import { ChatOpenAI } from '@langchain/openai';
+import { suggestionGeneratorPrompt } from '@/lib/prompts/suggestions';
-
-const suggestionGeneratorPrompt = `
-You are an AI suggestion generator for an AI powered search engine. You will be given a conversation below. You need to generate 4-5 suggestions based on the conversation. The suggestion should be relevant to the conversation that can be used by the user to ask the chat model for more information.
-You need to make sure the suggestions are relevant to the conversation and are helpful to the user. Keep a note that the user might use these suggestions to ask a chat model for more information. 
-Make sure the suggestions are medium in length and are informative and relevant to the conversation.
-
-Provide these suggestions separated by newlines between the XML tags <suggestions> and </suggestions>. For example:
-
-<suggestions>
-Tell me more about SpaceX and their recent projects
-What is the latest news on SpaceX?
-Who is the CEO of SpaceX?
-</suggestions>
-
-Conversation:
-{chat_history}
-`;
 
 type SuggestionGeneratorInput = {
-  chat_history: BaseMessage[];
+  chatHistory: BaseMessage[];
 };
 
 const outputParser = new ListLineOutputParser({
   key: 'suggestions',
 });
 
-const createSuggestionGeneratorChain = (llm: BaseChatModel) => {
+const generateSuggestions = async (
-  return RunnableSequence.from([
-    RunnableMap.from({
-      chat_history: (input: SuggestionGeneratorInput) =>
-        formatChatHistoryAsString(input.chat_history),
-    }),
-    PromptTemplate.fromTemplate(suggestionGeneratorPrompt),
-    llm,
-    outputParser,
-  ]);
-};
-
-const generateSuggestions = (
   input: SuggestionGeneratorInput,
   llm: BaseChatModel,
 ) => {
-  (llm as unknown as ChatOpenAI).temperature = 0;
+  const chatPrompt = await ChatPromptTemplate.fromMessages([
-  const suggestionGeneratorChain = createSuggestionGeneratorChain(llm);
+    new SystemMessage(suggestionGeneratorPrompt),
-  return suggestionGeneratorChain.invoke(input);
+    new HumanMessage(`<conversation>${formatChatHistoryAsString(input.chatHistory)}</conversation>`)
+  ]).formatMessages({})
+
+  const res = await llm.invoke(chatPrompt)
+
+  const suggestions = await outputParser.invoke(res)
+
+  return suggestions
 };
 
 export default generateSuggestions;

src/lib/prompts/suggestions/index.ts

@@ -0,0 +1,15 @@
+export const suggestionGeneratorPrompt = `
+You are an AI suggestion generator for an AI powered search engine. You will be given a conversation below. You need to generate 4-5 suggestions based on the conversation. The suggestion should be relevant to the conversation that can be used by the user to ask the chat model for more information.
+You need to make sure the suggestions are relevant to the conversation and are helpful to the user. Keep a note that the user might use these suggestions to ask a chat model for more information. 
+Make sure the suggestions are medium in length and are informative and relevant to the conversation.
+
+Provide these suggestions separated by newlines between the XML tags <suggestions> and </suggestions>. For example:
+
+<suggestions>
+Tell me more about SpaceX and their recent projects
+What is the latest news on SpaceX?
+Who is the CEO of SpaceX?
+</suggestions>
+
+Today's date is ${new Date().toISOString()}
+`;