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

haiku/rag/app.py

@@ -122,12 +122,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 +256,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/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/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/research/base.py

@@ -33,10 +33,18 @@ 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),
+            output_type=agent_output_type,
             system_prompt=self.get_system_prompt(),
         )
 
@@ -73,7 +81,7 @@ class BaseResearchAgent[T](ABC):
         pass
 
     async def run(
-        self, prompt: str, deps: ResearchDependencies, **kwargs
+        self, prompt: str, deps: "ResearchDependencies", **kwargs
     ) -> AgentRunResult[T]:
         """Execute the agent."""
         return await self._agent.run(prompt, deps=deps, **kwargs)

haiku/rag/research/evaluation_agent.py

@@ -11,7 +11,9 @@ class EvaluationResult(BaseModel):
         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
+        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)",

haiku/rag/research/orchestrator.py

@@ -12,6 +12,7 @@ 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
@@ -41,6 +42,9 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
         self.search_agent: SearchSpecialistAgent = SearchSpecialistAgent(
             provider, model
         )
+        self.presearch_agent: PresearchSurveyAgent = PresearchSurveyAgent(
+            provider, model
+        )
         self.evaluation_agent: AnalysisEvaluationAgent = AnalysisEvaluationAgent(
             provider, model
         )
@@ -61,7 +65,12 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
             "original_question": context.original_question,
             "unanswered_questions": context.sub_questions,
             "qa_responses": [
-                {"question": qa.query, "answer": qa.answer}
+                {
+                    "question": qa.query,
+                    "answer": qa.answer,
+                    "context_snippets": qa.context,
+                    "sources": qa.sources,
+                }
                 for qa in context.qa_responses
             ],
             "insights": context.insights,
@@ -99,12 +108,38 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
         # 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)
+
         # 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",
+        )
+
+        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}"
+        )
+
         plan_result: AgentRunResult[ResearchPlan] = await self.run(
-            f"Create a research plan for: {question}", deps=deps
+            plan_prompt, deps=deps
         )
 
         context.sub_questions = plan_result.output.sub_questions

haiku/rag/research/presearch_agent.py

@@ -0,0 +1,34 @@
+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]:
+        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."""
+            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/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/research/search_agent.py

@@ -42,6 +42,7 @@ class SearchSpecialistAgent(BaseResearchAgent[SearchAnswer]):
         ) -> 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

@@ -12,11 +12,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"

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

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

{haiku_rag-0.9.0.dist-info → haiku_rag-0.9.2.dist-info}/RECORD

@@ -1,5 +1,5 @@
 haiku/rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-haiku/rag/app.py,sha256=Smof7ZIe-oRGkDTap81VaKZGIeborD2X-oXKgBoJs9I,11763
+haiku/rag/app.py,sha256=o64L7aj5V8lYHxWhAKBNj1tGfXiN6xr0_Cc1dEYd3As,11483
 haiku/rag/chunker.py,sha256=PVe6ysv8UlacUd4Zb3_8RFWIaWDXnzBAy2VDJ4TaUsE,1555
 haiku/rag/cli.py,sha256=3nlzrT5FPCyfnu51KHchLG4Cj2eVv9YsuGHMShBnVb0,9845
 haiku/rag/client.py,sha256=NJVGXzVzpoVy1sttz_xEU7mXWtObKT8pGpvo5pZyzwc,21288
@@ -17,21 +17,22 @@ haiku/rag/embeddings/openai.py,sha256=fIFCk-jpUtaW0xsnrQnJ824O0UCjaGG2sgvBzREhil
 haiku/rag/embeddings/vllm.py,sha256=vhaUnCn6VMkfSluLhWKtSV-sekFaPsp4pKo2N7-SBCY,626
 haiku/rag/embeddings/voyageai.py,sha256=UW-MW4tJKnPB6Fs2P7A3yt-ZeRm46H9npckchSriPX8,661
 haiku/rag/qa/__init__.py,sha256=Sl7Kzrg9CuBOcMF01wc1NtQhUNWjJI0MhIHfCWrb8V4,434
-haiku/rag/qa/agent.py,sha256=f7hGWhjgzJKwa5BJkAO0KCxbgpwigPz5E9a26S9TUYI,2948
+haiku/rag/qa/agent.py,sha256=f4Keh-ESgctNbTg96QL95HYjINVLOcxa8t8crx92MMk,3081
 haiku/rag/qa/prompts.py,sha256=LhRfDtO8Pb06lpr4PpwEaKUYItZ5OiIkeqcCogcssHY,3347
 haiku/rag/reranking/__init__.py,sha256=IRXHs4qPu6VbGJQpzSwhgtVWWumURH_vEoVFE-extlo,894
 haiku/rag/reranking/base.py,sha256=LM9yUSSJ414UgBZhFTgxGprlRqzfTe4I1vgjricz2JY,405
 haiku/rag/reranking/cohere.py,sha256=1iTdiaa8vvb6oHVB2qpWzUOVkyfUcimVSZp6Qr4aq4c,1049
 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=hwCVV8fxnHTrLV2KCJ_Clqe_pPrCwTz-RW2b0BeGHeE,982
-haiku/rag/research/base.py,sha256=IW_VbeRlXTUfqh--jBS0dtIgSVXsbifPxZl8bfTLkDA,3686
+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=3YWAdfC6n27wAIdla7M72IE1aS4GqoL9DbnW4K1b35M,1357
-haiku/rag/research/orchestrator.py,sha256=AnDXCoJBbt4nYqaDKk5hiMi8WW1e8NwpRvzHLLnY3WY,10478
-haiku/rag/research/prompts.py,sha256=C_d9OGNgHfwSUY6n5L2c2J6OpCeBHwxtMjrLQOkdcxU,5221
-haiku/rag/research/search_agent.py,sha256=mYn3GlxoIPEle2NLkBqHI-VRV5PanoHxhjttVozsVis,2405
-haiku/rag/research/synthesis_agent.py,sha256=E7Iwfe0EAlmglIRMmRQ3kaNmEWIyEMpVFK3k4SPC5BM,1559
+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/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/store/__init__.py,sha256=hq0W0DAC7ysqhWSP2M2uHX8cbG6kbr-sWHxhq6qQcY0,103
 haiku/rag/store/engine.py,sha256=fNrykqMX7PRSCt4LSRfuJ66OLrb8BVYq2bpbfI2iaWU,8455
 haiku/rag/store/models/__init__.py,sha256=s0E72zneGlowvZrFWaNxHYjOAUjgWdLxzdYsnvNRVlY,88
@@ -42,8 +43,8 @@ haiku/rag/store/repositories/chunk.py,sha256=1RmPyEYRYOFbrALbmLOo62t3f-xO2KgxUjc
 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.0.dist-info/METADATA,sha256=ab5orVjoWGdapwaoPnwPdtuyetnErIxAvwDjl--9hfo,4681
-haiku_rag-0.9.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-haiku_rag-0.9.0.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
-haiku_rag-0.9.0.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
-haiku_rag-0.9.0.dist-info/RECORD,,
+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,,