feat(agents): update suggestion generator

src/app/api/suggestions/route.ts

@@ -19,19 +19,9 @@ export const POST = async (req: Request) => {
       body.chatModel.key,
     );
 
-    const chatHistory = body.chatHistory
-      .map((msg: any) => {
-        if (msg.role === 'user') {
-          return new HumanMessage(msg.content);
-        } else if (msg.role === 'assistant') {
-          return new AIMessage(msg.content);
-        }
-      })
-      .filter((msg) => msg !== undefined) as BaseMessage[];
-
     const suggestions = await generateSuggestions(
       {
-        chatHistory,
+        chatHistory: body.chatHistory,
       },
       llm,
     );

src/lib/agents/suggestions/index.ts

@@ -1,32 +1,39 @@
-import ListLineOutputParser from '@/lib/outputParsers/listLineOutputParser';
-import { ChatPromptTemplate, PromptTemplate } from '@langchain/core/prompts';
 import formatChatHistoryAsString from '@/lib/utils/formatHistory';
-import { BaseMessage, HumanMessage, SystemMessage } from '@langchain/core/messages';
-import { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import { suggestionGeneratorPrompt } from '@/lib/prompts/suggestions';
+import { ChatTurnMessage } from '@/lib/types';
+import z from 'zod';
+import BaseLLM from '@/lib/models/base/llm';
+import { i } from 'mathjs';
 
 type SuggestionGeneratorInput = {
-  chatHistory: BaseMessage[];
+  chatHistory: ChatTurnMessage[];
 };
 
-const outputParser = new ListLineOutputParser({
+const schema = z.object({
-  key: 'suggestions',
+  suggestions: z
+    .array(z.string())
+    .describe('List of suggested questions or prompts'),
 });
 
 const generateSuggestions = async (
   input: SuggestionGeneratorInput,
-  llm: BaseChatModel,
+  llm: BaseLLM<any>,
 ) => {
-  const chatPrompt = await ChatPromptTemplate.fromMessages([
+  const res = await llm.generateObject<z.infer<typeof schema>>({
-    new SystemMessage(suggestionGeneratorPrompt),
+    messages: [
-    new HumanMessage(`<conversation>${formatChatHistoryAsString(input.chatHistory)}</conversation>`)
+      {
-  ]).formatMessages({})
+        role: 'system',
+        content: suggestionGeneratorPrompt,
+      },
+      {
+        role: 'user',
+        content: `<chat_history>\n${formatChatHistoryAsString(input.chatHistory)}\n</chat_history>`,
+      },
+    ],
+    schema,
+  });
 
-  const res = await llm.invoke(chatPrompt)
+  return res.suggestions;
-
-  const suggestions = await outputParser.invoke(res)
-
-  return suggestions
 };
 
 export default generateSuggestions;

src/lib/prompts/suggestions/index.ts

@@ -3,13 +3,15 @@ You are an AI suggestion generator for an AI powered search engine. You will be
 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:
+Sample suggestions for a conversation about Elon Musk:
-
+{
-<suggestions>
+    "suggestions": [
-Tell me more about SpaceX and their recent projects
+        "What are Elon Musk's plans for SpaceX in the next decade?",
-What is the latest news on SpaceX?
+        "How has Tesla's stock performance been influenced by Elon Musk's leadership?",
-Who is the CEO of SpaceX?
+        "What are the key innovations introduced by Elon Musk in the electric vehicle industry?",
-</suggestions>
+        "How does Elon Musk's vision for renewable energy impact global sustainability efforts?"
+    ]
+}
 
 Today's date is ${new Date().toISOString()}
-`;
+`;