Skip to content

GitLab

Commit 0872bea7 authored by Jun Siang Cheah's avatar Jun Siang Cheah Committed by Timothy J. Baek
Browse files

feat: show RAG query results as citations

parent ba09fcd5
Show whitespace changes
Inline Side-by-side
Showing with 234 additions and 42 deletions
backend/apps/rag/utils.py
View file @ 0872bea7
......@@ -320,11 +320,19 @@ def rag_messages(
extracted_collections.extend(collection)
context_string = ""
citations = []
for context in relevant_contexts:
try:
if "documents" in context:
items = [item for item in context["documents"][0] if item is not None]
context_string += "\n\n".join(items)
if "metadatas" in context:
citations.append(
{
"document": context["documents"][0],
"metadata": context["metadatas"][0],
}
)
except Exception as e:
log.exception(e)
context_string = context_string.strip()
......@@ -355,7 +363,7 @@ def rag_messages(
messages[last_user_message_idx] = new_user_message
return messages
return messages, citations
def get_model_path(model: str, update_model: bool = False):
......
backend/main.py
View file @ 0872bea7
......@@ -15,7 +15,7 @@ from fastapi.middleware.wsgi import WSGIMiddleware
from fastapi.middleware.cors import CORSMiddleware
from starlette.exceptions import HTTPException as StarletteHTTPException
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import StreamingResponse
from apps.ollama.main import app as ollama_app
from apps.openai.main import app as openai_app
......@@ -102,6 +102,8 @@ origins = ["*"]
class RAGMiddleware(BaseHTTPMiddleware):
async def dispatch(self, request: Request, call_next):
return_citations = False
if request.method == "POST" and (
"/api/chat" in request.url.path or "/chat/completions" in request.url.path
):
......@@ -114,11 +116,15 @@ class RAGMiddleware(BaseHTTPMiddleware):
# Parse string to JSON
data = json.loads(body_str) if body_str else {}
return_citations = data.get("citations", False)
if "citations" in data:
del data["citations"]
# Example: Add a new key-value pair or modify existing ones
# data["modified"] = True # Example modification
if "docs" in data:
data = {**data}
data["messages"] = rag_messages(
data["messages"], citations = rag_messages(
docs=data["docs"],
messages=data["messages"],
template=rag_app.state.RAG_TEMPLATE,
......@@ -130,7 +136,9 @@ class RAGMiddleware(BaseHTTPMiddleware):
)
del data["docs"]
log.debug(f"data['messages']: {data['messages']}")
log.debug(
f"data['messages']: {data['messages']}, citations: {citations}"
)
modified_body_bytes = json.dumps(data).encode("utf-8")
......@@ -148,11 +156,36 @@ class RAGMiddleware(BaseHTTPMiddleware):
]
response = await call_next(request)
if return_citations:
# Inject the citations into the response
if isinstance(response, StreamingResponse):
# If it's a streaming response, inject it as SSE event or NDJSON line
content_type = response.headers.get("Content-Type")
if "text/event-stream" in content_type:
return StreamingResponse(
self.openai_stream_wrapper(response.body_iterator, citations),
)
if "application/x-ndjson" in content_type:
return StreamingResponse(
self.ollama_stream_wrapper(response.body_iterator, citations),
)
return response
async def _receive(self, body: bytes):
return {"type": "http.request", "body": body, "more_body": False}
async def openai_stream_wrapper(self, original_generator, citations):
yield f"data: {json.dumps({'citations': citations})}\n\n"
async for data in original_generator:
yield data
async def ollama_stream_wrapper(self, original_generator, citations):
yield f"{json.dumps({'citations': citations})}\n"
async for data in original_generator:
yield data
app.add_middleware(RAGMiddleware)
......
src/lib/apis/streaming/index.ts
View file @ 0872bea7
......@@ -4,6 +4,8 @@ import type { ParsedEvent } from 'eventsource-parser';
type TextStreamUpdate = {
done: boolean;
value: string;
// eslint-disable-next-line @typescript-eslint/no-explicit-any
citations?: any;
};
// createOpenAITextStream takes a responseBody with a SSE response,
......@@ -45,6 +47,11 @@ async function* openAIStreamToIterator(
const parsedData = JSON.parse(data);
console.log(parsedData);
if (parsedData.citations) {
yield { done: false, value: '', citations: parsedData.citations };
continue;
}
yield { done: false, value: parsedData.choices?.[0]?.delta?.content ?? '' };
} catch (e) {
console.error('Error extracting delta from SSE event:', e);
......@@ -62,6 +69,10 @@ async function* streamLargeDeltasAsRandomChunks(
yield textStreamUpdate;
return;
}
if (textStreamUpdate.citations) {
yield textStreamUpdate;
continue;
}
let content = textStreamUpdate.value;
if (content.length < 5) {
yield { done: false, value: content };
......
src/lib/components/chat/Messages/CitationsModal.svelte 0 → 100644
View file @ 0872bea7
<script lang="ts">
import { getContext, onMount, tick } from 'svelte';
import Modal from '$lib/components/common/Modal.svelte';
const i18n = getContext('i18n');
export let show = false;
export let citation: any[];
let mergedDocuments = [];
onMount(async () => {
console.log(citation);
// Merge the document with its metadata
mergedDocuments = citation.document?.map((c, i) => {
return {
document: c,
metadata: citation.metadata?.[i]
};
});
console.log(mergedDocuments);
});
</script>
<Modal size="lg" bind:show>
<div>
<div class=" flex justify-between dark:text-gray-300 px-5 py-4">
<div class=" text-lg font-medium self-center capitalize">
{$i18n.t('Citation')}
</div>
<button
class="self-center"
on:click={() => {
show = false;
}}
>
<svg
xmlns="http://www.w3.org/2000/svg"
viewBox="0 0 20 20"
fill="currentColor"
class="w-5 h-5"
>
<path
d="M6.28 5.22a.75.75 0 00-1.06 1.06L8.94 10l-3.72 3.72a.75.75 0 101.06 1.06L10 11.06l3.72 3.72a.75.75 0 101.06-1.06L11.06 10l3.72-3.72a.75.75 0 00-1.06-1.06L10 8.94 6.28 5.22z"
/>
</svg>
</button>
</div>
<hr class=" dark:border-gray-850" />
<div class="flex flex-col w-full px-5 py-4 dark:text-gray-200 overflow-y-scroll max-h-[22rem]">
{#each mergedDocuments as document}
<!-- Source from document.metadata.source -->
<div class="flex flex-col w-full">
<div class="text-lg font-medium dark:text-gray-300">
{$i18n.t('Source')}
</div>
<div class="text-sm dark:text-gray-400">
{document.metadata.source}
</div>
</div>
<!-- Content from document.document.content -->
<div class="flex flex-col w-full">
<div class="text-lg font-medium dark:text-gray-300">
{$i18n.t('Content')}
</div>
<pre class="text-sm dark:text-gray-400">
{document.document}
</pre>
</div>
<hr class=" dark:border-gray-850" />
{/each}
</div>
</div>
</Modal>
src/lib/components/chat/Messages/ResponseMessage.svelte
View file @ 0872bea7
......@@ -32,6 +32,7 @@
import { WEBUI_BASE_URL } from '$lib/constants';
import Tooltip from '$lib/components/common/Tooltip.svelte';
import RateComment from './RateComment.svelte';
import CitationsModal from '$lib/components/chat/Messages/CitationsModal.svelte';
export let modelfiles = [];
export let message;
......@@ -65,6 +66,8 @@
let showRateComment = false;
let showCitations = {};
$: tokens = marked.lexer(sanitizeResponseContent(message.content));
const renderer = new marked.Renderer();
......@@ -360,6 +363,48 @@
{/each}
</div>
{/if}
{#if message.citations}
<div class="my-2.5 w-full flex overflow-x-auto gap-2 flex-wrap">
{#each message.citations as citation}
<div>
<CitationsModal bind:show={showCitations[citation]} {citation} />
<button
class="h-16 w-[15rem] flex items-center space-x-3 px-2.5 dark:bg-gray-600 rounded-xl border border-gray-200 dark:border-none text-left"
type="button"
on:click={() => {
showCitations[citation] = !showCitations[citation];
}}
>
<div class="p-2.5 bg-red-400 text-white rounded-lg">
<svg
xmlns="http://www.w3.org/2000/svg"
viewBox="0 0 24 24"
fill="currentColor"
class="w-6 h-6"
>
<path
fill-rule="evenodd"
d="M5.625 1.5c-1.036 0-1.875.84-1.875 1.875v17.25c0 1.035.84 1.875 1.875 1.875h12.75c1.035 0 1.875-.84 1.875-1.875V12.75A3.75 3.75 0 0 0 16.5 9h-1.875a1.875 1.875 0 0 1-1.875-1.875V5.25A3.75 3.75 0 0 0 9 1.5H5.625ZM7.5 15a.75.75 0 0 1 .75-.75h7.5a.75.75 0 0 1 0 1.5h-7.5A.75.75 0 0 1 7.5 15Zm.75 2.25a.75.75 0 0 0 0 1.5H12a.75.75 0 0 0 0-1.5H8.25Z"
clip-rule="evenodd"
/>
<path
d="M12.971 1.816A5.23 5.23 0 0 1 14.25 5.25v1.875c0 .207.168.375.375.375H16.5a5.23 5.23 0 0 1 3.434 1.279 9.768 9.768 0 0 0-6.963-6.963Z"
/>
</svg>
</div>
<div class="flex flex-col justify-center -space-y-0.5">
<div class=" dark:text-gray-100 text-sm font-medium line-clamp-1">
{citation.metadata?.[0]?.source ?? 'N/A'}
</div>
<div class=" text-gray-500 text-sm">{$i18n.t('Document')}</div>
</div>
</button>
</div>
{/each}
</div>
{/if}
<div
class="prose chat-{message.role} w-full max-w-full dark:prose-invert prose-headings:my-0 prose-p:m-0 prose-p:-mb-6 prose-pre:my-0 prose-table:my-0 prose-blockquote:my-0 prose-img:my-0 prose-ul:-my-4 prose-ol:-my-4 prose-li:-my-3 prose-ul:-mb-6 prose-ol:-mb-8 prose-ol:p-0 prose-li:-mb-4 whitespace-pre-line"
......@@ -577,10 +622,11 @@
stroke-linejoin="round"
class="w-4 h-4"
xmlns="http://www.w3.org/2000/svg"
><path
d="M14 9V5a3 3 0 0 0-3-3l-4 9v11h11.28a2 2 0 0 0 2-1.7l1.38-9a2 2 0 0 0-2-2.3zM7 22H4a2 2 0 0 1-2-2v-7a2 2 0 0 1 2-2h3"
/></svg
>
<path
d="M14 9V5a3 3 0 0 0-3-3l-4 9v11h11.28a2 2 0 0 0 2-1.7l1.38-9a2 2 0 0 0-2-2.3zM7 22H4a2 2 0 0 1-2-2v-7a2 2 0 0 1 2-2h3"
/>
</svg>
</button>
</Tooltip>
......@@ -611,10 +657,11 @@
stroke-linejoin="round"
class="w-4 h-4"
xmlns="http://www.w3.org/2000/svg"
><path
d="M10 15v4a3 3 0 0 0 3 3l4-9V2H5.72a2 2 0 0 0-2 1.7l-1.38 9a2 2 0 0 0 2 2.3zm7-13h2.67A2.31 2.31 0 0 1 22 4v7a2.31 2.31 0 0 1-2.33 2H17"
/></svg
>
<path
d="M10 15v4a3 3 0 0 0 3 3l4-9V2H5.72a2 2 0 0 0-2 1.7l-1.38 9a2 2 0 0 0 2 2.3zm7-13h2.67A2.31 2.31 0 0 1 22 4v7a2.31 2.31 0 0 1-2.33 2H17"
/>
</svg>
</button>
</Tooltip>
{/if}
......@@ -637,35 +684,32 @@
fill="currentColor"
viewBox="0 0 24 24"
xmlns="http://www.w3.org/2000/svg"
><style>
>
<style>
.spinner_S1WN {
animation: spinner_MGfb 0.8s linear infinite;
animation-delay: -0.8s;
}
.spinner_Km9P {
animation-delay: -0.65s;
}
.spinner_JApP {
animation-delay: -0.5s;
}
@keyframes spinner_MGfb {
93.75%,
100% {
opacity: 0.2;
}
}
</style><circle class="spinner_S1WN" cx="4" cy="12" r="3" /><circle
class="spinner_S1WN spinner_Km9P"
cx="12"
cy="12"
r="3"
/><circle
class="spinner_S1WN spinner_JApP"
cx="20"
cy="12"
r="3"
/></svg
>
</style>
<circle class="spinner_S1WN" cx="4" cy="12" r="3" />
<circle class="spinner_S1WN spinner_Km9P" cx="12" cy="12" r="3" />
<circle class="spinner_S1WN spinner_JApP" cx="20" cy="12" r="3" />
</svg>
{:else if speaking}
<svg
xmlns="http://www.w3.org/2000/svg"
......@@ -718,35 +762,32 @@
fill="currentColor"
viewBox="0 0 24 24"
xmlns="http://www.w3.org/2000/svg"
><style>
>
<style>
.spinner_S1WN {
animation: spinner_MGfb 0.8s linear infinite;
animation-delay: -0.8s;
}
.spinner_Km9P {
animation-delay: -0.65s;
}
.spinner_JApP {
animation-delay: -0.5s;
}
@keyframes spinner_MGfb {
93.75%,
100% {
opacity: 0.2;
}
}
</style><circle class="spinner_S1WN" cx="4" cy="12" r="3" /><circle
class="spinner_S1WN spinner_Km9P"
cx="12"
cy="12"
r="3"
/><circle
class="spinner_S1WN spinner_JApP"
cx="20"
cy="12"
r="3"
/></svg
>
</style>
<circle class="spinner_S1WN" cx="4" cy="12" r="3" />
<circle class="spinner_S1WN spinner_Km9P" cx="12" cy="12" r="3" />
<circle class="spinner_S1WN spinner_JApP" cx="20" cy="12" r="3" />
</svg>
{:else}
<svg
xmlns="http://www.w3.org/2000/svg"
......
src/routes/(app)/+page.svelte
View file @ 0872bea7
......@@ -366,7 +366,8 @@
},
format: $settings.requestFormat ?? undefined,
keep_alive: $settings.keepAlive ?? undefined,
docs: docs.length > 0 ? docs : undefined
docs: docs.length > 0 ? docs : undefined,
citations: docs.length > 0
});
if (res && res.ok) {
......@@ -401,6 +402,11 @@
console.log(line);
let data = JSON.parse(line);
if ('citations' in data) {
responseMessage.citations = data.citations;
continue;
}
if ('detail' in data) {
throw data;
}
......@@ -598,7 +604,8 @@
num_ctx: $settings?.options?.num_ctx ?? undefined,
frequency_penalty: $settings?.options?.repeat_penalty ?? undefined,
max_tokens: $settings?.options?.num_predict ?? undefined,
docs: docs.length > 0 ? docs : undefined
docs: docs.length > 0 ? docs : undefined,
citations: docs.length > 0
},
model?.source?.toLowerCase() === 'litellm'
? `${LITELLM_API_BASE_URL}/v1`
......@@ -614,7 +621,7 @@
const textStream = await createOpenAITextStream(res.body, $settings.splitLargeChunks);
for await (const update of textStream) {
const { value, done } = update;
const { value, done, citations } = update;
if (done || stopResponseFlag || _chatId !== $chatId) {
responseMessage.done = true;
messages = messages;
......@@ -626,6 +633,11 @@
break;
}
if (citations) {
responseMessage.citations = citations;
continue;
}
if (responseMessage.content == '' && value == '\n') {
continue;
} else {
......
src/routes/(app)/c/[id]/+page.svelte
View file @ 0872bea7
......@@ -378,7 +378,8 @@
},
format: $settings.requestFormat ?? undefined,
keep_alive: $settings.keepAlive ?? undefined,
docs: docs.length > 0 ? docs : undefined
docs: docs.length > 0 ? docs : undefined,
citations: docs.length > 0
});
if (res && res.ok) {
......@@ -413,6 +414,11 @@
console.log(line);
let data = JSON.parse(line);
if ('citations' in data) {
responseMessage.citations = data.citations;
continue;
}
if ('detail' in data) {
throw data;
}
......@@ -610,7 +616,8 @@
num_ctx: $settings?.options?.num_ctx ?? undefined,
frequency_penalty: $settings?.options?.repeat_penalty ?? undefined,
max_tokens: $settings?.options?.num_predict ?? undefined,
docs: docs.length > 0 ? docs : undefined
docs: docs.length > 0 ? docs : undefined,
citations: docs.length > 0
},
model?.source?.toLowerCase() === 'litellm'
? `${LITELLM_API_BASE_URL}/v1`
......@@ -626,7 +633,7 @@
const textStream = await createOpenAITextStream(res.body, $settings.splitLargeChunks);
for await (const update of textStream) {
const { value, done } = update;
const { value, done, citations } = update;
if (done || stopResponseFlag || _chatId !== $chatId) {
responseMessage.done = true;
messages = messages;
......@@ -638,6 +645,11 @@
break;
}
if (citations) {
responseMessage.citations = citations;
continue;
}
if (responseMessage.content == '' && value == '\n') {
continue;
} else {
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment