Config research provider/model, CLI
This commit is contained in:
parent
86b44c47f9
commit
fe8e48df38
5 changed files with 132 additions and 4 deletions
|
|
@ -9,6 +9,7 @@ from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import Config
|
||||||
from haiku.rag.mcp import create_mcp_server
|
from haiku.rag.mcp import create_mcp_server
|
||||||
from haiku.rag.monitor import FileWatcher
|
from haiku.rag.monitor import FileWatcher
|
||||||
|
from haiku.rag.research.orchestrator import ResearchOrchestrator
|
||||||
from haiku.rag.store.models.chunk import Chunk
|
from haiku.rag.store.models.chunk import Chunk
|
||||||
from haiku.rag.store.models.document import Document
|
from haiku.rag.store.models.document import Document
|
||||||
|
|
||||||
|
|
@ -78,6 +79,83 @@ class HaikuRAGApp:
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self.console.print(f"[red]Error: {e}[/red]")
|
self.console.print(f"[red]Error: {e}[/red]")
|
||||||
|
|
||||||
|
async def research(
|
||||||
|
self, question: str, max_iterations: int = 3, verbose: bool = False
|
||||||
|
):
|
||||||
|
"""Run multi-agent research on a question."""
|
||||||
|
async with HaikuRAG(db_path=self.db_path) as client:
|
||||||
|
try:
|
||||||
|
# Create orchestrator with default config or fallback to QA
|
||||||
|
orchestrator = ResearchOrchestrator()
|
||||||
|
|
||||||
|
if verbose:
|
||||||
|
self.console.print(
|
||||||
|
f"[bold cyan]Starting research with {orchestrator.provider}:{orchestrator.model}[/bold cyan]"
|
||||||
|
)
|
||||||
|
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
|
||||||
|
self.console.print()
|
||||||
|
|
||||||
|
# Conduct research
|
||||||
|
report = await orchestrator.conduct_research(
|
||||||
|
question=question,
|
||||||
|
client=client,
|
||||||
|
max_iterations=max_iterations,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Display the report
|
||||||
|
self.console.print("[bold green]Research Report[/bold green]")
|
||||||
|
self.console.rule()
|
||||||
|
|
||||||
|
# Title and Executive Summary
|
||||||
|
self.console.print(f"[bold]{report.title}[/bold]")
|
||||||
|
self.console.print()
|
||||||
|
self.console.print("[bold cyan]Executive Summary:[/bold cyan]")
|
||||||
|
self.console.print(report.executive_summary)
|
||||||
|
self.console.print()
|
||||||
|
|
||||||
|
# Main Findings
|
||||||
|
if report.main_findings:
|
||||||
|
self.console.print("[bold cyan]Main Findings:[/bold cyan]")
|
||||||
|
for finding in report.main_findings:
|
||||||
|
self.console.print(f"• {finding}")
|
||||||
|
self.console.print()
|
||||||
|
|
||||||
|
# Themes
|
||||||
|
if report.themes:
|
||||||
|
self.console.print("[bold cyan]Key Themes:[/bold cyan]")
|
||||||
|
for theme, explanation in report.themes.items():
|
||||||
|
self.console.print(f"• [bold]{theme}[/bold]: {explanation}")
|
||||||
|
self.console.print()
|
||||||
|
|
||||||
|
# Conclusions
|
||||||
|
if report.conclusions:
|
||||||
|
self.console.print("[bold cyan]Conclusions:[/bold cyan]")
|
||||||
|
for conclusion in report.conclusions:
|
||||||
|
self.console.print(f"• {conclusion}")
|
||||||
|
self.console.print()
|
||||||
|
|
||||||
|
# Recommendations
|
||||||
|
if report.recommendations:
|
||||||
|
self.console.print("[bold cyan]Recommendations:[/bold cyan]")
|
||||||
|
for rec in report.recommendations:
|
||||||
|
self.console.print(f"• {rec}")
|
||||||
|
self.console.print()
|
||||||
|
|
||||||
|
# Limitations
|
||||||
|
if report.limitations:
|
||||||
|
self.console.print("[bold yellow]Limitations:[/bold yellow]")
|
||||||
|
for limitation in report.limitations:
|
||||||
|
self.console.print(f"• {limitation}")
|
||||||
|
self.console.print()
|
||||||
|
|
||||||
|
# Sources Summary
|
||||||
|
if report.sources_summary:
|
||||||
|
self.console.print("[bold cyan]Sources:[/bold cyan]")
|
||||||
|
self.console.print(report.sources_summary)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.console.print(f"[red]Error during research: {e}[/red]")
|
||||||
|
|
||||||
async def rebuild(self):
|
async def rebuild(self):
|
||||||
async with HaikuRAG(db_path=self.db_path, skip_validation=True) as client:
|
async with HaikuRAG(db_path=self.db_path, skip_validation=True) as client:
|
||||||
try:
|
try:
|
||||||
|
|
|
||||||
|
|
@ -235,6 +235,38 @@ def ask(
|
||||||
asyncio.run(app.ask(question=question, cite=cite))
|
asyncio.run(app.ask(question=question, cite=cite))
|
||||||
|
|
||||||
|
|
||||||
|
@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",
|
||||||
|
),
|
||||||
|
max_iterations: int = typer.Option(
|
||||||
|
3,
|
||||||
|
"--max-iterations",
|
||||||
|
"-n",
|
||||||
|
help="Maximum search/analyze iterations",
|
||||||
|
),
|
||||||
|
db: Path = typer.Option(
|
||||||
|
Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
|
||||||
|
"--db",
|
||||||
|
help="Path to the LanceDB database file",
|
||||||
|
),
|
||||||
|
verbose: bool = typer.Option(
|
||||||
|
False,
|
||||||
|
"--verbose",
|
||||||
|
help="Show verbose progress output",
|
||||||
|
),
|
||||||
|
):
|
||||||
|
app = HaikuRAGApp(db_path=db)
|
||||||
|
asyncio.run(
|
||||||
|
app.research(
|
||||||
|
question=question,
|
||||||
|
max_iterations=max_iterations,
|
||||||
|
verbose=verbose,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
@cli.command("settings", help="Display current configuration settings")
|
@cli.command("settings", help="Display current configuration settings")
|
||||||
def settings():
|
def settings():
|
||||||
app = HaikuRAGApp(db_path=Path()) # Don't need actual DB for settings
|
app = HaikuRAGApp(db_path=Path()) # Don't need actual DB for settings
|
||||||
|
|
|
||||||
|
|
@ -29,6 +29,10 @@ class AppConfig(BaseModel):
|
||||||
QA_PROVIDER: str = "ollama"
|
QA_PROVIDER: str = "ollama"
|
||||||
QA_MODEL: str = "qwen3"
|
QA_MODEL: str = "qwen3"
|
||||||
|
|
||||||
|
# Research defaults (fallback to QA if not provided via env)
|
||||||
|
RESEARCH_PROVIDER: str = ""
|
||||||
|
RESEARCH_MODEL: str = ""
|
||||||
|
|
||||||
CHUNK_SIZE: int = 256
|
CHUNK_SIZE: int = 256
|
||||||
CONTEXT_CHUNK_RADIUS: int = 0
|
CONTEXT_CHUNK_RADIUS: int = 0
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,10 +1,9 @@
|
||||||
"""Research orchestrator agent that coordinates specialized agents."""
|
|
||||||
|
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from pydantic import BaseModel, Field
|
from pydantic import BaseModel, Field
|
||||||
from pydantic_ai import RunContext
|
from pydantic_ai import RunContext
|
||||||
|
|
||||||
|
from haiku.rag.config import Config
|
||||||
from haiku.rag.research.analysis_agent import AnalysisAgent, AnalysisResult
|
from haiku.rag.research.analysis_agent import AnalysisAgent, AnalysisResult
|
||||||
from haiku.rag.research.base import BaseResearchAgent
|
from haiku.rag.research.base import BaseResearchAgent
|
||||||
from haiku.rag.research.clarification_agent import (
|
from haiku.rag.research.clarification_agent import (
|
||||||
|
|
@ -32,10 +31,11 @@ class ResearchPlan(BaseModel):
|
||||||
class ResearchOrchestrator(BaseResearchAgent):
|
class ResearchOrchestrator(BaseResearchAgent):
|
||||||
"""Orchestrator agent that coordinates the research workflow."""
|
"""Orchestrator agent that coordinates the research workflow."""
|
||||||
|
|
||||||
def __init__(self, provider: str, model: str):
|
def __init__(
|
||||||
|
self, provider: str = Config.RERANK_PROVIDER, model: str = Config.RERANK_MODEL
|
||||||
|
):
|
||||||
super().__init__(provider, model, output_type=ResearchPlan)
|
super().__init__(provider, model, output_type=ResearchPlan)
|
||||||
|
|
||||||
# Initialize specialized agents
|
|
||||||
self.search_agent = SearchSpecialistAgent(provider, model)
|
self.search_agent = SearchSpecialistAgent(provider, model)
|
||||||
self.analysis_agent = AnalysisAgent(provider, model)
|
self.analysis_agent = AnalysisAgent(provider, model)
|
||||||
self.clarification_agent = ClarificationAgent(provider, model)
|
self.clarification_agent = ClarificationAgent(provider, model)
|
||||||
|
|
|
||||||
|
|
@ -51,6 +51,20 @@ def create_mock_chunk(chunk_id: str, content: str, score: float = 0.8):
|
||||||
class TestResearchOrchestrator:
|
class TestResearchOrchestrator:
|
||||||
"""Test suite for ResearchOrchestrator."""
|
"""Test suite for ResearchOrchestrator."""
|
||||||
|
|
||||||
|
def test_orchestrator_uses_config_defaults(self):
|
||||||
|
"""Test that orchestrator uses config defaults when no args provided."""
|
||||||
|
orchestrator = ResearchOrchestrator()
|
||||||
|
|
||||||
|
# Should use RESEARCH_PROVIDER/MODEL if set, else QA_PROVIDER/MODEL
|
||||||
|
assert orchestrator.provider is not None
|
||||||
|
assert orchestrator.model is not None
|
||||||
|
|
||||||
|
# All agents should use the same provider/model
|
||||||
|
assert orchestrator.search_agent.provider == orchestrator.provider
|
||||||
|
assert orchestrator.search_agent.model == orchestrator.model
|
||||||
|
assert orchestrator.analysis_agent.provider == orchestrator.provider
|
||||||
|
assert orchestrator.analysis_agent.model == orchestrator.model
|
||||||
|
|
||||||
def test_orchestrator_initialization(self):
|
def test_orchestrator_initialization(self):
|
||||||
"""Test that orchestrator initializes all agents correctly."""
|
"""Test that orchestrator initializes all agents correctly."""
|
||||||
orchestrator = ResearchOrchestrator(provider="openai", model="gpt-4")
|
orchestrator = ResearchOrchestrator(provider="openai", model="gpt-4")
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue