haiku.rag 0.9.1__tar.gz → 0.9.3__tar.gz

{haiku_rag-0.9.1 → haiku_rag-0.9.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haiku.rag
-Version: 0.9.1
+Version: 0.9.3
 Summary: Agentic Retrieval Augmented Generation (RAG) with LanceDB
 Author-email: Yiorgis Gozadinos <ggozadinos@gmail.com>
 License: MIT

{haiku_rag-0.9.1 → haiku_rag-0.9.3}/docs/agents.md

@@ -43,6 +43,8 @@ The research workflow coordinates specialized agents to plan, search, analyze, a
 Components:
 
 - Orchestrator: Plans, coordinates, and loops until confidence is sufficient
+- Presearch Survey: Runs a quick KB scan and summarizes relevant chunk text to
+  ground the initial plan (plain-text summary; no URIs or scores)
 - Search Specialist: Performs targeted RAG searches and answers sub‑questions
 - Analysis & Evaluation: Extracts insights, identifies gaps, proposes new questions
 - Synthesis: Produces a final structured research report
@@ -68,14 +70,14 @@ from haiku.rag.client import HaikuRAG
 from haiku.rag.research import ResearchOrchestrator
 
 client = HaikuRAG(path_to_db)
-orchestrator = ResearchOrchestrator(provider="openai", model="gpt-4o-mini")
+orchestrator = ResearchOrchestrator(provider="ollama", model="gpt-oss")
 
 report = await orchestrator.conduct_research(
     question="What are the main drivers and recent trends of global temperature anomalies since 1990?",
     client=client,
     max_iterations=2,
     confidence_threshold=0.8,
-    verbose=False,
+    verbose=True,
 )
 
 print(report.title)

{haiku_rag-0.9.1 → haiku_rag-0.9.3}/pyproject.toml

@@ -1,7 +1,8 @@
 [project]
+
 name = "haiku.rag"
-version = "0.9.1"
 description = "Agentic Retrieval Augmented Generation (RAG) with LanceDB"
+version = "0.9.3"
 authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
 license = { text = "MIT" }
 readme = { file = "README.md", content-type = "text/markdown" }

{haiku_rag-0.9.1 → haiku_rag-0.9.3}/src/haiku/rag/app.py

@@ -101,7 +101,6 @@ class HaikuRAGApp:
                     client=client,
                     max_iterations=max_iterations,
                     verbose=verbose,
-                    console=self.console if verbose else None,
                 )
 
                 # Display the report
@@ -122,12 +121,7 @@ class HaikuRAGApp:
                         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()
+                # (Themes section removed)
 
                 # Conclusions
                 if report.conclusions:
@@ -261,7 +255,7 @@ class HaikuRAGApp:
                 elif transport == "sse":
                     await server.run_sse_async()
                 else:
-                    await server.run_http_async("streamable-http")
+                    await server.run_http_async(transport="streamable-http")
             except KeyboardInterrupt:
                 pass
             finally:

{haiku_rag-0.9.1 → haiku_rag-0.9.3}/src/haiku/rag/client.py

@@ -388,7 +388,7 @@ class HaikuRAG:
                 all_chunks = adjacent_chunks + [chunk]
 
                 # Get the range of orders for this expanded chunk
-                orders = [c.metadata.get("order", 0) for c in all_chunks]
+                orders = [c.order for c in all_chunks]
                 min_order = min(orders)
                 max_order = max(orders)
 
@@ -398,9 +398,7 @@ class HaikuRAG:
                         "score": score,
                         "min_order": min_order,
                         "max_order": max_order,
-                        "all_chunks": sorted(
-                            all_chunks, key=lambda c: c.metadata.get("order", 0)
-                        ),
+                        "all_chunks": sorted(all_chunks, key=lambda c: c.order),
                     }
                 )
 
@@ -459,7 +457,7 @@ class HaikuRAG:
                 # Merge all_chunks and deduplicate by order
                 all_chunks_dict = {}
                 for chunk in current["all_chunks"] + range_info["all_chunks"]:
-                    order = chunk.metadata.get("order", 0)
+                    order = chunk.order
                     all_chunks_dict[order] = chunk
                 current["all_chunks"] = [
                     all_chunks_dict[order] for order in sorted(all_chunks_dict.keys())

{haiku_rag-0.9.1 → haiku_rag-0.9.3}/src/haiku/rag/qa/agent.py

@@ -49,6 +49,9 @@ class QuestionAnswerAgent:
             limit: int = 3,
         ) -> list[SearchResult]:
             """Search the knowledge base for relevant documents."""
