feat: knowledge base functionality for the voice agent (#120)

* feat: upload file and store embedding

* feat: add documents in nodes

* feat: add openai embedding service

ui/src/components/ServiceConfiguration.tsx

@@ -14,7 +14,7 @@ import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
 import { VoiceSelector } from "@/components/VoiceSelector";
 import { useUserConfig } from "@/context/UserConfigContext";
 
-type ServiceSegment = "llm" | "tts" | "stt";
+type ServiceSegment = "llm" | "tts" | "stt" | "embeddings";
 
 interface SchemaProperty {
     type?: string;
@@ -41,6 +41,7 @@ const TAB_CONFIG: { key: ServiceSegment; label: string }[] = [
     { key: "llm", label: "LLM" },
     { key: "tts", label: "Voice" },
     { key: "stt", label: "Transcriber" },
+    { key: "embeddings", label: "Embedding" },
 ];
 
 // Display names for language codes (Deepgram + Sarvam)
@@ -109,12 +110,14 @@ export default function ServiceConfiguration() {
     const [schemas, setSchemas] = useState<Record<ServiceSegment, Record<string, ProviderSchema>>>({
         llm: {},
         tts: {},
-        stt: {}
+        stt: {},
+        embeddings: {}
     });
     const [serviceProviders, setServiceProviders] = useState<Record<ServiceSegment, string>>({
         llm: "",
         tts: "",
-        stt: ""
+        stt: "",
+        embeddings: ""
     });
     const [isManualModelInput, setIsManualModelInput] = useState(false);
     const [hasCheckedManualMode, setHasCheckedManualMode] = useState(false);
@@ -136,7 +139,8 @@ export default function ServiceConfiguration() {
                 setSchemas({
                     llm: response.data.llm as Record<string, ProviderSchema>,
                     tts: response.data.tts as Record<string, ProviderSchema>,
-                    stt: response.data.stt as Record<string, ProviderSchema>
+                    stt: response.data.stt as Record<string, ProviderSchema>,
+                    embeddings: response.data.embeddings as Record<string, ProviderSchema>
                 });
             } else {
                 console.error("Failed to fetch configurations");
@@ -147,7 +151,8 @@ export default function ServiceConfiguration() {
             const selectedProviders: Record<ServiceSegment, string> = {
                 llm: response.data.default_providers.llm,
                 tts: response.data.default_providers.tts,
-                stt: response.data.default_providers.stt
+                stt: response.data.default_providers.stt,
+                embeddings: response.data.default_providers.embeddings
             };
 
             const setServicePropertyValues = (service: ServiceSegment) => {
@@ -173,6 +178,7 @@ export default function ServiceConfiguration() {
             setServicePropertyValues("llm");
             setServicePropertyValues("tts");
             setServicePropertyValues("stt");
+            setServicePropertyValues("embeddings");
 
             // IMPORTANT: Reset form values BEFORE changing providers
             // Otherwise, Radix Select sees old values that don't match new provider's enum
@@ -246,7 +252,7 @@ export default function ServiceConfiguration() {
         setApiError(null);
         setIsSaving(true);
 
-        const userConfig = {
+        const userConfig: Record<ServiceSegment, Record<string, string | number>> = {
             llm: {
                 provider: serviceProviders.llm,
                 api_key: data.llm_api_key as string,
@@ -259,6 +265,11 @@ export default function ServiceConfiguration() {
             stt: {
                 provider: serviceProviders.stt,
                 api_key: data.stt_api_key as string
+            },
+            embeddings: {
+                provider: serviceProviders.embeddings,
+                api_key: data.embeddings_api_key as string,
+                model: data.embeddings_model as string
             }
         };
 
@@ -273,12 +284,25 @@ export default function ServiceConfiguration() {
             }
         });
 
+        // Build save config - only include embeddings if api_key is provided
+        const saveConfig: {
+            llm: Record<string, string | number>;
+            tts: Record<string, string | number>;
+            stt: Record<string, string | number>;
+            embeddings?: Record<string, string | number>;
+        } = {
+            llm: userConfig.llm,
+            tts: userConfig.tts,
+            stt: userConfig.stt
+        };
+
+        // Only include embeddings if user has configured it (has api_key)
+        if (userConfig.embeddings.api_key) {
+            saveConfig.embeddings = userConfig.embeddings;
+        }
+
         try {
-            await saveUserConfig({
-                llm: userConfig.llm,
-                tts: userConfig.tts,
-                stt: userConfig.stt
-            });
+            await saveUserConfig(saveConfig);
             setApiError(null);
         } catch (error: unknown) {
             if (error instanceof Error) {
@@ -543,7 +567,7 @@ export default function ServiceConfiguration() {
                 <Card>
                     <CardContent className="pt-6">
                         <Tabs defaultValue="llm" className="w-full">
-                            <TabsList className="grid w-full grid-cols-3 mb-6">
+                            <TabsList className="grid w-full grid-cols-4 mb-6">
                                 {TAB_CONFIG.map(({ key, label }) => (
                                     <TabsTrigger key={key} value={key}>
                                         {label}

ui/src/components/flow/DocumentBadges.tsx

@@ -0,0 +1,65 @@
+"use client";
+
+import { useCallback, useEffect, useState } from "react";
+
+import { useWorkflow } from "@/app/workflow/[workflowId]/contexts/WorkflowContext";
+import type { DocumentResponseSchema } from "@/client/types.gen";
+import { Badge } from "@/components/ui/badge";
+
+interface DocumentBadgesProps {
+    documentUuids: string[];
+    onStaleUuidsDetected?: (staleUuids: string[]) => void;
+}
+
+export const DocumentBadges = ({ documentUuids, onStaleUuidsDetected }: DocumentBadgesProps) => {
+    const { documents } = useWorkflow();
+    const [documentNames, setDocumentNames] = useState<Record<string, string>>({});
+
+    const processDocuments = useCallback((docs: DocumentResponseSchema[]) => {
+        const nameMap: Record<string, string> = {};
+        const validUuids = new Set<string>();
+
+        docs
+            .filter((doc) => documentUuids.includes(doc.document_uuid))
+            .forEach((doc) => {
+                nameMap[doc.document_uuid] = doc.filename;
+                validUuids.add(doc.document_uuid);
+            });
+        setDocumentNames(nameMap);
+
+        // Detect stale UUIDs - this only runs when we have loaded data (not undefined)
+        if (onStaleUuidsDetected) {
+            const staleUuids = documentUuids.filter(uuid => !validUuids.has(uuid));
+            if (staleUuids.length > 0) {
+                onStaleUuidsDetected(staleUuids);
+            }
+        }
+    }, [documentUuids, onStaleUuidsDetected]);
+
+    useEffect(() => {
+        if (documentUuids.length > 0 && documents !== undefined) {
+            processDocuments(documents);
+        } else if (documentUuids.length === 0) {
+            setDocumentNames({});
+        }
+    }, [documentUuids, documents, processDocuments]);
+
+    if (documentUuids.length === 0) {
+        return <></>;
+    }
+
+    // Show loading while data hasn't loaded yet
+    if (documents === undefined) {
+        return <Badge variant="outline">Loading...</Badge>;
+    }
+
+    return (
+        <>
+            {documentUuids.map((uuid) => (
+                <Badge key={uuid} variant="outline">
+                    {documentNames[uuid] || uuid}
+                </Badge>
+            ))}
+        </>
+    );
+};

ui/src/components/flow/DocumentSelector.tsx

@@ -0,0 +1,137 @@
+"use client";
+
+import { FileText } from "lucide-react";
+import Link from "next/link";
+import { useMemo } from "react";
+
+import type { DocumentResponseSchema } from "@/client/types.gen";
+import { Button } from "@/components/ui/button";
+import { Checkbox } from "@/components/ui/checkbox";
+import { Label } from "@/components/ui/label";
+
+interface DocumentSelectorProps {
+    value: string[];
+    onChange: (uuids: string[]) => void;
+    documents: DocumentResponseSchema[];
+    disabled?: boolean;
+    label?: string;
+    description?: string;
+    showLabel?: boolean;
+}
+
+export const DocumentSelector = ({
+    value,
+    onChange,
+    documents,
+    disabled = false,
+    label = "Knowledge Base Documents",
+    description = "Select documents that the agent can reference during conversations.",
+    showLabel = true,
+}: DocumentSelectorProps) => {
+    // Only show completed documents
+    const completedDocuments = useMemo(
+        () => documents.filter((doc) => doc.processing_status === "completed"),
+        [documents]
+    );
+
+    const handleToggle = (documentUuid: string, checked: boolean) => {
+        if (checked) {
+            onChange([...value, documentUuid]);
+        } else {
+            onChange(value.filter((uuid) => uuid !== documentUuid));
+        }
+    };
+
+    const formatFileSize = (bytes: number): string => {
+        if (bytes === 0) return "0 Bytes";
+        const k = 1024;
+        const sizes = ["Bytes", "KB", "MB", "GB"];
+        const i = Math.floor(Math.log(bytes) / Math.log(k));
+        return Math.round(bytes / Math.pow(k, i) * 100) / 100 + " " + sizes[i];
+    };
+
+    if (completedDocuments.length === 0) {
+        return (
+            <div className="space-y-2">
+                {showLabel && (
+                    <>
+                        <Label>{label}</Label>
+                        {description && (
+                            <Label className="text-xs text-muted-foreground">{description}</Label>
+                        )}
+                    </>
+                )}
+                <div className="border rounded-md p-4 space-y-3">
+                    <div className="text-sm text-muted-foreground text-center">
+                        No documents available. Upload documents to the knowledge base first.
+                    </div>
+                    <div className="flex justify-center">
+                        <Link href="/files">
+                            <Button variant="outline" size="sm">
+                                Upload Documents
+                            </Button>
+                        </Link>
+                    </div>
+                </div>
+            </div>
+        );
+    }
+
+    return (
+        <div className="space-y-2">
+            {showLabel && (
+                <>
+                    <Label>{label}</Label>
+                    {description && (
+                        <Label className="text-xs text-muted-foreground">{description}</Label>
+                    )}
+                </>
+            )}
+            <div className="border rounded-md max-h-[300px] overflow-y-auto">
+                <div className="divide-y">
+                    {completedDocuments.map((doc) => (
+                        <div
+                            key={doc.document_uuid}
+                            className="flex items-start gap-3 p-3 hover:bg-muted/50 transition-colors"
+                        >
+                            <Checkbox
+                                id={`doc-${doc.document_uuid}`}
+                                checked={value.includes(doc.document_uuid)}
+                                onCheckedChange={(checked) =>
+                                    handleToggle(doc.document_uuid, checked as boolean)
+                                }
+                                disabled={disabled}
+                            />
+                            <div className="flex-1 space-y-1">
+                                <label
+                                    htmlFor={`doc-${doc.document_uuid}`}
+                                    className="flex items-center gap-2 cursor-pointer"
+                                >
+                                    <div className="w-8 h-8 rounded-md bg-blue-500/10 flex items-center justify-center flex-shrink-0">
+                                        <FileText className="w-4 h-4 text-blue-500" />
+                                    </div>
+                                    <div className="flex-1 min-w-0">
+                                        <div className="text-sm font-medium truncate">
+                                            {doc.filename}
+                                        </div>
+                                        <div className="text-xs text-muted-foreground">
+                                            {formatFileSize(doc.file_size_bytes)} • {doc.total_chunks} chunks
+                                        </div>
+                                    </div>
+                                </label>
+                            </div>
+                        </div>
+                    ))}
+                </div>
+            </div>
+            <div className="flex items-center justify-between text-xs text-muted-foreground pt-1">
+                <span>
+                    {value.length} {value.length === 1 ? "document" : "documents"} selected
+                </span>
+                <Link href="/files" className="hover:underline">
+                    Manage Documents
+                </Link>
+            </div>
+        </div>
+    );
+};

ui/src/components/flow/ToolBadges.tsx

@@ -2,43 +2,43 @@
 
 import { useCallback, useEffect, useState } from "react";
 
-import { listToolsApiV1ToolsGet } from "@/client/sdk.gen";
+import { useWorkflow } from "@/app/workflow/[workflowId]/contexts/WorkflowContext";
 import type { ToolResponse } from "@/client/types.gen";
 import { Badge } from "@/components/ui/badge";
-import { useAuth } from "@/lib/auth";
 
 interface ToolBadgesProps {
     toolUuids: string[];
+    onStaleUuidsDetected?: (staleUuids: string[]) => void;
 }
 
-export function ToolBadges({ toolUuids }: ToolBadgesProps) {
-    const { getAccessToken } = useAuth();
-    const [tools, setTools] = useState<ToolResponse[]>([]);
+export function ToolBadges({ toolUuids, onStaleUuidsDetected }: ToolBadgesProps) {
+    const { tools } = useWorkflow();
+    const [selectedTools, setSelectedTools] = useState<ToolResponse[]>([]);
 
-    const fetchTools = useCallback(async () => {
-        try {
-            const accessToken = await getAccessToken();
-            const response = await listToolsApiV1ToolsGet({
-                headers: { Authorization: `Bearer ${accessToken}` },
-            });
-            if (response.data) {
-                setTools(response.data);
+    const processTools = useCallback((toolsData: ToolResponse[]) => {
+        const filtered = toolsData.filter(tool => toolUuids.includes(tool.tool_uuid));
+        setSelectedTools(filtered);
+
+        // Detect stale UUIDs - this only runs when we have loaded data (not undefined)
+        if (onStaleUuidsDetected) {
+            const validUuids = new Set(toolsData.map(tool => tool.tool_uuid));
+            const staleUuids = toolUuids.filter(uuid => !validUuids.has(uuid));
+            if (staleUuids.length > 0) {
+                onStaleUuidsDetected(staleUuids);
             }
-        } catch (error) {
-            console.error("Failed to fetch tools:", error);
         }
-    }, [getAccessToken]);
+    }, [toolUuids, onStaleUuidsDetected]);
 
     useEffect(() => {
-        if (toolUuids.length > 0) {
-            fetchTools();
+        if (toolUuids.length > 0 && tools !== undefined) {
+            processTools(tools);
+        } else if (toolUuids.length === 0) {
+            setSelectedTools([]);
         }
-    }, [toolUuids.length, fetchTools]);
+    }, [toolUuids, tools, processTools]);
 
-    const selectedTools = tools.filter((tool) => toolUuids.includes(tool.tool_uuid));
-
-    if (selectedTools.length === 0 && toolUuids.length > 0) {
-        // Still loading or tools not found
+    // Show loading while data hasn't loaded yet
+    if (tools === undefined && toolUuids.length > 0) {
         return (
             <div className="flex flex-wrap gap-1">
                 <Badge variant="outline" className="text-xs">

ui/src/components/flow/ToolSelector.tsx

@@ -1,20 +1,18 @@
 "use client";
 
-import { ExternalLink, Loader2 } from "lucide-react";
+import { ExternalLink } from "lucide-react";
 import Link from "next/link";
-import { useCallback, useEffect, useState } from "react";
 
 import { renderToolIcon } from "@/app/tools/config";
-import { listToolsApiV1ToolsGet } from "@/client/sdk.gen";
 import type { ToolResponse } from "@/client/types.gen";
 import { Button } from "@/components/ui/button";
 import { Checkbox } from "@/components/ui/checkbox";
 import { Label } from "@/components/ui/label";
-import { useAuth } from "@/lib/auth";
 
 interface ToolSelectorProps {
     value: string[];
     onChange: (uuids: string[]) => void;
+    tools: ToolResponse[];
     disabled?: boolean;
     label?: string;
     description?: string;
@@ -24,43 +22,14 @@ interface ToolSelectorProps {
 export function ToolSelector({
     value,
     onChange,
+    tools,
     disabled = false,
     label = "Tools",
     description = "Select tools that the agent can use during the conversation.",
     showLabel = true,
 }: ToolSelectorProps) {
-    const { getAccessToken } = useAuth();
-
-    const [tools, setTools] = useState<ToolResponse[]>([]);
-    const [loading, setLoading] = useState(false);
-
-    const fetchTools = useCallback(async () => {
-        setLoading(true);
-        try {
-            const accessToken = await getAccessToken();
-            const response = await listToolsApiV1ToolsGet({
-                headers: { Authorization: `Bearer ${accessToken}` },
-                query: { status: "active" },
-            });
-            if (response.error) {
-                console.error("Failed to fetch tools:", response.error);
-                setTools([]);
-                return;
-            }
-            if (response.data) {
-                setTools(response.data);
-            }
-        } catch (error) {
-            console.error("Failed to fetch tools:", error);
-            setTools([]);
-        } finally {
-            setLoading(false);
-        }
-    }, [getAccessToken]);
-
-    useEffect(() => {
-        fetchTools();
-    }, [fetchTools]);
+    // Filter to only show active tools
+    const activeTools = tools.filter((tool) => tool.status === "active");
 
     const handleToggle = (toolUuid: string, checked: boolean) => {
         if (checked) {
@@ -83,12 +52,7 @@ export function ToolSelector({
                 </>
             )}
 
-            {loading ? (
-                <div className="flex items-center gap-2 p-3 border rounded-md">
-                    <Loader2 className="h-4 w-4 animate-spin" />
-                    <span className="text-sm text-muted-foreground">Loading tools...</span>
-                </div>
-            ) : tools.length === 0 ? (
+            {activeTools.length === 0 ? (
                 <div className="p-4 border rounded-md text-center">
                     <p className="text-sm text-muted-foreground mb-2">
                         No tools available.
@@ -102,7 +66,7 @@ export function ToolSelector({
                 </div>
             ) : (
                 <div className="border rounded-md divide-y">
-                    {tools.map((tool) => {
+                    {activeTools.map((tool) => {
                         const isSelected = value.includes(tool.tool_uuid);
                         return (
                             <label

ui/src/components/flow/nodes/AgentNode.tsx

@@ -1,8 +1,11 @@
 import { NodeProps, NodeToolbar, Position } from "@xyflow/react";
-import { Edit, Headset, PlusIcon, Trash2Icon, Wrench } from "lucide-react";
-import { memo, useEffect, useMemo, useState } from "react";
+import { Edit, FileText, Headset, PlusIcon, Trash2Icon, Wrench } from "lucide-react";
+import { memo, useCallback, useEffect, useMemo, useState } from "react";
 
 import { useWorkflow } from "@/app/workflow/[workflowId]/contexts/WorkflowContext";
+import type { DocumentResponseSchema, ToolResponse } from "@/client/types.gen";
+import { DocumentBadges } from "@/components/flow/DocumentBadges";
+import { DocumentSelector } from "@/components/flow/DocumentSelector";
 import { ToolBadges } from "@/components/flow/ToolBadges";
 import { ToolSelector } from "@/components/flow/ToolSelector";
 import { ExtractionVariable, FlowNodeData } from "@/components/flow/types";
@@ -34,6 +37,10 @@ interface AgentNodeEditFormProps {
     setAddGlobalPrompt: (value: boolean) => void;
     toolUuids: string[];
     setToolUuids: (value: string[]) => void;
+    documentUuids: string[];
+    setDocumentUuids: (value: string[]) => void;
+    tools: ToolResponse[];
+    documents: DocumentResponseSchema[];
 }
 
 interface AgentNodeProps extends NodeProps {
@@ -42,7 +49,7 @@ interface AgentNodeProps extends NodeProps {
 
 export const AgentNode = memo(({ data, selected, id }: AgentNodeProps) => {
     const { open, setOpen, handleSaveNodeData, handleDeleteNode } = useNodeHandlers({ id });
-    const { saveWorkflow } = useWorkflow();
+    const { saveWorkflow, tools, documents } = useWorkflow();
 
     // Form state
     const [prompt, setPrompt] = useState(data.prompt);
@@ -55,6 +62,7 @@ export const AgentNode = memo(({ data, selected, id }: AgentNodeProps) => {
     const [variables, setVariables] = useState<ExtractionVariable[]>(data.extraction_variables ?? []);
     const [addGlobalPrompt, setAddGlobalPrompt] = useState(data.add_global_prompt ?? true);
     const [toolUuids, setToolUuids] = useState<string[]>(data.tool_uuids ?? []);
+    const [documentUuids, setDocumentUuids] = useState<string[]>(data.document_uuids ?? []);
 
     // Compute if form has unsaved changes (only check prompt, name)
     const isDirty = useMemo(() => {
@@ -75,6 +83,7 @@ export const AgentNode = memo(({ data, selected, id }: AgentNodeProps) => {
             extraction_variables: variables,
             add_global_prompt: addGlobalPrompt,
             tool_uuids: toolUuids.length > 0 ? toolUuids : undefined,
+            document_uuids: documentUuids.length > 0 ? documentUuids : undefined,
         });
         setOpen(false);
         // Save the workflow after updating node data with a small delay to ensure state is updated
@@ -94,6 +103,7 @@ export const AgentNode = memo(({ data, selected, id }: AgentNodeProps) => {
             setVariables(data.extraction_variables ?? []);
             setAddGlobalPrompt(data.add_global_prompt ?? true);
             setToolUuids(data.tool_uuids ?? []);
+            setDocumentUuids(data.document_uuids ?? []);
         }
         setOpen(newOpen);
     };
@@ -109,9 +119,34 @@ export const AgentNode = memo(({ data, selected, id }: AgentNodeProps) => {
             setVariables(data.extraction_variables ?? []);
             setAddGlobalPrompt(data.add_global_prompt ?? true);
             setToolUuids(data.tool_uuids ?? []);
+            setDocumentUuids(data.document_uuids ?? []);
         }
     }, [data, open]);
 
+    // Handle cleanup of stale document UUIDs
+    const handleStaleDocuments = useCallback((staleUuids: string[]) => {
+        const cleanedUuids = (data.document_uuids ?? []).filter(uuid => !staleUuids.includes(uuid));
+        handleSaveNodeData({
+            ...data,
+            document_uuids: cleanedUuids.length > 0 ? cleanedUuids : undefined,
+        });
+        setTimeout(async () => {
+            await saveWorkflow();
+        }, 100);
+    }, [data, handleSaveNodeData, saveWorkflow]);
+
+    // Handle cleanup of stale tool UUIDs
+    const handleStaleTools = useCallback((staleUuids: string[]) => {
+        const cleanedUuids = (data.tool_uuids ?? []).filter(uuid => !staleUuids.includes(uuid));
+        handleSaveNodeData({
+            ...data,
+            tool_uuids: cleanedUuids.length > 0 ? cleanedUuids : undefined,
+        });
+        setTimeout(async () => {
+            await saveWorkflow();
+        }, 100);
+    }, [data, handleSaveNodeData, saveWorkflow]);
+
     return (
         <>
             <NodeContent
@@ -136,7 +171,16 @@ export const AgentNode = memo(({ data, selected, id }: AgentNodeProps) => {
                             <Wrench className="h-3 w-3" />
                             <span>Tools:</span>
                         </div>
-                        <ToolBadges toolUuids={data.tool_uuids} />
+                        <ToolBadges toolUuids={data.tool_uuids} onStaleUuidsDetected={handleStaleTools} />
+                    </div>
+                )}
+                {data.document_uuids && data.document_uuids.length > 0 && (
+                    <div className="mt-3 pt-3 border-t border-border/50">
+                        <div className="flex items-center gap-1.5 text-xs text-muted-foreground mb-2">
+                            <FileText className="h-3 w-3" />
+                            <span>Documents:</span>
+                        </div>
+                        <DocumentBadges documentUuids={data.document_uuids} onStaleUuidsDetected={handleStaleDocuments} />
                     </div>
                 )}
             </NodeContent>
@@ -179,6 +223,10 @@ export const AgentNode = memo(({ data, selected, id }: AgentNodeProps) => {
                         setAddGlobalPrompt={setAddGlobalPrompt}
                         toolUuids={toolUuids}
                         setToolUuids={setToolUuids}
+                        documentUuids={documentUuids}
+                        setDocumentUuids={setDocumentUuids}
+                        tools={tools ?? []}
+                        documents={documents ?? []}
                     />
                 )}
             </NodeEditDialog>
@@ -203,6 +251,10 @@ const AgentNodeEditForm = ({
     setAddGlobalPrompt,
     toolUuids,
     setToolUuids,
+    documentUuids,
+    setDocumentUuids,
+    tools,
+    documents,
 }: AgentNodeEditFormProps) => {
     const handleVariableNameChange = (idx: number, value: string) => {
         const newVars = [...variables];
@@ -343,9 +395,20 @@ const AgentNodeEditForm = ({
                 <ToolSelector
                     value={toolUuids}
                     onChange={setToolUuids}
+                    tools={tools}
                     description="Select tools that the agent can invoke during this conversation step."
                 />
             </div>
+
+            {/* Documents Section */}
+            <div className="pt-4 border-t mt-4">
+                <DocumentSelector
+                    value={documentUuids}
+                    onChange={setDocumentUuids}
+                    documents={documents}
+                    description="Select documents from the knowledge base that the agent can reference during this conversation step."
+                />
+            </div>
         </div>
     );
 };

ui/src/components/flow/nodes/StartCall.tsx

@@ -1,8 +1,11 @@
 import { NodeProps, NodeToolbar, Position } from "@xyflow/react";
-import { Edit, Play, PlusIcon, Trash2Icon, Wrench } from "lucide-react";
-import { memo, useEffect, useMemo, useState } from "react";
+import { Edit, FileText, Play, PlusIcon, Trash2Icon, Wrench } from "lucide-react";
+import { memo, useCallback, useEffect, useMemo, useState } from "react";
 
 import { useWorkflow } from "@/app/workflow/[workflowId]/contexts/WorkflowContext";
+import type { DocumentResponseSchema, ToolResponse } from "@/client/types.gen";
+import { DocumentBadges } from "@/components/flow/DocumentBadges";
+import { DocumentSelector } from "@/components/flow/DocumentSelector";
 import { ToolBadges } from "@/components/flow/ToolBadges";
 import { ToolSelector } from "@/components/flow/ToolSelector";
 import { ExtractionVariable, FlowNodeData } from "@/components/flow/types";
@@ -41,6 +44,10 @@ interface StartCallEditFormProps {
     setVariables: (vars: ExtractionVariable[]) => void;
     toolUuids: string[];
     setToolUuids: (value: string[]) => void;
+    documentUuids: string[];
+    setDocumentUuids: (value: string[]) => void;
+    tools: ToolResponse[];
+    documents: DocumentResponseSchema[];
 }
 
 interface StartCallNodeProps extends NodeProps {
@@ -52,7 +59,7 @@ export const StartCall = memo(({ data, selected, id }: StartCallNodeProps) => {
         id,
         additionalData: { is_start: true }
     });
-    const { saveWorkflow } = useWorkflow();
+    const { saveWorkflow, tools, documents } = useWorkflow();
 
     // Form state
     const [prompt, setPrompt] = useState(data.prompt ?? "");
@@ -66,6 +73,7 @@ export const StartCall = memo(({ data, selected, id }: StartCallNodeProps) => {
     const [extractionPrompt, setExtractionPrompt] = useState(data.extraction_prompt ?? "");
     const [variables, setVariables] = useState<ExtractionVariable[]>(data.extraction_variables ?? []);
     const [toolUuids, setToolUuids] = useState<string[]>(data.tool_uuids ?? []);
+    const [documentUuids, setDocumentUuids] = useState<string[]>(data.document_uuids ?? []);
 
     // Compute if form has unsaved changes (only check prompt, name)
     const isDirty = useMemo(() => {
@@ -89,6 +97,7 @@ export const StartCall = memo(({ data, selected, id }: StartCallNodeProps) => {
             extraction_prompt: extractionPrompt,
             extraction_variables: variables,
             tool_uuids: toolUuids.length > 0 ? toolUuids : undefined,
+            document_uuids: documentUuids.length > 0 ? documentUuids : undefined,
         });
         setOpen(false);
         // Save the workflow after updating node data with a small delay to ensure state is updated
@@ -111,6 +120,7 @@ export const StartCall = memo(({ data, selected, id }: StartCallNodeProps) => {
             setExtractionPrompt(data.extraction_prompt ?? "");
             setVariables(data.extraction_variables ?? []);
             setToolUuids(data.tool_uuids ?? []);
+            setDocumentUuids(data.document_uuids ?? []);
         }
         setOpen(newOpen);
     };
@@ -129,9 +139,34 @@ export const StartCall = memo(({ data, selected, id }: StartCallNodeProps) => {
             setExtractionPrompt(data.extraction_prompt ?? "");
             setVariables(data.extraction_variables ?? []);
             setToolUuids(data.tool_uuids ?? []);
+            setDocumentUuids(data.document_uuids ?? []);
         }
     }, [data, open]);
 
+    // Handle cleanup of stale document UUIDs
+    const handleStaleDocuments = useCallback((staleUuids: string[]) => {
+        const cleanedUuids = (data.document_uuids ?? []).filter(uuid => !staleUuids.includes(uuid));
+        handleSaveNodeData({
+            ...data,
+            document_uuids: cleanedUuids.length > 0 ? cleanedUuids : undefined,
+        });
+        setTimeout(async () => {
+            await saveWorkflow();
+        }, 100);
+    }, [data, handleSaveNodeData, saveWorkflow]);
+
+    // Handle cleanup of stale tool UUIDs
+    const handleStaleTools = useCallback((staleUuids: string[]) => {
+        const cleanedUuids = (data.tool_uuids ?? []).filter(uuid => !staleUuids.includes(uuid));
+        handleSaveNodeData({
+            ...data,
+            tool_uuids: cleanedUuids.length > 0 ? cleanedUuids : undefined,
+        });
+        setTimeout(async () => {
+            await saveWorkflow();
+        }, 100);
+    }, [data, handleSaveNodeData, saveWorkflow]);
+
     return (
         <>
             <NodeContent
@@ -155,7 +190,16 @@ export const StartCall = memo(({ data, selected, id }: StartCallNodeProps) => {
                             <Wrench className="h-3 w-3" />
                             <span>Tools:</span>
                         </div>
-                        <ToolBadges toolUuids={data.tool_uuids} />
+                        <ToolBadges toolUuids={data.tool_uuids} onStaleUuidsDetected={handleStaleTools} />
+                    </div>
+                )}
+                {data.document_uuids && data.document_uuids.length > 0 && (
+                    <div className="mt-3 pt-3 border-t border-border/50">
+                        <div className="flex items-center gap-1.5 text-xs text-muted-foreground mb-2">
+                            <FileText className="h-3 w-3" />
+                            <span>Documents:</span>
+                        </div>
+                        <DocumentBadges documentUuids={data.document_uuids} onStaleUuidsDetected={handleStaleDocuments} />
                     </div>
                 )}
             </NodeContent>
@@ -199,6 +243,10 @@ export const StartCall = memo(({ data, selected, id }: StartCallNodeProps) => {
                         setVariables={setVariables}
                         toolUuids={toolUuids}
                         setToolUuids={setToolUuids}
+                        documentUuids={documentUuids}
+                        setDocumentUuids={setDocumentUuids}
+                        tools={tools ?? []}
+                        documents={documents ?? []}
                     />
                 )}
             </NodeEditDialog>
@@ -229,6 +277,10 @@ const StartCallEditForm = ({
     setVariables,
     toolUuids,
     setToolUuids,
+    documentUuids,
+    setDocumentUuids,
+    tools,
+    documents,
 }: StartCallEditFormProps) => {
     const handleVariableNameChange = (idx: number, value: string) => {
         const newVars = [...variables];
@@ -414,9 +466,20 @@ const StartCallEditForm = ({
                 <ToolSelector
                     value={toolUuids}
                     onChange={setToolUuids}
+                    tools={tools}
                     description="Select tools that the agent can invoke during this conversation step."
                 />
             </div>
+
+            {/* Documents Section */}
+            <div className="pt-4 border-t mt-4">
+                <DocumentSelector
+                    value={documentUuids}
+                    onChange={setDocumentUuids}
+                    documents={documents}
+                    description="Select documents from the knowledge base that the agent can reference during this conversation step."
+                />
+            </div>
         </div>
     );
 };

ui/src/components/flow/types.ts

@@ -42,6 +42,8 @@ export type FlowNodeData = {
     };
     // Tools - array of tool UUIDs that can be invoked by this node
     tool_uuids?: string[];
+    // Documents - array of knowledge base document UUIDs that can be referenced by this node
+    document_uuids?: string[];
 }
 
 export type FlowNode = {

ui/src/components/layout/AppSidebar.tsx

@@ -6,6 +6,7 @@ import {
   ChevronLeft,
   ChevronRight,
   CircleDollarSign,
+  Database,
   FileText,
   HelpCircle,
   Home,
@@ -114,6 +115,11 @@ export function AppSidebar() {
           url: "/tools",
           icon: Wrench,
         },
+        {
+          title: "Files",
+          url: "/files",
+          icon: Database,
+        },
         // {
         //   title: "Integrations",
         //   url: "/integrations",