haiku.rag 0.9.2__py3-none-any.whl → 0.9.3__py3-none-any.whl

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/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/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/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/research/evaluation_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 EVALUATION_AGENT_PROMPT
 
 
@@ -34,9 +39,47 @@ class AnalysisEvaluationAgent(BaseResearchAgent[EvaluationResult]):
     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
-
-    def register_tools(self) -> None:
-        """No additional tools needed - uses LLM capabilities directly."""
-        pass

haiku/rag/research/orchestrator.py

@@ -1,13 +1,15 @@
 from typing import Any
 
 from pydantic import BaseModel, Field
-from pydantic_ai.format_prompt import format_as_xml
 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.dependencies import (
+    ResearchContext,
+    ResearchDependencies,
+)
 from haiku.rag.research.evaluation_agent import (
     AnalysisEvaluationAgent,
     EvaluationResult,
@@ -31,7 +33,9 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
     """Orchestrator agent that coordinates the research workflow."""
 
     def __init__(
-        self, provider: str | None = Config.RESEARCH_PROVIDER, model: str | None = None
+        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
@@ -53,30 +57,15 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
     def get_system_prompt(self) -> str:
         return ORCHESTRATOR_PROMPT
 
-    def register_tools(self) -> None:
-        """Register orchestration tools."""
-        # Tools are no longer needed - orchestrator directly calls agents
-        pass
-
-    def _format_context_for_prompt(self, context: ResearchContext) -> str:
-        """Format the research context as XML for inclusion in prompts."""
+    def _should_stop_research(
+        self,
+        evaluation_result: AgentRunResult[EvaluationResult],
+        confidence_threshold: float,
+    ) -> bool:
+        """Determine if research should stop based on evaluation."""
 
-        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")
+        result = evaluation_result.output
+        return result.is_sufficient and result.confidence_score >= confidence_threshold
 
     async def conduct_research(
         self,
@@ -85,7 +74,6 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
         max_iterations: int = 3,
         confidence_threshold: float = 0.8,
         verbose: bool = False,
-        console: Console | None = None,
     ) -> ResearchReport:
         """Conduct comprehensive research on a question.
 
@@ -95,7 +83,6 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
             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
-            console: Optional Rich console for output
 
         Returns:
             ResearchReport with comprehensive findings
@@ -104,44 +91,27 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
         # Initialize context
         context = ResearchContext(original_question=question)
         deps = ResearchDependencies(client=client, context=context)
+        if verbose:
+            deps.console = Console()
 
-        # Use provided console or create a new one
-        console = console or Console() if verbose else None
-
-        # Run a simple presearch survey to summarize KB context
-        if console:
-            console.print(
-                "\n[bold cyan]🔎 Presearch: summarizing KB context...[/bold cyan]"
-            )
-
-        presearch_result = await self.presearch_agent.run(question, deps=deps)
-
+        console = deps.console
         # Create initial research plan
         if console:
             console.print("\n[bold cyan]📋 Creating research plan...[/bold cyan]")
 
-        # Include the presearch summary to ground the planning step.
-
-        planning_context_xml = format_as_xml(
-            {
-                "original_question": question,
-                "presearch_summary": presearch_result.output or "",
-            },
-            root_tag="planning_context",
-        )
-
+        # 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"{planning_context_xml}"
+            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:
@@ -152,7 +122,6 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
             console.print("   [bold]Sub-questions:[/bold]")
             for i, sq in enumerate(plan_result.output.sub_questions, 1):
                 console.print(f"      {i}. {sq}")
-            console.print()
 
         # Execute research iterations
         for iteration in range(max_iterations):
@@ -163,7 +132,6 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
 
             # Check if we have questions to search
             if not context.sub_questions:
-                # No more questions to explore
                 if console:
                     console.print(
                         "[yellow]No more questions to explore. Concluding research.[/yellow]"
@@ -171,90 +139,20 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
                 break
 
             # Use current sub-questions for this iteration
-            questions_to_search = context.sub_questions
+            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 i, q in enumerate(questions_to_search, 1):
-                    console.print(f"   {i}. {q}")
 
-            # Run searches for all questions and remove answered ones
-            answered_questions = []
             for search_question in questions_to_search:
-                try:
-                    await self.search_agent.run(search_question, deps=deps)
-                except Exception as e:  # pragma: no cover - defensive
-                    if console:
-                        console.print(
-                            f"\n   [red]×[/red] Omitting failed question: {search_question} ({e})"
-                        )
-                finally:
-                    answered_questions.append(search_question)
-
-                if console and context.qa_responses:
-                    # Show the last QA response (which should be for this question)
-                    latest_qa = context.qa_responses[-1]
-                    answer_preview = (
-                        latest_qa.answer[:150] + "..."
-                        if len(latest_qa.answer) > 150
-                        else latest_qa.answer
-                    )
-                    console.print(
-                        f"\n   [green]✓[/green] {search_question[:50]}..."
-                        if len(search_question) > 50
-                        else f"\n   [green]✓[/green] {search_question}"
-                    )
-                    console.print(f"      {answer_preview}")
-
-            # Remove answered questions from the list
-            for question in answered_questions:
-                if question in context.sub_questions:
-                    context.sub_questions.remove(question)
+                await self.search_agent.run(search_question, deps=deps)
 
             # Analysis and Evaluation phase
-            if console:
-                console.print(
-                    "\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]"
-                )
-
-            # Format context for the evaluation agent
-            context_xml = self._format_context_for_prompt(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."""
-
-            evaluation_result = await self.evaluation_agent.run(
-                evaluation_prompt,
-                deps=deps,
-            )
-
-            if console and evaluation_result.output:
-                output = evaluation_result.output
-                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}")
 
-            # Store insights
-            for insight in evaluation_result.output.key_insights:
-                context.add_insight(insight)
-
-            # Add new questions to the sub-questions list
-            for new_q in evaluation_result.output.new_questions:
-                if new_q not in context.sub_questions:
-                    context.sub_questions.append(new_q)
+            evaluation_result = await self.evaluation_agent.run("", deps=deps)
 
             # Check if research is sufficient
             if self._should_stop_research(evaluation_result, confidence_threshold):
@@ -265,36 +163,8 @@ Evaluate the research progress for the original question and identify any remain
                 break
 
         # Generate final report
-        if console:
-            console.print(
-                "\n[bold cyan]📝 Generating final research report...[/bold cyan]"
-            )
-
-        # Format context for the synthesis agent
-        final_context_xml = self._format_context_for_prompt(context)
-        synthesis_prompt = f"""Generate a comprehensive research report based on all gathered information.
-
-{final_context_xml}
-
-Create a detailed report that synthesizes all findings into a coherent response."""
-
         report_result: AgentRunResult[ResearchReport] = await self.synthesis_agent.run(
-            synthesis_prompt, deps=deps
+            "", deps=deps
         )
 
-        if console:
-            console.print("[bold green]✅ Research complete![/bold green]")
-
         return report_result.output
-
-    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
-
-        # Stop if the agent indicates sufficient information AND confidence exceeds threshold
-        return result.is_sufficient and result.confidence_score >= confidence_threshold

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/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/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/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/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

haiku/rag/store/repositories/chunk.py

@@ -28,7 +28,9 @@ class ChunkRepository:
     def _ensure_fts_index(self) -> None:
         """Ensure FTS index exists on the content column."""
         try:
-            self.store.chunks_table.create_fts_index("content", replace=True)
+            self.store.chunks_table.create_fts_index(
+                "content", replace=True, with_position=True, remove_stop_words=False
+            )
         except Exception as e:
             # Log the error but don't fail - FTS might already exist
             logger.debug(f"FTS index creation skipped: {e}")
@@ -59,11 +61,16 @@ class ChunkRepository:
             embedding = entity.embedding
         else:
             embedding = await self.embedder.embed(entity.content)
+        order_val = int(entity.order)
+
         chunk_record = self.store.ChunkRecord(
             id=chunk_id,
             document_id=entity.document_id,
             content=entity.content,
-            metadata=json.dumps(entity.metadata),
+            metadata=json.dumps(
+                {k: v for k, v in entity.metadata.items() if k != "order"}
+            ),
+            order=order_val,
             vector=embedding,
         )
 
@@ -90,11 +97,13 @@ class ChunkRepository:
             return None
 
         chunk_record = results[0]
+        md = json.loads(chunk_record.metadata)
         return Chunk(
             id=chunk_record.id,
             document_id=chunk_record.document_id,
             content=chunk_record.content,
-            metadata=json.loads(chunk_record.metadata) if chunk_record.metadata else {},
+            metadata=md,
+            order=chunk_record.order,
         )
 
     async def update(self, entity: Chunk) -> Chunk:
@@ -102,13 +111,17 @@ class ChunkRepository:
         assert entity.id, "Chunk ID is required for update"
 
         embedding = await self.embedder.embed(entity.content)
+        order_val = int(entity.order)
 
         self.store.chunks_table.update(
             where=f"id = '{entity.id}'",
             values={
                 "document_id": entity.document_id,
                 "content": entity.content,
-                "metadata": json.dumps(entity.metadata),
+                "metadata": json.dumps(
+                    {k: v for k, v in entity.metadata.items() if k != "order"}
+                ),
+                "order": order_val,
                 "vector": embedding,
             },
         )
@@ -140,15 +153,19 @@ class ChunkRepository:
 
         results = list(query.to_pydantic(self.store.ChunkRecord))
 
-        return [
-            Chunk(
-                id=chunk.id,
-                document_id=chunk.document_id,
-                content=chunk.content,
-                metadata=json.loads(chunk.metadata) if chunk.metadata else {},
+        chunks: list[Chunk] = []
+        for rec in results:
+            md = json.loads(rec.metadata)
+            chunks.append(
+                Chunk(
+                    id=rec.id,
+                    document_id=rec.document_id,
+                    content=rec.content,
+                    metadata=md,
+                    order=rec.order,
+                )
             )
-            for chunk in results
-        ]
+        return chunks
 
     async def create_chunks_for_document(
         self, document_id: str, document: DoclingDocument
@@ -191,7 +208,8 @@ class ChunkRepository:
                 id=chunk_id,
                 document_id=document_id,
                 content=chunk_text,
-                metadata=json.dumps({"order": order}),
+                metadata=json.dumps({}),
+                order=order,
                 vector=embedding,
             )
             chunk_records.append(chunk_record)
@@ -200,7 +218,8 @@ class ChunkRepository:
                 id=chunk_id,
                 document_id=document_id,
                 content=chunk_text,
-                metadata={"order": order},
+                metadata={},
+                order=order,
             )
             created_chunks.append(chunk)
 
@@ -219,8 +238,10 @@ class ChunkRepository:
         self.store.chunks_table = self.store.db.create_table(
             "chunks", schema=self.store.ChunkRecord
         )
-        # Create FTS index on the new table
-        self.store.chunks_table.create_fts_index("content", replace=True)
+        # Create FTS index on the new table with phrase query support
+        self.store.chunks_table.create_fts_index(
+            "content", replace=True, with_position=True, remove_stop_words=False
+        )
 
     async def delete_by_document_id(self, document_id: str) -> bool:
         """Delete all chunks for a document."""
@@ -298,37 +319,36 @@ class ChunkRepository:
         doc_uri = doc_results[0].uri if doc_results else None
         doc_meta = doc_results[0].metadata if doc_results else "{}"
 
-        # Sort by order in metadata
-        chunks = [
-            Chunk(
-                id=chunk.id,
-                document_id=chunk.document_id,
-                content=chunk.content,
-                metadata=json.loads(chunk.metadata) if chunk.metadata else {},
-                document_uri=doc_uri,
-                document_meta=json.loads(doc_meta) if doc_meta else {},
+        chunks: list[Chunk] = []
+        for rec in results:
+            md = json.loads(rec.metadata)
+            chunks.append(
+                Chunk(
+                    id=rec.id,
+                    document_id=rec.document_id,
+                    content=rec.content,
+                    metadata=md,
+                    order=rec.order,
+                    document_uri=doc_uri,
+                    document_meta=json.loads(doc_meta),
+                )
             )
-            for chunk in results
-        ]
 
-        chunks.sort(key=lambda c: c.metadata.get("order", 0))
+        chunks.sort(key=lambda c: c.order)
         return chunks
 
     async def get_adjacent_chunks(self, chunk: Chunk, num_adjacent: int) -> list[Chunk]:
         """Get adjacent chunks before and after the given chunk within the same document."""
         assert chunk.document_id, "Document id is required for adjacent chunk finding"
 
-        chunk_order = chunk.metadata.get("order")
-        if chunk_order is None:
-            return []
+        chunk_order = chunk.order
 
-        # Get all chunks for the document
+        # Fetch chunks for the same document and filter by order proximity
         all_chunks = await self.get_by_document_id(chunk.document_id)
 
-        # Filter to adjacent chunks
-        adjacent_chunks = []
+        adjacent_chunks: list[Chunk] = []
         for c in all_chunks:
-            c_order = c.metadata.get("order", 0)
+            c_order = c.order
             if c.id != chunk.id and abs(c_order - chunk_order) <= num_adjacent:
                 adjacent_chunks.append(c)
 
@@ -380,15 +400,16 @@ class ChunkRepository:
             doc_uri = doc.uri if doc else None
             doc_meta = doc.metadata if doc else "{}"
 
+            md = json.loads(chunk_record.metadata)
+
             chunk = Chunk(
                 id=chunk_record.id,
                 document_id=chunk_record.document_id,
                 content=chunk_record.content,
-                metadata=json.loads(chunk_record.metadata)
-                if chunk_record.metadata
-                else {},
+                metadata=md,
+                order=chunk_record.order,
                 document_uri=doc_uri,
-                document_meta=json.loads(doc_meta) if doc_meta else {},
+                document_meta=json.loads(doc_meta),
             )
 
             # Get score from arrow result

haiku/rag/store/repositories/document.py

@@ -34,7 +34,7 @@ class DocumentRepository:
             id=record.id,
             content=record.content,
             uri=record.uri,
-            metadata=json.loads(record.metadata) if record.metadata else {},
+            metadata=json.loads(record.metadata),
             created_at=datetime.fromisoformat(record.created_at)
             if record.created_at
             else datetime.now(),
@@ -194,7 +194,7 @@ class DocumentRepository:
                 )
                 for order, chunk in enumerate(chunks):
                     chunk.document_id = created_doc.id
-                    chunk.metadata["order"] = order
+                    chunk.order = order
                     await self.chunk_repository.create(chunk)
 
             return created_doc

haiku/rag/store/repositories/settings.py

@@ -84,11 +84,18 @@ class SettingsRepository:
         )
 
         if existing:
-            # Only update when configuration actually changed to avoid needless new versions
-            existing_payload = (
-                json.loads(existing[0].settings) if existing[0].settings else {}
-            )
-            if existing_payload != current_config:
+            # Preserve existing version if present to avoid interfering with upgrade flow
+            try:
+                existing_settings = (
+                    json.loads(existing[0].settings) if existing[0].settings else {}
+                )
+            except Exception:
+                existing_settings = {}
+            if "version" in existing_settings:
+                current_config["version"] = existing_settings["version"]
+
+            # Update existing settings
+            if existing_settings != current_config:
                 self.store.settings_table.update(
                     where="id = 'settings'",
                     values={"settings": json.dumps(current_config)},

haiku/rag/store/upgrades/__init__.py

@@ -1 +1,60 @@
-upgrades = []
+import logging
+from collections.abc import Callable
+from dataclasses import dataclass
+
+from packaging.version import Version, parse
+
+from haiku.rag.store.engine import Store
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Upgrade:
+    """Represents a database upgrade step."""
+
+    version: str
+    apply: Callable[[Store], None]
+    description: str = ""
+
+
+# Registry of upgrade steps (ordered by version)
+upgrades: list[Upgrade] = []
+
+
+def run_pending_upgrades(store: Store, from_version: str, to_version: str) -> None:
+    """Run upgrades where from_version < step.version <= to_version."""
+    v_from: Version = parse(from_version)
+    v_to: Version = parse(to_version)
+
+    # Ensure that tests/development run available code upgrades even if the
+    # installed package version hasn't been bumped to include them yet.
+    if upgrades:
+        highest_step_version: Version = max(parse(u.version) for u in upgrades)
+        if highest_step_version > v_to:
+            v_to = highest_step_version
+
+    # Determine applicable steps
+    sorted_steps = sorted(upgrades, key=lambda u: parse(u.version))
+    applicable = [s for s in sorted_steps if v_from < parse(s.version) <= v_to]
+    if applicable:
+        logger.info("%d upgrade step(s) pending", len(applicable))
+
+    # Apply in ascending order
+    for idx, step in enumerate(applicable, start=1):
+        logger.info(
+            "Applying upgrade %s: %s (%d/%d)",
+            step.version,
+            step.description or "",
+            idx,
+            len(applicable),
+        )
+        step.apply(store)
+        logger.info("Completed upgrade %s", step.version)
+
+
+from .v0_9_3 import upgrade_fts_phrase as upgrade_0_9_3_fts  # noqa: E402
+from .v0_9_3 import upgrade_order as upgrade_0_9_3_order  # noqa: E402
+
+upgrades.append(upgrade_0_9_3_order)
+upgrades.append(upgrade_0_9_3_fts)

haiku/rag/store/upgrades/v0_9_3.py

@@ -0,0 +1,112 @@
+import json
+
+from lancedb.pydantic import LanceModel, Vector
+from pydantic import Field
+
+from haiku.rag.store.engine import Store
+from haiku.rag.store.upgrades import Upgrade
+
+
+def _infer_vector_dim(store: Store) -> int:
+    """Infer vector dimension from existing data; fallback to embedder config."""
+    try:
+        arrow = store.chunks_table.search().limit(1).to_arrow()
+        rows = arrow.to_pylist()
+        if rows:
+            vec = rows[0].get("vector")
+            if isinstance(vec, list) and vec:
+                return len(vec)
+    except Exception:
+        pass
+    # Fallback to configured embedder vector dim
+    return getattr(store.embedder, "_vector_dim", 1024)
+
+
+def _apply_chunk_order(store: Store) -> None:
+    """Add integer 'order' column to chunks and backfill from metadata."""
+
+    vector_dim = _infer_vector_dim(store)
+
+    class ChunkRecordV2(LanceModel):
+        id: str
+        document_id: str
+        content: str
+        metadata: str = Field(default="{}")
+        order: int = Field(default=0)
+        vector: Vector(vector_dim) = Field(  # type: ignore
+            default_factory=lambda: [0.0] * vector_dim
+        )
+
+    # Read existing chunks
+    try:
+        chunks_arrow = store.chunks_table.search().to_arrow()
+        rows = chunks_arrow.to_pylist()
+    except Exception:
+        rows = []
+
+    new_chunk_records: list[ChunkRecordV2] = []
+    for row in rows:
+        md_raw = row.get("metadata") or "{}"
+        try:
+            md = json.loads(md_raw) if isinstance(md_raw, str) else md_raw
+        except Exception:
+            md = {}
+        # Extract and normalize order
+        order_val = 0
+        try:
+            if isinstance(md, dict) and "order" in md:
+                order_val = int(md["order"])  # type: ignore[arg-type]
+        except Exception:
+            order_val = 0
+
+        if isinstance(md, dict) and "order" in md:
+            md = {k: v for k, v in md.items() if k != "order"}
+
+        vec = row.get("vector") or [0.0] * vector_dim
+
+        new_chunk_records.append(
+            ChunkRecordV2(
+                id=row.get("id"),
+                document_id=row.get("document_id"),
+                content=row.get("content", ""),
+                metadata=json.dumps(md),
+                order=order_val,
+                vector=vec,
+            )
+        )
+
+    # Recreate chunks table with new schema
+    try:
+        store.db.drop_table("chunks")
+    except Exception:
+        pass
+
+    store.chunks_table = store.db.create_table("chunks", schema=ChunkRecordV2)
+    store.chunks_table.create_fts_index("content", replace=True)
+
+    if new_chunk_records:
+        store.chunks_table.add(new_chunk_records)
+
+
+upgrade_order = Upgrade(
+    version="0.9.3",
+    apply=_apply_chunk_order,
+    description="Add 'order' column to chunks and backfill from metadata",
+)
+
+
+def _apply_fts_phrase_support(store: Store) -> None:
+    """Recreate FTS index with phrase query support and no stop-word removal."""
+    try:
+        store.chunks_table.create_fts_index(
+            "content", replace=True, with_position=True, remove_stop_words=False
+        )
+    except Exception:
+        pass
+
+
+upgrade_fts_phrase = Upgrade(
+    version="0.9.3",
+    apply=_apply_fts_phrase_support,
+    description="Enable FTS phrase queries (with positions) and keep stop-words",
+)

{haiku_rag-0.9.2.dist-info → haiku_rag-0.9.3.dist-info}/METADATA

@@ -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.dist-info → haiku_rag-0.9.3.dist-info}/RECORD

