feat: knowledge file download API and file citations (#1836)

Co-authored-by: Ryan Hopper-Lowe <[email protected]>

Author: iwilltry42

pkg/api/handlers/agent.go

@@ -442,6 +442,15 @@ func (a *AgentHandler) ListKnowledgeFiles(req api.Context) error {
 	return listKnowledgeFiles(req, agentName, "", knowledgeSetNames[0], knowledgeSource)
 }
 
+func (a *AgentHandler) GetKnowledgeFile(req api.Context) error {
+	knowledgeSetNames, _, err := a.getKnowledgeSetsAndName(req, req.PathValue("agent_id"))
+	if err != nil {
+		return err
+	}
+
+	return getKnowledgeFileFromAllowedSets(req, a.gptscript, knowledgeSetNames, req.PathValue("file"))
+}
+
 func (a *AgentHandler) UploadKnowledgeFile(req api.Context) error {
 	knowledgeSetNames, agentName, err := a.getKnowledgeSetsAndName(req, req.PathValue("id"))
 	if err != nil {

pkg/api/handlers/assistants.go

@@ -363,6 +363,14 @@ func (a *AssistantHandler) Knowledge(req api.Context) error {
 	return listKnowledgeFiles(req, "", thread.Name, thread.Status.KnowledgeSetNames[0], nil)
 }
 
+func (a *AssistantHandler) GetKnowledgeFile(req api.Context) error {
+	thread, err := getThreadForScope(req)
+	if err != nil {
+		return err
+	}
+	return getKnowledgeFile(req, a.gptScript, thread, nil, req.PathValue("file"))
+}
+
 func (a *AssistantHandler) UploadKnowledge(req api.Context) error {
 	thread, err := getThreadForScope(req)
 	if err != nil {

pkg/api/handlers/files.go

@@ -4,7 +4,9 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"net/url"
 	"path/filepath"
+	"slices"
 	"strings"
 
 	"github.com/gptscript-ai/go-gptscript"
@@ -52,6 +54,63 @@ func getWorkspaceFromKnowledgeSet(req api.Context, knowledgeSetName string) (*v1
 	return &ws, req.Get(&ws, knowledgeSet.Status.WorkspaceName)
 }
 
+// getKnowledgeFile retrieves a knowledge file from the workspace associated with the knowledge set.
+// It works for both thread and agent knowledge sets. If the knowledge set is not found in the thread, it will be looked up in the agent.
+func getKnowledgeFile(req api.Context, gClient *gptscript.GPTScript, thread *v1.Thread, agent *v1.Agent, fileRef string) error {
+	var err error
+
+	// make sure that the selected knowledge set belongs either to the thread or to the agent
+	var knowledgeSetNames []string
+	if thread != nil {
+		knowledgeSetNames = thread.Status.KnowledgeSetNames
+		if agent == nil {
+			agent, err = getAssistant(req, thread.Spec.AgentName)
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	if agent != nil {
+		knowledgeSetNames = append(knowledgeSetNames, agent.Status.KnowledgeSetNames...)
+	}
+
+	return getKnowledgeFileFromAllowedSets(req, gClient, knowledgeSetNames, fileRef)
+}
+
+// getKnowledgeFileFromAllowedSets retrieves a knowledge file from the workspace associated with the knowledge set, if the knowledge set is in the list of allowed knowledge sets.
+// The fileRef is expected to be in the URL-encoded format [<knowledgeSet.Namespace>/]<knowledgeSet.Name>::<filename>.
+func getKnowledgeFileFromAllowedSets(req api.Context, gClient *gptscript.GPTScript, knowledgeSetNames []string, fileRef string) error {
+	var knowledgeSetName string
+
+	file, err := url.PathUnescape(fileRef)
+	if err != nil {
+		return types.NewErrBadRequest("invalid knowledgeFile reference")
+	}
+
+	parts := strings.Split(file, "::")
+	if len(parts) != 2 {
+		return types.NewErrBadRequest("invalid knowledgeFile path")
+	}
+	knowledgeSetName, file = parts[0], parts[1]
+
+	if parts := strings.Split(knowledgeSetName, "/"); len(parts) > 1 {
+		knowledgeSetName = parts[1] // may come in as <namespace>/<knowledgeset>, we don't care about the namespace right now
+	}
+
+	if !slices.Contains(knowledgeSetNames, knowledgeSetName) {
+		return types.NewErrNotFound("knowledge set %q not accessible", knowledgeSetName)
+	}
+
+	ws, err := getWorkspaceFromKnowledgeSet(req, knowledgeSetName)
+	if err != nil {
+		return err
+	}
+
+	req.SetPathValue("file", file)
+	return getFileInWorkspace(req.Context(), req, gClient, ws.Status.WorkspaceID, "") // knowledge files are stored in the root of the workspace (we have one workspace per knowledge set)
+}
+
 func listKnowledgeFiles(req api.Context, agentName, threadName, knowledgeSetName string, knowledgeSource *v1.KnowledgeSource) error {
 	var (
 		files               v1.KnowledgeFileList
@@ -177,6 +236,7 @@ func getFileInWorkspace(ctx context.Context, req api.Context, gClient *gptscript
 	}
 
 	req.ResponseWriter.Header().Set("Content-Type", "application/octet-stream")
+	req.ResponseWriter.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%q", file)) // make sure the file is downloaded with only the filename, not e.g. the dataset prefix
 	_, err = req.ResponseWriter.Write(data)
 	return err
 }

pkg/api/handlers/threads.go

@@ -339,6 +339,18 @@ func (a *ThreadHandler) Knowledge(req api.Context) error {
 	return listKnowledgeFiles(req, "", thread.Name, thread.Status.KnowledgeSetNames[0], nil)
 }
 
+func (a *ThreadHandler) GetKnowledgeFile(req api.Context) error {
+	var (
+		threadID = req.PathValue("id")
+	)
+
+	var thread v1.Thread
+	if err := req.Get(&thread, threadID); err != nil {
+		return err
+	}
+	return getKnowledgeFile(req, a.gptscript, &thread, nil, req.PathValue("file"))
+}
+
 func (a *ThreadHandler) UploadKnowledge(req api.Context) error {
 	var (
 		threadID = req.PathValue("id")

pkg/api/router/router.go

@@ -92,6 +92,7 @@ func Router(services *services.Services) (http.Handler, error) {
 	mux.HandleFunc("DELETE /api/assistants/{assistant_id}/projects/{project_id}/files/{file...}", assistants.DeleteFile)
 	// Assistant knowledge files
 	mux.HandleFunc("GET /api/assistants/{assistant_id}/projects/{project_id}/knowledge", assistants.Knowledge)
+	mux.HandleFunc("GET /api/assistants/{assistant_id}/projects/{project_id}/knowledge/{file}", assistants.GetKnowledgeFile)
 	mux.HandleFunc("POST /api/assistants/{assistant_id}/projects/{project_id}/knowledge/{file}", assistants.UploadKnowledge)
 	mux.HandleFunc("DELETE /api/assistants/{assistant_id}/projects/{project_id}/knowledge/{file...}", assistants.DeleteKnowledge)
 	// Env
@@ -198,6 +199,7 @@ func Router(services *services.Services) (http.Handler, error) {
 
 	// Agent knowledge files
 	mux.HandleFunc("GET /api/agents/{agent_id}/knowledge-files", agents.ListKnowledgeFiles)
+	mux.HandleFunc("GET /api/agents/{agent_id}/knowledge-files/{file}", agents.GetKnowledgeFile)
 	mux.HandleFunc("POST /api/agents/{id}/knowledge-files/{file...}", agents.UploadKnowledgeFile)
 	mux.HandleFunc("DELETE /api/agents/{id}/knowledge-files/{file...}", agents.DeleteKnowledgeFile)
 	mux.HandleFunc("POST /api/agents/{agent_id}/knowledge-files/{file_id}/ingest", agents.ReIngestKnowledgeFile)
@@ -239,6 +241,7 @@ func Router(services *services.Services) (http.Handler, error) {
 
 	// Thread knowledge files
 	mux.HandleFunc("GET /api/threads/{id}/knowledge-files", threads.Knowledge)
+	mux.HandleFunc("GET /api/threads/{id}/knowledge-files/{file}", threads.GetKnowledgeFile)
 	mux.HandleFunc("POST /api/threads/{id}/knowledge-files/{file}", threads.UploadKnowledge)
 	mux.HandleFunc("DELETE /api/threads/{id}/knowledge-files/{file...}", threads.DeleteKnowledge)
 

ui/admin/app/components/chat/Message.tsx

@@ -1,11 +1,13 @@
 import "@radix-ui/react-tooltip";
-import { AlertCircleIcon, WrenchIcon } from "lucide-react";
+import { AlertCircleIcon, BrainIcon, WrenchIcon } from "lucide-react";
 import React, { useDeferredValue, useMemo, useState } from "react";
 import { useForm } from "react-hook-form";
 
 import { AgentIcons } from "~/lib/model/agents";
 import { AuthPrompt } from "~/lib/model/chatEvents";
+import { KnowledgeFileNamespace } from "~/lib/model/knowledge";
 import { Message as MessageType } from "~/lib/model/messages";
+import { ApiRoutes } from "~/lib/routers/apiRoutes";
 import { PromptApiService } from "~/lib/service/api/PromptApi";
 import { cn, formatTime } from "~/lib/utils";
 
@@ -14,7 +16,7 @@ import { ToolCallInfo } from "~/components/chat/ToolCallInfo";
 import { ControlledInput } from "~/components/form/controlledInputs";
 import { ToolIcon } from "~/components/tools/ToolIcon";
 import { Avatar, AvatarFallback, AvatarImage } from "~/components/ui/avatar";
-import { Button } from "~/components/ui/button";
+import { Button, buttonVariants } from "~/components/ui/button";
 import {
 	Dialog,
 	DialogContent,
@@ -26,6 +28,11 @@ import { Form } from "~/components/ui/form";
 import { Link } from "~/components/ui/link";
 import { Markdown } from "~/components/ui/markdown";
 import { ScrollArea } from "~/components/ui/scroll-area";
+import {
+	Tooltip,
+	TooltipContent,
+	TooltipTrigger,
+} from "~/components/ui/tooltip";
 import { useAnimatedText } from "~/hooks/messages/useAnimatedText";
 import { useAsync } from "~/hooks/useAsync";
 
@@ -369,44 +376,80 @@ export function SourceCitations({
 	if (!message.knowledgeSources || !show) return null;
 
 	const formatUrl = (url: string) => {
-		return url.replace(/(https?:\/\/)?(www\.)?/, "");
+		return url.replace(/^(https?:\/\/)?(knowledge:\/\/)?/, "");
 	};
 
-	const citations = new Map(
-		message.knowledgeSources
-			.filter((s) => !!s.url)
-			.map((s) => {
-				return [
-					formatUrl(s.url as string),
-					{ ...s, url: new URL(s.url as string) },
-				];
-			})
-	);
-
 	return (
 		<div className="flex flex-col gap-2 pt-4">
 			<h4>Sources</h4>
 			<div className="flex flex-wrap gap-2 pb-3">
-				{Array.from(citations.entries()).map(([key, { url }]) => (
+				{message.knowledgeSources.map((s) => {
+					if (!s.url) return null;
+
+					if (s.url.startsWith("knowledge://"))
+						return renderKnowledgeSource(s.url);
+
+					return renderLink(s.url);
+				})}
+			</div>
+		</div>
+	);
+
+	function renderKnowledgeSource(url: string) {
+		if (!message.threadId) return null;
+
+		const filePath = url.replace("knowledge://", "");
+		const [_, fileName] = decodeURIComponent(filePath).split("::");
+
+		return (
+			<Tooltip>
+				<TooltipContent>{decodeURIComponent(url)}</TooltipContent>
+
+				<TooltipTrigger asChild>
+					<a
+						className={buttonVariants({ variant: "secondary", size: "sm" })}
+						key={url}
+						href={
+							ApiRoutes.knowledgeFiles.getKnowledgeFileById(
+								KnowledgeFileNamespace.Threads,
+								message.threadId,
+								filePath
+							).url
+						}
+						download={fileName}
+					>
+						<BrainIcon /> {fileName}
+					</a>
+				</TooltipTrigger>
+			</Tooltip>
+		);
+	}
+
+	function renderLink(url: string) {
+		const urlObj = new URL(url);
+		const formatted = formatUrl(url);
+		return (
+			<Tooltip>
+				<TooltipContent>{url}</TooltipContent>
+				<TooltipTrigger asChild>
 					<Link
 						as="button"
 						variant="secondary"
 						size="sm"
-						to={url.toString()}
+						to={url}
 						target="_blank"
 						rel="noreferrer"
-						key={key}
 					>
 						<img
-							src={`${url.origin}/favicon.ico`}
+							src={`${urlObj.origin}/favicon.ico`}
 							alt="Favicon"
 							onError={(e) => (e.currentTarget.src = "/favicon.ico")}
 							className="size-4"
 						/>
-						{key.length > 25 ? key.slice(0, 25) + "..." : key}
+						{formatted.length > 25 ? formatted.slice(0, 25) + "..." : formatted}
 					</Link>
-				))}
-			</div>
-		</div>
-	);
+				</TooltipTrigger>
+			</Tooltip>
+		);
+	}
 }

ui/admin/app/lib/model/messages.ts

@@ -1,6 +1,5 @@
 import {
 	AuthPrompt,
-	ChatEvent,
 	KnowledgeToolOutput,
 	ToolCall,
 } from "~/lib/model/chatEvents";
@@ -20,6 +19,7 @@ export interface Message {
 	contentID?: string;
 	time?: Date | string;
 	knowledgeSources?: KnowledgeToolOutput;
+	threadId?: string;
 }
 
 export const runsToMessages = (runs: Run[]) => {
@@ -48,56 +48,8 @@ export const toolCallMessage = (toolCall: ToolCall): Message => ({
 	tools: [toolCall],
 });
 
-export const promptMessage = (prompt: AuthPrompt, runID: string): Message => ({
+export const promptMessage = (prompt: AuthPrompt): Message => ({
 	sender: "agent",
 	text: prompt.message,
 	prompt,
-	runId: runID,
 });
-
-export const chatEventsToMessages = (events: ChatEvent[]) => {
-	const messages: Message[] = [];
-
-	for (const event of events) {
-		const { content, input, toolCall, runID, error, prompt } = event;
-
-		if (error) {
-			messages.push({
-				sender: "agent",
-				text: `Error: ${error}`,
-				runId: runID,
-				error: true,
-			});
-			continue;
-		}
-
-		if (input) {
-			messages.push({
-				sender: "user",
-				text: input,
-				runId: runID,
-			});
-			continue;
-		}
-
-		if (toolCall) {
-			messages.push(toolCallMessage(toolCall));
-			continue;
-		}
-
-		if (prompt) {
-			messages.push(promptMessage(prompt, runID));
-			continue;
-		}
-
-		if (content) {
-			messages.push({
-				sender: "agent",
-				text: content,
-				runId: runID,
-			});
-		}
-	}
-
-	return messages;
-};

ui/admin/app/lib/routers/apiRoutes.ts

@@ -103,6 +103,11 @@ export const ApiRoutes = {
 	knowledgeFiles: {
 		getKnowledgeFiles: (namespace: KnowledgeFileNamespace, entityId: string) =>
 			buildUrl(`/${namespace}/${entityId}/knowledge-files`),
+		getKnowledgeFileById: (
+			namespace: KnowledgeFileNamespace,
+			entityId: string,
+			fileName: string
+		) => buildUrl(`/${namespace}/${entityId}/knowledge-files/${fileName}`),
 		addKnowledgeFile: (
 			namespace: KnowledgeFileNamespace,
 			entityId: string,