[feat] Add support for endpoint web sources (#1544)

* [feat] Add support for endpoint web sources

* chores: lint

---------

Co-authored-by: Nathan Sarrazin <[email protected]>

src/lib/components/chat/ChatMessage.svelte

@@ -23,10 +23,12 @@
 
 	import OpenWebSearchResults from "../OpenWebSearchResults.svelte";
 	import {
+		MessageUpdateType,
 		MessageWebSearchUpdateType,
 		type MessageToolUpdate,
 		type MessageWebSearchSourcesUpdate,
 		type MessageWebSearchUpdate,
+		type MessageFinalAnswerUpdate,
 	} from "$lib/types/MessageUpdate";
 	import { base } from "$app/paths";
 	import { useConvTreeStore } from "$lib/stores/convTree";
@@ -149,6 +151,10 @@
 	$: searchUpdates = (message.updates?.filter(({ type }) => type === "webSearch") ??
 		[]) as MessageWebSearchUpdate[];
 
+	$: messageFinalAnswer = message.updates?.find(
+		({ type }) => type === MessageUpdateType.FinalAnswer
+	) as MessageFinalAnswerUpdate;
+
 	// filter all updates with type === "tool" then group them by uuid field
 
 	$: toolUpdates = message.updates
@@ -314,7 +320,29 @@
 					{/each}
 				</div>
 			{/if}
+
+			<!-- Endpoint web sources -->
+			{#if messageFinalAnswer?.webSources && messageFinalAnswer.webSources.length}
+				<div class="mt-4 flex flex-wrap items-center gap-x-2 gap-y-1.5 text-sm">
+					<div class="text-gray-400">Sources:</div>
+					{#each messageFinalAnswer.webSources as { uri, title }}
+						<a
+							class="flex items-center gap-2 whitespace-nowrap rounded-lg border bg-white px-2 py-1.5 leading-none hover:border-gray-300 dark:border-gray-800 dark:bg-gray-900 dark:hover:border-gray-700"
+							href={uri}
+							target="_blank"
+						>
+							<img
+								class="h-3.5 w-3.5 rounded"
+								src="https://www.google.com/s2/favicons?sz=64&domain_url={new URL(uri).hostname}"
+								alt="{title} favicon"
+							/>
+							<div>{title}</div>
+						</a>
+					{/each}
+				</div>
+			{/if}
 		</div>
+
 		{#if !loading && (message.content || toolUpdates)}
 			<div
 				class="absolute -bottom-4 right-0 flex max-md:transition-all md:group-hover:visible md:group-hover:opacity-100

src/lib/server/endpoints/endpoints.ts

@@ -48,13 +48,14 @@ export interface EndpointParameters {
 interface CommonEndpoint {
 	weight: number;
 }
-type TextGenerationStreamOutputWithTools = TextGenerationStreamOutput & {
+export type TextGenerationStreamOutputWithToolsAndWebSources = TextGenerationStreamOutput & {
 	token: TextGenerationStreamToken & { toolCalls?: ToolCall[] };
+	webSources?: { uri: string; title: string }[];
 };
 // type signature for the endpoint
 export type Endpoint = (
 	params: EndpointParameters
-) => Promise<AsyncGenerator<TextGenerationStreamOutputWithTools, void, void>>;
+) => Promise<AsyncGenerator<TextGenerationStreamOutputWithToolsAndWebSources, void, void>>;
 
 // generator function that takes in parameters for defining the endpoint and return the endpoint
 export type EndpointGenerator<T extends CommonEndpoint> = (parameters: T) => Endpoint;

src/lib/server/endpoints/google/endpointVertex.ts

@@ -5,10 +5,9 @@ import {
 	type Content,
 	type TextPart,
 } from "@google-cloud/vertexai";
-import type { Endpoint } from "../endpoints";
+import type { Endpoint, TextGenerationStreamOutputWithToolsAndWebSources } from "../endpoints";
 import { z } from "zod";
 import type { Message } from "$lib/types/Message";
-import type { TextGenerationStreamOutput } from "@huggingface/inference";
 import { createImageProcessorOptionsValidator, makeImageProcessor } from "../images";
 import { createDocumentProcessorOptionsValidator, makeDocumentProcessor } from "../document";
 
@@ -170,6 +169,8 @@ export function endpointVertex(input: z.input<typeof endpointVertexParametersSch
 		return (async function* () {
 			let generatedText = "";
 
+			const webSources = [];
+
 			for await (const data of result.stream) {
 				if (!data?.candidates?.length) break; // Handle case where no candidates are present
 
@@ -183,9 +184,29 @@ export function endpointVertex(input: z.input<typeof endpointVertexParametersSch
 
 				const isLastChunk = !!candidate.finishReason;
 
+				const candidateWebSources = candidate.groundingMetadata?.groundingChunks
+					?.map((chunk) => {
+						const uri = chunk.web?.uri ?? chunk.retrievedContext?.uri;
+						const title = chunk.web?.title ?? chunk.retrievedContext?.title;
+
+						if (!uri || !title) {
+							return null;
+						}
+
+						return {
+							uri,
+							title,
+						};
+					})
+					.filter((source) => source !== null);
+
+				if (candidateWebSources) {
+					webSources.push(...candidateWebSources);
+				}
+
 				const content = firstPart.text;
 				generatedText += content;
-				const output: TextGenerationStreamOutput = {
+				const output: TextGenerationStreamOutputWithToolsAndWebSources = {
 					token: {
 						id: tokenId++,
 						text: content,
@@ -194,6 +215,7 @@ export function endpointVertex(input: z.input<typeof endpointVertexParametersSch
 					},
 					generated_text: isLastChunk ? generatedText : null,
 					details: null,
+					webSources,
 				};
 				yield output;
 

src/lib/server/textGeneration/generate.ts

@@ -33,7 +33,12 @@ export async function* generate(
 				text = text.slice(0, text.length - stopToken.length);
 			}
 
-			yield { type: MessageUpdateType.FinalAnswer, text, interrupted };
+			yield {
+				type: MessageUpdateType.FinalAnswer,
+				text,
+				interrupted,
+				webSources: output.webSources,
+			};
 			continue;
 		}
 

src/lib/types/MessageUpdate.ts

@@ -124,4 +124,5 @@ export interface MessageFinalAnswerUpdate {
 	type: MessageUpdateType.FinalAnswer;
 	text: string;
 	interrupted: boolean;
+	webSources?: { uri: string; title: string }[];
 }