Merge pull request #31 from ChiragAgg5k/chore-improve-system-prompt

chore: shifted system prompt to prefixMessages

Author: loks0n

src/documents.js

@@ -3,6 +3,9 @@ import { readFile } from "fs/promises";
 import { Document } from "langchain/document";
 import { MarkdownTextSplitter } from "langchain/text_splitter";
 
+/**
+ * @returns {Promise<Document[]>}
+ */
 const getDocumentation = async () => {
   const filenames = await glob([
     "./sources/website/src/routes/docs/**/*.markdoc",
@@ -36,6 +39,9 @@ const getDocumentation = async () => {
   );
 };
 
+/**
+ * @returns {Promise<Document[]>} Array of Document objects containing processed references
+ */
 const getReferences = async () => {
   const filenames = await glob(["./sources/references/**/*.md"]);
 
@@ -64,7 +70,7 @@ export const getDocuments = async () => {
   return await splitDocuments([...documentation, ...references]);
 };
 
-/**x
+/**
  * @param {Document[]} documents
  * @returns {Promise<Document<Record<string, any>>[]>}
  */
@@ -80,6 +86,10 @@ async function splitDocuments(documents) {
   return await splitter.createDocuments(texts, metadatas);
 }
 
+/**
+ * @param {string} contents
+ * @returns {Object.<string, string>}
+ */
 function parseMarkdownFrontmatter(contents) {
   const raw = contents.match(/^---\n([\s\S]*?)\n---/);
   if (!raw) {
@@ -94,6 +104,10 @@ function parseMarkdownFrontmatter(contents) {
   return frontmatter;
 }
 
+/**
+ * @param {string} filename
+ * @returns {{sdk: string, service: string}}
+ */
 function parseReferenceData(filename) {
   const [sdk, service] = filename
     .replace("sources/references/", "")

src/embeddings.js

@@ -4,6 +4,9 @@ import { OpenAIChat } from "langchain/llms/openai";
 import { loadQAStuffChain } from "langchain/chains";
 import { getDocuments } from "./documents.js";
 
+/**
+ * @returns {Promise<VectorStoreRetriever<HNSWLib>>}
+ */
 export const intializeDocumentRetriever = async () => {
   const embeddings = new OpenAIEmbeddings({
     openAIApiKey: process.env._APP_ASSISTANT_OPENAI_API_KEY,
@@ -15,7 +18,11 @@ export const intializeDocumentRetriever = async () => {
   return vectorStore.asRetriever(5);
 };
 
-export const getOpenAIChat = async (onToken) =>
+/**
+ * @param {function} onToken
+ * @param {string} systemPrompt
+ */
+export const getOpenAIChat = async (onToken, systemPrompt) =>
   new OpenAIChat({
     modelName: "gpt-4o",
     openAIApiKey: process.env._APP_ASSISTANT_OPENAI_API_KEY,
@@ -27,8 +34,18 @@ export const getOpenAIChat = async (onToken) =>
         handleLLMNewToken: onToken,
       },
     ],
+    prefixMessages: [
+      {
+        role: "system",
+        content: systemPrompt,
+      },
+    ],
   });
 
-export const getRagChain = async (onToken) => {
-  return loadQAStuffChain(await getOpenAIChat(onToken));
+/**
+ * @param {function} onToken
+ * @param {string} systemPrompt
+ */
+export const getRagChain = async (onToken, systemPrompt) => {
+  return loadQAStuffChain(await getOpenAIChat(onToken, systemPrompt));
 };

src/main.js

@@ -30,7 +30,6 @@ app.post("/v1/models/assistant/prompt", async (req, res) => {
     return;
   }
 
-  // raw to text
   const decoder = new TextDecoder();
   const text = decoder.decode(req.body);
 
@@ -41,11 +40,11 @@ app.post("/v1/models/assistant/prompt", async (req, res) => {
 
   const chain = await getRagChain((token) => {
     res.write(token);
-  });
+  }, systemPrompt);
 
   await chain.call({
     input_documents: relevantDocuments,
-    question: `${systemPrompt}\n\n${prompt}`,
+    question: prompt,
   });
 
   const sources = new Set(
@@ -71,14 +70,11 @@ app.post("/v1/models/generic/prompt", async (req, res) => {
   let { prompt, systemPrompt } = JSON.parse(text);
   systemPrompt ??= SYSTEM_PROMPT;
 
-  const chain = await getOpenAIChat((token) => {
+  const chat = await getOpenAIChat((token) => {
     res.write(token);
-  });
+  }, systemPrompt);
 
-  await chain.invoke([
-    ["system", systemPrompt],
-    ["human", prompt],
-  ])
+  await chat.call(prompt);
 
   res.end();
 });