feat(app): migrate suggestion chain

This commit is contained in:
ItzCrazyKns
2025-11-13 11:47:28 +05:30
parent 1a8889c71c
commit 41fe009847
3 changed files with 38 additions and 47 deletions

View File

@@ -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,
);

View File

@@ -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 <suggestions> and </suggestions>. For example:
<suggestions>
Tell me more about SpaceX and their recent projects
What is the latest news on SpaceX?
Who is the CEO of SpaceX?
</suggestions>
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(`<conversation>${formatChatHistoryAsString(input.chatHistory)}</conversation>`)
]).formatMessages({})
const res = await llm.invoke(chatPrompt)
const suggestions = await outputParser.invoke(res)
return suggestions
};
export default generateSuggestions;

View File

@@ -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 <suggestions> and </suggestions>. For example:
<suggestions>
Tell me more about SpaceX and their recent projects
What is the latest news on SpaceX?
Who is the CEO of SpaceX?
</suggestions>
Today's date is ${new Date().toISOString()}
`;