@@ -1,8 +1,8 @@
 haiku/rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-haiku/rag/app.py,sha256=o64L7aj5V8lYHxWhAKBNj1tGfXiN6xr0_Cc1dEYd3As,11483
+haiku/rag/app.py,sha256=nkud-OHic3HIgEEiNOKVvhmW98DPpDe6HokBSz-xV7w,11420
 haiku/rag/chunker.py,sha256=PVe6ysv8UlacUd4Zb3_8RFWIaWDXnzBAy2VDJ4TaUsE,1555
 haiku/rag/cli.py,sha256=3nlzrT5FPCyfnu51KHchLG4Cj2eVv9YsuGHMShBnVb0,9845
-haiku/rag/client.py,sha256=NJVGXzVzpoVy1sttz_xEU7mXWtObKT8pGpvo5pZyzwc,21288
+haiku/rag/client.py,sha256=QgJQu7g7JjAzWN6R10NeDqpFf89Dml_LiWce4QRHLHc,21177
 haiku/rag/config.py,sha256=SPEIv2IElZmZh4Wsp8gk7ViRW5ZzD-UGmIqRAXscDdI,2134
 haiku/rag/logging.py,sha256=dm65AwADpcQsH5OAPtRA-4hsw0w5DK-sGOvzYkj6jzw,1720
 haiku/rag/mcp.py,sha256=bR9Y-Nz-hvjiql20Y0KE0hwNGwyjmPGX8K9d-qmXptY,4683
@@ -25,26 +25,27 @@ haiku/rag/reranking/cohere.py,sha256=1iTdiaa8vvb6oHVB2qpWzUOVkyfUcimVSZp6Qr4aq4c
 haiku/rag/reranking/mxbai.py,sha256=46sVTsTIkzIX9THgM3u8HaEmgY7evvEyB-N54JTHvK8,867
 haiku/rag/reranking/vllm.py,sha256=xVGH9ss-ISWdJ5SKUUHUbTqBo7PIEmA_SQv0ScdJ6XA,1479
 haiku/rag/research/__init__.py,sha256=qLF41YayAxW_VeHhuTceVuz9hw1FNbuRV9VMhonUMW0,1078
-haiku/rag/research/base.py,sha256=ZUvqh-IxU8r4mOPRKjwZcHciKcIfrTnP6Q_9jVElelQ,4041
-haiku/rag/research/dependencies.py,sha256=vZctKC5donqhm8LFO6hQdXZZXzjdW1__4eUlhyZn058,1573
-haiku/rag/research/evaluation_agent.py,sha256=yyBobKr8MRwiox59I2Jqycp02ju9EGVaI9FceRGL188,1386
-haiku/rag/research/orchestrator.py,sha256=LrxRG74BWun0T4uOxhc9AuitxbPioS_eG_nE098ftyY,11765
-haiku/rag/research/presearch_agent.py,sha256=vf-WlY46g5tuuLKMBuPXXYYffynsBw7KVLr8LoTNHnU,1292
+haiku/rag/research/base.py,sha256=X5n6myUG_Oz4i8WGfyKZ39YzK13rOkyvwGKwSBfL50k,4043
+haiku/rag/research/dependencies.py,sha256=N7mnFwa_uyWYH0NtbEHp5JJvNGN64Q8HHfY41E8Irx0,2362
+haiku/rag/research/evaluation_agent.py,sha256=VMegemd9Vln3jfZbeHzMfb7rUPFNzNxi5Y_l1zrddl8,2994
+haiku/rag/research/orchestrator.py,sha256=nvSRdIs77kSb1CZaQUYZM_Zl5xLP8K6noVgnixpeLJI,6329
+haiku/rag/research/presearch_agent.py,sha256=MpakZ9HSynv73EnWakwUuytfKpiN_8lEqZlVc3zZjGU,1427
 haiku/rag/research/prompts.py,sha256=pVRB7_b_p3JaLF1bC3ANTbSFY78ypSjDhoq6peoU6jo,5685
-haiku/rag/research/search_agent.py,sha256=0iK7vCd9w7h8pWJgB6VUSPOdjlzB8peboNSXxuEGBK0,2464
-haiku/rag/research/synthesis_agent.py,sha256=jo5rg7aL4zGXLQP105cANqRPIiwJLqYe2unO5BQkNvE,1511
+haiku/rag/research/search_agent.py,sha256=xn2MlEyL9te_dtZqTzW81lGw7fYmyUzn26mvzX52hNA,2599
+haiku/rag/research/synthesis_agent.py,sha256=FQCt8wbaaKOwgIOQazTNAmohBMZRUDoVzHkByYhbGg8,2182
 haiku/rag/store/__init__.py,sha256=hq0W0DAC7ysqhWSP2M2uHX8cbG6kbr-sWHxhq6qQcY0,103
-haiku/rag/store/engine.py,sha256=fNrykqMX7PRSCt4LSRfuJ66OLrb8BVYq2bpbfI2iaWU,8455
+haiku/rag/store/engine.py,sha256=-3MZJYft2XTWaLuyKha8DKhWQeU5E5CBeskXXF5fXso,9555
 haiku/rag/store/models/__init__.py,sha256=s0E72zneGlowvZrFWaNxHYjOAUjgWdLxzdYsnvNRVlY,88
