diff --git a/src/app/api/videos/route.ts b/src/app/api/videos/route.ts index 0ace57f..1417226 100644 --- a/src/app/api/videos/route.ts +++ b/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, diff --git a/src/lib/agents/media/video.ts b/src/lib/agents/media/video.ts index 1bf7f75..60fc04f 100644 --- a/src/lib/agents/media/video.ts +++ b/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 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', - '\n\n\nHow does a car work?\n', - ], - ['assistant', 'How does a car work?'], - [ - 'user', - '\n\n\nWhat is the theory of relativity?\n', - ], - ['assistant', 'Theory of relativity'], - [ - 'user', - '\n\n\nHow does an AC work?\n', - ], - ['assistant', 'AC working'], - [ - 'user', - '{chat_history}\n\n{query}\n', - ], - ]), - 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(`${formatChatHistoryAsString(input.chatHistory)}\n\n\n${input.query}\n`) + ]).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; diff --git a/src/lib/prompts/media/videos.ts b/src/lib/prompts/media/videos.ts new file mode 100644 index 0000000..b4a0d55 --- /dev/null +++ b/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 element. Do not include any explanation or additional text. +`; + +export const videoSearchFewShots: BaseMessageLike[] = [ + [ + 'user', + '\n\n\nHow does a car work?\n', + ], + ['assistant', 'How does a car work?'], + [ + 'user', + '\n\n\nWhat is the theory of relativity?\n', + ], + ['assistant', 'Theory of relativity'], + [ + 'user', + '\n\n\nHow does an AC work?\n', + ], + ['assistant', 'AC working'], +] \ No newline at end of file