Interactive research agent through AGUI client-side tool calls in CLI
This commit is contained in:
parent
b60c5583aa
commit
218126de8d
11 changed files with 809 additions and 50 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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).
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
465
haiku_rag_slim/haiku/rag/cli_chat.py
Normal file
465
haiku_rag_slim/haiku/rag/cli_chat.py
Normal file
|
|
@ -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))
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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] = {}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Reference in a new issue