haiku.rag/examples/ag-ui-research/frontend/components/Agent.tsx

449 lines
9.6 KiB
TypeScript

"use client";
import {
CopilotKit,
useCoAgent,
useCopilotAction,
useCopilotContext,
} from "@copilotkit/react-core";
import { CopilotChat } from "@copilotkit/react-ui";
import "@copilotkit/react-ui/styles.css";
import { useState } from "react";
import DocumentSelector from "./DocumentSelector";
import StateDisplay from "./StateDisplay";
interface SearchAnswer {
query: string;
answer: string;
confidence: number;
cited_chunks: string[];
citations: {
document_id: string;
chunk_id: string;
document_uri: string;
document_title?: string;
page_numbers: number[];
headings?: string[];
content: string;
}[];
}
interface ResearchState {
context: {
original_question: string;
sub_questions: string[];
qa_responses: SearchAnswer[];
};
iterations: number;
max_iterations: number;
confidence_threshold: number;
max_concurrency: number;
last_eval: {
new_questions: string[];
confidence_score: number;
is_sufficient: boolean;
reasoning: string;
} | null;
result?: {
title: string;
executive_summary: string;
main_findings: string[];
conclusions: string[];
limitations: string[];
recommendations: string[];
sources_summary: string;
};
current_activity?: string;
current_activity_message?: string;
documentFilter?: string[];
}
interface DecisionArgs {
original_question: string;
sub_questions: string[];
qa_responses: SearchAnswer[];
}
type DecisionAction = "search" | "synthesize" | "modify_questions";
interface DecisionResult {
action: DecisionAction;
questions?: string[];
}
function DecisionUI({
args,
onResolve,
}: {
args: DecisionArgs;
onResolve: (result: DecisionResult) => void | Promise<void>;
}) {
const [editableQuestions, setEditableQuestions] = useState<string[]>(
args.sub_questions || [],
);
const [newQuestion, setNewQuestion] = useState("");
const [submitting, setSubmitting] = useState(false);
const qaCount = args.qa_responses?.length || 0;
const hasQuestions = editableQuestions.length > 0;
const canSearch = hasQuestions && !submitting;
const canSynthesize = qaCount > 0 && !submitting;
const questionsModified =
editableQuestions.length !== args.sub_questions.length ||
editableQuestions.some((q, i) => q !== args.sub_questions[i]);
const handleSubmit = (action: DecisionAction, questions?: string[]) => {
setSubmitting(true);
onResolve({ action, questions });
};
const handleSearch = () => {
handleSubmit(
questionsModified ? "modify_questions" : "search",
editableQuestions,
);
};
const handleSynthesize = () => {
handleSubmit("synthesize");
};
const handleRemoveQuestion = (index: number) => {
if (submitting) return;
setEditableQuestions(editableQuestions.filter((_, i) => i !== index));
};
const handleAddQuestion = () => {
if (submitting || !newQuestion.trim()) return;
setEditableQuestions([...editableQuestions, newQuestion.trim()]);
setNewQuestion("");
};
if (submitting) {
return null;
}
return (
<div
style={{
marginBottom: "1rem",
background: "#f0f9ff",
border: "2px solid #0ea5e9",
borderRadius: "8px",
padding: "1rem",
}}
>
<div
style={{
fontWeight: "bold",
color: "#0369a1",
marginBottom: "0.75rem",
fontSize: "1rem",
}}
>
Research Decision Point
</div>
<div
style={{
fontSize: "0.85rem",
color: "#64748b",
marginBottom: "0.75rem",
}}
>
{qaCount} answers collected
</div>
<div style={{ marginBottom: "0.75rem" }}>
<div
style={{
fontSize: "0.8rem",
color: "#475569",
marginBottom: "0.5rem",
}}
>
Pending Questions ({editableQuestions.length}):
</div>
{editableQuestions.map((q, idx) => (
<div
key={`question-${idx}`}
style={{
display: "flex",
alignItems: "center",
gap: "0.5rem",
padding: "0.375rem 0.5rem",
background: "white",
borderRadius: "4px",
marginBottom: "0.25rem",
fontSize: "0.85rem",
}}
>
<span style={{ flex: 1 }}>{q}</span>
<button
type="button"
onClick={() => handleRemoveQuestion(idx)}
style={{
background: "#ef4444",
color: "white",
border: "none",
borderRadius: "4px",
padding: "0.25rem 0.5rem",
cursor: "pointer",
fontSize: "0.75rem",
}}
>
Remove
</button>
</div>
))}
</div>
<div style={{ display: "flex", gap: "0.5rem", marginBottom: "1rem" }}>
<input
type="text"
value={newQuestion}
onChange={(e) => setNewQuestion(e.target.value)}
placeholder="Add a new question..."
style={{
flex: 1,
padding: "0.5rem",
border: "1px solid #cbd5e1",
borderRadius: "4px",
fontSize: "0.85rem",
}}
onKeyDown={(e) => {
if (e.key === "Enter") handleAddQuestion();
}}
/>
<button
type="button"
onClick={handleAddQuestion}
style={{
background: "#22c55e",
color: "white",
border: "none",
borderRadius: "4px",
padding: "0.5rem 1rem",
cursor: "pointer",
fontSize: "0.85rem",
}}
>
Add
</button>
</div>
<div style={{ display: "flex", gap: "0.5rem" }}>
<button
type="button"
onClick={handleSearch}
disabled={!canSearch}
style={{
flex: 1,
background: canSearch ? "#0ea5e9" : "#94a3b8",
color: "white",
border: "none",
borderRadius: "4px",
padding: "0.75rem",
cursor: canSearch ? "pointer" : "not-allowed",
fontWeight: "bold",
fontSize: "0.9rem",
}}
>
Search ({editableQuestions.length})
</button>
<button
type="button"
onClick={handleSynthesize}
disabled={!canSynthesize}
style={{
flex: 1,
background: canSynthesize ? "#8b5cf6" : "#94a3b8",
color: "white",
border: "none",
borderRadius: "4px",
padding: "0.75rem",
cursor: canSynthesize ? "pointer" : "not-allowed",
fontWeight: "bold",
fontSize: "0.9rem",
}}
>
Generate Report
</button>
</div>
</div>
);
}
const BACKEND_URL =
process.env.NEXT_PUBLIC_BACKEND_URL || "http://localhost:8000";
function AgentContent() {
const { state, setState, running } = useCoAgent<ResearchState>({
name: "research_agent",
initialState: {
context: {
original_question: "",
sub_questions: [],
qa_responses: [],
},
iterations: 0,
max_iterations: 3,
confidence_threshold: 0.8,
max_concurrency: 1,
last_eval: null,
documentFilter: [],
},
});
const { threadId } = useCopilotContext();
const handleDocumentFilterChange = (ids: string[]) => {
setState({ ...state, documentFilter: ids });
};
const sendToolResult = async (result: DecisionResult) => {
if (!threadId) {
console.error("No threadId available to send tool result");
return;
}
try {
const response = await fetch(`${BACKEND_URL}/v1/research/stream`, {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
threadId,
messages: [
{
id: crypto.randomUUID(),
role: "tool",
content: JSON.stringify(result),
},
],
}),
});
if (!response.ok) {
console.error("Failed to send tool result:", response.status);
}
} catch (error) {
console.error("Error sending tool result:", error);
}
};
useCopilotAction({
name: "human_decision",
description: "Pause for human decision on research direction",
parameters: [
{
name: "original_question",
type: "string",
description: "The original research question",
},
{
name: "sub_questions",
type: "string[]",
description: "Pending sub-questions to search",
},
{
name: "qa_responses",
type: "object[]",
description: "Answers collected so far",
},
],
renderAndWaitForResponse: ({ args, status }) => {
if (status === "complete") {
return null;
}
return (
<DecisionUI
args={args as unknown as DecisionArgs}
onResolve={sendToolResult}
/>
);
},
});
return (
<>
<style>{`
.chat-container {
width: 50%;
height: 100vh;
border-right: 1px solid #e2e8f0;
display: flex;
flex-direction: column;
}
.chat-container > * {
flex: 1;
min-height: 0;
}
`}</style>
<div style={{ display: "flex", height: "100vh" }}>
<div className="chat-container">
<CopilotChat
labels={{
title: "Research Assistant",
initial:
"Hello! I can help you conduct deep research on complex questions using the haiku.rag knowledge base. Ask me anything!",
}}
/>
</div>
<div
style={{
width: "50%",
height: "100vh",
overflow: "auto",
background: "#f7fafc",
}}
>
<div style={{ padding: "2rem" }}>
<header style={{ marginBottom: "2rem" }}>
<h1
style={{
fontSize: "2rem",
fontWeight: "bold",
marginBottom: "0.5rem",
color: "#1a202c",
}}
>
Research State
</h1>
<p
style={{
fontSize: "0.875rem",
color: "#4a5568",
lineHeight: "1.6",
}}
>
Live updates from the research agent
</p>
</header>
{!running && (
<div style={{ marginBottom: "1rem" }}>
<DocumentSelector
selected={state.documentFilter || []}
onChange={handleDocumentFilterChange}
/>
</div>
)}
<StateDisplay state={state} />
</div>
</div>
</div>
</>
);
}
export default function Agent() {
return (
<CopilotKit runtimeUrl="/api/copilotkit" agent="research_agent">
<AgentContent />
</CopilotKit>
);
}