feat(app): migrate image search chain

src/app/api/images/route.ts

@@ -1,4 +1,4 @@
-import handleImageSearch from '@/lib/agents/media/image';
+import searchImages from '@/lib/agents/media/image';
 import ModelRegistry from '@/lib/models/registry';
 import { ModelWithProvider } from '@/lib/models/types';
 import { AIMessage, BaseMessage, HumanMessage } from '@langchain/core/messages';
@@ -13,6 +13,13 @@ export const POST = async (req: Request) => {
   try {
     const body: ImageSearchBody = 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 images = await handleImageSearch(
+    const images = await searchImages(
       {
-        chat_history: chatHistory,
+        chatHistory: chatHistory,
         query: body.query,
       },
       llm,

src/lib/agents/media/image.ts

@@ -7,101 +7,59 @@ import {
 } from '@langchain/core/runnables';
 import { ChatPromptTemplate } 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 { StringOutputParser } from '@langchain/core/output_parsers';
 import { searchSearxng } from '@/lib/searxng';
 import type { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import LineOutputParser from '@/lib/outputParsers/lineOutputParser';
-
-const imageSearchChainPrompt = `
-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 the web for images.
-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 { imageSearchFewShots, imageSearchPrompt } from '@/lib/prompts/media/image';
 
 type ImageSearchChainInput = {
-  chat_history: BaseMessage[];
+  chatHistory: BaseMessage[];
   query: string;
 };
 
-interface ImageSearchResult {
+type ImageSearchResult = {
   img_src: string;
   url: string;
   title: string;
 }
 
-const strParser = new StringOutputParser();
+const outputParser = new LineOutputParser({
+  key: 'query',
+})
 
-const createImageSearchChain = (llm: BaseChatModel) => {
-  return RunnableSequence.from([
-    RunnableMap.from({
-      chat_history: (input: ImageSearchChainInput) => {
-        return formatChatHistoryAsString(input.chat_history);
-      },
-      query: (input: ImageSearchChainInput) => {
-        return input.query;
-      },
-    }),
-    ChatPromptTemplate.fromMessages([
-      ['system', imageSearchChainPrompt],
-      [
-        'user',
-        '<conversation>\n</conversation>\n<follow_up>\nWhat is a cat?\n</follow_up>',
-      ],
-      ['assistant', '<query>A cat</query>'],
-
-      [
-        'user',
-        '<conversation>\n</conversation>\n<follow_up>\nWhat is a car? How does it work?\n</follow_up>',
-      ],
-      ['assistant', '<query>Car working</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: ['bing images', 'google images'],
-      });
-
-      const images: ImageSearchResult[] = [];
-
-      res.results.forEach((result) => {
-        if (result.img_src && result.url && result.title) {
-          images.push({
-            img_src: result.img_src,
-            url: result.url,
-            title: result.title,
-          });
-        }
-      });
-
-      return images.slice(0, 10);
-    }),
-  ]);
-};
-
-const handleImageSearch = (
+const searchImages = async (
   input: ImageSearchChainInput,
   llm: BaseChatModel,
 ) => {
-  const imageSearchChain = createImageSearchChain(llm);
-  return imageSearchChain.invoke(input);
+  const chatPrompt = await ChatPromptTemplate.fromMessages([
+    new SystemMessage(imageSearchPrompt),
+    ...imageSearchFewShots,
+    new HumanMessage(`<conversation>\n${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: ['bing images', 'google images'],
+  });
+
+  const images: ImageSearchResult[] = [];
+
+  searchRes.results.forEach((result) => {
+    if (result.img_src && result.url && result.title) {
+      images.push({
+        img_src: result.img_src,
+        url: result.url,
+        title: result.title,
+      });
+    }
+  });
+
+  return images.slice(0, 10);
 };
 
-export default handleImageSearch;
+export default searchImages;

src/lib/prompts/media/image.ts

@@ -0,0 +1,26 @@
+import { BaseMessageLike } from "@langchain/core/messages";
+
+export const imageSearchPrompt = `
+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 the web for images.
+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 imageSearchFewShots: BaseMessageLike[] = [
+    [
+        'user',
+        '<conversation>\n</conversation>\n<follow_up>\nWhat is a cat?\n</follow_up>',
+    ],
+    ['assistant', '<query>A cat</query>'],
+
+    [
+        'user',
+        '<conversation>\n</conversation>\n<follow_up>\nWhat is a car? How does it work?\n</follow_up>',
+    ],
+    ['assistant', '<query>Car working</query>'],
+    [
+        'user',
+        '<conversation>\n</conversation>\n<follow_up>\nHow does an AC work?\n</follow_up>',
+    ],
+    ['assistant', '<query>AC working</query>']
+]