haiku.rag 0.9.2__tar.gz → 0.9.3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haiku.rag
-Version: 0.9.2
+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.2 → haiku_rag-0.9.3}/docs/agents.md

@@ -70,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.2 → haiku_rag-0.9.3}/pyproject.toml

@@ -1,7 +1,8 @@
 [project]
+
 name = "haiku.rag"
-version = "0.9.2"
 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.2 → 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

{haiku_rag-0.9.2 → 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.2 → haiku_rag-0.9.3}/src/haiku/rag/research/base.py

@@ -45,7 +45,8 @@ class BaseResearchAgent[T](ABC):
             model=model_obj,
             deps_type=ResearchDependencies,
             output_type=agent_output_type,
-            system_prompt=self.get_system_prompt(),
+            instructions=self.get_system_prompt(),
+            retries=3,
         )
 
         # Register tools
@@ -75,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.2 → 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.2 → haiku_rag-0.9.3}/src/haiku/rag/research/presearch_agent.py

@@ -15,6 +15,12 @@ class PresearchSurveyAgent(BaseResearchAgent[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:
@@ -28,7 +34,6 @@ class PresearchSurveyAgent(BaseResearchAgent[str]):
             limit: int = 6,
         ) -> str:
             """Return verbatim concatenation of relevant chunk texts."""
-            query = query.replace('"', "")
             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.2 → 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,9 +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
-            # XXX: Investigate how to do that with 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.2 → 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
 
 
@@ -30,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

{haiku_rag-0.9.2 → haiku_rag-0.9.3}/src/haiku/rag/store/engine.py

@@ -35,6 +35,7 @@ def create_chunk_model(vector_dim: int):
         document_id: str
         content: str
         metadata: str = Field(default="{}")
+        order: int = Field(default=0)
         vector: Vector(vector_dim) = Field(default_factory=lambda: [0.0] * vector_dim)  # type: ignore
 
     return ChunkRecord
@@ -117,8 +118,10 @@ class Store:
             self.chunks_table = self.db.open_table("chunks")
         else:
             self.chunks_table = self.db.create_table("chunks", schema=self.ChunkRecord)
-            # Create FTS index on the new table
-            self.chunks_table.create_fts_index("content", replace=True)
+            # Create FTS index on the new table with phrase query support
+            self.chunks_table.create_fts_index(
+                "content", replace=True, with_position=True, remove_stop_words=False
+            )
 
         # Create or get settings table
         if "settings" in existing_tables:
@@ -133,21 +136,41 @@ class Store:
                 [SettingsRecord(id="settings", settings=json.dumps(settings_data))]
             )
 
-        # Set current version in settings
-        current_version = metadata.version("haiku.rag")
-        self.set_haiku_version(current_version)
-
-        # Check if we need to perform upgrades
+        # Run pending upgrades based on stored version and package version
         try:
-            existing_settings = list(
-                self.settings_table.search().limit(1).to_pydantic(SettingsRecord)
+            from haiku.rag.store.upgrades import run_pending_upgrades
+
+            current_version = metadata.version("haiku.rag")
+            db_version = self.get_haiku_version()
+
+            run_pending_upgrades(self, db_version, current_version)
+
+            # After upgrades complete (or if none), set stored version
+            # to the greater of the installed package version and the
+            # highest available upgrade step version in code.
+            try:
+                from packaging.version import parse as _v
+
+                from haiku.rag.store.upgrades import upgrades as _steps
+
+                highest_step = max((_v(u.version) for u in _steps), default=None)
+                effective_version = (
+                    str(max(_v(current_version), highest_step))
+                    if highest_step is not None
+                    else current_version
+                )
+            except Exception:
+                effective_version = current_version
+
+            self.set_haiku_version(effective_version)
+        except Exception as e:
+            # Avoid hard failure on initial connection; log and continue so CLI remains usable.
+            logger.warning(
+                "Skipping upgrade due to error (db=%s -> pkg=%s): %s",
+                self.get_haiku_version(),
+                metadata.version("haiku.rag") if hasattr(metadata, "version") else "",
+                e,
             )
-            if existing_settings:
-                db_version = self.get_haiku_version()  # noqa: F841
-                # TODO: Add upgrade logic here similar to SQLite version when needed
-        except Exception:
-            # Settings table might not exist yet in fresh databases
-            pass
 
     def get_haiku_version(self) -> str:
         """Returns the user version stored in settings."""
@@ -201,8 +224,10 @@ class Store:
         self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
         self.chunks_table = self.db.create_table("chunks", schema=self.ChunkRecord)
 
-        # Create FTS index on the new table
-        self.chunks_table.create_fts_index("content", replace=True)
+        # Create FTS index on the new table with phrase query support
+        self.chunks_table.create_fts_index(
+            "content", replace=True, with_position=True, remove_stop_words=False
+        )
 
     def close(self):
         """Close the database connection."""

{haiku_rag-0.9.2 → haiku_rag-0.9.3}/src/haiku/rag/store/models/chunk.py

@@ -10,6 +10,7 @@ class Chunk(BaseModel):
     document_id: str | None = None
     content: str
     metadata: dict = {}
+    order: int = 0
     document_uri: str | None = None
     document_meta: dict = {}
     embedding: list[float] | None = None