+
+            # Remove quotes from queries as this requires positional indexing in lancedb
+            query = query.replace('"', "")
             search_results = await ctx.deps.client.search(query, limit=limit)
             expanded_results = await ctx.deps.client.expand_context(search_results)
 

{haiku_rag-0.9.1 → haiku_rag-0.9.3}/src/haiku/rag/research/__init__.py

@@ -12,6 +12,7 @@ from haiku.rag.research.evaluation_agent import (
     EvaluationResult,
 )
 from haiku.rag.research.orchestrator import ResearchOrchestrator, ResearchPlan
+from haiku.rag.research.presearch_agent import PresearchSurveyAgent
 from haiku.rag.research.search_agent import SearchSpecialistAgent
 from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent
 
@@ -25,6 +26,7 @@ __all__ = [
     # Specialized agents
     "SearchAnswer",
     "SearchSpecialistAgent",
+    "PresearchSurveyAgent",
     "AnalysisEvaluationAgent",
     "EvaluationResult",
     "SynthesisAgent",

{haiku_rag-0.9.1 → haiku_rag-0.9.3}/src/haiku/rag/research/base.py

@@ -33,11 +33,20 @@ class BaseResearchAgent[T](ABC):
         # Import deps type lazily to avoid circular import during module load
         from haiku.rag.research.dependencies import ResearchDependencies
 
+        # If the agent is expected to return plain text, pass `str` directly.
+        # Otherwise, wrap the model with ToolOutput for robust tool-handling retries.
+        agent_output_type: Any
+        if self.output_type is str:  # plain text output
+            agent_output_type = str
+        else:
+            agent_output_type = ToolOutput(self.output_type, max_retries=3)
+
         self._agent = Agent(
             model=model_obj,
             deps_type=ResearchDependencies,
-            output_type=ToolOutput(self.output_type, max_retries=3),
-            system_prompt=self.get_system_prompt(),
+            output_type=agent_output_type,
+            instructions=self.get_system_prompt(),
+            retries=3,
         )
 
         # Register tools
@@ -67,7 +76,6 @@ class BaseResearchAgent[T](ABC):
         """Return the system prompt for this agent."""
         pass
 
-    @abstractmethod
     def register_tools(self) -> None:
         """Register agent-specific tools."""
         pass

{haiku_rag-0.9.1 → haiku_rag-0.9.3}/src/haiku/rag/research/dependencies.py

@@ -1,4 +1,6 @@
 from pydantic import BaseModel, Field
+from pydantic_ai import format_as_xml
+from rich.console import Console
 
 from haiku.rag.client import HaikuRAG
 from haiku.rag.research.base import SearchAnswer
@@ -43,3 +45,25 @@ class ResearchDependencies(BaseModel):
 
     client: HaikuRAG = Field(description="RAG client for document operations")
     context: ResearchContext = Field(description="Shared research context")
+    console: Console | None = None
+
+
+def _format_context_for_prompt(context: ResearchContext) -> str:
+    """Format the research context as XML for inclusion in prompts."""
+
+    context_data = {
+        "original_question": context.original_question,
+        "unanswered_questions": context.sub_questions,
+        "qa_responses": [
+            {
+                "question": qa.query,
+                "answer": qa.answer,
+                "context_snippets": qa.context,
+                "sources": qa.sources,
+            }
+            for qa in context.qa_responses
+        ],
+        "insights": context.insights,
+        "gaps": context.gaps,
+    }
+    return format_as_xml(context_data, root_tag="research_context")

haiku_rag-0.9.3/src/haiku/rag/research/evaluation_agent.py

@@ -0,0 +1,85 @@
+from pydantic import BaseModel, Field
+from pydantic_ai.run import AgentRunResult
+
+from haiku.rag.research.base import BaseResearchAgent
+from haiku.rag.research.dependencies import (
+    ResearchDependencies,
+    _format_context_for_prompt,
+)
+from haiku.rag.research.prompts import EVALUATION_AGENT_PROMPT
+
+
+class EvaluationResult(BaseModel):
+    """Result of analysis and evaluation."""
+
+    key_insights: list[str] = Field(
+        description="Main insights extracted from the research so far"
+    )
+    new_questions: list[str] = Field(
+        description="New sub-questions to add to the research (max 3)",
+        max_length=3,
+        default=[],
+    )
+    confidence_score: float = Field(
+        description="Confidence level in the completeness of research (0-1)",
+        ge=0.0,
+        le=1.0,
+    )
+    is_sufficient: bool = Field(
+        description="Whether the research is sufficient to answer the original question"
+    )
+    reasoning: str = Field(
+        description="Explanation of why the research is or isn't complete"
+    )
+
+
+class AnalysisEvaluationAgent(BaseResearchAgent[EvaluationResult]):
+    """Agent that analyzes findings and evaluates research completeness."""
+
+    def __init__(self, provider: str, model: str) -> None:
+        super().__init__(provider, model, output_type=EvaluationResult)
+
+    async def run(
+        self, prompt: str, deps: ResearchDependencies, **kwargs
+    ) -> AgentRunResult[EvaluationResult]:
+        console = deps.console
+        if console:
+            console.print(
+                "\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]"
+            )
+
+        # Format context for the evaluation agent
+        context_xml = _format_context_for_prompt(deps.context)
+        evaluation_prompt = f"""Analyze all gathered information and evaluate the completeness of research.
+
+{context_xml}
+
+Evaluate the research progress for the original question and identify any remaining gaps."""
+
+        result = await super().run(evaluation_prompt, deps, **kwargs)
+        output = result.output
+
+        # Store insights
+        for insight in output.key_insights:
+            deps.context.add_insight(insight)
+
+        # Add new questions to the sub-questions list
+        for new_q in output.new_questions:
+            if new_q not in deps.context.sub_questions:
+                deps.context.sub_questions.append(new_q)
+
+        if console:
+            if output.key_insights:
+                console.print("   [bold]Key insights:[/bold]")
+                for insight in output.key_insights:
+                    console.print(f"   • {insight}")
+            console.print(
+                f"   Confidence: [yellow]{output.confidence_score:.1%}[/yellow]"
+            )
+            status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
+            console.print(f"   Sufficient: {status}")
+
+        return result
+
+    def get_system_prompt(self) -> str:
+        return EVALUATION_AGENT_PROMPT

haiku_rag-0.9.3/src/haiku/rag/research/orchestrator.py

@@ -0,0 +1,170 @@
+from typing import Any
+
+from pydantic import BaseModel, Field
+from pydantic_ai.run import AgentRunResult
+from rich.console import Console
+
+from haiku.rag.config import Config
+from haiku.rag.research.base import BaseResearchAgent
+from haiku.rag.research.dependencies import (
+    ResearchContext,
+    ResearchDependencies,
+)
+from haiku.rag.research.evaluation_agent import (
+    AnalysisEvaluationAgent,
+    EvaluationResult,
+)
+from haiku.rag.research.presearch_agent import PresearchSurveyAgent
+from haiku.rag.research.prompts import ORCHESTRATOR_PROMPT
+from haiku.rag.research.search_agent import SearchSpecialistAgent
+from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent
+
+
+class ResearchPlan(BaseModel):
+    """Research execution plan."""
+
+    main_question: str = Field(description="The main research question")
+    sub_questions: list[str] = Field(
+        description="Decomposed sub-questions to investigate (max 3)", max_length=3
+    )
+
+
+class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
+    """Orchestrator agent that coordinates the research workflow."""
+
+    def __init__(
+        self,
+        provider: str | None = Config.RESEARCH_PROVIDER,
+        model: str | None = None,
+    ):
+        # Use provided values or fall back to config defaults
+        provider = provider or Config.RESEARCH_PROVIDER or Config.QA_PROVIDER
+        model = model or Config.RESEARCH_MODEL or Config.QA_MODEL
+
+        super().__init__(provider, model, output_type=ResearchPlan)
+
+        self.search_agent: SearchSpecialistAgent = SearchSpecialistAgent(
+            provider, model
+        )
+        self.presearch_agent: PresearchSurveyAgent = PresearchSurveyAgent(
+            provider, model
+        )
+        self.evaluation_agent: AnalysisEvaluationAgent = AnalysisEvaluationAgent(
+            provider, model
+        )
+        self.synthesis_agent: SynthesisAgent = SynthesisAgent(provider, model)
+
+    def get_system_prompt(self) -> str:
+        return ORCHESTRATOR_PROMPT
+
+    def _should_stop_research(
+        self,
+        evaluation_result: AgentRunResult[EvaluationResult],
+        confidence_threshold: float,
+    ) -> bool:
+        """Determine if research should stop based on evaluation."""
+
+        result = evaluation_result.output
+        return result.is_sufficient and result.confidence_score >= confidence_threshold
+
+    async def conduct_research(
+        self,
+        question: str,
+        client: Any,
+        max_iterations: int = 3,
+        confidence_threshold: float = 0.8,
+        verbose: bool = False,
+    ) -> ResearchReport:
+        """Conduct comprehensive research on a question.
+
+        Args:
+            question: The research question to investigate
+            client: HaikuRAG client for document operations
+            max_iterations: Maximum number of search-analyze-clarify cycles
+            confidence_threshold: Minimum confidence level to stop research (0-1)
+            verbose: If True, print progress and intermediate results
+
+        Returns:
+            ResearchReport with comprehensive findings
+        """
+
+        # Initialize context
+        context = ResearchContext(original_question=question)
+        deps = ResearchDependencies(client=client, context=context)
+        if verbose:
+            deps.console = Console()
+
+        console = deps.console
+        # Create initial research plan
+        if console:
+            console.print("\n[bold cyan]📋 Creating research plan...[/bold cyan]")
+
+        # Run a simple presearch survey to summarize KB context
+        presearch_result = await self.presearch_agent.run(question, deps=deps)
+        plan_prompt = (
+            "Create a research plan for the main question below.\n\n"
+            f"Main question: {question}\n\n"
+            "Use this brief presearch summary to inform the plan. Focus the 3 sub-questions "
+            "on the most important aspects not already obvious from the current KB context.\n\n"
+            f"{presearch_result.output}"
+        )
+
+        plan_result: AgentRunResult[ResearchPlan] = await self.run(
+            plan_prompt, deps=deps
+        )
+        context.sub_questions = plan_result.output.sub_questions
+
+        if console:
+            console.print("\n[bold green]✅ Research Plan Created:[/bold green]")
+            console.print(
+                f"   [bold]Main Question:[/bold] {plan_result.output.main_question}"
+            )
+            console.print("   [bold]Sub-questions:[/bold]")
+            for i, sq in enumerate(plan_result.output.sub_questions, 1):
+                console.print(f"      {i}. {sq}")
+
+        # Execute research iterations
+        for iteration in range(max_iterations):
+            if console:
+                console.rule(
+                    f"[bold yellow]🔄 Iteration {iteration + 1}/{max_iterations}[/bold yellow]"
+                )
+
+            # Check if we have questions to search
+            if not context.sub_questions:
+                if console:
+                    console.print(
+                        "[yellow]No more questions to explore. Concluding research.[/yellow]"
+                    )
+                break
+
+            # Use current sub-questions for this iteration
+            questions_to_search = context.sub_questions[:]
+
+            # Search phase - answer all questions in this iteration
+            if console:
+                console.print(
+                    f"\n[bold cyan]🔍 Searching & Answering {len(questions_to_search)} questions:[/bold cyan]"
+                )
+
+            for search_question in questions_to_search:
+                await self.search_agent.run(search_question, deps=deps)
+
+            # Analysis and Evaluation phase
+
+            evaluation_result = await self.evaluation_agent.run("", deps=deps)
+
+            # Check if research is sufficient
+            if self._should_stop_research(evaluation_result, confidence_threshold):
+                if console:
+                    console.print(
+                        f"\n[bold green]✅ Stopping research:[/bold green] {evaluation_result.output.reasoning}"
+                    )
+                break
+
+        # Generate final report
+        report_result: AgentRunResult[ResearchReport] = await self.synthesis_agent.run(
+            "", deps=deps
+        )
+
+        return report_result.output

haiku_rag-0.9.3/src/haiku/rag/research/presearch_agent.py

@@ -0,0 +1,39 @@
+from pydantic_ai import RunContext
+from pydantic_ai.run import AgentRunResult
+
+from haiku.rag.research.base import BaseResearchAgent
+from haiku.rag.research.dependencies import ResearchDependencies
+from haiku.rag.research.prompts import PRESEARCH_AGENT_PROMPT
+
+
+class PresearchSurveyAgent(BaseResearchAgent[str]):
+    """Presearch agent that gathers verbatim context and summarizes it."""
+
+    def __init__(self, provider: str, model: str) -> None:
+        super().__init__(provider, model, str)
+
+    async def run(
+        self, prompt: str, deps: ResearchDependencies, **kwargs
+    ) -> AgentRunResult[str]:
+        console = deps.console
+        if console:
+            console.print(
+                "\n[bold cyan]🔎 Presearch: summarizing KB context...[/bold cyan]"
+            )
+
+        return await super().run(prompt, deps, **kwargs)
+
+    def get_system_prompt(self) -> str:
+        return PRESEARCH_AGENT_PROMPT
+
+    def register_tools(self) -> None:
+        @self.agent.tool
+        async def gather_context(
+            ctx: RunContext[ResearchDependencies],
+            query: str,
+            limit: int = 6,
+        ) -> str:
+            """Return verbatim concatenation of relevant chunk texts."""
+            results = await ctx.deps.client.search(query, limit=limit)
+            expanded = await ctx.deps.client.expand_context(results)
+            return "\n\n".join(chunk.content for chunk, _ in expanded)

{haiku_rag-0.9.1 → haiku_rag-0.9.3}/src/haiku/rag/research/prompts.py

@@ -114,3 +114,16 @@ Focus on creating a report that provides clear value to the reader by:
 - Highlighting the most important findings
 - Explaining the implications of the research
 - Suggesting concrete next steps"""
+
+PRESEARCH_AGENT_PROMPT = """You are a rapid research surveyor.
+
+Task:
+- Call the gather_context tool once with the main question to obtain a
+  relevant texts from the Knowledge Base (KB).
+- Read that context and produce a brief natural-language summary describing
+  what the KB appears to contain relative to the question.
+
+Rules:
+- Base the summary strictly on the provided text; do not invent.
+- Output only the summary as plain text (one short paragraph).
+"""

{haiku_rag-0.9.1 → haiku_rag-0.9.3}/src/haiku/rag/research/search_agent.py

@@ -21,10 +21,17 @@ class SearchSpecialistAgent(BaseResearchAgent[SearchAnswer]):
         Pydantic AI enforces `SearchAnswer` as the output model; we just store
         the QA response with the last search results as sources.
         """
-        result = await super().run(prompt, deps, **kwargs)
+        console = deps.console
+        if console:
+            console.print(f"\t{prompt}")
 
-        if result.output:
-            deps.context.add_qa_response(result.output)
+        result = await super().run(prompt, deps, **kwargs)
+        deps.context.add_qa_response(result.output)
+        deps.context.sub_questions.remove(prompt)
+        if console:
+            answer = result.output.answer
+            answer_preview = answer[:150] + "…" if len(answer) > 150 else answer
+            console.log(f"\n   [green]✓[/green] {answer_preview}")
 
         return result
 
@@ -41,8 +48,6 @@ class SearchSpecialistAgent(BaseResearchAgent[SearchAnswer]):
             limit: int = 5,
         ) -> str:
             """Search the KB and return a concise context pack."""
