From fe8e48df3863fcbf32d52d51238617d686301d36 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 15 Sep 2025 14:24:18 +0300 Subject: [PATCH] Config research provider/model, CLI --- src/haiku/rag/app.py | 78 ++++++++++++++++++++++++++ src/haiku/rag/cli.py | 32 +++++++++++ src/haiku/rag/config.py | 4 ++ src/haiku/rag/research/orchestrator.py | 8 +-- tests/research/test_orchestrator.py | 14 +++++ 5 files changed, 132 insertions(+), 4 deletions(-) diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index 334f0cc2..c39f9a3e 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -9,6 +9,7 @@ from haiku.rag.client import HaikuRAG from haiku.rag.config import Config from haiku.rag.mcp import create_mcp_server 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.document import Document @@ -78,6 +79,83 @@ class HaikuRAGApp: except Exception as e: 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 with HaikuRAG(db_path=self.db_path, skip_validation=True) as client: try: diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py index f64e63a5..ec8e0cf3 100644 --- a/src/haiku/rag/cli.py +++ b/src/haiku/rag/cli.py @@ -235,6 +235,38 @@ def ask( 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") def settings(): app = HaikuRAGApp(db_path=Path()) # Don't need actual DB for settings diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index 5a58beec..1b044c65 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -29,6 +29,10 @@ class AppConfig(BaseModel): QA_PROVIDER: str = "ollama" QA_MODEL: str = "qwen3" + # Research defaults (fallback to QA if not provided via env) + RESEARCH_PROVIDER: str = "" + RESEARCH_MODEL: str = "" + CHUNK_SIZE: int = 256 CONTEXT_CHUNK_RADIUS: int = 0 diff --git a/src/haiku/rag/research/orchestrator.py b/src/haiku/rag/research/orchestrator.py index cdb748fa..127d421d 100644 --- a/src/haiku/rag/research/orchestrator.py +++ b/src/haiku/rag/research/orchestrator.py @@ -1,10 +1,9 @@ -"""Research orchestrator agent that coordinates specialized agents.""" - from typing import Any from pydantic import BaseModel, Field from pydantic_ai import RunContext +from haiku.rag.config import Config from haiku.rag.research.analysis_agent import AnalysisAgent, AnalysisResult from haiku.rag.research.base import BaseResearchAgent from haiku.rag.research.clarification_agent import ( @@ -32,10 +31,11 @@ class ResearchPlan(BaseModel): class ResearchOrchestrator(BaseResearchAgent): """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) - # Initialize specialized agents self.search_agent = SearchSpecialistAgent(provider, model) self.analysis_agent = AnalysisAgent(provider, model) self.clarification_agent = ClarificationAgent(provider, model) diff --git a/tests/research/test_orchestrator.py b/tests/research/test_orchestrator.py index 1eccffa3..6dce457e 100644 --- a/tests/research/test_orchestrator.py +++ b/tests/research/test_orchestrator.py @@ -51,6 +51,20 @@ def create_mock_chunk(chunk_id: str, content: str, score: float = 0.8): class TestResearchOrchestrator: """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): """Test that orchestrator initializes all agents correctly.""" orchestrator = ResearchOrchestrator(provider="openai", model="gpt-4")