From 4fa2b76f3d79d8d6efeab6a29d4438b1df0c0a36 Mon Sep 17 00:00:00 2001
From: Thuc Pham <51660321+thucpn@users.noreply.github.com>
Date: Wed, 28 Aug 2024 16:47:00 +0700
Subject: [PATCH] feat: implement citation for TS (#257)
---
.changeset/wise-adults-behave.md | 5 ++++
helpers/env-variables.ts | 7 +----
.../engines/typescript/chat/chat.ts | 12 ++++++++-
.../typescript/chat/nodePostprocessors.ts | 26 +++++++++++++++++++
4 files changed, 43 insertions(+), 7 deletions(-)
create mode 100644 .changeset/wise-adults-behave.md
create mode 100644 templates/components/engines/typescript/chat/nodePostprocessors.ts
diff --git a/.changeset/wise-adults-behave.md b/.changeset/wise-adults-behave.md
new file mode 100644
index 00000000..1c0a1ba5
--- /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 77b60ac4..783eb43e 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 c0841aa5..fa2e9edc 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 00000000..66065f5e
--- /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();
--
GitLab