From 7062e637e86589facac93221e0ae90d4bb0b2bf8 Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Thu, 21 Nov 2024 18:26:38 -0800
Subject: [PATCH] refac: inline citations

---
 backend/open_webui/config.py                     | 10 ++++++----
 backend/open_webui/main.py                       | 16 ++++++++++++++--
 .../components/chat/Messages/Citations.svelte    |  4 ++++
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py
index 57b9fbcfd..3b12b2f94 100644
--- a/backend/open_webui/config.py
+++ b/backend/open_webui/config.py
@@ -1182,7 +1182,7 @@ CHUNK_OVERLAP = PersistentConfig(
 )
 
 DEFAULT_RAG_TEMPLATE = """### Task:
-Respond to the user query using the provided context, incorporating inline citations in the format [source_id].
+Respond to the user query using the provided context, incorporating inline citations in the format [source_id] **only when the <source_id> tag is explicitly provided** in the context.
 
 ### Guidelines:
 - If you don't know the answer, clearly state that.
@@ -1190,16 +1190,18 @@ Respond to the user query using the provided context, incorporating inline citat
 - Respond in the same language as the user's query.
 - If the context is unreadable or of poor quality, inform the user and provide the best possible answer.
 - If the answer isn't present in the context but you possess the knowledge, explain this to the user and provide the answer using your own understanding.
-- Include inline citations using [source_id] corresponding to the sources listed in the context.
+- **Only include inline citations using [source_id] when a <source_id> tag is explicitly provided in the context.**  
+- Do not cite if the <source_id> tag is not provided in the context.  
 - Do not use XML tags in your response.
 - Ensure citations are concise and directly related to the information provided.
 
 ### Example of Citation:
-If the user asks about a specific topic and the information is found in "whitepaper.pdf", the response should include the citation like so:  
+If the user asks about a specific topic and the information is found in "whitepaper.pdf" with a provided <source_id>, the response should include the citation like so:  
 * "According to the study, the proposed method increases efficiency by 20% [whitepaper.pdf]."
+If no <source_id> is present, the response should omit the citation.
 
 ### Output:
-Provide a clear and direct response to the user's query, including inline citations in the format [source_id] where relevant.
+Provide a clear and direct response to the user's query, including inline citations in the format [source_id] only when the <source_id> tag is present in the context.
 
 <context>
 {{CONTEXT}}
diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py
index 0279cee05..d62ef158e 100644
--- a/backend/open_webui/main.py
+++ b/backend/open_webui/main.py
@@ -463,6 +463,8 @@ async def chat_completion_tools_handler(
             except Exception as e:
                 tool_output = str(e)
 
+            print(tools[tool_function_name]["citation"])
+
             if tools[tool_function_name]["citation"]:
                 citations.append(
                     {
@@ -473,6 +475,9 @@ async def chat_completion_tools_handler(
                         "metadata": [{"source": tool_function_name}],
                     }
                 )
+            else:
+                citations.append({})
+
             if tools[tool_function_name]["file_handler"]:
                 skip_files = True
 
@@ -485,7 +490,7 @@ async def chat_completion_tools_handler(
         log.exception(f"Error: {e}")
         content = None
 
-    log.debug(f"tool_contexts: {contexts}")
+    log.debug(f"tool_contexts: {contexts} {citations}")
 
     if skip_files and "files" in body.get("metadata", {}):
         del body["metadata"]["files"]
@@ -683,7 +688,14 @@ class ChatCompletionMiddleware(BaseHTTPMiddleware):
             for context_idx, context in enumerate(contexts):
                 print(context)
                 source_id = citations[context_idx].get("source", {}).get("name", "")
-                context_string += f"<source><source_id>{source_id}</source_id><source_context>{context}</source_context></source>\n"
+
+                print(f"\n\n\n\n{source_id}\n\n\n\n")
+                if source_id:
+                    context_string += f"<source><source_id>{source_id}</source_id><source_context>{context}</source_context></source>\n"
+                else:
+                    context_string += (
+                        f"<source><source_context>{context}</source_context></source>\n"
+                    )
 
             context_string = context_string.strip()
             prompt = get_last_user_message(body["messages"])
diff --git a/src/lib/components/chat/Messages/Citations.svelte b/src/lib/components/chat/Messages/Citations.svelte
index 80d0e0279..1e24518cd 100644
--- a/src/lib/components/chat/Messages/Citations.svelte
+++ b/src/lib/components/chat/Messages/Citations.svelte
@@ -43,6 +43,10 @@
 
 	$: {
 		_citations = citations.reduce((acc, citation) => {
+			if (Object.keys(citation).length === 0) {
+				return acc;
+			}
+
 			citation.document.forEach((document, index) => {
 				const metadata = citation.metadata?.[index];
 				const distance = citation.distances?.[index];