feat: add documents in nodes

This commit is contained in:
Abhishek Kumar 2026-01-16 17:23:24 +05:30
parent ec1417da87
commit eb41285204
5 changed files with 312 additions and 2 deletions

View file

@ -0,0 +1,67 @@
"use client";
import { Badge } from "@/components/ui/badge";
import { listDocumentsApiV1KnowledgeBaseDocumentsGet } from "@/client/sdk.gen";
import { useAuth } from "@/lib/auth";
import { useCallback, useEffect, useState } from "react";
interface DocumentBadgesProps {
documentUuids: string[];
}
export const DocumentBadges = ({ documentUuids }: DocumentBadgesProps) => {
const { getAccessToken } = useAuth();
const [documentNames, setDocumentNames] = useState<Record<string, string>>({});
const [loading, setLoading] = useState(false);
const fetchDocuments = useCallback(async () => {
if (documentUuids.length === 0) return;
setLoading(true);
try {
const accessToken = await getAccessToken();
const response = await listDocumentsApiV1KnowledgeBaseDocumentsGet({
headers: { Authorization: `Bearer ${accessToken}` },
query: {
limit: 100,
},
});
if (response.data) {
const nameMap: Record<string, string> = {};
response.data.documents
.filter((doc) => documentUuids.includes(doc.document_uuid))
.forEach((doc) => {
nameMap[doc.document_uuid] = doc.filename;
});
setDocumentNames(nameMap);
}
} catch (error) {
console.error("Failed to fetch documents:", error);
} finally {
setLoading(false);
}
}, [documentUuids, getAccessToken]);
useEffect(() => {
fetchDocuments();
}, [fetchDocuments]);
if (documentUuids.length === 0) {
return <></>;
}
if (loading) {
return <Badge variant="outline">Loading...</Badge>;
}
return (
<>
{documentUuids.map((uuid) => (
<Badge key={uuid} variant="outline">
{documentNames[uuid] || uuid}
</Badge>
))}
</>
);
};

View file

@ -0,0 +1,181 @@
"use client";
import { Button } from "@/components/ui/button";
import { Checkbox } from "@/components/ui/checkbox";
import { Label } from "@/components/ui/label";
import { listDocumentsApiV1KnowledgeBaseDocumentsGet } from "@/client/sdk.gen";
import type { DocumentResponseSchema } from "@/client/types.gen";
import { useAuth } from "@/lib/auth";
import { FileText } from "lucide-react";
import Link from "next/link";
import { useCallback, useEffect, useState } from "react";
interface DocumentSelectorProps {
value: string[];
onChange: (uuids: string[]) => void;
disabled?: boolean;
label?: string;
description?: string;
showLabel?: boolean;
}
export const DocumentSelector = ({
value,
onChange,
disabled = false,
label = "Knowledge Base Documents",
description = "Select documents that the agent can reference during conversations.",
showLabel = true,
}: DocumentSelectorProps) => {
const { getAccessToken } = useAuth();
const [documents, setDocuments] = useState<DocumentResponseSchema[]>([]);
const [loading, setLoading] = useState(false);
const fetchDocuments = useCallback(async () => {
setLoading(true);
try {
const accessToken = await getAccessToken();
const response = await listDocumentsApiV1KnowledgeBaseDocumentsGet({
headers: { Authorization: `Bearer ${accessToken}` },
query: {
limit: 100,
},
});
if (response.data) {
// Only show completed documents
const completedDocs = response.data.documents.filter(
(doc) => doc.processing_status === "completed"
);
setDocuments(completedDocs);
}
} catch (error) {
console.error("Failed to fetch documents:", error);
} finally {
setLoading(false);
}
}, [getAccessToken]);
useEffect(() => {
fetchDocuments();
}, [fetchDocuments]);
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 (loading) {
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 text-sm text-muted-foreground text-center">
Loading documents...
</div>
</div>
);
}
if (documents.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">
{documents.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>
);
};

View file

@ -1,8 +1,10 @@
import { NodeProps, NodeToolbar, Position } from "@xyflow/react";
import { Edit, Headset, PlusIcon, Trash2Icon, Wrench } from "lucide-react";
import { Edit, FileText, Headset, PlusIcon, Trash2Icon, Wrench } from "lucide-react";
import { memo, useEffect, useMemo, useState } from "react";
import { useWorkflow } from "@/app/workflow/[workflowId]/contexts/WorkflowContext";
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 +36,8 @@ interface AgentNodeEditFormProps {
setAddGlobalPrompt: (value: boolean) => void;
toolUuids: string[];
setToolUuids: (value: string[]) => void;
documentUuids: string[];
setDocumentUuids: (value: string[]) => void;
}
interface AgentNodeProps extends NodeProps {
@ -55,6 +59,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 +80,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 +100,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,6 +116,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 ?? []);
}
}, [data, open]);
@ -139,6 +147,15 @@ export const AgentNode = memo(({ data, selected, id }: AgentNodeProps) => {
<ToolBadges toolUuids={data.tool_uuids} />
</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} />
</div>
)}
</NodeContent>
<NodeToolbar isVisible={selected} position={Position.Right}>
@ -179,6 +196,8 @@ export const AgentNode = memo(({ data, selected, id }: AgentNodeProps) => {
setAddGlobalPrompt={setAddGlobalPrompt}
toolUuids={toolUuids}
setToolUuids={setToolUuids}
documentUuids={documentUuids}
setDocumentUuids={setDocumentUuids}
/>
)}
</NodeEditDialog>
@ -203,6 +222,8 @@ const AgentNodeEditForm = ({
setAddGlobalPrompt,
toolUuids,
setToolUuids,
documentUuids,
setDocumentUuids,
}: AgentNodeEditFormProps) => {
const handleVariableNameChange = (idx: number, value: string) => {
const newVars = [...variables];
@ -346,6 +367,15 @@ const AgentNodeEditForm = ({
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}
description="Select documents from the knowledge base that the agent can reference during this conversation step."
/>
</div>
</div>
);
};

View file

@ -1,8 +1,10 @@
import { NodeProps, NodeToolbar, Position } from "@xyflow/react";
import { Edit, Play, PlusIcon, Trash2Icon, Wrench } from "lucide-react";
import { Edit, FileText, Play, PlusIcon, Trash2Icon, Wrench } from "lucide-react";
import { memo, useEffect, useMemo, useState } from "react";
import { useWorkflow } from "@/app/workflow/[workflowId]/contexts/WorkflowContext";
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 +43,8 @@ interface StartCallEditFormProps {
setVariables: (vars: ExtractionVariable[]) => void;
toolUuids: string[];
setToolUuids: (value: string[]) => void;
documentUuids: string[];
setDocumentUuids: (value: string[]) => void;
}
interface StartCallNodeProps extends NodeProps {
@ -66,6 +70,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 +94,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 +117,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,6 +136,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 ?? []);
}
}, [data, open]);
@ -158,6 +166,15 @@ export const StartCall = memo(({ data, selected, id }: StartCallNodeProps) => {
<ToolBadges toolUuids={data.tool_uuids} />
</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} />
</div>
)}
</NodeContent>
<NodeToolbar isVisible={selected} position={Position.Right}>
@ -199,6 +216,8 @@ export const StartCall = memo(({ data, selected, id }: StartCallNodeProps) => {
setVariables={setVariables}
toolUuids={toolUuids}
setToolUuids={setToolUuids}
documentUuids={documentUuids}
setDocumentUuids={setDocumentUuids}
/>
)}
</NodeEditDialog>
@ -229,6 +248,8 @@ const StartCallEditForm = ({
setVariables,
toolUuids,
setToolUuids,
documentUuids,
setDocumentUuids,
}: StartCallEditFormProps) => {
const handleVariableNameChange = (idx: number, value: string) => {
const newVars = [...variables];
@ -417,6 +438,15 @@ const StartCallEditForm = ({
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}
description="Select documents from the knowledge base that the agent can reference during this conversation step."
/>
</div>
</div>
);
};

View file

@ -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 = {