feat: implement citation for TS (#257)

.changeset/wise-adults-behave.md

+---
+"create-llama": patch
+---
+
+feat: implement citation for TS

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.

templates/components/engines/typescript/chat/chat.ts

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

templates/components/engines/typescript/chat/nodePostprocessors.ts

+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();