diff --git a/CHANGELOG.md b/CHANGELOG.md index e986e0b0..acb50d7a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,19 @@ ### Added +- **Interactive Research Mode**: Human-in-the-loop research using graph-based decision nodes + - `haiku-rag research --interactive` starts conversational CLI chat + - Natural language interpretation for user commands (search, modify questions, synthesize) + - Chat with assistant before starting research, and during decision points + - Review collected answers and pending questions at each decision point + - Add, remove, or modify sub-questions through natural conversation + - New `human_decide` graph node emits AG-UI tool calls (`TOOL_CALL_START/ARGS/END`) for frontend integration + - New `emit_tool_call_start()`, `emit_tool_call_args()`, `emit_tool_call_end()` AG-UI event helpers + - New `AGUIEmitter.emit()` method for direct event emission +- **AG-UI Research Example**: Human-in-the-loop research with client-side tool calling + - Frontend handles `human_decision` tool calls via AG-UI `TOOL_CALL_*` events + - Tool results sent directly to backend `/v1/research/stream` endpoint + - Backend queues decisions and continues the research graph - **HotpotQA Evaluation**: Added HotpotQA dataset adapter for multi-hop QA benchmarks - Extracts unique documents from validation set context paragraphs - Uses MAP for retrieval evaluation (multiple supporting documents per question) diff --git a/docs/agents.md b/docs/agents.md index 59c48482..ff11b1bf 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -196,3 +196,35 @@ report = await graph.run(state=state, deps=deps) ``` The filter applies to all search operations in the graph. See [Filtering Search Results](python.md#filtering-search-results) for available filter columns and syntax. + +### Interactive Research Mode + +Interactive mode provides human-in-the-loop control over the research process through a conversational interface. + +**CLI usage:** + +```bash +# Start interactive research mode +haiku-rag research --interactive + +# Start with a specific question +haiku-rag research --interactive "How does X work?" + +# With document filter +haiku-rag research --interactive --filter "uri LIKE '%report%'" +``` + +In interactive mode, you can: + +- Chat with the assistant before starting research +- Review the generated sub-questions after planning +- Add, remove, or modify questions through natural conversation +- Execute searches and review collected answers +- Continue researching or synthesize when ready + +For a web-based interactive experience, see the [AG-UI Research Example](https://github.com/ggozad/haiku.rag/tree/main/examples/ag-ui-research). The example demonstrates AG-UI client-side tool calling: + +- Frontend handles `human_decision` tool calls via AG-UI `TOOL_CALL_*` events +- Decision UI rendered inline in the chat at each decision point +- Question editing (add/remove) and action buttons (Search, Generate Report) +- Tool results sent directly to the backend endpoint which queues decisions and continues the graph diff --git a/docs/cli.md b/docs/cli.md index 3e0edf0f..ad59fcd9 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -183,10 +183,24 @@ Filter to specific documents: haiku-rag research "What are the key findings?" --filter "uri LIKE '%paper%'" ``` +Interactive mode with human-in-the-loop: + +```bash +# Start interactive research mode +haiku-rag research --interactive + +# Start with a specific question +haiku-rag research --interactive "How does haiku.rag work?" + +# With document filter +haiku-rag research --interactive --filter "uri LIKE '%docs%'" +``` + Flags: - `--verbose`: Show planning, searching previews, evaluation summary, and stop reason - `--filter`: SQL WHERE clause to filter documents (see [Filtering Search Results](python.md#filtering-search-results)) +- `--interactive` / `-i`: Start interactive research mode with human-in-the-loop decision points Research parameters like `max_iterations`, `confidence_threshold`, and `max_concurrency` are configured in your [configuration file](configuration/index.md) under the `research` section. diff --git a/examples/a2a-server/haiku_rag_a2a/a2a/__init__.py b/examples/a2a-server/haiku_rag_a2a/a2a/__init__.py index 745c9903..11ab5cac 100644 --- a/examples/a2a-server/haiku_rag_a2a/a2a/__init__.py +++ b/examples/a2a-server/haiku_rag_a2a/a2a/__init__.py @@ -6,7 +6,7 @@ import logfire from pydantic_ai import Agent, RunContext from haiku.rag.config import AppConfig, Config -from haiku.rag.graph.common import get_model +from haiku.rag.utils import get_model from .context import load_message_history, save_message_history from .models import A2AConfig, AgentDependencies, SearchResult diff --git a/examples/ag-ui-research/README.md b/examples/ag-ui-research/README.md index df6850dc..3925fba1 100644 --- a/examples/ag-ui-research/README.md +++ b/examples/ag-ui-research/README.md @@ -1,13 +1,13 @@ # Interactive Research Assistant -Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/), [Pydantic Graph](https://ai.pydantic.dev/graph/), and [AG-UI](https://docs.ag-ui.com/). Ask complex questions and watch the research process unfold in real-time. +Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/), [Pydantic Graph](https://ai.pydantic.dev/graph/), and [AG-UI](https://docs.ag-ui.com/). Ask complex questions and watch the research process unfold in real-time with human-in-the-loop control. [Watch demo video](https://vimeo.com/1128874386) ## Features -- **Multi-iteration research graph**: Automated question decomposition and search -- **Intelligent evaluation**: Confidence-based decision making with automatic iteration until sufficient information is gathered +- **Human-in-the-loop research**: Review and modify questions at decision points, then continue searching or generate report +- **Multi-iteration research graph**: Automated question decomposition and parallel search - **Live state synchronization**: Real-time delta updates of research progress via AG-UI protocol - **Rich reporting**: Generates comprehensive research reports with findings, conclusions, and sources @@ -25,9 +25,7 @@ Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/), **Option A: Create a new database** ```bash - mkdir -p data - haiku-rag add "Your documents here" --db data/haiku_rag.lancedb - # Or add from files + haiku-rag init --db data/haiku_rag.lancedb haiku-rag add-src document.pdf --db data/haiku_rag.lancedb ``` @@ -63,27 +61,29 @@ Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/), DB_PATH=/path/to/your/existing/haiku_rag.lancedb # If using an existing db. ``` -1. **Start the application** +4. **Start the application** ```bash docker compose up --build ``` -2. **Access the interface** +5. **Access the interface** - Frontend: http://localhost:3000 - Backend health: http://localhost:8000/health ## How It Works 1. **Ask a question**: Type your research question in the chat -2. **Plan phase**: The research graph automatically: - - Decomposes your question into targeted sub-questions - - Gathers initial context about the topic -3. **Research iterations**: The graph autonomously: - - Searches the knowledge base for each sub-question in parallel - - Assesses confidence in gathered information - - Generates new follow-up questions if needed - - Iterates until confidence threshold is met or max iterations reached -4. **Synthesis**: Generates a comprehensive research report with: +2. **Plan phase**: The research graph decomposes your question into targeted sub-questions +3. **Decision point**: Review the proposed questions in the right panel + - Add new questions using the input field + - Remove questions you don't need + - Click **Search** to execute searches for pending questions + - Click **Generate Report** to skip to synthesis (when you have enough answers) +4. **Research iterations**: After each search cycle, you return to a decision point where you can: + - Review collected answers + - Add follow-up questions based on findings + - Continue searching or generate the final report +5. **Synthesis**: Generates a comprehensive research report with: - Executive summary - Main findings with supporting evidence - Conclusions and recommendations @@ -93,38 +93,42 @@ Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/), ### Agent + Graph Pattern -This example demonstrates the **agent+graph** architecture pattern: +This example demonstrates the **agent+graph** architecture with AG-UI client-side tool calls: 1. **Conversational Agent** (`agent.py`): - Pydantic AI agent handles user conversations - Decides when to invoke the research tool based on user intent - Responds directly to greetings/casual chat without tools - - Formats research results for the user -2. **Research Graph** (haiku.rag): +2. **Interactive Research Graph** (haiku.rag): - Multi-step research workflow invoked by the agent's tool - - Autonomous execution with plan → search → analyze → decide → synthesize flow - - Emits AG-UI events for real-time progress tracking + - At decision points, emits AG-UI `TOOL_CALL_START/ARGS/END` events for `human_decision` + - Waits for tool result via async queue before continuing -3. **Shared Event Stream**: +3. **Client-Side Tool Handling** (AG-UI pattern): + - Frontend listens for `human_decision` tool calls via AG-UI events + - Renders decision UI inline in chat when tool call is received + - User decision sent directly to backend `/v1/research/stream` endpoint + - Backend extracts tool result from messages and routes to waiting graph via async queue + +4. **Shared Event Stream**: - `AGUIEmitter` is shared between agent and graph - Events from both flow through a single stream to the frontend - - Custom streaming endpoint (`main.py`) uses anyio memory streams for proper async handling + - `STATE_DELTA` events sync research state to frontend in real-time ### Components - **Backend** (Python): - Uses published `ghcr.io/ggozad/haiku.rag:latest` Docker image as base - - `agent.py`: Pydantic AI agent with `run_research` tool - - `main.py`: Custom AG-UI streaming endpoint with anyio memory object streams + - `agent.py`: Pydantic AI agent with `run_research` tool, manages `ActiveResearch` registry + - `main.py`: Custom AG-UI streaming endpoint, extracts tool results from messages - Real-time event forwarding from emitter to SSE stream - - Filters out `ACTIVITY_SNAPSHOT` events (not yet supported by CopilotKit) - **Frontend** (Next.js/React): - - CopilotKit for AG-UI protocol integration + - AG-UI protocol integration for real-time streaming + - Handles `human_decision` tool calls with inline decision UI - Split-pane UI: chat on left, live research state on right - - Real-time state synchronization via Server-Sent Events (SSE) - - `StateDisplay` component with collapsible sections for questions and report + - Tool results sent directly to backend endpoint ## Configuration diff --git a/examples/ag-ui-research/backend/agent.py b/examples/ag-ui-research/backend/agent.py index e697145f..c30e17c2 100644 --- a/examples/ag-ui-research/backend/agent.py +++ b/examples/ag-ui-research/backend/agent.py @@ -1,6 +1,7 @@ """Research assistant agent with graph integration.""" -from dataclasses import dataclass +import asyncio +from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING @@ -9,10 +10,10 @@ from pydantic_ai import Agent, RunContext from haiku.rag.client import HaikuRAG from haiku.rag.config import load_yaml_config from haiku.rag.config.models import AppConfig -from haiku.rag.graph.common import get_model from haiku.rag.graph.research.dependencies import ResearchContext from haiku.rag.graph.research.graph import build_research_graph -from haiku.rag.graph.research.state import ResearchDeps, ResearchState +from haiku.rag.graph.research.state import HumanDecision, ResearchDeps, ResearchState +from haiku.rag.utils import get_model if TYPE_CHECKING: from haiku.rag.graph.agui.emitter import AGUIEmitter @@ -27,6 +28,20 @@ Config = ( ) +@dataclass +class ActiveResearch: + """Tracks state for active research awaiting human decision.""" + + queue: asyncio.Queue[HumanDecision] + sub_questions: list[str] = field(default_factory=list) + qa_responses: list[dict] = field(default_factory=list) + original_question: str = "" + + +# Global registry of active research by thread_id +_active_research: dict[str, ActiveResearch] = {} + + @dataclass class AgentDeps: """Dependencies for research agent.""" @@ -34,6 +49,8 @@ class AgentDeps: client: HaikuRAG agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None search_filter: str | None = None + thread_id: str | None = None + research_result: "ResearchReport | None" = None model = get_model(Config.research.model, Config) @@ -50,10 +67,10 @@ CRITICAL RULES: 4. NEVER answer substantive questions from your own knowledge - always use the tool How to decide: -- "Hi" / "Hello" / "How are you?" → Respond directly, NO tools -- "What can you do?" → Respond directly, NO tools -- "How does X work in the codebase?" → Use run_research tool -- "Tell me about Y" → Use run_research tool +- "Hi" / "Hello" / "How are you?" -> Respond directly, NO tools +- "What can you do?" -> Respond directly, NO tools +- "How does X work in the codebase?" -> Use run_research tool +- "Tell me about Y" -> Use run_research tool When you use run_research, the graph will decompose questions, search the knowledge base, and generate a comprehensive report. @@ -70,23 +87,41 @@ async def run_research(ctx: RunContext[AgentDeps], question: str) -> str: DO NOT use for greetings or casual conversation. """ if ctx.deps.agui_emitter: - ctx.deps.agui_emitter.log(f"🔍 Starting research on: {question}") + ctx.deps.agui_emitter.log(f"Starting research on: {question}") - graph = build_research_graph(Config) + # Create queue for human decisions + queue: asyncio.Queue[HumanDecision] = asyncio.Queue() + + # Build interactive graph + graph = build_research_graph(Config, interactive=True) context = ResearchContext(original_question=question) state = ResearchState.from_config(context=context, config=Config) state.search_filter = ctx.deps.search_filter + # Register active research for decision endpoint to find + thread_id = ctx.deps.thread_id + if thread_id: + _active_research[thread_id] = ActiveResearch( + queue=queue, + sub_questions=[], + qa_responses=[], + original_question=question, + ) + graph_deps = ResearchDeps( client=ctx.deps.client, agui_emitter=ctx.deps.agui_emitter, + human_input_queue=queue, + interactive=True, ) try: result = await graph.run(state=state, deps=graph_deps) if ctx.deps.agui_emitter: - ctx.deps.agui_emitter.log("✅ Research complete!") + ctx.deps.agui_emitter.log("Research complete!") + # Store result for main.py to emit RUN_FINISHED after agent completes + ctx.deps.research_result = result return f"""Research completed successfully! @@ -108,5 +143,9 @@ The full research report with all citations has been provided to the user. except Exception as e: if ctx.deps.agui_emitter: - ctx.deps.agui_emitter.log(f"❌ Research error: {str(e)}") + ctx.deps.agui_emitter.log(f"Research error: {str(e)}") return f"I encountered an error while researching: {str(e)}" + finally: + # Cleanup + if thread_id and thread_id in _active_research: + del _active_research[thread_id] diff --git a/examples/ag-ui-research/backend/main.py b/examples/ag-ui-research/backend/main.py index d432a8ba..a5211c4b 100644 --- a/examples/ag-ui-research/backend/main.py +++ b/examples/ag-ui-research/backend/main.py @@ -1,8 +1,9 @@ +import json import logging import os from pathlib import Path -from agent import AgentDeps, agent +from agent import AgentDeps, _active_research, agent from anyio import create_memory_object_stream, create_task_group from anyio.streams.memory import MemoryObjectSendStream from starlette.applications import Starlette @@ -19,7 +20,7 @@ from haiku.rag.graph.agui.emitter import AGUIEmitter from haiku.rag.graph.agui.server import RunAgentInput, format_sse_event from haiku.rag.graph.research.dependencies import ResearchContext from haiku.rag.graph.research.models import ResearchReport -from haiku.rag.graph.research.state import ResearchState +from haiku.rag.graph.research.state import HumanDecision, ResearchState logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" @@ -62,11 +63,54 @@ def get_client(effective_db_path: Path) -> HaikuRAG: return _client_cache[path_key] +def extract_tool_result(messages: list[dict]) -> dict | None: + """Extract human_decision tool result from messages if present.""" + for msg in reversed(messages): + # Check for tool result message (CopilotKit sends role="tool") + if msg.get("role") == "tool": + content = msg.get("content") + # Content may be a string (JSON) or dict + if isinstance(content, str): + try: + content = json.loads(content) + except json.JSONDecodeError: + continue + if isinstance(content, dict) and "action" in content: + return content + return None + + async def stream_research_agent(request: Request) -> StreamingResponse: """Agent streaming endpoint with research graph integration.""" body = await request.json() + logger.info(f"Received request body keys: {list(body.keys())}") + if "tools" in body: + logger.info(f"Frontend tools received: {body['tools']}") input_data = RunAgentInput(**body) + thread_id = input_data.thread_id + active_research = _active_research.get(thread_id) if thread_id else None + + # Check if this is a tool result for active research + if active_research and input_data.messages: + tool_result = extract_tool_result(input_data.messages) + if tool_result: + logger.info(f"Received tool result: {tool_result}") + action = tool_result.get("action", "search") + questions = tool_result.get("questions") + + decision = HumanDecision( + action=action, + questions=questions, + ) + await active_research.queue.put(decision) + + # Return acknowledgment - the original stream will continue + return StreamingResponse( + iter([format_sse_event({"type": "TOOL_RESULT_RECEIVED"})]), + media_type="text/event-stream", + ) + user_message = "" if input_data.messages: user_message = input_data.messages[-1].get("content", "") @@ -79,11 +123,11 @@ async def stream_research_agent(request: Request) -> StreamingResponse: """Execute agent and forward emitter events to memory stream.""" async with send_stream: try: - # Create shared emitter + # Create shared emitter (use_deltas=True for CopilotKit compatibility) emitter: AGUIEmitter[ResearchState, ResearchReport] = AGUIEmitter( thread_id=input_data.thread_id, run_id=input_data.run_id, - use_deltas=False, + use_deltas=True, ) # Get client @@ -104,6 +148,7 @@ async def stream_research_agent(request: Request) -> StreamingResponse: client=client, agui_emitter=emitter, search_filter=search_filter, + thread_id=thread_id, ) # Start run with empty initial state @@ -117,9 +162,12 @@ async def stream_research_agent(request: Request) -> StreamingResponse: # Forward emitter events to stream async def forward_events(): async for event in emitter: - # Log events for debugging - logger.info(f"AG-UI Event: {event}") event_type = event.get("type") + logger.info(f"AG-UI event: {event_type}") + + # Log tool call events for debugging + if event_type and event_type.startswith("TOOL_CALL"): + logger.info(f"Tool call event: {event}") # Convert ACTIVITY_SNAPSHOT to STATE_DELTA for CopilotKit # As CopilotKit does not handle ACTIVITY_SNAPSHOT events @@ -129,7 +177,6 @@ async def stream_research_agent(request: Request) -> StreamingResponse: message = content.get("message", "") # Emit STATE_DELTA to patch activity info into state - # Use "add" op which creates or replaces the value delta_event = { "type": "STATE_DELTA", "delta": [ @@ -148,6 +195,21 @@ async def stream_research_agent(request: Request) -> StreamingResponse: await send_stream.send(format_sse_event(delta_event)) continue + # Sync state to ActiveResearch when human_decision tool call + if event_type == "TOOL_CALL_ARGS" and thread_id: + delta = event.get("delta", "{}") + args = ( + json.loads(delta) if isinstance(delta, str) else delta + ) + active = _active_research.get(thread_id) + if active: + active.sub_questions = list( + args.get("sub_questions", []) + ) + active.qa_responses = list(args.get("qa_responses", [])) + if "original_question" in args: + active.original_question = args["original_question"] + await send_stream.send(format_sse_event(event)) # Run agent and event forwarding concurrently @@ -156,6 +218,9 @@ async def stream_research_agent(request: Request) -> StreamingResponse: result = await agent.run(user_message, deps=agent_deps) emitter.log(result.output) + # Emit RUN_FINISHED with research result if available + if agent_deps.research_result: + emitter.finish_run(agent_deps.research_result) await emitter.close() except Exception as e: diff --git a/examples/ag-ui-research/docker-compose.yml b/examples/ag-ui-research/docker-compose.yml index 5ea19039..5d83bfbc 100644 --- a/examples/ag-ui-research/docker-compose.yml +++ b/examples/ag-ui-research/docker-compose.yml @@ -32,9 +32,13 @@ services: - OLLAMA_BASE_URL=${OLLAMA_BASE_URL} # Prevent Python bytecode caching for development - PYTHONDONTWRITEBYTECODE=1 + # Use local haiku_rag_slim for development + - PYTHONPATH=/app/haiku_rag_slim volumes: - ${DB_PATH}:/app/data/haiku.rag.lancedb - ./haiku.rag.yaml:/app/haiku.rag.yaml:ro + # Mount local haiku_rag_slim for development + - ../../haiku_rag_slim:/app/haiku_rag_slim:ro networks: - ag-ui-network extra_hosts: diff --git a/examples/ag-ui-research/frontend/Dockerfile b/examples/ag-ui-research/frontend/Dockerfile index b5da8b25..0959f3ef 100644 --- a/examples/ag-ui-research/frontend/Dockerfile +++ b/examples/ag-ui-research/frontend/Dockerfile @@ -4,7 +4,7 @@ FROM node:22-alpine WORKDIR /app COPY package.json package-lock.json* ./ -RUN npm ci +RUN npm install --legacy-peer-deps COPY . . EXPOSE 3000 diff --git a/examples/ag-ui-research/frontend/components/Agent.tsx b/examples/ag-ui-research/frontend/components/Agent.tsx index 75429a91..2f1ebc92 100644 --- a/examples/ag-ui-research/frontend/components/Agent.tsx +++ b/examples/ag-ui-research/frontend/components/Agent.tsx @@ -1,64 +1,282 @@ "use client"; -import { CopilotKit, useCoAgent } from "@copilotkit/react-core"; +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 Citation { - document_id: string; - chunk_id: string; - document_uri: string; - document_title?: string; - page_numbers: number[]; - headings?: string[]; - content: string; -} - interface SearchAnswer { query: string; answer: string; confidence: number; cited_chunks: string[]; - citations: Citation[]; + citations: { + document_id: string; + chunk_id: string; + document_uri: string; + document_title?: string; + page_numbers: number[]; + headings?: string[]; + content: string; + }[]; } -interface ResearchContext { +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[]; } -interface EvaluationResult { - new_questions: string[]; - confidence_score: number; - is_sufficient: boolean; - reasoning: string; +type DecisionAction = "search" | "synthesize" | "modify_questions"; + +interface DecisionResult { + action: DecisionAction; + questions?: string[]; } -interface ResearchReport { - title: string; - executive_summary: string; - main_findings: string[]; - conclusions: string[]; - limitations: string[]; - recommendations: string[]; - sources_summary: string; +function DecisionUI({ + args, + onResolve, +}: { + args: DecisionArgs; + onResolve: (result: DecisionResult) => void | Promise; +}) { + const [editableQuestions, setEditableQuestions] = useState( + 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 ( +
+
+ Research Decision Point +
+ +
+ {qaCount} answers collected +
+ +
+
+ Pending Questions ({editableQuestions.length}): +
+ {editableQuestions.map((q, idx) => ( +
+ {q} + +
+ ))} +
+ +
+ 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(); + }} + /> + +
+ +
+ + +
+
+ ); } -interface ResearchState { - context: ResearchContext; - iterations: number; - max_iterations: number; - confidence_threshold: number; - max_concurrency: number; - last_eval: EvaluationResult | null; - result?: ResearchReport; - current_activity?: string; - current_activity_message?: string; - documentFilter?: string[]; -} +const BACKEND_URL = + process.env.NEXT_PUBLIC_BACKEND_URL || "http://localhost:8000"; function AgentContent() { const { state, setState, running } = useCoAgent({ @@ -78,10 +296,76 @@ function AgentContent() { }, }); + 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 ( + + ); + }, + }); + return ( <>
- {/* Chat on the left */}
- {/* State display on the right */}
-
- -
+ {!running && ( +
+ +
+ )}
diff --git a/examples/ag-ui-research/frontend/components/DocumentSelector.tsx b/examples/ag-ui-research/frontend/components/DocumentSelector.tsx index 33060026..7c4c93bd 100644 --- a/examples/ag-ui-research/frontend/components/DocumentSelector.tsx +++ b/examples/ag-ui-research/frontend/components/DocumentSelector.tsx @@ -298,22 +298,29 @@ export default function DocumentSelector({ }} />
+ {doc.title && ( +
+ {doc.title} +
+ )}
- {doc.title || "Untitled"} -
-
{/* Question */} - {state.context.original_question && ( + {state.context?.original_question && (
)} - {/* Research Progress - only show when research has started */} - {(state.iterations > 0 || (state.current_activity && !state.result)) && ( + {/* Research Progress - only show when research is in progress (not when complete) */} + {(state.iterations > 0 || state.current_activity) && !state.result && (
- {/* Current Activity - hide when complete */} - {state.current_activity && !state.result && ( + {/* Current Activity */} + {state.current_activity && (
)} - {/* Sub-Questions and QA Responses */} - {(state.context.sub_questions.length > 0 || - state.context.qa_responses.length > 0) && ( + {/* Answers */} + {state.context?.qa_responses && state.context.qa_responses.length > 0 && (
- - Sub-Questions ({state.context.sub_questions.length}) • Answers ( - {state.context.qa_responses.length}) - + Answers ({state.context.qa_responses.length}) {expandedSections.questions ? "▼" : "▶"} {expandedSections.questions && ( @@ -408,38 +404,6 @@ export default function StateDisplay({ state }: StateDisplayProps) { borderRadius: "0 0 4px 4px", }} > - {/* Show pending sub_questions */} - {state.context.sub_questions.map((question, idx) => ( -
-
- ⏳ -
-
- -
-
- ))} - {/* Show all qa_responses (each has query + answer) */} {state.context.qa_responses.map((qaResponse, idx) => { const questionId = `q-${idx}`; diff --git a/examples/ag-ui-research/frontend/package.json b/examples/ag-ui-research/frontend/package.json index d2e2c811..85e66050 100644 --- a/examples/ag-ui-research/frontend/package.json +++ b/examples/ag-ui-research/frontend/package.json @@ -12,9 +12,9 @@ }, "dependencies": { "@ag-ui/client": "^0.0.42", - "@copilotkit/react-core": "^1.10.6", - "@copilotkit/react-ui": "^1.10.6", - "@copilotkit/runtime": "^1.10.6", + "@copilotkit/react-core": "^1.50.0", + "@copilotkit/react-ui": "^1.50.0", + "@copilotkit/runtime": "^1.50.0", "next": "15.5.5", "react": "^19.0.0", "react-dom": "^19.0.0" diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index 42f8fbd9..bd1af435 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -321,7 +321,8 @@ def ask( @cli.command("research", help="Run multi-agent research and output a concise report") def research( question: str = typer.Argument( - help="The research question to investigate", + None, + help="The research question to investigate (required unless --interactive)", ), db: Path | None = typer.Option( None, @@ -339,9 +340,34 @@ def research( "-f", help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")", ), + interactive: bool = typer.Option( + False, + "--interactive", + "-i", + help="Start interactive research mode with human-in-the-loop", + ), ): app = create_app(db) - asyncio.run(app.research(question=question, verbose=verbose, filter=filter)) + + if interactive: + from haiku.rag.cli_chat import interactive_research + from haiku.rag.client import HaikuRAG + + client = HaikuRAG(db_path=app.db_path, config=app.config) + try: + interactive_research( + client=client, + config=app.config, + search_filter=filter, + question=question, + ) + finally: + client.close() + else: + if question is None: + typer.echo("Error: Question is required unless using --interactive mode") + raise typer.Exit(1) + asyncio.run(app.research(question=question, verbose=verbose, filter=filter)) @cli.command("settings", help="Display current configuration settings") diff --git a/haiku_rag_slim/haiku/rag/cli_chat.py b/haiku_rag_slim/haiku/rag/cli_chat.py new file mode 100644 index 00000000..68882cf3 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/cli_chat.py @@ -0,0 +1,489 @@ +"""Interactive CLI chat loop for research graph with human-in-the-loop.""" + +import asyncio +import json + +from pydantic_ai import Agent +from rich.console import Console +from rich.markdown import Markdown +from rich.panel import Panel +from rich.prompt import Prompt + +from haiku.rag.client import HaikuRAG +from haiku.rag.config import get_config +from haiku.rag.config.models import AppConfig +from haiku.rag.graph.agui.emitter import AGUIEmitter +from haiku.rag.graph.research.dependencies import ResearchContext +from haiku.rag.graph.research.graph import build_research_graph +from haiku.rag.graph.research.models import ResearchReport +from haiku.rag.graph.research.state import HumanDecision, ResearchDeps, ResearchState +from haiku.rag.utils import get_model + +INITIAL_CHAT_PROMPT = """You are a research assistant. The user hasn't started a research task yet. + +You can: +1. Chat with the user - greet them, answer questions about what you can do +2. Detect when they want to research something + +## Actions: +- "chat": User is chatting, greeting, or asking questions (set message with your response) +- "research": User wants to research a topic (extract the research question into research_question) + +## Guidelines: +- If the user provides a clear research question or topic, set action="research" and extract the question +- If the user is just chatting or asking what you can do, set action="chat" and respond helpfully +- Be friendly and explain you can help them research topics by searching a knowledge base + +Examples: +- "hi" → action="chat", message="Hello! I'm a research assistant. I can help you research topics by searching through documents and synthesizing findings. What would you like to explore?" +- "what can you do?" → action="chat", message="I help you conduct research! Give me a question or topic, and I'll break it into sub-questions, search for answers, and synthesize a report. What are you curious about?" +- "tell me about Python's memory management" → action="research", research_question="How does Python's memory management work?" +- "I want to understand how RAG systems work" → action="research", research_question="How do RAG (Retrieval-Augmented Generation) systems work?" +""" + +RESEARCH_ASSISTANT_PROMPT = """You are a research assistant helping the user conduct research on a topic. + +You are at a decision point in the research workflow. You can: +1. Chat with the user - answer questions, discuss the research, make suggestions +2. Take workflow actions when the user requests them + +## Workflow Actions (set in the action field): +- "search": Search the pending questions (user says: "go", "search", "yes", "continue", "looks good") +- "synthesize": Generate final report (user says: "done", "finish", "synthesize", "generate report") +- "add_questions": Add NEW research questions to the existing list +- "modify_questions": REPLACE all pending questions with a new list (use when user wants to remove, keep only certain questions, or change the questions) +- "chat": Have a conversation without modifying questions + +## IMPORTANT - Modifying Questions: +- "use only the first question" → action="modify_questions", questions=[first question from the list] +- "drop questions 2 and 3" → action="modify_questions", questions=[remaining questions] +- "keep only questions about X" → action="modify_questions", questions=[filtered list] +- "remove the duplicate" → action="modify_questions", questions=[deduplicated list] +- When user wants to reduce/filter/keep-only, use "modify_questions" NOT "chat" + +## Guidelines: +- If the user wants to modify the question list in ANY way (remove, keep only, filter), use "modify_questions" +- For "modify_questions", include ALL questions that should remain in the questions field +- You can combine "chat" with a message to explain what you're doing +- If just chatting without changes, set action="chat" and provide helpful response in message +""" + + +async def initial_chat( + user_message: str, + config: AppConfig, +) -> HumanDecision: + """Handle initial conversation before research starts. + + Args: + user_message: The user's message + config: Application configuration + + Returns: + HumanDecision with chat response or research question + """ + agent: Agent[None, HumanDecision] = Agent( + model=get_model(config.research.model, config), + output_type=HumanDecision, + instructions=INITIAL_CHAT_PROMPT, + retries=2, + ) + + result = await agent.run(user_message) + return result.output + + +async def interpret_user_decision( + user_message: str, + sub_questions: list[str], + qa_responses: list[dict], + config: AppConfig, +) -> HumanDecision: + """Interpret a natural language user message into a HumanDecision. + + Args: + user_message: The user's natural language input + sub_questions: Current sub-questions pending search + qa_responses: Answers already collected + config: Application configuration + + Returns: + HumanDecision with the interpreted action, questions, and/or message + """ + agent: Agent[None, HumanDecision] = Agent( + model=get_model(config.research.model, config), + output_type=HumanDecision, + instructions=RESEARCH_ASSISTANT_PROMPT, + retries=2, + ) + + # Build context with full research state + answers_summary = "" + if qa_responses: + answers_parts = [] + for qa in qa_responses: + conf = f"{qa['confidence']:.0%}" if qa.get("confidence") else "N/A" + answers_parts.append( + f"Q: {qa['query']}\nA: {qa['answer'][:300]}... (confidence: {conf})" + ) + answers_summary = "\n\n".join(answers_parts) + + context = f"""Current research state: +- Answers collected: {len(qa_responses)} +- Pending questions to search: {len(sub_questions)} + +Pending questions: +{chr(10).join(f"- {q}" for q in sub_questions) if sub_questions else "(none)"} + +{f"Collected answers:{chr(10)}{answers_summary}" if answers_summary else ""} + +User message: {user_message}""" + + result = await agent.run(context) + return result.output + + +async def run_interactive_research( + question: str, + client: HaikuRAG, + config: AppConfig | None = None, + search_filter: str | None = None, +) -> ResearchReport: + """Run interactive research with human-in-the-loop decision points. + + Args: + question: The research question + client: HaikuRAG client for document operations + config: Application configuration (uses global config if None) + search_filter: Optional SQL WHERE clause to filter documents + + Returns: + ResearchReport with the final synthesis + """ + config = config or get_config() + console = Console() + + # Build interactive graph + graph = build_research_graph(config=config, include_plan=True, interactive=True) + + # Create async queue for human input + human_input_queue: asyncio.Queue[HumanDecision] = asyncio.Queue() + + # Create emitter + emitter: AGUIEmitter[ResearchState, ResearchReport] = AGUIEmitter() + + # Create deps with queue + deps = ResearchDeps( + client=client, + agui_emitter=emitter, + human_input_queue=human_input_queue, + interactive=True, + ) + + # Create initial state + context = ResearchContext(original_question=question) + state = ResearchState.from_config(context=context, config=config) + state.search_filter = search_filter + + # Start the run + emitter.start_run(state) + + # Run graph in background task + async def run_graph() -> ResearchReport: + try: + result = await graph.run(state=state, deps=deps) + emitter.finish_run(result) + return result + except Exception as e: + emitter.error(e) + raise + + graph_task = asyncio.create_task(run_graph()) + + # Process events and handle human decision points + try: + async for event in emitter: + event_type = event.get("type") + + if event_type == "STEP_STARTED": + step_name = event.get("stepName", "") + if step_name == "plan": + console.print("[dim]Planning research...[/dim]") + elif step_name.startswith("search:"): + query = step_name.replace("search: ", "") + console.print(f"[dim]Searching: {query}[/dim]") + elif step_name == "synthesize": + console.print("[dim]Synthesizing report...[/dim]") + + elif event_type == "STATE_SNAPSHOT" or event_type == "STATE_DELTA": + # State updated, could show progress + pass + + elif event_type == "TOOL_CALL_START": + tool_name = event.get("toolCallName") + if tool_name == "human_decision": + # Will get args in next event + pass + + elif event_type == "TOOL_CALL_ARGS": + delta = event.get("delta", "{}") + args = json.loads(delta) if isinstance(delta, str) else delta + original_question = args.get("original_question", "") + sub_questions = list(args.get("sub_questions", [])) + qa_responses = args.get("qa_responses", []) + iterations = args.get("iterations", 0) + + # Loop for modifications until user wants to proceed + while True: + # Show research state + console.print() + console.print( + Panel( + f"[bold]{original_question}[/bold]", + title="Research Question", + border_style="blue", + ) + ) + + # Show collected answers + if qa_responses: + answers_text = [] + for i, qa in enumerate(qa_responses, 1): + conf = ( + f"{qa['confidence']:.0%}" + if qa.get("confidence") + else "N/A" + ) + answer_preview = ( + qa["answer"][:200] + "..." + if len(qa["answer"]) > 200 + else qa["answer"] + ) + answers_text.append( + f"[cyan]{i}. {qa['query']}[/cyan]\n" + f" [dim]Confidence: {conf} | Citations: {qa.get('citations_count', 0)}[/dim]\n" + f" {answer_preview}" + ) + console.print( + Panel( + "\n\n".join(answers_text), + title=f"Answers Collected ({len(qa_responses)})", + border_style="green", + ) + ) + + # Show pending questions + if sub_questions: + console.print( + Panel( + "\n".join( + f"{i + 1}. {q}" for i, q in enumerate(sub_questions) + ), + title="Pending Questions to Search", + border_style="cyan", + ) + ) + else: + console.print("[dim]No pending questions.[/dim]") + + if iterations > 0: + console.print(f"[dim]Iteration: {iterations}[/dim]") + + # Prompt user with context-aware hints + console.print() + hints = [] + if sub_questions: + hints.append("search questions") + hints.append("modify questions") + if qa_responses: + hints.append("generate report") + hint_text = f" [dim]({', '.join(hints)})[/dim]" if hints else "" + user_input = Prompt.ask( + f"[bold]What would you like to do?[/bold]{hint_text}" + ) + + # Chat with research assistant + console.print("[dim]Thinking...[/dim]") + decision = await interpret_user_decision( + user_message=user_input, + sub_questions=sub_questions, + qa_responses=qa_responses, + config=config, + ) + + # Handle modifications and chat locally, continue loop + if decision.action == "chat": + if decision.message: + console.print( + f"\n[bold cyan]Assistant:[/bold cyan] {decision.message}" + ) + continue + elif decision.action == "add_questions" and decision.questions: + sub_questions.extend(decision.questions) + console.print( + f"[green]Added {len(decision.questions)} question(s)[/green]" + ) + continue + elif decision.action == "modify_questions" and decision.questions: + sub_questions = list(decision.questions) + console.print( + f"[green]Replaced with {len(decision.questions)} question(s)[/green]" + ) + continue + + # User wants to proceed - send final decision + action_display = { + "search": "Searching questions", + "synthesize": "Generating report", + } + console.print( + f"[dim]→ {action_display.get(decision.action, decision.action)}[/dim]" + ) + + # Include any accumulated question changes + if decision.action == "search": + decision = HumanDecision( + action="modify_questions", questions=sub_questions + ) + + await human_input_queue.put(decision) + break + + elif event_type == "TEXT_MESSAGE_CHUNK": + # Log message from graph + message = event.get("delta", "") + if message: + console.print(f"[dim]{message}[/dim]") + + elif event_type == "RUN_FINISHED": + break + + elif event_type == "RUN_ERROR": + error_msg = event.get("message", "Unknown error") + console.print(f"[red]Error: {error_msg}[/red]") + break + + # Wait for graph to complete + report = await graph_task + return report + + except Exception as e: + graph_task.cancel() + raise e + finally: + await emitter.close() + + +async def run_chat_loop( + client: HaikuRAG, + config: AppConfig | None = None, + search_filter: str | None = None, + question: str | None = None, +) -> None: + """Run an interactive chat loop for research. + + Args: + client: HaikuRAG client for document operations + config: Application configuration (uses global config if None) + search_filter: Optional SQL WHERE clause to filter documents + question: Optional initial research question (skips initial chat if provided) + """ + config = config or get_config() + console = Console() + + console.print( + Panel( + "[bold cyan]Interactive Research Mode[/bold cyan]\n\n" + "Chat with me or tell me what you'd like to research.\n" + "Type [green]exit[/green] or [green]quit[/green] to end the session.", + title="haiku.rag Research Assistant", + border_style="cyan", + ) + ) + + while True: + try: + # Use provided question or get one through conversation + if question: + research_question = question + console.print(f"[dim]Starting research: {research_question}[/dim]") + question = None # Clear so subsequent loops go through chat + else: + # Initial conversation loop - chat until user wants to research + research_question = None + while research_question is None: + user_input = Prompt.ask("\n[bold blue]You[/bold blue]") + + if not user_input.strip(): + continue + + if user_input.lower().strip() in ("exit", "quit", "q"): + console.print("[dim]Goodbye![/dim]") + return + + console.print("[dim]Thinking...[/dim]") + decision = await initial_chat(user_input, config) + + if decision.action == "research" and decision.research_question: + research_question = decision.research_question + console.print( + f"[dim]Starting research: {research_question}[/dim]" + ) + elif decision.action == "chat" and decision.message: + console.print( + f"\n[bold cyan]Assistant:[/bold cyan] {decision.message}" + ) + else: + # Fallback - treat as research question + research_question = user_input + + console.print() + report = await run_interactive_research( + question=research_question, + client=client, + config=config, + search_filter=search_filter, + ) + + # Display final report + console.print() + console.print( + Panel( + Markdown(f"## {report.title}\n\n{report.executive_summary}"), + title="Research Report", + border_style="green", + ) + ) + + if report.main_findings: + findings = "\n".join(f"- {f}" for f in report.main_findings[:5]) + console.print(Markdown(f"**Key Findings:**\n{findings}")) + + if report.conclusions: + conclusions = "\n".join(f"- {c}" for c in report.conclusions[:3]) + console.print(Markdown(f"**Conclusions:**\n{conclusions}")) + + console.print(Markdown(f"**Sources:** {report.sources_summary}")) + + except KeyboardInterrupt: + console.print("\n[dim]Goodbye![/dim]") + return + except Exception as e: + console.print(f"[red]Error: {e}[/red]") + + +def interactive_research( + client: HaikuRAG, + config: AppConfig | None = None, + search_filter: str | None = None, + question: str | None = None, +) -> None: + """Entry point for interactive research mode. + + Args: + client: HaikuRAG client for document operations + config: Application configuration (uses global config if None) + search_filter: Optional SQL WHERE clause to filter documents + question: Optional initial research question (skips initial chat if provided) + """ + asyncio.run(run_chat_loop(client, config, search_filter, question)) diff --git a/haiku_rag_slim/haiku/rag/graph/agui/emitter.py b/haiku_rag_slim/haiku/rag/graph/agui/emitter.py index b2cf5e3e..7df847f5 100644 --- a/haiku_rag_slim/haiku/rag/graph/agui/emitter.py +++ b/haiku_rag_slim/haiku/rag/graph/agui/emitter.py @@ -80,8 +80,8 @@ class AGUIEmitter[StateT: BaseModel, ResultT]: self._thread_id = self._generate_thread_id(state_json) # RunStarted (state snapshot follows immediately with full state) - self._emit(emit_run_started(self._thread_id, self._run_id)) - self._emit(emit_state_snapshot(initial_state)) + self.emit(emit_run_started(self._thread_id, self._run_id)) + self.emit(emit_state_snapshot(initial_state)) # Store a deep copy to detect future changes self._last_state = initial_state.model_copy(deep=True) @@ -92,12 +92,12 @@ class AGUIEmitter[StateT: BaseModel, ResultT]: step_name: Name of the step being started """ self._current_step = step_name - self._emit(emit_step_started(step_name)) + self.emit(emit_step_started(step_name)) def finish_step(self) -> None: """Emit StepFinished event for the current step.""" if self._current_step: - self._emit(emit_step_finished(self._current_step)) + self.emit(emit_step_finished(self._current_step)) self._current_step = None def log(self, message: str, role: str = "assistant") -> None: @@ -107,7 +107,7 @@ class AGUIEmitter[StateT: BaseModel, ResultT]: message: The message content role: The role of the sender (default: assistant) """ - self._emit(emit_text_message(message, role)) + self.emit(emit_text_message(message, role)) def update_state(self, new_state: StateT) -> None: """Emit StateDelta or StateSnapshot for state change. @@ -117,10 +117,10 @@ class AGUIEmitter[StateT: BaseModel, ResultT]: """ if self._use_deltas and self._last_state is not None: # Emit delta for incremental updates - self._emit(emit_state_delta(self._last_state, new_state)) + self.emit(emit_state_delta(self._last_state, new_state)) else: # Emit full snapshot for initial state or when deltas disabled - self._emit(emit_state_snapshot(new_state)) + self.emit(emit_state_snapshot(new_state)) # Store a deep copy to detect future changes self._last_state = new_state.model_copy(deep=True) @@ -139,7 +139,7 @@ class AGUIEmitter[StateT: BaseModel, ResultT]: """ if message_id is None: message_id = str(uuid4()) - self._emit(emit_activity(message_id, activity_type, content)) + self.emit(emit_activity(message_id, activity_type, content)) def finish_run(self, result: ResultT) -> None: """Emit RunFinished event. @@ -147,7 +147,7 @@ class AGUIEmitter[StateT: BaseModel, ResultT]: Args: result: The final result from the graph """ - self._emit(emit_run_finished(self._thread_id, self._run_id, result)) + self.emit(emit_run_finished(self._thread_id, self._run_id, result)) def error(self, error: Exception, code: str | None = None) -> None: """Emit RunError event. @@ -156,9 +156,9 @@ class AGUIEmitter[StateT: BaseModel, ResultT]: error: The exception that occurred code: Optional error code """ - self._emit(emit_run_error(str(error), code)) + self.emit(emit_run_error(str(error), code)) - def _emit(self, event: AGUIEvent) -> None: + def emit(self, event: AGUIEvent) -> None: """Put event in queue. Args: diff --git a/haiku_rag_slim/haiku/rag/graph/agui/events.py b/haiku_rag_slim/haiku/rag/graph/agui/events.py index eec980f2..5cb1728b 100644 --- a/haiku_rag_slim/haiku/rag/graph/agui/events.py +++ b/haiku_rag_slim/haiku/rag/graph/agui/events.py @@ -252,3 +252,62 @@ def emit_activity_delta( "activityType": activity_type, "patch": patch, } + + +def emit_tool_call_start( + tool_call_id: str, + tool_name: str, + parent_message_id: str | None = None, +) -> dict[str, Any]: + """Create a ToolCallStart event. + + Args: + tool_call_id: Unique identifier for this tool call + tool_name: Name of the tool being called + parent_message_id: Optional parent message ID + + Returns: + ToolCallStart event dict + """ + event: dict[str, Any] = { + "type": "TOOL_CALL_START", + "toolCallId": tool_call_id, + "toolCallName": tool_name, + } + if parent_message_id: + event["parentMessageId"] = parent_message_id + return event + + +def emit_tool_call_args(tool_call_id: str, args: dict[str, Any]) -> dict[str, Any]: + """Create a ToolCallArgs event. + + Args: + tool_call_id: Identifier for the tool call + args: Tool arguments + + Returns: + ToolCallArgs event dict + """ + import json + + return { + "type": "TOOL_CALL_ARGS", + "toolCallId": tool_call_id, + "delta": json.dumps(args), + } + + +def emit_tool_call_end(tool_call_id: str) -> dict[str, Any]: + """Create a ToolCallEnd event. + + Args: + tool_call_id: Identifier for the tool call being completed + + Returns: + ToolCallEnd event dict + """ + return { + "type": "TOOL_CALL_END", + "toolCallId": tool_call_id, + } diff --git a/haiku_rag_slim/haiku/rag/graph/agui/server.py b/haiku_rag_slim/haiku/rag/graph/agui/server.py index 1ccba823..26b463e6 100644 --- a/haiku_rag_slim/haiku/rag/graph/agui/server.py +++ b/haiku_rag_slim/haiku/rag/graph/agui/server.py @@ -166,7 +166,10 @@ def create_agui_server( # pragma: no cover from haiku.rag.client import HaikuRAG from haiku.rag.graph.research.dependencies import ResearchContext from haiku.rag.graph.research.graph import build_research_graph - from haiku.rag.graph.research.state import ResearchDeps, ResearchState + from haiku.rag.graph.research.state import ( + ResearchDeps, + ResearchState, + ) # Store client reference for proper lifecycle management _client_cache: dict[str, HaikuRAG] = {} diff --git a/haiku_rag_slim/haiku/rag/graph/research/graph.py b/haiku_rag_slim/haiku/rag/graph/research/graph.py index 4e190665..e3c8befa 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/graph.py +++ b/haiku_rag_slim/haiku/rag/graph/research/graph.py @@ -1,4 +1,6 @@ import asyncio +from typing import Literal +from uuid import uuid4 from pydantic_ai import Agent, RunContext, format_as_xml from pydantic_ai.output import ToolOutput @@ -7,6 +9,13 @@ from pydantic_graph.beta.join import reduce_list_append from haiku.rag.config import Config from haiku.rag.config.models import AppConfig +from haiku.rag.graph.agui.events import ( + emit_text_message_end, + emit_text_message_start, + emit_tool_call_args, + emit_tool_call_end, + emit_tool_call_start, +) from haiku.rag.graph.research.dependencies import ResearchContext, ResearchDependencies from haiku.rag.graph.research.models import ( EvaluationResult, @@ -54,12 +63,14 @@ def format_context_for_prompt(context: ResearchContext) -> str: def build_research_graph( config: AppConfig = Config, include_plan: bool = True, + interactive: bool = False, ) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]: """Build the Research graph. Args: config: AppConfig object (uses config.research for provider, model, and graph parameters) include_plan: Whether to include the planning step (False for execute-only mode) + interactive: Whether to include human decision nodes for HIL Returns: Configured Research graph @@ -240,7 +251,7 @@ def build_research_graph( @g.step async def get_batch( - ctx: StepContext[ResearchState, ResearchDeps, None | bool], + ctx: StepContext[ResearchState, ResearchDeps, None | bool | str], ) -> list[str] | None: """Get all remaining questions for this iteration.""" state = ctx.state @@ -302,9 +313,16 @@ def build_research_graph( state.last_eval = output state.iterations += 1 + # Get already-answered questions to avoid duplicates + answered_queries = {qa.query.lower() for qa in state.context.qa_responses} + for new_q in output.new_questions: - if new_q not in state.context.sub_questions: - state.context.sub_questions.append(new_q) + # Skip if already in pending or already answered + if new_q in state.context.sub_questions: + continue + if new_q.lower() in answered_queries: + continue + state.context.sub_questions.append(new_q) if deps.agui_emitter: deps.agui_emitter.update_state(state) @@ -329,9 +347,82 @@ def build_research_graph( if deps.agui_emitter: deps.agui_emitter.finish_step() + @g.step + async def human_decide( + ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer] | None | bool], + ) -> Literal["search", "synthesize"]: + """Wait for human decision on whether to continue searching or synthesize.""" + state = ctx.state + deps = ctx.deps + + if deps.agui_emitter: + deps.agui_emitter.start_step("human_decide") + deps.agui_emitter.update_state(state) + + try: + # Emit tool call for human input wrapped in a message context + # This makes the tool call appear as if emitted by the LLM + message_id = str(uuid4()) + tool_call_id = str(uuid4()) + + if deps.agui_emitter: + # Start an assistant message to contain the tool call + deps.agui_emitter.emit(emit_text_message_start(message_id)) + # Emit tool call with parent message reference + deps.agui_emitter.emit( + emit_tool_call_start(tool_call_id, "human_decision", message_id) + ) + # Include full state for display + qa_responses = [ + { + "query": qa.query, + "answer": qa.answer, + "confidence": qa.confidence, + "citations_count": len(qa.citations), + } + for qa in state.context.qa_responses + ] + deps.agui_emitter.emit( + emit_tool_call_args( + tool_call_id, + { + "original_question": state.context.original_question, + "sub_questions": list(state.context.sub_questions), + "qa_responses": qa_responses, + "iterations": state.iterations, + }, + ) + ) + deps.agui_emitter.emit(emit_tool_call_end(tool_call_id)) + # End the message after tool call + deps.agui_emitter.emit(emit_text_message_end(message_id)) + + # Wait for human input + if deps.human_input_queue is None: + raise RuntimeError("human_input_queue is required for interactive mode") + + decision = await deps.human_input_queue.get() + + # Process decision + if decision.action == "modify_questions" and decision.questions: + state.context.sub_questions = list(decision.questions) + elif decision.action == "add_questions" and decision.questions: + state.context.sub_questions.extend(decision.questions) + + if deps.agui_emitter: + deps.agui_emitter.update_state(state) + + if decision.action in ("search", "modify_questions", "add_questions"): + return "search" + else: + return "synthesize" + finally: + if deps.agui_emitter: + deps.agui_emitter.finish_step() + @g.step async def synthesize( - ctx: StepContext[ResearchState, ResearchDeps, None | bool], + ctx: StepContext[ResearchState, ResearchDeps, None | bool | str], ) -> ResearchReport: """Generate final research report.""" state = ctx.state @@ -375,39 +466,76 @@ def build_research_graph( initial_factory=list[SearchAnswer], ) - if include_plan: + if interactive: + # Interactive mode: human decides after plan and after evaluation + if include_plan: + g.add( + g.edge_from(g.start_node).to(plan), + g.edge_from(plan).to(human_decide), + ) + else: + g.add(g.edge_from(g.start_node).to(human_decide)) + g.add( - g.edge_from(g.start_node).to(plan), - g.edge_from(plan).to(get_batch), + g.edge_from(human_decide).to( + g.decision() + .branch( + g.match(str, matches=lambda x: x == "search") + .label("Search") + .to(get_batch) + ) + .branch( + g.match(str, matches=lambda x: x == "synthesize") + .label("Synthesize") + .to(synthesize) + ) + ), + g.edge_from(get_batch).to( + g.decision() + .branch(g.match(list).label("Has questions").map().to(search_one)) + .branch(g.match(type(None)).label("No questions").to(human_decide)) + ), + g.edge_from(search_one).to(collect_answers), + # After search, evaluate to suggest new questions, then human decides + g.edge_from(collect_answers).to(decide), + g.edge_from(decide).to(human_decide), + g.edge_from(synthesize).to(g.end_node), ) else: - g.add(g.edge_from(g.start_node).to(get_batch)) - - g.add( - g.edge_from(get_batch).to( - g.decision() - .branch(g.match(list).label("Has questions").map().to(search_one)) - .branch(g.match(type(None)).label("No questions").to(synthesize)) - ), - g.edge_from(search_one).to(collect_answers), - g.edge_from(collect_answers).to(decide), - ) - - g.add( - g.edge_from(decide).to( - g.decision() - .branch( - g.match(bool, matches=lambda x: x) - .label("Continue research") - .to(get_batch) + # Non-interactive mode: automatic decision based on confidence/iterations + if include_plan: + g.add( + g.edge_from(g.start_node).to(plan), + g.edge_from(plan).to(get_batch), ) - .branch( - g.match(bool, matches=lambda x: not x) - .label("Done researching") - .to(synthesize) - ) - ), - g.edge_from(synthesize).to(g.end_node), - ) + else: + g.add(g.edge_from(g.start_node).to(get_batch)) + + g.add( + g.edge_from(get_batch).to( + g.decision() + .branch(g.match(list).label("Has questions").map().to(search_one)) + .branch(g.match(type(None)).label("No questions").to(synthesize)) + ), + g.edge_from(search_one).to(collect_answers), + g.edge_from(collect_answers).to(decide), + ) + + g.add( + g.edge_from(decide).to( + g.decision() + .branch( + g.match(bool, matches=lambda x: x) + .label("Continue research") + .to(get_batch) + ) + .branch( + g.match(bool, matches=lambda x: not x) + .label("Done researching") + .to(synthesize) + ) + ), + g.edge_from(synthesize).to(g.end_node), + ) return g.build() diff --git a/haiku_rag_slim/haiku/rag/graph/research/state.py b/haiku_rag_slim/haiku/rag/graph/research/state.py index fbf4fc39..93448a4c 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/state.py +++ b/haiku_rag_slim/haiku/rag/graph/research/state.py @@ -1,6 +1,6 @@ import asyncio from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal from pydantic import BaseModel, Field @@ -13,6 +13,17 @@ if TYPE_CHECKING: from haiku.rag.graph.agui.emitter import AGUIEmitter +class HumanDecision(BaseModel): + """Human decision input for interactive research.""" + + action: Literal[ + "search", "synthesize", "modify_questions", "add_questions", "chat", "research" + ] + questions: list[str] | None = None + message: str | None = None + research_question: str | None = None + + @dataclass class ResearchDeps: """Dependencies for research graph execution.""" @@ -20,6 +31,8 @@ class ResearchDeps: client: HaikuRAG agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None semaphore: asyncio.Semaphore | None = None + human_input_queue: asyncio.Queue[HumanDecision] | None = None + interactive: bool = False def emit_log(self, message: str, state: "ResearchState | None" = None) -> None: """Emit a log message through AG-UI events.""" diff --git a/tests/graph/agui/test_events.py b/tests/graph/agui/test_events.py index 6b807e32..5c8795d2 100644 --- a/tests/graph/agui/test_events.py +++ b/tests/graph/agui/test_events.py @@ -11,6 +11,9 @@ from haiku.rag.graph.agui.events import ( emit_step_finished, emit_step_started, emit_text_message, + emit_tool_call_args, + emit_tool_call_end, + emit_tool_call_start, ) @@ -135,6 +138,46 @@ def test_emit_activity(): assert event["content"] == {"message": "Working on task"} +def test_emit_tool_call_start(): + """Test TOOL_CALL_START event creation.""" + event = emit_tool_call_start("call-1", "search_documents") + + assert event["type"] == "TOOL_CALL_START" + assert event["toolCallId"] == "call-1" + assert event["toolCallName"] == "search_documents" + assert "parentMessageId" not in event + + +def test_emit_tool_call_start_with_parent(): + """Test TOOL_CALL_START event with parent message ID.""" + event = emit_tool_call_start("call-1", "search", parent_message_id="msg-1") + + assert event["type"] == "TOOL_CALL_START" + assert event["toolCallId"] == "call-1" + assert event["toolCallName"] == "search" + assert event["parentMessageId"] == "msg-1" + + +def test_emit_tool_call_args(): + """Test TOOL_CALL_ARGS event creation.""" + import json + + args = {"query": "test query", "limit": 10} + event = emit_tool_call_args("call-1", args) + + assert event["type"] == "TOOL_CALL_ARGS" + assert event["toolCallId"] == "call-1" + assert event["delta"] == json.dumps(args) + + +def test_emit_tool_call_end(): + """Test TOOL_CALL_END event creation.""" + event = emit_tool_call_end("call-1") + + assert event["type"] == "TOOL_CALL_END" + assert event["toolCallId"] == "call-1" + + def test_event_structure_consistency(): """Test that all events have consistent structure.""" events = [ @@ -146,6 +189,9 @@ def test_event_structure_consistency(): emit_text_message("text"), emit_state_snapshot(TestState(value=1)), emit_activity("m1", "type", {"content": "value"}), + emit_tool_call_start("c1", "tool"), + emit_tool_call_args("c1", {"arg": "value"}), + emit_tool_call_end("c1"), ] for event in events: diff --git a/tests/graph/test_research_graph.py b/tests/graph/test_research_graph.py index 66f7482e..79a20453 100644 --- a/tests/graph/test_research_graph.py +++ b/tests/graph/test_research_graph.py @@ -1,3 +1,5 @@ +import asyncio + import pytest from pydantic_ai.models.test import TestModel @@ -5,7 +7,7 @@ from haiku.rag.client import HaikuRAG from haiku.rag.graph.agui.stream import stream_graph from haiku.rag.graph.research.dependencies import ResearchContext from haiku.rag.graph.research.graph import build_research_graph -from haiku.rag.graph.research.state import ResearchDeps, ResearchState +from haiku.rag.graph.research.state import HumanDecision, ResearchDeps, ResearchState @pytest.mark.asyncio @@ -61,3 +63,83 @@ async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path): assert "STEP_STARTED" in event_types client.close() + + +@pytest.mark.asyncio +async def test_interactive_graph_with_human_decision(monkeypatch, temp_db_path): + """Test interactive research graph pauses and resumes with human decisions.""" + + # Mock get_model to return TestModel + def test_model_factory(_provider, _model, _config=None): + return TestModel() + + monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory) + monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory) + + # Build interactive graph + graph = build_research_graph(interactive=True) + + state = ResearchState( + context=ResearchContext(original_question="What is haiku.rag?"), + max_iterations=1, + confidence_threshold=0.5, + max_concurrency=2, + ) + + # Create human input queue + human_input_queue: asyncio.Queue[HumanDecision] = asyncio.Queue() + + client = HaikuRAG(temp_db_path, create=True) + deps = ResearchDeps( + client=client, + human_input_queue=human_input_queue, + interactive=True, + ) + + events = [] + tool_call_received = asyncio.Event() + result = None + + async def run_graph(): + nonlocal result + async for event in stream_graph(graph, state, deps): + events.append(event) + if event["type"] == "TOOL_CALL_START": + tool_name = event.get("toolCallName") + if tool_name == "human_decision": + tool_call_received.set() + elif event["type"] == "RUN_FINISHED": + result = event["result"] + elif event["type"] == "RUN_ERROR": + pytest.fail(f"Graph execution failed: {event['message']}") + + async def send_decisions(): + # Wait for first tool call (after planning) + await asyncio.wait_for(tool_call_received.wait(), timeout=30) + tool_call_received.clear() + + # Send search decision + await human_input_queue.put(HumanDecision(action="search")) + + # Wait for second tool call (after search cycle) + await asyncio.wait_for(tool_call_received.wait(), timeout=30) + + # Send synthesize decision + await human_input_queue.put(HumanDecision(action="synthesize")) + + # Run graph and decision sender concurrently + await asyncio.gather(run_graph(), send_decisions()) + + # Verify result + assert result is not None, ( + f"No result. Events collected: {[e['type'] for e in events]}" + ) + assert isinstance(result, dict) + assert "title" in result + + # Verify human_decision tool calls were emitted + event_types = [e["type"] for e in events] + assert "TOOL_CALL_START" in event_types + assert "TOOL_CALL_END" in event_types + + client.close()