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: