From 218126de8d377cf14731f5780888b5e2c5654c84 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 16 Dec 2025 14:10:31 +0200 Subject: [PATCH 01/10] Interactive research agent through AGUI client-side tool calls in CLI --- CHANGELOG.md | 9 + docs/agents.md | 24 + examples/ag-ui-research/backend/agent.py | 2 +- haiku_rag_slim/haiku/rag/cli.py | 29 +- haiku_rag_slim/haiku/rag/cli_chat.py | 465 ++++++++++++++++++ .../haiku/rag/graph/agui/emitter.py | 22 +- haiku_rag_slim/haiku/rag/graph/agui/events.py | 57 +++ haiku_rag_slim/haiku/rag/graph/agui/server.py | 5 +- .../haiku/rag/graph/research/graph.py | 187 +++++-- .../haiku/rag/graph/research/state.py | 15 +- tests/graph/agui/test_events.py | 44 ++ 11 files changed, 809 insertions(+), 50 deletions(-) create mode 100644 haiku_rag_slim/haiku/rag/cli_chat.py diff --git a/CHANGELOG.md b/CHANGELOG.md index e986e0b0..94e0f6d0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,15 @@ ### 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 - **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..103afa9c 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -196,3 +196,27 @@ 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 + +# 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). diff --git a/examples/ag-ui-research/backend/agent.py b/examples/ag-ui-research/backend/agent.py index e697145f..9537a4c3 100644 --- a/examples/ag-ui-research/backend/agent.py +++ b/examples/ag-ui-research/backend/agent.py @@ -9,10 +9,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.utils import get_model if TYPE_CHECKING: from haiku.rag.graph.agui.emitter import AGUIEmitter diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index 42f8fbd9..cc8be9f6 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,33 @@ 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, + ) + 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..bf67f300 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/cli_chat.py @@ -0,0 +1,465 @@ +"""Interactive CLI chat loop for research graph with human-in-the-loop.""" + +import asyncio + +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": + args = event.get("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 for natural language input + console.print() + user_input = Prompt.ask("[bold]What would you like to do?[/bold]") + + # 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, +) -> 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 + """ + 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: + # 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]Interrupted. Type 'exit' to quit.[/dim]") + 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, +) -> 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 + """ + asyncio.run(run_chat_loop(client, config, search_filter)) 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..49710126 100644 --- a/haiku_rag_slim/haiku/rag/graph/agui/events.py +++ b/haiku_rag_slim/haiku/rag/graph/agui/events.py @@ -252,3 +252,60 @@ 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 + """ + return { + "type": "TOOL_CALL_ARGS", + "toolCallId": tool_call_id, + "delta": 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..8d996607 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,11 @@ 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_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 +61,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 +249,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 +311,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 +345,75 @@ 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 + tool_call_id = str(uuid4()) + + if deps.agui_emitter: + deps.agui_emitter.emit( + emit_tool_call_start(tool_call_id, "human_decision") + ) + # 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)) + + # 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 +457,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..297229e6 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,44 @@ 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.""" + 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"] == 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 +187,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: From 9818cac4ebae66857069969156047e5ad7a296eb Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 16 Dec 2025 15:32:10 +0200 Subject: [PATCH 02/10] Fix agui tool call args being a json string --- haiku_rag_slim/haiku/rag/cli_chat.py | 4 +++- haiku_rag_slim/haiku/rag/graph/agui/events.py | 4 +++- tests/graph/agui/test_events.py | 4 +++- 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/cli_chat.py b/haiku_rag_slim/haiku/rag/cli_chat.py index bf67f300..35070708 100644 --- a/haiku_rag_slim/haiku/rag/cli_chat.py +++ b/haiku_rag_slim/haiku/rag/cli_chat.py @@ -1,6 +1,7 @@ """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 @@ -225,7 +226,8 @@ async def run_interactive_research( pass elif event_type == "TOOL_CALL_ARGS": - args = event.get("delta", {}) + 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", []) diff --git a/haiku_rag_slim/haiku/rag/graph/agui/events.py b/haiku_rag_slim/haiku/rag/graph/agui/events.py index 49710126..5cb1728b 100644 --- a/haiku_rag_slim/haiku/rag/graph/agui/events.py +++ b/haiku_rag_slim/haiku/rag/graph/agui/events.py @@ -289,10 +289,12 @@ def emit_tool_call_args(tool_call_id: str, args: dict[str, Any]) -> dict[str, An Returns: ToolCallArgs event dict """ + import json + return { "type": "TOOL_CALL_ARGS", "toolCallId": tool_call_id, - "delta": args, + "delta": json.dumps(args), } diff --git a/tests/graph/agui/test_events.py b/tests/graph/agui/test_events.py index 297229e6..5c8795d2 100644 --- a/tests/graph/agui/test_events.py +++ b/tests/graph/agui/test_events.py @@ -160,12 +160,14 @@ def test_emit_tool_call_start_with_parent(): 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"] == args + assert event["delta"] == json.dumps(args) def test_emit_tool_call_end(): From ff13be1210b28d45b9dd38fe13b3a7ea49cd5c7c Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 17 Dec 2025 12:03:42 +0200 Subject: [PATCH 03/10] Human-in-the-loop in ag-ui-example --- CHANGELOG.md | 2 + docs/agents.md | 6 +- examples/ag-ui-research/README.md | 55 ++--- examples/ag-ui-research/backend/agent.py | 56 ++++- examples/ag-ui-research/backend/main.py | 69 +++++- .../frontend/components/Agent.tsx | 213 +++++++++++++++++- .../frontend/components/StateDisplay.tsx | 50 +--- tests/graph/test_research_graph.py | 84 ++++++- 8 files changed, 443 insertions(+), 92 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 94e0f6d0..0ffb6b76 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,8 @@ - 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**: Updated with interactive decision UI + - Decision panel with question editing (add/remove) at each decision point - **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 103afa9c..69bc0e82 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -219,4 +219,8 @@ In interactive mode, you can: - 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). +For a web-based interactive experience with visual decision UI, see the [AG-UI Research Example](https://github.com/ggozad/haiku.rag/tree/main/examples/ag-ui-research). The web interface provides: + +- Question editing panel to add/remove sub-questions at decision points +- Search and Generate Report buttons for controlling research flow +- Live state display showing answers, confidence, and progress diff --git a/examples/ag-ui-research/README.md b/examples/ag-ui-research/README.md index df6850dc..1fa202c9 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 @@ -99,32 +99,35 @@ This example demonstrates the **agent+graph** architecture pattern: - 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 + - Pauses at decision points waiting for human input via async queue - Emits AG-UI events for real-time progress tracking -3. **Shared Event Stream**: +3. **Decision Endpoint** (`main.py`): + - `/v1/research/decide` receives human decisions from frontend + - Forwards decisions to the waiting graph via `HumanDecision` queue + - Supports actions: `search`, `synthesize`, `modify_questions` + +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, decision endpoint for human input - 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 - Split-pane UI: chat on left, live research state on right + - Decision UI: question editor with add/remove, search and generate report buttons - Real-time state synchronization via Server-Sent Events (SSE) - - `StateDisplay` component with collapsible sections for questions and report ## Configuration diff --git a/examples/ag-ui-research/backend/agent.py b/examples/ag-ui-research/backend/agent.py index 9537a4c3..fe15557d 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 @@ -11,7 +12,7 @@ from haiku.rag.config import load_yaml_config from haiku.rag.config.models import AppConfig 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: @@ -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,7 @@ class AgentDeps: client: HaikuRAG agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None search_filter: str | None = None + thread_id: str | None = None model = get_model(Config.research.model, Config) @@ -50,10 +66,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 +86,39 @@ 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!") return f"""Research completed successfully! @@ -108,5 +140,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..6254c53e 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" @@ -79,11 +80,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 @@ -99,11 +100,14 @@ async def stream_research_agent(request: Request) -> StreamingResponse: ids_str = ", ".join(f"'{id}'" for id in document_ids) search_filter = f"id IN ({ids_str})" + thread_id = input_data.thread_id + # Create agent dependencies with shared emitter agent_deps = AgentDeps( client=client, agui_emitter=emitter, search_filter=search_filter, + thread_id=thread_id, ) # Start run with empty initial state @@ -148,6 +152,37 @@ async def stream_research_agent(request: Request) -> StreamingResponse: await send_stream.send(format_sse_event(delta_event)) continue + # When human_decision tool starts, set awaiting_decision flag + if event_type == "TOOL_CALL_START": + tool_name = event.get("toolCallName") + if tool_name == "human_decision": + delta_event = { + "type": "STATE_DELTA", + "delta": [ + { + "op": "add", + "path": "/awaiting_decision", + "value": True, + } + ], + } + await send_stream.send(format_sse_event(delta_event)) + + # 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 @@ -248,10 +283,36 @@ async def visualize_chunk(request: Request) -> JSONResponse: ) +async def research_decide(request: Request) -> JSONResponse: + """Endpoint to receive human decisions for active research.""" + body = await request.json() + action = body.get("action", "search") + questions = body.get("questions", []) + + # Get first active research (single-user example) + active = next(iter(_active_research.values()), None) + if not active: + return JSONResponse({"error": "No active research found"}, status_code=404) + + # When "search" action is sent with questions, use "modify_questions" to update them + effective_action = ( + "modify_questions" if action == "search" and questions else action + ) + + decision = HumanDecision( + action=effective_action, + questions=questions or None, + ) + await active.queue.put(decision) + + return JSONResponse({"status": "ok", "action": effective_action}) + + # Create Starlette app app = Starlette( routes=[ Route("/v1/research/stream", stream_research_agent, methods=["POST"]), + Route("/v1/research/decide", research_decide, methods=["POST"]), Route("/api/documents", list_documents, methods=["GET"]), Route("/api/visualize/{chunk_id}", visualize_chunk, methods=["GET"]), Route("/health", health_check, methods=["GET"]), diff --git a/examples/ag-ui-research/frontend/components/Agent.tsx b/examples/ag-ui-research/frontend/components/Agent.tsx index 75429a91..4cc1c731 100644 --- a/examples/ag-ui-research/frontend/components/Agent.tsx +++ b/examples/ag-ui-research/frontend/components/Agent.tsx @@ -3,6 +3,7 @@ import { CopilotKit, useCoAgent } from "@copilotkit/react-core"; import { CopilotChat } from "@copilotkit/react-ui"; import "@copilotkit/react-ui/styles.css"; +import { useState, useEffect } from "react"; import DocumentSelector from "./DocumentSelector"; import StateDisplay from "./StateDisplay"; @@ -58,6 +59,7 @@ interface ResearchState { current_activity?: string; current_activity_message?: string; documentFilter?: string[]; + awaiting_decision?: boolean; } function AgentContent() { @@ -75,13 +77,61 @@ function AgentContent() { max_concurrency: 1, last_eval: null, documentFilter: [], + awaiting_decision: false, }, }); + const [editableQuestions, setEditableQuestions] = useState([]); + const [newQuestion, setNewQuestion] = useState(""); + const [submitting, setSubmitting] = useState(false); + + // Sync editable questions when state changes + useEffect(() => { + if (state.awaiting_decision && state.context.sub_questions) { + setEditableQuestions([...state.context.sub_questions]); + } + }, [state.awaiting_decision, state.context.sub_questions]); + const handleDocumentFilterChange = (ids: string[]) => { setState({ ...state, documentFilter: ids }); }; + const handleRemoveQuestion = (index: number) => { + setEditableQuestions(editableQuestions.filter((_, i) => i !== index)); + }; + + const handleAddQuestion = () => { + if (newQuestion.trim()) { + setEditableQuestions([...editableQuestions, newQuestion.trim()]); + setNewQuestion(""); + } + }; + + const handleDecision = async (action: "search" | "synthesize") => { + setSubmitting(true); + try { + const response = await fetch( + `${process.env.NEXT_PUBLIC_API_URL || "http://localhost:8000"}/v1/research/decide`, + { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + thread_id: state.context.original_question, // Use as identifier + action, + questions: editableQuestions, + }), + } + ); + if (response.ok) { + setState({ ...state, awaiting_decision: false }); + } + } catch (error) { + console.error("Failed to send decision:", error); + } finally { + setSubmitting(false); + } + }; + return ( <>
- {/* Chat on the left */}
- {/* State display on the right */}
- {/* Document filter - hidden when research is running */} {!running && (
)} - {/* Decision UI when awaiting human input - hidden when report exists, submitting, or not running */} - {state.awaiting_decision && !submitting && !state.result && running && ( -
-
- Research Decision Point -
- -
- {state.context.qa_responses?.length || 0} answers collected | Iteration {state.iterations || 0} -
- - {/* Questions list */} -
-
- Pending Questions ({editableQuestions.length}): -
- {editableQuestions.map((q, idx) => ( -
- {q} - -
- ))} -
- - {/* Add question input */} -
- setNewQuestion(e.target.value)} - placeholder="Add a new question..." - disabled={submitting} - style={{ - flex: 1, - padding: "0.5rem", - border: "1px solid #cbd5e1", - borderRadius: "4px", - fontSize: "0.85rem", - }} - onKeyDown={(e) => { - if (e.key === "Enter") { - handleAddQuestion(); - } - }} - /> - -
- - {/* Action buttons */} -
- - -
-
- )} -
diff --git a/examples/ag-ui-research/frontend/components/StateDisplay.tsx b/examples/ag-ui-research/frontend/components/StateDisplay.tsx index aa7ee112..3c9e129f 100644 --- a/examples/ag-ui-research/frontend/components/StateDisplay.tsx +++ b/examples/ag-ui-research/frontend/components/StateDisplay.tsx @@ -162,7 +162,7 @@ export default function StateDisplay({ state }: StateDisplayProps) { }} > {/* Question */} - {state.context.original_question && ( + {state.context?.original_question && (
0 && ( + {state.context?.qa_responses && state.context.qa_responses.length > 0 && (
Date: Wed, 17 Dec 2025 17:55:40 +0200 Subject: [PATCH 07/10] QOL fixes for interactive cli research --- haiku_rag_slim/haiku/rag/cli.py | 1 + haiku_rag_slim/haiku/rag/cli_chat.py | 59 +++++++++++++++++----------- 2 files changed, 37 insertions(+), 23 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index cc8be9f6..bd1af435 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -359,6 +359,7 @@ def research( client=client, config=app.config, search_filter=filter, + question=question, ) finally: client.close() diff --git a/haiku_rag_slim/haiku/rag/cli_chat.py b/haiku_rag_slim/haiku/rag/cli_chat.py index 35070708..db79c758 100644 --- a/haiku_rag_slim/haiku/rag/cli_chat.py +++ b/haiku_rag_slim/haiku/rag/cli_chat.py @@ -369,6 +369,7 @@ 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. @@ -376,6 +377,7 @@ async def run_chat_loop( 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() @@ -392,31 +394,39 @@ async def run_chat_loop( while True: try: - # 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]") + # 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 not user_input.strip(): + continue - if user_input.lower().strip() in ("exit", "quit", "q"): - console.print("[dim]Goodbye![/dim]") - return + 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) + 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 + 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( @@ -447,7 +457,8 @@ async def run_chat_loop( console.print(Markdown(f"**Sources:** {report.sources_summary}")) except KeyboardInterrupt: - console.print("\n[dim]Interrupted. Type 'exit' to quit.[/dim]") + console.print("\n[dim]Goodbye![/dim]") + return except Exception as e: console.print(f"[red]Error: {e}[/red]") @@ -456,6 +467,7 @@ 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. @@ -463,5 +475,6 @@ def interactive_research( 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)) + asyncio.run(run_chat_loop(client, config, search_filter, question)) From 3ec17c65e146e99da6a0ccc6c4e69bc298ac6ebe Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 17 Dec 2025 17:58:03 +0200 Subject: [PATCH 08/10] Update docs for interactive cli --- docs/agents.md | 3 +++ docs/cli.md | 14 ++++++++++++++ 2 files changed, 17 insertions(+) diff --git a/docs/agents.md b/docs/agents.md index e5caf19f..ff11b1bf 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -207,6 +207,9 @@ Interactive mode provides human-in-the-loop control over the research process th # 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%'" ``` 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. From bf5711f7ab13b331d4148bd2e8433437dc485b05 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 17 Dec 2025 18:02:17 +0200 Subject: [PATCH 09/10] Give hints in interactive cli --- haiku_rag_slim/haiku/rag/cli_chat.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/cli_chat.py b/haiku_rag_slim/haiku/rag/cli_chat.py index db79c758..68882cf3 100644 --- a/haiku_rag_slim/haiku/rag/cli_chat.py +++ b/haiku_rag_slim/haiku/rag/cli_chat.py @@ -289,9 +289,18 @@ async def run_interactive_research( if iterations > 0: console.print(f"[dim]Iteration: {iterations}[/dim]") - # Prompt user for natural language input + # Prompt user with context-aware hints console.print() - user_input = Prompt.ask("[bold]What would you like to do?[/bold]") + 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]") From 440ec123c6438c59f929b66a880a5fec5c0563d3 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 18 Dec 2025 09:42:55 +0200 Subject: [PATCH 10/10] If documents have no title display uri alone, ag-ui-example --- .../frontend/components/DocumentSelector.tsx | 35 +++++++++++-------- 1 file changed, 21 insertions(+), 14 deletions(-) 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"} -
-