-haiku/rag/store/models/chunk.py,sha256=ZNyTfO6lh3rXWLVYO3TZcitbL4LSUGr42fR6jQQ5iQc,364
+haiku/rag/store/models/chunk.py,sha256=Ww_hj3DMwJLNM33l1GvIP84yzDFc6cxfiWcotUfWSYg,383
 haiku/rag/store/models/document.py,sha256=zSSpt6pyrMJAIXGQvIcqojcqUzwZnhp3WxVokaWxNRc,396
 haiku/rag/store/repositories/__init__.py,sha256=Olv5dLfBQINRV3HrsfUpjzkZ7Qm7goEYyMNykgo_DaY,291
-haiku/rag/store/repositories/chunk.py,sha256=1RmPyEYRYOFbrALbmLOo62t3f-xO2KgxUjcvPdrRZlc,14467
-haiku/rag/store/repositories/document.py,sha256=XoLCrMrZqs0iCZoHlDOfRDaVUux77Vdu5iZczduF1rY,7812
-haiku/rag/store/repositories/settings.py,sha256=wx3fuP_5CpPflZHRrIkeoer6ml-iD0qXERh5k6MQRzI,5291
-haiku/rag/store/upgrades/__init__.py,sha256=wUiEoSiHTahvuagx93E4FB07v123AhdbOjwUkPusiIg,14
-haiku_rag-0.9.2.dist-info/METADATA,sha256=IM9tGvye83CRTj2wOFtPP7oD9KtJvp3RXh4QdCFknD4,4681
-haiku_rag-0.9.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-haiku_rag-0.9.2.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
-haiku_rag-0.9.2.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
-haiku_rag-0.9.2.dist-info/RECORD,,
+haiku/rag/store/repositories/chunk.py,sha256=O2SEhQy3ZptWjwwpxS-L8KNq2tEqEBqheHfLw-M_FqA,15012
+haiku/rag/store/repositories/document.py,sha256=m11SamQoGYs5ODfmarJGU1yIcqtgmnba-5bGOPQuYrI,7773
+haiku/rag/store/repositories/settings.py,sha256=7XMBMavU8zRgdBoQzQg0Obfa7UKjuVnBugidTC6sEW0,5548
+haiku/rag/store/upgrades/__init__.py,sha256=gDOxiq3wdZPr3JoenjNYxx0cpgZJhbaFKNX2fzXRq1Q,1852
+haiku/rag/store/upgrades/v0_9_3.py,sha256=NrjNilQSgDtFWRbL3ZUtzQzJ8tf9u0dDRJtnDFwwbdw,3322
+haiku_rag-0.9.3.dist-info/METADATA,sha256=iCae4YtZ0meIQTZLUNree_-74F3irXvArPvdSxVz8ZM,4681
+haiku_rag-0.9.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+haiku_rag-0.9.3.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
+haiku_rag-0.9.3.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
+haiku_rag-0.9.3.dist-info/RECORD,,