diff --git a/.changeset/wise-adults-behave.md b/.changeset/wise-adults-behave.md new file mode 100644 index 0000000000000000000000000000000000000000..1c0a1ba53504b424f80892d7ce153ebe85030967 --- /dev/null +++ b/.changeset/wise-adults-behave.md @@ -0,0 +1,5 @@ +--- +"create-llama": patch +--- + +feat: implement citation for TS diff --git a/helpers/env-variables.ts b/helpers/env-variables.ts index 77b60ac46aee36e8e0fdb120500c8f376005641b..783eb43e8d9c2de42dfd644036ed91b8ac07fe77 100644 --- a/helpers/env-variables.ts +++ b/helpers/env-variables.ts @@ -454,12 +454,7 @@ const getSystemPromptEnv = ( }, ]; - // Citation only works with FastAPI along with the chat engine and data source provided for now. - if ( - framework === "fastapi" && - tools?.length == 0 && - (dataSources?.length ?? 0 > 0) - ) { + if (tools?.length == 0 && (dataSources?.length ?? 0 > 0)) { const citationPrompt = `'You have provided information from a knowledge base that has been passed to you in nodes of information. Each node has useful metadata such as node ID, file name, page, etc. Please add the citation to the data node for each sentence or paragraph that you reference in the provided information. diff --git a/templates/components/engines/typescript/chat/chat.ts b/templates/components/engines/typescript/chat/chat.ts index c0841aa5bdd1cc609303322cdfdf633dfea3df7b..fa2e9edc4583f19701f0813449ca4087c8886581 100644 --- a/templates/components/engines/typescript/chat/chat.ts +++ b/templates/components/engines/typescript/chat/chat.ts @@ -1,5 +1,6 @@ import { ContextChatEngine, Settings } from "llamaindex"; import { getDataSource } from "./index"; +import { nodeCitationProcessor } from "./nodePostprocessors"; import { generateFilters } from "./queryFilter"; export async function createChatEngine(documentIds?: string[], params?: any) { @@ -14,9 +15,18 @@ export async function createChatEngine(documentIds?: string[], params?: any) { filters: generateFilters(documentIds || []), }); + const systemPrompt = process.env.SYSTEM_PROMPT; + const citationPrompt = process.env.SYSTEM_CITATION_PROMPT; + const prompt = + [systemPrompt, citationPrompt].filter((p) => p).join("\n") || undefined; + const nodePostprocessors = citationPrompt + ? [nodeCitationProcessor] + : undefined; + return new ContextChatEngine({ chatModel: Settings.llm, retriever, - systemPrompt: process.env.SYSTEM_PROMPT, + systemPrompt: prompt, + nodePostprocessors, }); } diff --git a/templates/components/engines/typescript/chat/nodePostprocessors.ts b/templates/components/engines/typescript/chat/nodePostprocessors.ts new file mode 100644 index 0000000000000000000000000000000000000000..66065f5e967af5c51854ea20a1a3c7f2dd98b6f9 --- /dev/null +++ b/templates/components/engines/typescript/chat/nodePostprocessors.ts @@ -0,0 +1,26 @@ +import { + BaseNodePostprocessor, + MessageContent, + NodeWithScore, +} from "llamaindex"; + +class NodeCitationProcessor implements BaseNodePostprocessor { + /** + * Append node_id into metadata for citation purpose. + * Config SYSTEM_CITATION_PROMPT in your runtime environment variable to enable this feature. + */ + async postprocessNodes( + nodes: NodeWithScore[], + query?: MessageContent, + ): Promise<NodeWithScore[]> { + for (const nodeScore of nodes) { + if (!nodeScore.node || !nodeScore.node.metadata) { + continue; // Skip nodes with missing properties + } + nodeScore.node.metadata["node_id"] = nodeScore.node.id_; + } + return nodes; + } +} + +export const nodeCitationProcessor = new NodeCitationProcessor();