haiku.rag-slim 0.16.0__py3-none-any.whl → 0.24.0__py3-none-any.whl

haiku/rag/cli_chat.py

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