feat(app): migrate video search chain

src/app/api/videos/route.ts

@@ -13,6 +13,13 @@ export const POST = async (req: Request) => {
   try {
     const body: VideoSearchBody = 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 +30,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 videos = await handleVideoSearch(
       {
-        chat_history: chatHistory,
+        chatHistory: chatHistory,
         query: body.query,
       },
       llm,

src/lib/agents/media/video.ts

@@ -1,110 +1,65 @@
-import {
-  RunnableSequence,
-  RunnableMap,
-  RunnableLambda,
-} from '@langchain/core/runnables';
 import { ChatPromptTemplate } from '@langchain/core/prompts';
 import formatChatHistoryAsString from '@/lib/utils/formatHistory';
-import { BaseMessage } from '@langchain/core/messages';
-import { StringOutputParser } from '@langchain/core/output_parsers';
+import { BaseMessage, HumanMessage, SystemMessage } from '@langchain/core/messages';
 import { searchSearxng } from '@/lib/searxng';
 import type { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import LineOutputParser from '@/lib/outputParsers/lineOutputParser';
-
-const videoSearchChainPrompt = `
-You will be given a conversation below and a follow up question. You need to rephrase the follow-up question so it is a standalone question that can be used by the LLM to search Youtube for videos.
-You need to make sure the rephrased question agrees with the conversation and is relevant to the conversation.
-Output only the rephrased query wrapped in an XML <query> element. Do not include any explanation or additional text.
-`;
+import { videoSearchFewShots, videoSearchPrompt } from '@/lib/prompts/media/videos';
 
 type VideoSearchChainInput = {
-  chat_history: BaseMessage[];
+  chatHistory: BaseMessage[];
   query: string;
 };
 
-interface VideoSearchResult {
+type VideoSearchResult = {
   img_src: string;
   url: string;
   title: string;
   iframe_src: string;
 }
 
-const strParser = new StringOutputParser();
+const outputParser = new LineOutputParser({
+  key: 'query',
+});
 
-const createVideoSearchChain = (llm: BaseChatModel) => {
-  return RunnableSequence.from([
-    RunnableMap.from({
-      chat_history: (input: VideoSearchChainInput) => {
-        return formatChatHistoryAsString(input.chat_history);
-      },
-      query: (input: VideoSearchChainInput) => {
-        return input.query;
-      },
-    }),
-    ChatPromptTemplate.fromMessages([
-      ['system', videoSearchChainPrompt],
-      [
-        'user',
-        '<conversation>\n</conversation>\n<follow_up>\nHow does a car work?\n</follow_up>',
-      ],
-      ['assistant', '<query>How does a car work?</query>'],
-      [
-        'user',
-        '<conversation>\n</conversation>\n<follow_up>\nWhat is the theory of relativity?\n</follow_up>',
-      ],
-      ['assistant', '<query>Theory of relativity</query>'],
-      [
-        'user',
-        '<conversation>\n</conversation>\n<follow_up>\nHow does an AC work?\n</follow_up>',
-      ],
-      ['assistant', '<query>AC working</query>'],
-      [
-        'user',
-        '<conversation>{chat_history}</conversation>\n<follow_up>\n{query}\n</follow_up>',
-      ],
-    ]),
-    llm,
-    strParser,
-    RunnableLambda.from(async (input: string) => {
-      const queryParser = new LineOutputParser({
-        key: 'query',
-      });
-      return await queryParser.parse(input);
-    }),
-    RunnableLambda.from(async (input: string) => {
-      const res = await searchSearxng(input, {
-        engines: ['youtube'],
-      });
-
-      const videos: VideoSearchResult[] = [];
-
-      res.results.forEach((result) => {
-        if (
-          result.thumbnail &&
-          result.url &&
-          result.title &&
-          result.iframe_src
-        ) {
-          videos.push({
-            img_src: result.thumbnail,
-            url: result.url,
-            title: result.title,
-            iframe_src: result.iframe_src,
-          });
-        }
-      });
-
-      return videos.slice(0, 10);
-    }),
-  ]);
-};
-
-const handleVideoSearch = (
+const searchVideos = async (
   input: VideoSearchChainInput,
   llm: BaseChatModel,
 ) => {
-  const videoSearchChain = createVideoSearchChain(llm);
-  return videoSearchChain.invoke(input);
+  const chatPrompt = await ChatPromptTemplate.fromMessages([
+    new SystemMessage(videoSearchPrompt),
+    ...videoSearchFewShots,
+    new HumanMessage(`<conversation>${formatChatHistoryAsString(input.chatHistory)}\n</conversation>\n<follow_up>\n${input.query}\n</follow_up>`)
+  ]).formatMessages({})
+
+  const res = await llm.invoke(chatPrompt)
+
+  const query = await outputParser.invoke(res)
+
+  const searchRes = await searchSearxng(query!, {
+    engines: ['youtube'],
+  });
+
+  const videos: VideoSearchResult[] = [];
+
+  searchRes.results.forEach((result) => {
+    if (
+      result.thumbnail &&
+      result.url &&
+      result.title &&
+      result.iframe_src
+    ) {
+      videos.push({
+        img_src: result.thumbnail,
+        url: result.url,
+        title: result.title,
+        iframe_src: result.iframe_src,
+      });
+    }
+  });
+
+  return videos.slice(0, 10);
+
 };
 
-export default handleVideoSearch;
+export default searchVideos;

src/lib/prompts/media/videos.ts

@@ -0,0 +1,25 @@
+import { BaseMessageLike } from "@langchain/core/messages";
+
+export const videoSearchPrompt = `
+You will be given a conversation below and a follow up question. You need to rephrase the follow-up question so it is a standalone question that can be used by the LLM to search Youtube for videos.
+You need to make sure the rephrased question agrees with the conversation and is relevant to the conversation.
+Output only the rephrased query wrapped in an XML <query> element. Do not include any explanation or additional text.
+`;
+
+export const videoSearchFewShots: BaseMessageLike[] = [
+    [
+        'user',
+        '<conversation>\n</conversation>\n<follow_up>\nHow does a car work?\n</follow_up>',
+    ],
+    ['assistant', '<query>How does a car work?</query>'],
+    [
+        'user',
+        '<conversation>\n</conversation>\n<follow_up>\nWhat is the theory of relativity?\n</follow_up>',
+    ],
+    ['assistant', '<query>Theory of relativity</query>'],
+    [
+        'user',
+        '<conversation>\n</conversation>\n<follow_up>\nHow does an AC work?\n</follow_up>',
+    ],
+    ['assistant', '<query>AC working</query>'],
+]