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