Refactor ag-ui-research. Drop human-in-the-loop, use MemoryObjectSendStream to merge the graph and agent streams together
This commit is contained in:
Yiorgis Gozadinos
parent
4a623934b2
commit
2c43f034de
16 changed files with 12671 additions and 14635 deletions
|
|
@ -21,6 +21,12 @@
|
|||
- **CLI AG-UI Flag**: New `--agui` flag for `serve` command to start AG-UI server
|
||||
- **Graph Module**: New unified `haiku.rag.graph` module containing all graph-related functionality
|
||||
- **Common Graph Nodes**: New factory functions (`create_plan_node`, `create_search_node`) in `haiku.rag.graph.common.nodes` for reusable graph components
|
||||
- **AG-UI Research Example**: New full-stack example (`examples/ag-ui-research`) demonstrating agent+graph architecture with CopilotKit frontend
|
||||
- Pydantic AI agent with research tool that invokes the research graph
|
||||
- Custom AG-UI streaming endpoint with anyio memory streams
|
||||
- React/Next.js frontend with split-pane UI showing live research state
|
||||
- Real-time progress tracking of questions, answers, insights, and gaps
|
||||
- Docker Compose setup for easy local development
|
||||
|
||||
### Changed
|
||||
|
||||
|
|
|
|||
|
|
@ -3,6 +3,11 @@
|
|||
# Must be an absolute path to an existing database created with haiku-rag
|
||||
DB_PATH=/absolute/path/to/your/haiku.rag.lancedb
|
||||
|
||||
# Ollama API base URL (if using Ollama for local models)
|
||||
# If running Ollama on your host machine, use your machine's IP address
|
||||
# that the Docker container can reach (not localhost)
|
||||
OLLAMA_BASE_URL=http://host.docker.internal:11434
|
||||
|
||||
# API keys (set as needed for your QA provider)
|
||||
# OPENAI_API_KEY=your-key-here
|
||||
# ANTHROPIC_API_KEY=your-key-here
|
||||
|
|
|
|||
|
|
@ -1,16 +1,16 @@
|
|||
# Interactive Research Assistant
|
||||
|
||||
Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/), [Pydantic AI](https://ai.pydantic.dev/), and [AG-UI](https://docs.ag-ui.com/). Ask complex questions and watch the research process unfold in real-time.
|
||||
Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/), [Pydantic Graph](https://ai.pydantic.dev/graph/), and [AG-UI](https://docs.ag-ui.com/). Ask complex questions and watch the research process unfold in real-time.
|
||||
|
||||
[Watch demo video](https://vimeo.com/1128874386)
|
||||
|
||||
## Features
|
||||
|
||||
- **Multi-step research workflow**: Question decomposition, search, analysis, and synthesis
|
||||
- **Human-in-the-loop**: Approve or revise research plans before execution
|
||||
- **Live state synchronization**: Real-time updates of research progress between backend and frontend
|
||||
- **Context expansion**: Automatically expands top search results for better context
|
||||
- **Rich reporting**: Generates structured reports with findings, conclusions, and citations
|
||||
- **Multi-iteration research graph**: Automated question decomposition, search, insight extraction, and gap analysis
|
||||
- **Intelligent evaluation**: Confidence-based decision making with automatic iteration until sufficient information is gathered
|
||||
- **Live state synchronization**: Real-time delta updates of research progress via AG-UI protocol
|
||||
- **Insight & gap tracking**: Structured insights with provenance and automatic gap identification
|
||||
- **Rich reporting**: Generates comprehensive research reports with findings, conclusions, and sources
|
||||
|
||||
## Quick Start
|
||||
|
||||
|
|
@ -57,32 +57,67 @@ Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/),
|
|||
## How It Works
|
||||
|
||||
1. **Ask a question**: Type your research question in the chat
|
||||
2. **Review the plan**: The agent decomposes your question into 3 sub-questions
|
||||
3. **Approve or revise**: Choose to approve the plan or request changes
|
||||
4. **Watch it work**: The agent automatically:
|
||||
- Searches the knowledge base for each sub-question
|
||||
- Extracts key insights from search results
|
||||
- Evaluates overall confidence in findings
|
||||
5. **Get your report**: Receive a structured research report with citations
|
||||
2. **Plan phase**: The research graph automatically:
|
||||
- Decomposes your question into targeted sub-questions
|
||||
- Gathers initial context about the topic
|
||||
3. **Research iterations**: The graph autonomously:
|
||||
- Searches the knowledge base for each sub-question in parallel
|
||||
- Extracts structured insights with source provenance
|
||||
- Identifies information gaps and assesses confidence
|
||||
- Generates new follow-up questions for gaps
|
||||
- Iterates until confidence threshold is met or max iterations reached
|
||||
4. **Synthesis**: Generates a comprehensive research report with:
|
||||
- Executive summary
|
||||
- Main findings with supporting evidence
|
||||
- Conclusions and recommendations
|
||||
- Source citations
|
||||
|
||||
## Architecture
|
||||
|
||||
- **Backend** (Python): Pydantic AI agent with haiku.rag integration
|
||||
- Uses published `ghcr.io/ggozad/haiku.rag:latest` Docker image as base
|
||||
- `agent.py`: Research agent with tool definitions
|
||||
- `main.py`: Starlette app serving AG-UI protocol
|
||||
### Agent + Graph Pattern
|
||||
|
||||
- **Frontend** (Next.js): CopilotKit/AG-UI interface
|
||||
- Real-time state synchronization with backend
|
||||
- Interactive approval workflow
|
||||
- Collapsible research plan and insights display
|
||||
This example demonstrates the **agent+graph** architecture pattern:
|
||||
|
||||
1. **Conversational Agent** (`agent.py`):
|
||||
- Pydantic AI agent handles user conversations
|
||||
- Decides when to invoke the research tool based on user intent
|
||||
- Responds directly to greetings/casual chat without tools
|
||||
- Formats research results for the user
|
||||
|
||||
2. **Research Graph** (haiku.rag):
|
||||
- Multi-step research workflow invoked by the agent's tool
|
||||
- Autonomous execution with plan → search → analyze → decide → synthesize flow
|
||||
- Emits AG-UI events for real-time progress tracking
|
||||
|
||||
3. **Shared Event Stream**:
|
||||
- `AGUIEmitter` is shared between agent and graph
|
||||
- Events from both flow through a single stream to the frontend
|
||||
- Custom streaming endpoint (`main.py`) uses anyio memory streams for proper async handling
|
||||
|
||||
### Components
|
||||
|
||||
- **Backend** (Python):
|
||||
- Uses published `ghcr.io/ggozad/haiku.rag:latest` Docker image as base
|
||||
- `agent.py`: Pydantic AI agent with `run_research` tool
|
||||
- `main.py`: Custom AG-UI streaming endpoint with anyio memory object streams
|
||||
- Real-time event forwarding from emitter to SSE stream
|
||||
- Filters out `ACTIVITY_SNAPSHOT` events (not yet supported by CopilotKit)
|
||||
|
||||
- **Frontend** (Next.js/React):
|
||||
- CopilotKit for AG-UI protocol integration
|
||||
- Split-pane UI: chat on left, live research state on right
|
||||
- Real-time state synchronization via Server-Sent Events (SSE)
|
||||
- `StateDisplay` component with collapsible sections for questions, insights, and gaps
|
||||
|
||||
## Configuration
|
||||
|
||||
Configuration is done through `haiku.rag.yaml` (see `haiku.rag.yaml.example`):
|
||||
|
||||
- `qa.provider`: LLM provider (default: `ollama`)
|
||||
- `qa.model`: Model name (default: `gpt-oss:latest`)
|
||||
- `research.provider`: LLM provider (default: `ollama`)
|
||||
- `research.model`: Model name (default: `gpt-oss:latest`)
|
||||
- `research.max_iterations`: Maximum research iterations (default: `3`)
|
||||
- `research.confidence_threshold`: Confidence threshold for completion (default: `0.8`)
|
||||
- `research.max_concurrency`: Parallel sub-question processing (default: `1`)
|
||||
- `providers.ollama.base_url`: Ollama endpoint (default: `http://host.docker.internal:11434`)
|
||||
|
||||
Environment variables (see `.env.example`):
|
||||
|
|
|
|||
|
|
@ -3,15 +3,13 @@ FROM ghcr.io/ggozad/haiku.rag:latest
|
|||
WORKDIR /app
|
||||
|
||||
# Copy backend application files
|
||||
COPY agent.py main.py ./
|
||||
COPY pyproject.toml ./
|
||||
COPY main.py agent.py ./
|
||||
|
||||
# Install backend dependencies
|
||||
# Note: haiku-rag is already installed in the base image
|
||||
# Install additional dependencies for the example
|
||||
# Note: haiku-rag-slim is already installed in the base image
|
||||
RUN pip install --no-cache-dir \
|
||||
starlette>=0.45.2 \
|
||||
uvicorn[standard]>=0.34.2 \
|
||||
pydantic-ai-slim[ag-ui,openai]>=1.1.0 \
|
||||
python-dotenv>=1.0.1
|
||||
|
||||
EXPOSE 8000
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
# Haiku.rag Research Assistant Backend
|
||||
|
||||
FastAPI backend for the haiku.rag interactive research assistant, using Pydantic AI with AG-UI protocol support.
|
||||
Starlette backend for the haiku.rag interactive research assistant, using the research graph with AG-UI protocol support.
|
||||
|
||||
## Setup
|
||||
|
||||
|
|
@ -11,7 +11,17 @@ uv run python main.py
|
|||
|
||||
The server starts on `http://localhost:8000` and uses [haiku.rag configuration](https://ggozad.github.io/haiku.rag/configuration/).
|
||||
|
||||
## Architecture
|
||||
|
||||
The backend uses `create_agui_server()` from `haiku.rag.graph.agui.server` which provides:
|
||||
|
||||
- **Research graph execution**: Multi-iteration research workflow with insight/gap tracking
|
||||
- **AG-UI protocol**: Server-Sent Events (SSE) streaming for real-time state updates
|
||||
- **Delta state updates**: Efficient incremental state synchronization using JSON Patch operations
|
||||
- **Both research and deep_qa endpoints**: `/agent/research` and `/agent/deep_qa`
|
||||
|
||||
## Endpoints
|
||||
|
||||
- `GET /health` - Health check
|
||||
- `POST /agent` - AG-UI protocol endpoint
|
||||
- `GET /health` - Health check with configuration info
|
||||
- `POST /agent/research/stream` - Research graph streaming endpoint (AG-UI protocol)
|
||||
- `POST /agent/deep_qa/stream` - Deep QA graph streaming endpoint (AG-UI protocol)
|
||||
|
|
|
|||
|
|
@ -1,432 +1,111 @@
|
|||
import json
|
||||
from dataclasses import dataclass
|
||||
"""Research assistant agent with graph integration."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from ag_ui.core import EventType, StateSnapshotEvent
|
||||
from pydantic import BaseModel
|
||||
from pydantic_ai import Agent, RunContext
|
||||
from pydantic_ai.ag_ui import StateDeps
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import Config
|
||||
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
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.graph.agui.emitter import AGUIEmitter
|
||||
from haiku.rag.graph.research.models import ResearchReport
|
||||
|
||||
class ResearchState(BaseModel):
|
||||
"""Shared state between research agent and frontend."""
|
||||
|
||||
question: str = ""
|
||||
phase: str = "idle"
|
||||
status: str = ""
|
||||
plan: list[dict] = []
|
||||
current_question_index: int = 0
|
||||
insights: list[dict] = []
|
||||
document_registry: dict[str, dict] = {}
|
||||
current_document: dict | None = None
|
||||
confidence: float = 0.0
|
||||
final_report: dict | None = None
|
||||
# Load config
|
||||
config_path = Path("/app/haiku.rag.yaml")
|
||||
Config = (
|
||||
AppConfig.model_validate(load_yaml_config(config_path))
|
||||
if config_path.exists()
|
||||
else AppConfig()
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResearchDeps(StateDeps[ResearchState]):
|
||||
"""Dependencies for the research agent with HaikuRAG client."""
|
||||
class AgentDeps:
|
||||
"""Dependencies for research agent."""
|
||||
|
||||
client: HaikuRAG
|
||||
agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None
|
||||
|
||||
|
||||
def _as_state_snapshot(ctx: RunContext[ResearchDeps]) -> StateSnapshotEvent:
|
||||
return StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=ctx.deps.state)
|
||||
model = get_model(Config.research.provider, Config.research.model)
|
||||
|
||||
|
||||
def create_agent(
|
||||
qa_provider: str = Config.qa.provider, qa_model: str = Config.qa.model
|
||||
) -> Agent[ResearchDeps, str]:
|
||||
"""Create and configure the research agent.
|
||||
|
||||
Args:
|
||||
qa_provider: QA provider for the agent (default: from Config.qa.provider)
|
||||
qa_model: Model name to use (default: from Config.qa.model)
|
||||
"""
|
||||
print(f"[AGENT SETUP] Creating agent with provider={qa_provider}, model={qa_model}")
|
||||
agent = Agent(
|
||||
model=get_model(qa_provider, qa_model),
|
||||
deps_type=ResearchDeps,
|
||||
instructions="""You are a research co-pilot powered by haiku.rag.
|
||||
|
||||
Your workflow MUST follow these exact steps in order:
|
||||
1. Call propose_research_plan with the user's question
|
||||
2. After propose_research_plan completes, IMMEDIATELY call approve_research_plan (with no arguments)
|
||||
3. WAIT for approve_research_plan to return:
|
||||
- If it returns "APPROVED", proceed to step 4
|
||||
- If it returns "REVISE", ask the user "How would you like me to revise the research plan?" and wait for their response
|
||||
- Once you receive their revision feedback, revise the plan and go back to step 1
|
||||
4. Once approved, process questions ONE AT A TIME:
|
||||
- Call search_question(question_id=0) and WAIT for it to complete
|
||||
- Then call extract_insights_from_results(question_id=0) and WAIT for it to complete
|
||||
- Then call search_question(question_id=1) and WAIT for it to complete
|
||||
- Then call extract_insights_from_results(question_id=1) and WAIT for it to complete
|
||||
- Then call search_question(question_id=2) and WAIT for it to complete
|
||||
- Then call extract_insights_from_results(question_id=2) and WAIT for it to complete
|
||||
5. After all questions are processed, call evaluate_research_confidence
|
||||
6. Ask user if they want to finalize or continue researching
|
||||
7. When user approves, call synthesize_final_report
|
||||
agent = Agent(
|
||||
model,
|
||||
deps_type=AgentDeps,
|
||||
system_prompt="""You are an advanced research assistant powered by haiku.rag.
|
||||
|
||||
CRITICAL RULES:
|
||||
- MANDATORY: Call approve_research_plan immediately after propose_research_plan - NO EXCEPTIONS
|
||||
- If approve_research_plan returns "REVISE", ask the user for revision feedback naturally in chat
|
||||
- Call ONE tool at a time - wait for each tool to return before calling the next
|
||||
- NEVER call extract_insights_from_results until search_question has completed and returned results
|
||||
- DO NOT explain what you're about to do - just call the tool
|
||||
- The state updates will show the user what's happening - you don't need to narrate
|
||||
- Process all 3 questions automatically without asking for approval between them
|
||||
1. For greetings (hi, hello, hey, etc) or casual chat: respond directly WITHOUT using any tools
|
||||
2. For questions about yourself or the system: respond directly WITHOUT using any tools
|
||||
3. For substantive questions requiring information: ALWAYS use the run_research tool
|
||||
4. NEVER answer substantive questions from your own knowledge - always use the tool
|
||||
|
||||
Document Viewing:
|
||||
- When user asks to "show document X", call get_full_document with the document_uri
|
||||
How to decide:
|
||||
- "Hi" / "Hello" / "How are you?" → Respond directly, NO tools
|
||||
- "What can you do?" → Respond directly, NO tools
|
||||
- "How does X work in the codebase?" → Use run_research tool
|
||||
- "Tell me about Y" → Use run_research tool
|
||||
|
||||
Remember: Call tools ONE AT A TIME in sequence. Each tool must complete before calling the next.
|
||||
""",
|
||||
When you use run_research, the graph will decompose questions, search the knowledge base,
|
||||
extract insights, and generate a comprehensive report.
|
||||
|
||||
Be friendly and conversational in all responses.""",
|
||||
)
|
||||
|
||||
|
||||
@agent.tool
|
||||
async def run_research(ctx: RunContext[AgentDeps], question: str) -> str:
|
||||
"""Execute research graph on a substantive question.
|
||||
|
||||
Use for questions requiring knowledge base search.
|
||||
DO NOT use for greetings or casual conversation.
|
||||
"""
|
||||
if ctx.deps.agui_emitter:
|
||||
ctx.deps.agui_emitter.log(f"🔍 Starting research on: {question}")
|
||||
|
||||
graph = build_research_graph(Config)
|
||||
context = ResearchContext(original_question=question)
|
||||
state = ResearchState.from_config(context=context, config=Config)
|
||||
|
||||
graph_deps = ResearchDeps(
|
||||
client=ctx.deps.client,
|
||||
agui_emitter=ctx.deps.agui_emitter,
|
||||
)
|
||||
|
||||
@agent.tool
|
||||
async def propose_research_plan(
|
||||
ctx: RunContext[ResearchDeps], question: str
|
||||
) -> StateSnapshotEvent:
|
||||
"""Propose a research plan by decomposing the question into sub-questions."""
|
||||
ctx.deps.state.question = question
|
||||
ctx.deps.state.phase = "planning"
|
||||
ctx.deps.state.status = "Decomposing question into sub-questions..."
|
||||
try:
|
||||
result = await graph.run(state=state, deps=graph_deps)
|
||||
|
||||
decompose_prompt = f"""Break down this research question into exactly 3 specific sub-questions that would help answer it comprehensively.
|
||||
if ctx.deps.agui_emitter:
|
||||
ctx.deps.agui_emitter.log("✅ Research complete!")
|
||||
|
||||
Research Question: {question}
|
||||
return f"""Research completed successfully!
|
||||
|
||||
Return ONLY a JSON array of sub-questions, like: ["Question 1?", "Question 2?", "Question 3?"]"""
|
||||
Question: {question}
|
||||
|
||||
response = await ctx.deps.client.ask(decompose_prompt)
|
||||
Executive Summary: {result.executive_summary}
|
||||
|
||||
try:
|
||||
sub_questions = json.loads(response)
|
||||
except json.JSONDecodeError:
|
||||
sub_questions = [
|
||||
q.strip().lstrip("0123456789.-) ")
|
||||
for q in response.split("\n")
|
||||
if q.strip()
|
||||
][:3]
|
||||
Main Findings:
|
||||
{chr(10).join(f"- {finding}" for finding in result.main_findings[:3])}
|
||||
|
||||
plan = [
|
||||
{"id": i, "question": q, "status": "pending"}
|
||||
for i, q in enumerate(sub_questions)
|
||||
]
|
||||
Conclusions:
|
||||
{chr(10).join(f"- {conclusion}" for conclusion in result.conclusions[:2])}
|
||||
|
||||
ctx.deps.state.plan = plan
|
||||
ctx.deps.state.current_question_index = 0
|
||||
ctx.deps.state.status = f"Proposed plan with {len(plan)} sub-questions"
|
||||
Total insights gathered: {len(state.context.insights)}
|
||||
Confidence: {f"{state.last_eval.confidence_score:.0%}" if state.last_eval else "N/A"}
|
||||
Iterations completed: {state.iterations}
|
||||
|
||||
return _as_state_snapshot(ctx)
|
||||
The full research report with all citations has been provided to the user.
|
||||
"""
|
||||
|
||||
@agent.tool
|
||||
async def search_question(
|
||||
ctx: RunContext[ResearchDeps],
|
||||
question_id: int,
|
||||
search_type: str = "hybrid",
|
||||
) -> StateSnapshotEvent:
|
||||
"""Execute search for a specific sub-question."""
|
||||
plan = ctx.deps.state.plan
|
||||
if question_id >= len(plan):
|
||||
raise ValueError(f"Question ID {question_id} not found in plan")
|
||||
|
||||
question = plan[question_id]["question"]
|
||||
ctx.deps.state.phase = "searching"
|
||||
ctx.deps.state.current_question_index = question_id
|
||||
ctx.deps.state.status = f"Searching: {question}"
|
||||
plan[question_id]["status"] = "searching"
|
||||
|
||||
search_results = await ctx.deps.client.search(
|
||||
question, limit=5, search_type=search_type
|
||||
)
|
||||
|
||||
expanded_map = {}
|
||||
if search_results:
|
||||
expanded_results = await ctx.deps.client.expand_context(
|
||||
search_results[:3], radius=2
|
||||
)
|
||||
expanded_map = {
|
||||
chunk.id: (chunk, score) for chunk, score in expanded_results
|
||||
}
|
||||
|
||||
results = []
|
||||
for chunk, score in search_results:
|
||||
doc_uri = chunk.document_uri or "unknown"
|
||||
doc_title = chunk.document_title or chunk.document_uri or "Unknown"
|
||||
|
||||
if doc_uri not in ctx.deps.state.document_registry:
|
||||
ctx.deps.state.document_registry[doc_uri] = {
|
||||
"title": doc_title,
|
||||
"chunks_referenced": [],
|
||||
}
|
||||
|
||||
if (
|
||||
chunk.id
|
||||
not in ctx.deps.state.document_registry[doc_uri]["chunks_referenced"]
|
||||
):
|
||||
ctx.deps.state.document_registry[doc_uri]["chunks_referenced"].append(
|
||||
chunk.id
|
||||
)
|
||||
|
||||
expanded_chunk, _ = (
|
||||
expanded_map[chunk.id] if chunk.id in expanded_map else (chunk, score)
|
||||
)
|
||||
result_data = {
|
||||
"chunk": expanded_chunk.content[:500],
|
||||
"chunk_id": chunk.id,
|
||||
"document_uri": doc_uri,
|
||||
"document_title": doc_title,
|
||||
"chunk_position": chunk.order,
|
||||
"full_chunk_content": expanded_chunk.content,
|
||||
"score": round(score, 3),
|
||||
"expanded": chunk.id in expanded_map,
|
||||
}
|
||||
results.append(result_data)
|
||||
|
||||
plan[question_id]["search_results"] = {
|
||||
"type": search_type,
|
||||
"results": results,
|
||||
}
|
||||
plan[question_id]["status"] = "searched"
|
||||
ctx.deps.state.status = f"Found {len(results)} results"
|
||||
|
||||
return _as_state_snapshot(ctx)
|
||||
|
||||
@agent.tool
|
||||
async def extract_insights_from_results(
|
||||
ctx: RunContext[ResearchDeps],
|
||||
question_id: int,
|
||||
) -> StateSnapshotEvent:
|
||||
"""Extract key insights from search results for a specific question."""
|
||||
plan = ctx.deps.state.plan
|
||||
if question_id >= len(plan):
|
||||
raise ValueError(f"Question ID {question_id} not found in plan")
|
||||
|
||||
question_item = plan[question_id]
|
||||
if "search_results" not in question_item:
|
||||
raise ValueError(
|
||||
f"No search results found for question ID {question_id}. "
|
||||
f"You must call search_question(question_id={question_id}) first."
|
||||
)
|
||||
|
||||
search_results = question_item["search_results"]
|
||||
ctx.deps.state.phase = "analyzing"
|
||||
ctx.deps.state.status = "Extracting insights from results..."
|
||||
|
||||
context_parts = [
|
||||
f"[Result {idx}] [Source: {r['document_title']}] {r['full_chunk_content']}"
|
||||
for idx, r in enumerate(search_results["results"])
|
||||
]
|
||||
context = "\n\n".join(context_parts)
|
||||
|
||||
class InsightResult(BaseModel):
|
||||
summary: str
|
||||
confidence: float
|
||||
result_indices: list[int]
|
||||
|
||||
class InsightsList(BaseModel):
|
||||
insights: list[InsightResult]
|
||||
|
||||
question_text = question_item["question"]
|
||||
extract_prompt = f"""Analyze these search results and extract 1-3 key insights that help answer the question: "{question_text}"
|
||||
|
||||
Search Results:
|
||||
{context}
|
||||
|
||||
For each insight, reference which result numbers (0, 1, 2, etc.) support it."""
|
||||
|
||||
insight_agent: Agent[None, InsightsList] = Agent(
|
||||
ctx.model,
|
||||
output_type=InsightsList,
|
||||
retries=3,
|
||||
)
|
||||
|
||||
result = await insight_agent.run(extract_prompt)
|
||||
raw_insights = [
|
||||
{
|
||||
"summary": insight.summary,
|
||||
"confidence": insight.confidence,
|
||||
"result_indices": insight.result_indices,
|
||||
}
|
||||
for insight in result.output.insights
|
||||
]
|
||||
|
||||
new_insights = []
|
||||
for insight in raw_insights:
|
||||
source_refs = []
|
||||
for idx in insight.get("result_indices", []):
|
||||
if 0 <= idx < len(search_results["results"]):
|
||||
result = search_results["results"][idx]
|
||||
source_refs.append(
|
||||
{
|
||||
"chunk_id": result["chunk_id"],
|
||||
"document_uri": result["document_uri"],
|
||||
"document_title": result["document_title"],
|
||||
"chunk_position": result["chunk_position"],
|
||||
}
|
||||
)
|
||||
|
||||
new_insights.append(
|
||||
{
|
||||
"summary": insight["summary"],
|
||||
"confidence": insight.get("confidence", 0.7),
|
||||
"source_refs": source_refs,
|
||||
}
|
||||
)
|
||||
|
||||
ctx.deps.state.insights.extend(new_insights)
|
||||
plan[question_id]["status"] = "done"
|
||||
ctx.deps.state.status = f"Extracted {len(new_insights)} insights"
|
||||
|
||||
return _as_state_snapshot(ctx)
|
||||
|
||||
@agent.tool
|
||||
async def evaluate_research_confidence(
|
||||
ctx: RunContext[ResearchDeps],
|
||||
) -> StateSnapshotEvent:
|
||||
"""Evaluate overall confidence in the research findings."""
|
||||
insights = ctx.deps.state.insights
|
||||
if not insights:
|
||||
raise ValueError("No insights collected yet")
|
||||
|
||||
ctx.deps.state.phase = "evaluating"
|
||||
ctx.deps.state.status = "Evaluating research confidence..."
|
||||
|
||||
confidences = [i.get("confidence", 0.5) for i in insights]
|
||||
overall_confidence = sum(confidences) / len(confidences) if confidences else 0
|
||||
|
||||
eval_prompt = f"""Evaluate if these insights provide a confident answer to: "{ctx.deps.state.question}"
|
||||
|
||||
Insights collected:
|
||||
{chr(10).join([f"- {i['summary']}" for i in insights])}
|
||||
|
||||
Assess:
|
||||
1. Do we have enough information to answer the question?
|
||||
2. What gaps remain?
|
||||
3. Overall confidence (0.0-1.0)
|
||||
|
||||
Return JSON: {{"confidence": 0.0-1.0, "gaps": ["gap1", "gap2"], "recommendation": "continue" or "finalize"}}"""
|
||||
|
||||
response = await ctx.deps.client.ask(eval_prompt)
|
||||
|
||||
try:
|
||||
evaluation = json.loads(response)
|
||||
overall_confidence = evaluation.get("confidence", overall_confidence)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
ctx.deps.state.confidence = overall_confidence
|
||||
ctx.deps.state.status = f"Confidence: {overall_confidence:.0%}"
|
||||
|
||||
return _as_state_snapshot(ctx)
|
||||
|
||||
@agent.tool
|
||||
async def synthesize_final_report(
|
||||
ctx: RunContext[ResearchDeps],
|
||||
) -> StateSnapshotEvent:
|
||||
"""Generate final research report with citations."""
|
||||
insights = ctx.deps.state.insights
|
||||
if not insights:
|
||||
raise ValueError("No insights to synthesize")
|
||||
|
||||
ctx.deps.state.phase = "synthesizing"
|
||||
ctx.deps.state.status = "Generating final report..."
|
||||
|
||||
insights_summary = []
|
||||
for i in insights:
|
||||
source_titles = [ref["document_title"] for ref in i.get("source_refs", [])]
|
||||
unique_sources = list(dict.fromkeys(source_titles))
|
||||
insights_summary.append(
|
||||
f"- {i['summary']} (sources: {', '.join(unique_sources[:2])})"
|
||||
)
|
||||
|
||||
report_prompt = f"""Generate a comprehensive research report answering: "{ctx.deps.state.question}"
|
||||
|
||||
Based on these insights:
|
||||
{chr(10).join(insights_summary)}
|
||||
|
||||
Create a structured report with:
|
||||
- Executive Summary (2-3 sentences)
|
||||
- Main Findings (bullet points)
|
||||
- Conclusions
|
||||
- Sources (list the document titles mentioned above)
|
||||
|
||||
Return JSON with format:
|
||||
{{
|
||||
"title": "...",
|
||||
"summary": "...",
|
||||
"findings": ["finding1", "finding2", ...],
|
||||
"conclusions": ["conclusion1", ...],
|
||||
"sources": ["source1", "source2", ...]
|
||||
}}"""
|
||||
|
||||
response = await ctx.deps.client.ask(report_prompt)
|
||||
|
||||
try:
|
||||
report = json.loads(response)
|
||||
except json.JSONDecodeError:
|
||||
report = {
|
||||
"title": ctx.deps.state.question,
|
||||
"summary": response[:300],
|
||||
"findings": [i["summary"] for i in insights],
|
||||
"conclusions": ["See findings above"],
|
||||
"sources": [],
|
||||
}
|
||||
|
||||
citations = [
|
||||
{
|
||||
"document_uri": doc_uri,
|
||||
"document_title": doc_info["title"],
|
||||
"chunk_ids": doc_info["chunks_referenced"],
|
||||
}
|
||||
for doc_uri, doc_info in ctx.deps.state.document_registry.items()
|
||||
]
|
||||
report["citations"] = citations
|
||||
|
||||
ctx.deps.state.final_report = report
|
||||
ctx.deps.state.phase = "done"
|
||||
ctx.deps.state.status = "Research complete"
|
||||
|
||||
return _as_state_snapshot(ctx)
|
||||
|
||||
@agent.tool
|
||||
async def get_full_document(
|
||||
ctx: RunContext[ResearchDeps],
|
||||
document_uri: str,
|
||||
) -> StateSnapshotEvent:
|
||||
"""Retrieve and display the full content of a document by its URI."""
|
||||
ctx.deps.state.status = f"Retrieving document: {document_uri}"
|
||||
document = await ctx.deps.client.get_document_by_uri(document_uri)
|
||||
|
||||
if document is None:
|
||||
ctx.deps.state.status = f"Document not found: {document_uri}"
|
||||
ctx.deps.state.current_document = {
|
||||
"uri": document_uri,
|
||||
"title": "Not Found",
|
||||
"content": f"Document with URI '{document_uri}' was not found.",
|
||||
"total_chunks": 0,
|
||||
}
|
||||
else:
|
||||
all_chunks = await ctx.deps.client.search(
|
||||
query="", limit=1000, search_type="fts"
|
||||
)
|
||||
chunks_for_doc = [
|
||||
c for c, _ in all_chunks if c.document_uri == document_uri
|
||||
]
|
||||
|
||||
ctx.deps.state.current_document = {
|
||||
"uri": document.uri or document_uri,
|
||||
"title": document.title or "Untitled",
|
||||
"content": document.content,
|
||||
"total_chunks": len(chunks_for_doc),
|
||||
"metadata": document.metadata,
|
||||
}
|
||||
ctx.deps.state.status = f"Retrieved: {document.title or document_uri}"
|
||||
|
||||
return _as_state_snapshot(ctx)
|
||||
|
||||
return agent
|
||||
except Exception as e:
|
||||
if ctx.deps.agui_emitter:
|
||||
ctx.deps.agui_emitter.log(f"❌ Research error: {str(e)}")
|
||||
return f"I encountered an error while researching: {str(e)}"
|
||||
|
|
|
|||
|
|
@ -1,87 +1,171 @@
|
|||
import logging
|
||||
import os
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
|
||||
from agent import ResearchDeps, ResearchState, create_agent
|
||||
from agent import AgentDeps, agent
|
||||
from anyio import create_memory_object_stream, create_task_group
|
||||
from anyio.streams.memory import MemoryObjectSendStream
|
||||
from starlette.applications import Starlette
|
||||
from starlette.middleware import Middleware
|
||||
from starlette.middleware.cors import CORSMiddleware
|
||||
from starlette.responses import JSONResponse
|
||||
from starlette.routing import Mount, Route
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import JSONResponse, StreamingResponse
|
||||
from starlette.routing import Route
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.config import load_yaml_config
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.graph.agui.emitter import AGUIEmitter
|
||||
from haiku.rag.graph.agui.server import RunAgentInput, format_sse_event
|
||||
from haiku.rag.graph.research.dependencies import ResearchContext
|
||||
from haiku.rag.graph.research.models import ResearchReport
|
||||
from haiku.rag.graph.research.state import ResearchState
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
client: HaikuRAG | None = None
|
||||
ag_ui_app = None
|
||||
# Load config from mounted haiku.rag.yaml
|
||||
config_path = Path("/app/haiku.rag.yaml")
|
||||
if config_path.exists():
|
||||
yaml_data = load_yaml_config(config_path)
|
||||
Config = AppConfig.model_validate(yaml_data)
|
||||
else:
|
||||
# Fallback to default config
|
||||
Config = AppConfig()
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app):
|
||||
global client
|
||||
db_path_str = os.getenv("DB_PATH", "haiku_rag.lancedb")
|
||||
db_path = Path(db_path_str)
|
||||
# Get DB path from environment
|
||||
db_path_str = os.getenv("DB_PATH", "haiku_rag.lancedb")
|
||||
db_path = Path(db_path_str)
|
||||
|
||||
if not db_path.exists():
|
||||
logger.error(f"Database not found at {db_path}")
|
||||
logger.error("Run: haiku-rag add <path-to-documents>")
|
||||
raise RuntimeError(f"Database not found: {db_path}")
|
||||
if not db_path.exists():
|
||||
logger.error(f"Database not found at {db_path}")
|
||||
logger.error("Run: haiku-rag add <path-to-documents>")
|
||||
raise RuntimeError(f"Database not found: {db_path}")
|
||||
|
||||
logger.info(f"Initializing HaikuRAG client with database: {db_path}")
|
||||
client = HaikuRAG(db_path)
|
||||
logger.info("Research assistant backend ready")
|
||||
logger.info(f"QA Provider: {Config.qa.provider}, Model: {Config.qa.model}")
|
||||
logger.info(f"Initializing research assistant with database: {db_path}")
|
||||
logger.info(
|
||||
f"Research Provider: {Config.research.provider}, Model: {Config.research.model}"
|
||||
)
|
||||
|
||||
yield
|
||||
|
||||
if client:
|
||||
logger.info("Closing HaikuRAG client")
|
||||
client.close()
|
||||
# Store client reference for proper lifecycle management
|
||||
_client_cache: dict[str, HaikuRAG] = {}
|
||||
|
||||
|
||||
agent = create_agent()
|
||||
def get_client(effective_db_path: Path) -> HaikuRAG:
|
||||
"""Get or create cached client."""
|
||||
path_key = str(effective_db_path)
|
||||
if path_key not in _client_cache:
|
||||
_client_cache[path_key] = HaikuRAG(db_path=effective_db_path, config=Config)
|
||||
return _client_cache[path_key]
|
||||
|
||||
|
||||
async def health(request):
|
||||
db_path_str = os.getenv("DB_PATH", "haiku_rag.lancedb")
|
||||
async def stream_research_agent(request: Request) -> StreamingResponse:
|
||||
"""Agent streaming endpoint with research graph integration."""
|
||||
body = await request.json()
|
||||
input_data = RunAgentInput(**body)
|
||||
|
||||
user_message = ""
|
||||
if input_data.messages:
|
||||
user_message = input_data.messages[-1].get("content", "")
|
||||
|
||||
send_stream, receive_stream = create_memory_object_stream[str]()
|
||||
|
||||
async def run_agent_with_streaming(
|
||||
send_stream: MemoryObjectSendStream[str],
|
||||
) -> None:
|
||||
"""Execute agent and forward emitter events to memory stream."""
|
||||
async with send_stream:
|
||||
try:
|
||||
# Create shared emitter
|
||||
emitter: AGUIEmitter[ResearchState, ResearchReport] = AGUIEmitter(
|
||||
thread_id=input_data.thread_id,
|
||||
run_id=input_data.run_id,
|
||||
use_deltas=False,
|
||||
)
|
||||
|
||||
# Get client
|
||||
effective_db_path = input_data.config.get("db_path") or db_path
|
||||
if isinstance(effective_db_path, str):
|
||||
effective_db_path = Path(effective_db_path)
|
||||
client = get_client(effective_db_path)
|
||||
|
||||
# Create agent dependencies with shared emitter
|
||||
agent_deps = AgentDeps(client=client, agui_emitter=emitter)
|
||||
|
||||
# Start run with empty initial state
|
||||
emitter.start_run(
|
||||
initial_state=ResearchState.from_config(
|
||||
context=ResearchContext(original_question=""),
|
||||
config=Config,
|
||||
)
|
||||
)
|
||||
|
||||
# Forward emitter events to stream
|
||||
async def forward_events():
|
||||
async for event in emitter:
|
||||
# Filter out ACTIVITY_SNAPSHOT - not supported by CopilotKit
|
||||
if event.get("type") == "ACTIVITY_SNAPSHOT":
|
||||
continue
|
||||
await send_stream.send(format_sse_event(event))
|
||||
|
||||
# Run agent and event forwarding concurrently
|
||||
async with create_task_group() as tg:
|
||||
tg.start_soon(forward_events)
|
||||
|
||||
result = await agent.run(user_message, deps=agent_deps)
|
||||
emitter.log(result.output)
|
||||
await emitter.close()
|
||||
|
||||
except Exception as e:
|
||||
logger.exception("Error executing agent")
|
||||
try:
|
||||
await send_stream.send(
|
||||
format_sse_event({"type": "error", "error": str(e)})
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def event_generator():
|
||||
"""Generate SSE events from memory stream."""
|
||||
async with create_task_group() as tg:
|
||||
tg.start_soon(run_agent_with_streaming, send_stream)
|
||||
async with receive_stream:
|
||||
async for event_str in receive_stream:
|
||||
yield event_str
|
||||
|
||||
return StreamingResponse(
|
||||
event_generator(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
"X-Accel-Buffering": "no",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def health_check(_: Request) -> JSONResponse:
|
||||
"""Health check endpoint with configuration info."""
|
||||
return JSONResponse(
|
||||
{
|
||||
"status": "healthy",
|
||||
"agent_model": str(agent.model),
|
||||
"qa_provider": Config.qa.provider,
|
||||
"qa_model": Config.qa.model,
|
||||
"ollama_base_url": Config.providers.ollama.base_url,
|
||||
"db_path": db_path_str,
|
||||
"db_exists": Path(db_path_str).exists(),
|
||||
"research_provider": Config.research.provider,
|
||||
"research_model": Config.research.model,
|
||||
"db_path": str(db_path),
|
||||
"db_exists": db_path.exists(),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def get_ag_ui_app():
|
||||
global ag_ui_app
|
||||
if ag_ui_app is None and client is not None:
|
||||
research_deps = ResearchDeps(client=client, state=ResearchState())
|
||||
logger.info("Creating AG-UI app")
|
||||
ag_ui_app = agent.to_ag_ui(deps=research_deps)
|
||||
return ag_ui_app
|
||||
|
||||
|
||||
async def agent_endpoint(scope, receive, send):
|
||||
app = get_ag_ui_app()
|
||||
if app is None:
|
||||
response = JSONResponse({"error": "Client not initialized"}, status_code=503)
|
||||
await response(scope, receive, send)
|
||||
return
|
||||
await app(scope, receive, send)
|
||||
|
||||
|
||||
# Create Starlette app
|
||||
app = Starlette(
|
||||
routes=[
|
||||
Route("/health", health),
|
||||
Mount("/agent", agent_endpoint),
|
||||
Route("/v1/research/stream", stream_research_agent, methods=["POST"]),
|
||||
Route("/health", health_check, methods=["GET"]),
|
||||
],
|
||||
middleware=[
|
||||
Middleware(
|
||||
|
|
@ -92,17 +176,11 @@ app = Starlette(
|
|||
allow_headers=["*"],
|
||||
)
|
||||
],
|
||||
lifespan=lifespan,
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
|
||||
print("Starting haiku.rag research assistant backend...")
|
||||
print(f"Agent model: {agent.model}")
|
||||
print(f"QA provider: {Config.qa.provider}")
|
||||
print(f"QA model: {Config.qa.model}")
|
||||
|
||||
uvicorn.run(
|
||||
"main:app",
|
||||
host="0.0.0.0",
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ dependencies = [
|
|||
"uvicorn[standard]>=0.34.2",
|
||||
"pydantic-ai-slim[ag-ui,openai]>=1.1.0",
|
||||
"python-dotenv>=1.0.1",
|
||||
"haiku-rag>=0.12.1",
|
||||
"haiku-rag-slim @ file:///Users/ggozad/dev/open-source/haiku.rag-agui/haiku_rag_slim",
|
||||
]
|
||||
|
||||
[dependency-groups]
|
||||
|
|
@ -18,6 +18,9 @@ dev = [
|
|||
"ruff>=0.13.0",
|
||||
]
|
||||
|
||||
[tool.hatch.metadata]
|
||||
allow-direct-references = true
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["."]
|
||||
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -10,11 +10,16 @@ services:
|
|||
# API keys (set these in your shell or .env file)
|
||||
- OPENAI_API_KEY=${OPENAI_API_KEY}
|
||||
- ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY}
|
||||
# Ollama connection (use value from .env)
|
||||
- OLLAMA_BASE_URL=${OLLAMA_BASE_URL}
|
||||
# Prevent Python bytecode caching for development
|
||||
- PYTHONDONTWRITEBYTECODE=1
|
||||
volumes:
|
||||
- ./backend:/app
|
||||
- /app/.venv
|
||||
- ${DB_PATH}:/app/data/haiku.rag.lancedb
|
||||
- ./haiku.rag.yaml:/app/haiku.rag.yaml:ro
|
||||
- ./backend/main.py:/app/main.py
|
||||
- ./backend/agent.py:/app/agent.py
|
||||
- ../../haiku_rag_slim/haiku:/app/.venv/lib/python3.13/site-packages/haiku
|
||||
networks:
|
||||
- ag-ui-network
|
||||
extra_hosts:
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ const runtime = new CopilotRuntime({
|
|||
agents: {
|
||||
// "research_agent" maps to the agent name used in useCoAgent() on the frontend
|
||||
research_agent: new HttpAgent({
|
||||
url: `${process.env.BACKEND_URL || "http://backend:8000"}/agent`,
|
||||
url: `${process.env.BACKEND_URL || "http://backend:8000"}/v1/research/stream`,
|
||||
}),
|
||||
},
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,208 +1,94 @@
|
|||
"use client";
|
||||
|
||||
import {
|
||||
CopilotKit,
|
||||
useCoAgent,
|
||||
useCoAgentStateRender,
|
||||
useCopilotAction,
|
||||
} from "@copilotkit/react-core";
|
||||
import { CopilotKit, useCoAgent } from "@copilotkit/react-core";
|
||||
import { CopilotChat } from "@copilotkit/react-ui";
|
||||
import "@copilotkit/react-ui/styles.css";
|
||||
import StateDisplay from "./StateDisplay";
|
||||
|
||||
interface SourceRef {
|
||||
chunk_id: string;
|
||||
document_uri: string;
|
||||
document_title: string;
|
||||
chunk_position: number;
|
||||
interface InsightRecord {
|
||||
id: string;
|
||||
summary: string;
|
||||
status: string;
|
||||
notes?: string;
|
||||
supporting_sources: string[];
|
||||
originating_questions: string[];
|
||||
}
|
||||
|
||||
interface GapRecord {
|
||||
id: string;
|
||||
description: string;
|
||||
severity: string;
|
||||
blocking: boolean;
|
||||
resolved: boolean;
|
||||
notes?: string;
|
||||
supporting_sources: string[];
|
||||
resolved_by: string[];
|
||||
}
|
||||
|
||||
interface SearchAnswer {
|
||||
query: string;
|
||||
answer: string;
|
||||
confidence: number;
|
||||
context: string[];
|
||||
sources: string[];
|
||||
}
|
||||
|
||||
interface ResearchContext {
|
||||
original_question: string;
|
||||
sub_questions: string[];
|
||||
qa_responses: SearchAnswer[];
|
||||
insights: InsightRecord[];
|
||||
gaps: GapRecord[];
|
||||
}
|
||||
|
||||
interface EvaluationResult {
|
||||
confidence: number;
|
||||
reasoning: string;
|
||||
should_continue: boolean;
|
||||
gaps_identified: string[];
|
||||
follow_up_questions: string[];
|
||||
}
|
||||
|
||||
interface ResearchReport {
|
||||
question: string;
|
||||
summary: string;
|
||||
findings: string[];
|
||||
conclusions: string[];
|
||||
insights_used: string[];
|
||||
methodology: string;
|
||||
}
|
||||
|
||||
interface ResearchState {
|
||||
question: string;
|
||||
phase: string;
|
||||
status: string;
|
||||
plan: Array<{
|
||||
id: number;
|
||||
question: string;
|
||||
status: string;
|
||||
search_results?: {
|
||||
type: string;
|
||||
results: Array<{
|
||||
chunk: string;
|
||||
chunk_id: string;
|
||||
document_uri: string;
|
||||
document_title: string;
|
||||
chunk_position: number;
|
||||
full_chunk_content: string;
|
||||
score: number;
|
||||
expanded: boolean;
|
||||
}>;
|
||||
};
|
||||
}>;
|
||||
current_question_index: number;
|
||||
insights: Array<{
|
||||
summary: string;
|
||||
confidence: number;
|
||||
source_refs: SourceRef[];
|
||||
}>;
|
||||
document_registry: Record<
|
||||
string,
|
||||
{
|
||||
title: string;
|
||||
chunks_referenced: string[];
|
||||
}
|
||||
>;
|
||||
current_document: {
|
||||
uri: string;
|
||||
title: string;
|
||||
content: string;
|
||||
total_chunks: number;
|
||||
metadata?: Record<string, unknown>;
|
||||
} | null;
|
||||
confidence: number;
|
||||
final_report: {
|
||||
title: string;
|
||||
summary: string;
|
||||
findings: string[];
|
||||
conclusions: string[];
|
||||
sources: string[];
|
||||
citations: Array<{
|
||||
document_uri: string;
|
||||
document_title: string;
|
||||
chunk_ids: string[];
|
||||
}>;
|
||||
context: ResearchContext;
|
||||
iterations: number;
|
||||
max_iterations: number;
|
||||
confidence_threshold: number;
|
||||
max_concurrency: number;
|
||||
last_eval: EvaluationResult | null;
|
||||
last_analysis: {
|
||||
insights_extracted: InsightRecord[];
|
||||
gaps_identified: GapRecord[];
|
||||
} | null;
|
||||
result?: ResearchReport;
|
||||
}
|
||||
|
||||
function AgentContent() {
|
||||
const { state } = useCoAgent<ResearchState>({
|
||||
name: "research_agent",
|
||||
initialState: {
|
||||
question: "",
|
||||
phase: "idle",
|
||||
status: "",
|
||||
plan: [],
|
||||
current_question_index: 0,
|
||||
insights: [],
|
||||
document_registry: {},
|
||||
current_document: null,
|
||||
confidence: 0.0,
|
||||
final_report: null,
|
||||
},
|
||||
});
|
||||
|
||||
useCopilotAction({
|
||||
name: "approve_research_plan",
|
||||
description:
|
||||
"Request user approval for the research plan. Returns 'APPROVED' if approved or 'REVISE' if user wants to revise.",
|
||||
parameters: [],
|
||||
renderAndWaitForResponse: ({ respond, status }) => (
|
||||
<div
|
||||
style={{
|
||||
padding: "1.5rem",
|
||||
background: "white",
|
||||
borderRadius: "8px",
|
||||
border: "2px solid #4299e1",
|
||||
marginBottom: "1rem",
|
||||
boxShadow: "0 2px 8px rgba(0,0,0,0.1)",
|
||||
}}
|
||||
>
|
||||
<h3
|
||||
style={{
|
||||
fontSize: "1.25rem",
|
||||
fontWeight: "bold",
|
||||
marginBottom: "1rem",
|
||||
color: "#2d3748",
|
||||
}}
|
||||
>
|
||||
Research Plan Approval
|
||||
</h3>
|
||||
<p
|
||||
style={{
|
||||
fontSize: "0.875rem",
|
||||
color: "#4a5568",
|
||||
marginBottom: "1rem",
|
||||
}}
|
||||
>
|
||||
Please review the research plan in the right pane.
|
||||
</p>
|
||||
|
||||
<div
|
||||
style={{
|
||||
display: "flex",
|
||||
gap: "1rem",
|
||||
}}
|
||||
className={status !== "executing" ? "hidden" : ""}
|
||||
>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => respond?.("REVISE")}
|
||||
disabled={status !== "executing"}
|
||||
style={{
|
||||
flex: 1,
|
||||
padding: "0.75rem",
|
||||
background: "white",
|
||||
border: "2px solid #e2e8f0",
|
||||
borderRadius: "6px",
|
||||
fontSize: "0.875rem",
|
||||
fontWeight: "600",
|
||||
cursor: status === "executing" ? "pointer" : "not-allowed",
|
||||
opacity: status === "executing" ? 1 : 0.5,
|
||||
}}
|
||||
>
|
||||
Revise Plan
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => respond?.("APPROVED")}
|
||||
disabled={status !== "executing"}
|
||||
style={{
|
||||
flex: 1,
|
||||
padding: "0.75rem",
|
||||
background: "#4299e1",
|
||||
color: "white",
|
||||
border: "none",
|
||||
borderRadius: "6px",
|
||||
fontSize: "0.875rem",
|
||||
fontWeight: "600",
|
||||
cursor: status === "executing" ? "pointer" : "not-allowed",
|
||||
opacity: status === "executing" ? 1 : 0.5,
|
||||
}}
|
||||
>
|
||||
Approve & Start Research
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
),
|
||||
});
|
||||
|
||||
useCoAgentStateRender<ResearchState>({
|
||||
name: "research_agent",
|
||||
render: ({ state: newState }) => {
|
||||
const phaseMessages: Record<string, string> = {
|
||||
planning: "Planning research...",
|
||||
searching: "Searching...",
|
||||
analyzing: "Extracting insights...",
|
||||
evaluating: `Evaluating confidence: ${(newState.confidence * 100).toFixed(0)}%`,
|
||||
synthesizing: "Generating final report...",
|
||||
done: "Research complete!",
|
||||
};
|
||||
const phaseMessage =
|
||||
phaseMessages[newState.phase] || newState.status || "Ready";
|
||||
|
||||
return (
|
||||
<div
|
||||
style={{
|
||||
padding: "1rem",
|
||||
background: "#e6f7ff",
|
||||
borderRadius: "4px",
|
||||
marginBottom: "0.5rem",
|
||||
border: "1px solid #91d5ff",
|
||||
}}
|
||||
>
|
||||
<strong>Research Update:</strong> {phaseMessage}
|
||||
</div>
|
||||
);
|
||||
context: {
|
||||
original_question: "",
|
||||
sub_questions: [],
|
||||
qa_responses: [],
|
||||
insights: [],
|
||||
gaps: [],
|
||||
},
|
||||
iterations: 0,
|
||||
max_iterations: 3,
|
||||
confidence_threshold: 0.8,
|
||||
max_concurrency: 1,
|
||||
last_eval: null,
|
||||
last_analysis: null,
|
||||
},
|
||||
});
|
||||
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
23368
examples/ag-ui-research/frontend/package-lock.json
generated
23368
examples/ag-ui-research/frontend/package-lock.json
generated
File diff suppressed because it is too large
Load diff
|
|
@ -1,20 +1,23 @@
|
|||
# haiku.rag configuration for ag-ui-research example
|
||||
# Copy to haiku.rag.yaml and customize
|
||||
|
||||
qa:
|
||||
research:
|
||||
provider: ollama
|
||||
model: gpt-oss:latest
|
||||
max_iterations: 3
|
||||
confidence_threshold: 0.8
|
||||
max_concurrency: 1
|
||||
|
||||
providers:
|
||||
ollama:
|
||||
base_url: http://host.docker.internal:11434
|
||||
|
||||
# For OpenAI:
|
||||
# qa:
|
||||
# research:
|
||||
# provider: openai
|
||||
# model: gpt-4o-mini
|
||||
|
||||
# For Anthropic:
|
||||
# qa:
|
||||
# research:
|
||||
# provider: anthropic
|
||||
# model: claude-3-5-haiku-20241022
|
||||
|
|
|
|||
|
|
@ -156,8 +156,11 @@ def create_search_node[AgentDepsT: GraphAgentDeps](
|
|||
deps: GraphDeps = ctx.deps # type: ignore[assignment]
|
||||
sub_q = ctx.inputs
|
||||
|
||||
# Create unique step name from question text
|
||||
step_name = f"search: {sub_q}"
|
||||
|
||||
if deps.agui_emitter and with_step_wrapper:
|
||||
deps.agui_emitter.start_step("search_one")
|
||||
deps.agui_emitter.start_step(step_name)
|
||||
|
||||
try:
|
||||
# Create semaphore if not already provided
|
||||
|
|
|
|||
Loading…
Reference in a new issue