-            # Remove quotes from queries as this requires positional indexing in lancedb
-            query = query.replace('"', "")
             search_results = await ctx.deps.client.search(query, limit=limit)
             expanded = await ctx.deps.client.expand_context(search_results)
 

{haiku_rag-0.9.1 → haiku_rag-0.9.3}/src/haiku/rag/research/synthesis_agent.py

@@ -1,6 +1,11 @@
 from pydantic import BaseModel, Field
+from pydantic_ai.run import AgentRunResult
 
 from haiku.rag.research.base import BaseResearchAgent
+from haiku.rag.research.dependencies import (
+    ResearchDependencies,
+    _format_context_for_prompt,
+)
 from haiku.rag.research.prompts import SYNTHESIS_AGENT_PROMPT
 
 
@@ -12,11 +17,12 @@ class ResearchReport(BaseModel):
     main_findings: list[str] = Field(
         description="Primary research findings with supporting evidence"
     )
-    themes: dict[str, str] = Field(description="Major themes and their explanations")
     conclusions: list[str] = Field(description="Evidence-based conclusions")
-    limitations: list[str] = Field(description="Limitations of the current research")
+    limitations: list[str] = Field(
+        description="Limitations of the current research", default=[]
+    )
     recommendations: list[str] = Field(
-        description="Actionable recommendations based on findings"
+        description="Actionable recommendations based on findings", default=[]
     )
     sources_summary: str = Field(
         description="Summary of sources used and their reliability"
@@ -29,11 +35,26 @@ class SynthesisAgent(BaseResearchAgent[ResearchReport]):
     def __init__(self, provider: str, model: str) -> None:
         super().__init__(provider, model, output_type=ResearchReport)
 
+    async def run(
+        self, prompt: str, deps: ResearchDependencies, **kwargs
+    ) -> AgentRunResult[ResearchReport]:
+        console = deps.console
+        if console:
+            console.print(
+                "\n[bold cyan]📝 Generating final research report...[/bold cyan]"
+            )
+
+        context_xml = _format_context_for_prompt(deps.context)
+        synthesis_prompt = f"""Generate a comprehensive research report based on all gathered information.
+
+{context_xml}
+
+Create a detailed report that synthesizes all findings into a coherent response."""
+        result = await super().run(synthesis_prompt, deps, **kwargs)
+        if console:
+            console.print("[bold green]✅ Research complete![/bold green]")
+
+        return result
+
     def get_system_prompt(self) -> str:
         return SYNTHESIS_AGENT_PROMPT
-
-    def register_tools(self) -> None:
-        """Register synthesis-specific tools."""
-        # The agent will use its LLM capabilities directly for synthesis
-        # The structured output will guide the report generation
-        pass