"""Interactive CLI chat loop for research graph with human-in-the-loop.""" import asyncio import json from pydantic_ai import Agent from rich.console import Console from rich.markdown import Markdown from rich.panel import Panel from rich.prompt import Prompt from haiku.rag.client import HaikuRAG from haiku.rag.config import get_config from haiku.rag.config.models import AppConfig from haiku.rag.graph.agui.emitter import AGUIEmitter from haiku.rag.graph.research.dependencies import ResearchContext from haiku.rag.graph.research.graph import build_research_graph from haiku.rag.graph.research.models import ResearchReport from haiku.rag.graph.research.state import HumanDecision, ResearchDeps, ResearchState from haiku.rag.utils import get_model INITIAL_CHAT_PROMPT = """You are a research assistant. The user hasn't started a research task yet. You can: 1. Chat with the user - greet them, answer questions about what you can do 2. Detect when they want to research something ## Actions: - "chat": User is chatting, greeting, or asking questions (set message with your response) - "research": User wants to research a topic (extract the research question into research_question) ## Guidelines: - If the user provides a clear research question or topic, set action="research" and extract the question - If the user is just chatting or asking what you can do, set action="chat" and respond helpfully - Be friendly and explain you can help them research topics by searching a knowledge base Examples: - "hi" → action="chat", message="Hello! I'm a research assistant. I can help you research topics by searching through documents and synthesizing findings. What would you like to explore?" - "what can you do?" → action="chat", message="I help you conduct research! Give me a question or topic, and I'll break it into sub-questions, search for answers, and synthesize a report. What are you curious about?" - "tell me about Python's memory management" → action="research", research_question="How does Python's memory management work?" - "I want to understand how RAG systems work" → action="research", research_question="How do RAG (Retrieval-Augmented Generation) systems work?" """ RESEARCH_ASSISTANT_PROMPT = """You are a research assistant helping the user conduct research on a topic. You are at a decision point in the research workflow. You can: 1. Chat with the user - answer questions, discuss the research, make suggestions 2. Take workflow actions when the user requests them ## Workflow Actions (set in the action field): - "search": Search the pending questions (user says: "go", "search", "yes", "continue", "looks good") - "synthesize": Generate final report (user says: "done", "finish", "synthesize", "generate report") - "add_questions": Add NEW research questions to the existing list - "modify_questions": REPLACE all pending questions with a new list (use when user wants to remove, keep only certain questions, or change the questions) - "chat": Have a conversation without modifying questions ## IMPORTANT - Modifying Questions: - "use only the first question" → action="modify_questions", questions=[first question from the list] - "drop questions 2 and 3" → action="modify_questions", questions=[remaining questions] - "keep only questions about X" → action="modify_questions", questions=[filtered list] - "remove the duplicate" → action="modify_questions", questions=[deduplicated list] - When user wants to reduce/filter/keep-only, use "modify_questions" NOT "chat" ## Guidelines: - If the user wants to modify the question list in ANY way (remove, keep only, filter), use "modify_questions" - For "modify_questions", include ALL questions that should remain in the questions field - You can combine "chat" with a message to explain what you're doing - If just chatting without changes, set action="chat" and provide helpful response in message """ async def initial_chat( user_message: str, config: AppConfig, ) -> HumanDecision: """Handle initial conversation before research starts. Args: user_message: The user's message config: Application configuration Returns: HumanDecision with chat response or research question """ agent: Agent[None, HumanDecision] = Agent( model=get_model(config.research.model, config), output_type=HumanDecision, instructions=INITIAL_CHAT_PROMPT, retries=2, ) result = await agent.run(user_message) return result.output async def interpret_user_decision( user_message: str, sub_questions: list[str], qa_responses: list[dict], config: AppConfig, ) -> HumanDecision: """Interpret a natural language user message into a HumanDecision. Args: user_message: The user's natural language input sub_questions: Current sub-questions pending search qa_responses: Answers already collected config: Application configuration Returns: HumanDecision with the interpreted action, questions, and/or message """ agent: Agent[None, HumanDecision] = Agent( model=get_model(config.research.model, config), output_type=HumanDecision, instructions=RESEARCH_ASSISTANT_PROMPT, retries=2, ) # Build context with full research state answers_summary = "" if qa_responses: answers_parts = [] for qa in qa_responses: conf = f"{qa['confidence']:.0%}" if qa.get("confidence") else "N/A" answers_parts.append( f"Q: {qa['query']}\nA: {qa['answer'][:300]}... (confidence: {conf})" ) answers_summary = "\n\n".join(answers_parts) context = f"""Current research state: - Answers collected: {len(qa_responses)} - Pending questions to search: {len(sub_questions)} Pending questions: {chr(10).join(f"- {q}" for q in sub_questions) if sub_questions else "(none)"} {f"Collected answers:{chr(10)}{answers_summary}" if answers_summary else ""} User message: {user_message}""" result = await agent.run(context) return result.output async def run_interactive_research( question: str, client: HaikuRAG, config: AppConfig | None = None, search_filter: str | None = None, ) -> ResearchReport: """Run interactive research with human-in-the-loop decision points. Args: question: The research question client: HaikuRAG client for document operations config: Application configuration (uses global config if None) search_filter: Optional SQL WHERE clause to filter documents Returns: ResearchReport with the final synthesis """ config = config or get_config() console = Console() # Build interactive graph graph = build_research_graph(config=config, include_plan=True, interactive=True) # Create async queue for human input human_input_queue: asyncio.Queue[HumanDecision] = asyncio.Queue() # Create emitter emitter: AGUIEmitter[ResearchState, ResearchReport] = AGUIEmitter() # Create deps with queue deps = ResearchDeps( client=client, agui_emitter=emitter, human_input_queue=human_input_queue, interactive=True, ) # Create initial state context = ResearchContext(original_question=question) state = ResearchState.from_config(context=context, config=config) state.search_filter = search_filter # Start the run emitter.start_run(state) # Run graph in background task async def run_graph() -> ResearchReport: try: result = await graph.run(state=state, deps=deps) emitter.finish_run(result) return result except Exception as e: emitter.error(e) raise graph_task = asyncio.create_task(run_graph()) # Process events and handle human decision points try: async for event in emitter: event_type = event.get("type") if event_type == "STEP_STARTED": step_name = event.get("stepName", "") if step_name == "plan": console.print("[dim]Planning research...[/dim]") elif step_name.startswith("search:"): query = step_name.replace("search: ", "") console.print(f"[dim]Searching: {query}[/dim]") elif step_name == "synthesize": console.print("[dim]Synthesizing report...[/dim]") elif event_type == "STATE_SNAPSHOT" or event_type == "STATE_DELTA": # State updated, could show progress pass elif event_type == "TOOL_CALL_START": tool_name = event.get("toolCallName") if tool_name == "human_decision": # Will get args in next event pass elif event_type == "TOOL_CALL_ARGS": delta = event.get("delta", "{}") args = json.loads(delta) if isinstance(delta, str) else delta original_question = args.get("original_question", "") sub_questions = list(args.get("sub_questions", [])) qa_responses = args.get("qa_responses", []) iterations = args.get("iterations", 0) # Loop for modifications until user wants to proceed while True: # Show research state console.print() console.print( Panel( f"[bold]{original_question}[/bold]", title="Research Question", border_style="blue", ) ) # Show collected answers if qa_responses: answers_text = [] for i, qa in enumerate(qa_responses, 1): conf = ( f"{qa['confidence']:.0%}" if qa.get("confidence") else "N/A" ) answer_preview = ( qa["answer"][:200] + "..." if len(qa["answer"]) > 200 else qa["answer"] ) answers_text.append( f"[cyan]{i}. {qa['query']}[/cyan]\n" f" [dim]Confidence: {conf} | Citations: {qa.get('citations_count', 0)}[/dim]\n" f" {answer_preview}" ) console.print( Panel( "\n\n".join(answers_text), title=f"Answers Collected ({len(qa_responses)})", border_style="green", ) ) # Show pending questions if sub_questions: console.print( Panel( "\n".join( f"{i + 1}. {q}" for i, q in enumerate(sub_questions) ), title="Pending Questions to Search", border_style="cyan", ) ) else: console.print("[dim]No pending questions.[/dim]") if iterations > 0: console.print(f"[dim]Iteration: {iterations}[/dim]") # Prompt user with context-aware hints console.print() hints = [] if sub_questions: hints.append("search questions") hints.append("modify questions") if qa_responses: hints.append("generate report") hint_text = f" [dim]({', '.join(hints)})[/dim]" if hints else "" user_input = Prompt.ask( f"[bold]What would you like to do?[/bold]{hint_text}" ) # Chat with research assistant console.print("[dim]Thinking...[/dim]") decision = await interpret_user_decision( user_message=user_input, sub_questions=sub_questions, qa_responses=qa_responses, config=config, ) # Handle modifications and chat locally, continue loop if decision.action == "chat": if decision.message: console.print( f"\n[bold cyan]Assistant:[/bold cyan] {decision.message}" ) continue elif decision.action == "add_questions" and decision.questions: sub_questions.extend(decision.questions) console.print( f"[green]Added {len(decision.questions)} question(s)[/green]" ) continue elif decision.action == "modify_questions" and decision.questions: sub_questions = list(decision.questions) console.print( f"[green]Replaced with {len(decision.questions)} question(s)[/green]" ) continue # User wants to proceed - send final decision action_display = { "search": "Searching questions", "synthesize": "Generating report", } console.print( f"[dim]→ {action_display.get(decision.action, decision.action)}[/dim]" ) # Include any accumulated question changes if decision.action == "search": decision = HumanDecision( action="modify_questions", questions=sub_questions ) await human_input_queue.put(decision) break elif event_type == "TEXT_MESSAGE_CHUNK": # Log message from graph message = event.get("delta", "") if message: console.print(f"[dim]{message}[/dim]") elif event_type == "RUN_FINISHED": break elif event_type == "RUN_ERROR": error_msg = event.get("message", "Unknown error") console.print(f"[red]Error: {error_msg}[/red]") break # Wait for graph to complete report = await graph_task return report except Exception as e: graph_task.cancel() raise e finally: await emitter.close() async def run_chat_loop( client: HaikuRAG, config: AppConfig | None = None, search_filter: str | None = None, question: str | None = None, ) -> None: """Run an interactive chat loop for research. Args: client: HaikuRAG client for document operations config: Application configuration (uses global config if None) search_filter: Optional SQL WHERE clause to filter documents question: Optional initial research question (skips initial chat if provided) """ config = config or get_config() console = Console() console.print( Panel( "[bold cyan]Interactive Research Mode[/bold cyan]\n\n" "Chat with me or tell me what you'd like to research.\n" "Type [green]exit[/green] or [green]quit[/green] to end the session.", title="haiku.rag Research Assistant", border_style="cyan", ) ) while True: try: # Use provided question or get one through conversation if question: research_question = question console.print(f"[dim]Starting research: {research_question}[/dim]") question = None # Clear so subsequent loops go through chat else: # Initial conversation loop - chat until user wants to research research_question = None while research_question is None: user_input = Prompt.ask("\n[bold blue]You[/bold blue]") if not user_input.strip(): continue if user_input.lower().strip() in ("exit", "quit", "q"): console.print("[dim]Goodbye![/dim]") return console.print("[dim]Thinking...[/dim]") decision = await initial_chat(user_input, config) if decision.action == "research" and decision.research_question: research_question = decision.research_question console.print( f"[dim]Starting research: {research_question}[/dim]" ) elif decision.action == "chat" and decision.message: console.print( f"\n[bold cyan]Assistant:[/bold cyan] {decision.message}" ) else: # Fallback - treat as research question research_question = user_input console.print() report = await run_interactive_research( question=research_question, client=client, config=config, search_filter=search_filter, ) # Display final report console.print() console.print( Panel( Markdown(f"## {report.title}\n\n{report.executive_summary}"), title="Research Report", border_style="green", ) ) if report.main_findings: findings = "\n".join(f"- {f}" for f in report.main_findings[:5]) console.print(Markdown(f"**Key Findings:**\n{findings}")) if report.conclusions: conclusions = "\n".join(f"- {c}" for c in report.conclusions[:3]) console.print(Markdown(f"**Conclusions:**\n{conclusions}")) console.print(Markdown(f"**Sources:** {report.sources_summary}")) except KeyboardInterrupt: console.print("\n[dim]Goodbye![/dim]") return except Exception as e: console.print(f"[red]Error: {e}[/red]") def interactive_research( client: HaikuRAG, config: AppConfig | None = None, search_filter: str | None = None, question: str | None = None, ) -> None: """Entry point for interactive research mode. Args: client: HaikuRAG client for document operations config: Application configuration (uses global config if None) search_filter: Optional SQL WHERE clause to filter documents question: Optional initial research question (skips initial chat if provided) """ asyncio.run(run_chat_loop(client, config, search_filter, question))