diff --git a/src/app/api/suggestions/route.ts b/src/app/api/suggestions/route.ts index f942758..2dc7248 100644 --- a/src/app/api/suggestions/route.ts +++ b/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, ); diff --git a/src/lib/agents/suggestions/index.ts b/src/lib/agents/suggestions/index.ts index 3beb1bd..03302ac 100644 --- a/src/lib/agents/suggestions/index.ts +++ b/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'; - -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 and . For example: - - -Tell me more about SpaceX and their recent projects -What is the latest news on SpaceX? -Who is the CEO of SpaceX? - - -Conversation: -{chat_history} -`; +import { suggestionGeneratorPrompt } from '@/lib/prompts/suggestions'; type SuggestionGeneratorInput = { - chat_history: BaseMessage[]; + chatHistory: BaseMessage[]; }; const outputParser = new ListLineOutputParser({ key: 'suggestions', }); -const createSuggestionGeneratorChain = (llm: BaseChatModel) => { - return RunnableSequence.from([ - RunnableMap.from({ - chat_history: (input: SuggestionGeneratorInput) => - formatChatHistoryAsString(input.chat_history), - }), - PromptTemplate.fromTemplate(suggestionGeneratorPrompt), - llm, - outputParser, - ]); -}; - -const generateSuggestions = ( +const generateSuggestions = async ( input: SuggestionGeneratorInput, llm: BaseChatModel, ) => { - (llm as unknown as ChatOpenAI).temperature = 0; - const suggestionGeneratorChain = createSuggestionGeneratorChain(llm); - return suggestionGeneratorChain.invoke(input); + const chatPrompt = await ChatPromptTemplate.fromMessages([ + new SystemMessage(suggestionGeneratorPrompt), + new HumanMessage(`${formatChatHistoryAsString(input.chatHistory)}`) + ]).formatMessages({}) + + const res = await llm.invoke(chatPrompt) + + const suggestions = await outputParser.invoke(res) + + return suggestions }; export default generateSuggestions; diff --git a/src/lib/prompts/suggestions/index.ts b/src/lib/prompts/suggestions/index.ts new file mode 100644 index 0000000..daa99d4 --- /dev/null +++ b/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 and . For example: + + +Tell me more about SpaceX and their recent projects +What is the latest news on SpaceX? +Who is the CEO of SpaceX? + + +Today's date is ${new Date().toISOString()} +`; \ No newline at end of file