diff --git a/src/documents.js b/src/documents.js index 1bdfdde..23d5ff5 100644 --- a/src/documents.js +++ b/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} + */ const getDocumentation = async () => { const filenames = await glob([ "./sources/website/src/routes/docs/**/*.markdoc", @@ -36,6 +39,9 @@ const getDocumentation = async () => { ); }; +/** + * @returns {Promise} 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>[]>} */ @@ -80,6 +86,10 @@ async function splitDocuments(documents) { return await splitter.createDocuments(texts, metadatas); } +/** + * @param {string} contents + * @returns {Object.} + */ 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/", "") diff --git a/src/embeddings.js b/src/embeddings.js index 78c6625..efd9fb3 100644 --- a/src/embeddings.js +++ b/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>} + */ 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)); }; diff --git a/src/main.js b/src/main.js index a0d4c77..6f0bd13 100644 --- a/src/main.js +++ b/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(); });