haiku.rag 0.10.2__tar.gz → 0.11.0__tar.gz

{haiku_rag-0.10.2 → haiku_rag-0.11.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haiku.rag
-Version: 0.10.2
+Version: 0.11.0
 Summary: Agentic Retrieval Augmented Generation (RAG) with LanceDB
 Author-email: Yiorgis Gozadinos <ggozadinos@gmail.com>
 License: MIT
@@ -102,11 +102,12 @@ haiku-rag serve
 ```python
 from haiku.rag.client import HaikuRAG
 from haiku.rag.research import (
+    PlanNode,
     ResearchContext,
     ResearchDeps,
     ResearchState,
     build_research_graph,
-    PlanNode,
+    stream_research_graph,
 )
 
 async with HaikuRAG("database.lancedb") as client:
@@ -128,22 +129,40 @@ async with HaikuRAG("database.lancedb") as client:
 
     # Multi‑agent research pipeline (Plan → Search → Evaluate → Synthesize)
     graph = build_research_graph()
+    question = (
+        "What are the main drivers and trends of global temperature "
+        "anomalies since 1990?"
+    )
     state = ResearchState(
-        question=(
-            "What are the main drivers and trends of global temperature "
-            "anomalies since 1990?"
-        ),
-        context=ResearchContext(original_question="…"),
+        context=ResearchContext(original_question=question),
         max_iterations=2,
         confidence_threshold=0.8,
-        max_concurrency=3,
+        max_concurrency=2,
     )
     deps = ResearchDeps(client=client)
-    start = PlanNode(provider=None, model=None)
-    result = await graph.run(start, state=state, deps=deps)
-    report = result.output
-    print(report.title)
-    print(report.executive_summary)
+
+    # Blocking run (final result only)
+    result = await graph.run(
+        PlanNode(provider="openai", model="gpt-4o-mini"),
+        state=state,
+        deps=deps,
+    )
+    print(result.output.title)
+
+    # Streaming progress (log/report/error events)
+    async for event in stream_research_graph(
+        graph,
+        PlanNode(provider="openai", model="gpt-4o-mini"),
+        state,
+        deps,
+    ):
+        if event.type == "log":
+            iteration = event.state.iterations if event.state else state.iterations
+            print(f"[{iteration}] {event.message}")
+        elif event.type == "report":
+            print("\nResearch complete!\n")
+            print(event.report.title)
+            print(event.report.executive_summary)
 ```
 
 ## MCP Server

{haiku_rag-0.10.2 → haiku_rag-0.11.0}/README.md

@@ -64,11 +64,12 @@ haiku-rag serve
 ```python
 from haiku.rag.client import HaikuRAG
 from haiku.rag.research import (
+    PlanNode,
     ResearchContext,
     ResearchDeps,
     ResearchState,
     build_research_graph,
-    PlanNode,
+    stream_research_graph,
 )
 
 async with HaikuRAG("database.lancedb") as client:
@@ -90,22 +91,40 @@ async with HaikuRAG("database.lancedb") as client:
 
     # Multi‑agent research pipeline (Plan → Search → Evaluate → Synthesize)
     graph = build_research_graph()
+    question = (
+        "What are the main drivers and trends of global temperature "
+        "anomalies since 1990?"
+    )
     state = ResearchState(
-        question=(
-            "What are the main drivers and trends of global temperature "
-            "anomalies since 1990?"
-        ),
-        context=ResearchContext(original_question="…"),
+        context=ResearchContext(original_question=question),
         max_iterations=2,
         confidence_threshold=0.8,
-        max_concurrency=3,
+        max_concurrency=2,
     )
     deps = ResearchDeps(client=client)
-    start = PlanNode(provider=None, model=None)
-    result = await graph.run(start, state=state, deps=deps)
-    report = result.output
-    print(report.title)
-    print(report.executive_summary)
+
+    # Blocking run (final result only)
+    result = await graph.run(
+        PlanNode(provider="openai", model="gpt-4o-mini"),
+        state=state,
+        deps=deps,
+    )
+    print(result.output.title)
+
+    # Streaming progress (log/report/error events)
+    async for event in stream_research_graph(
+        graph,
+        PlanNode(provider="openai", model="gpt-4o-mini"),
+        state,
+        deps,
+    ):
+        if event.type == "log":
+            iteration = event.state.iterations if event.state else state.iterations
+            print(f"[{iteration}] {event.message}")
+        elif event.type == "report":
+            print("\nResearch complete!\n")
+            print(event.report.title)
+            print(event.report.executive_summary)
 ```
 
 ## MCP Server

{haiku_rag-0.10.2 → haiku_rag-0.11.0}/docs/agents.md

@@ -47,9 +47,10 @@ title: Research graph
 ---
 stateDiagram-v2
   PlanNode --> SearchDispatchNode
-  SearchDispatchNode --> EvaluateNode
-  EvaluateNode --> SearchDispatchNode
-  EvaluateNode --> SynthesizeNode
+  SearchDispatchNode --> AnalyzeInsightsNode
+  AnalyzeInsightsNode --> DecisionNode
+  DecisionNode --> SearchDispatchNode
+  DecisionNode --> SynthesizeNode
   SynthesizeNode --> [*]
 ```
 
@@ -57,12 +58,15 @@ Key nodes:
 
 - Plan: builds up to 3 standalone sub‑questions (uses an internal presearch tool)
 - Search (batched): answers sub‑questions using the KB with minimal, verbatim context
-- Evaluate: extracts insights, proposes new questions, and checks sufficiency/confidence
+- Analyze: aggregates fresh insights, updates gaps, and suggests new sub-questions
+- Decision: checks sufficiency/confidence thresholds and chooses whether to iterate
 - Synthesize: generates a final structured report
 
 Primary models:
 
 - `SearchAnswer` — one per sub‑question (query, answer, context, sources)
+- `InsightRecord` / `GapRecord` — structured tracking of findings and open issues
+- `InsightAnalysis` — output of the analysis stage (insights, gaps, commentary)
 - `EvaluationResult` — insights, new questions, sufficiency, confidence
 - `ResearchReport` — final report (title, executive summary, findings, conclusions, …)
 
@@ -76,30 +80,75 @@ haiku-rag research "How does haiku.rag organize and query documents?" \
   --verbose
 ```
 
-Python usage:
+Python usage (blocking result):
 
 ```python
 from haiku.rag.client import HaikuRAG
 from haiku.rag.research import (
+    PlanNode,
     ResearchContext,
     ResearchDeps,
     ResearchState,
     build_research_graph,
-    PlanNode,
 )
 
 async with HaikuRAG(path_to_db) as client:
     graph = build_research_graph()
+    question = "What are the main drivers and trends of global temperature anomalies since 1990?"
     state = ResearchState(
-        question="What are the main drivers and trends of global temperature anomalies since 1990?",
-        context=ResearchContext(original_question=... ),
+        context=ResearchContext(original_question=question),
         max_iterations=2,
         confidence_threshold=0.8,
-        max_concurrency=3,
+        max_concurrency=2,
     )
     deps = ResearchDeps(client=client)
-    result = await graph.run(PlanNode(provider=None, model=None), state=state, deps=deps)
+
+    result = await graph.run(
+        PlanNode(provider="openai", model="gpt-4o-mini"),
+        state=state,
+        deps=deps,
+    )
+
     report = result.output
     print(report.title)
     print(report.executive_summary)
 ```
+
+Python usage (streamed events):
+
+```python
+from haiku.rag.client import HaikuRAG
+from haiku.rag.research import (
+    PlanNode,
+    ResearchContext,
+    ResearchDeps,
+    ResearchState,
+    build_research_graph,
+    stream_research_graph,
+)
+
+async with HaikuRAG(path_to_db) as client:
+    graph = build_research_graph()
+    question = "What are the main drivers and trends of global temperature anomalies since 1990?"
+    state = ResearchState(
+        context=ResearchContext(original_question=question),
+        max_iterations=2,
+        confidence_threshold=0.8,
+        max_concurrency=2,
+    )
+    deps = ResearchDeps(client=client)
+
+    async for event in stream_research_graph(
+        graph,
+        PlanNode(provider="openai", model="gpt-4o-mini"),
+        state,
+        deps,
+    ):
+        if event.type == "log":
+            iteration = event.state.iterations if event.state else state.iterations
+            print(f"[{iteration}] {event.message}")
+        elif event.type == "report":
+            print("\nResearch complete!\n")
+            print(event.report.title)
+            print(event.report.executive_summary)
+```

{haiku_rag-0.10.2 → haiku_rag-0.11.0}/docs/cli.md

@@ -113,6 +113,8 @@ Flags:
 - `--max-concurrency`: number of sub-questions searched in parallel each iteration (default: 3)
 - `--verbose`: show planning, searching previews, evaluation summary, and stop reason
 
+When `--verbose` is set the CLI also consumes the internal research stream, printing every `log` event as agents progress through planning, search, evaluation, and synthesis. If you build your own integration, call `stream_research_graph` to access the same `log`, `report`, and `error` events and render them however you like while the graph is running.
+
 ## Server
 
 Start the MCP server:

{haiku_rag-0.10.2 → haiku_rag-0.11.0}/pyproject.toml

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

{haiku_rag-0.10.2 → haiku_rag-0.11.0}/src/haiku/rag/app.py

@@ -18,6 +18,7 @@ from haiku.rag.research.graph import (
     ResearchState,
     build_research_graph,
 )
+from haiku.rag.research.stream import stream_research_graph
 from haiku.rag.store.models.chunk import Chunk
 from haiku.rag.store.models.document import Document
 
@@ -221,9 +222,9 @@ class HaikuRAGApp:
                     self.console.print()
 
                 graph = build_research_graph()
+                context = ResearchContext(original_question=question)
                 state = ResearchState(
-                    question=question,
-                    context=ResearchContext(original_question=question),
+                    context=context,
                     max_iterations=max_iterations,
                     confidence_threshold=confidence_threshold,
                     max_concurrency=max_concurrency,
@@ -236,22 +237,20 @@ class HaikuRAGApp:
                     provider=Config.RESEARCH_PROVIDER or Config.QA_PROVIDER,
                     model=Config.RESEARCH_MODEL or Config.QA_MODEL,
                 )
-                # Prefer graph.run; fall back to iter if unavailable
                 report = None
-                try:
-                    result = await graph.run(start, state=state, deps=deps)
-                    report = result.output
-                except Exception:
-                    from pydantic_graph import End
-
-                    async with graph.iter(start, state=state, deps=deps) as run:
-                        node = run.next_node
-                        while not isinstance(node, End):
-                            node = await run.next(node)
-                        if run.result:
-                            report = run.result.output
+                async for event in stream_research_graph(graph, start, state, deps):
+                    if event.type == "report":
+                        report = event.report
+                        break
+                    if event.type == "error":
+                        self.console.print(
+                            f"[red]Error during research: {event.message}[/red]"
+                        )
+                        return
+
                 if report is None:
-                    raise RuntimeError("Graph did not produce a report")
+                    self.console.print("[red]Research did not produce a report.[/red]")
+                    return
 
                 # Display the report
                 self.console.print("[bold green]Research Report[/bold green]")

{haiku_rag-0.10.2 → haiku_rag-0.11.0}/src/haiku/rag/research/__init__.py

@@ -6,6 +6,11 @@ from haiku.rag.research.graph import (
     build_research_graph,
 )
 from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer
+from haiku.rag.research.stream import (
+    ResearchStateSnapshot,
+    ResearchStreamEvent,
+    stream_research_graph,
+)
 
 __all__ = [
     "ResearchDependencies",
@@ -17,4 +22,7 @@ __all__ = [
     "ResearchState",
     "PlanNode",
     "build_research_graph",
+    "stream_research_graph",
+    "ResearchStreamEvent",
+    "ResearchStateSnapshot",
 ]

haiku_rag-0.11.0/src/haiku/rag/research/common.py

@@ -0,0 +1,118 @@
+from typing import TYPE_CHECKING, Any
+
+from pydantic_ai import format_as_xml
+from pydantic_ai.models.openai import OpenAIChatModel
+from pydantic_ai.providers.ollama import OllamaProvider
+from pydantic_ai.providers.openai import OpenAIProvider
+
+from haiku.rag.config import Config
+from haiku.rag.research.dependencies import ResearchContext
+from haiku.rag.research.models import InsightAnalysis
+
+if TYPE_CHECKING:  # pragma: no cover
+    from haiku.rag.research.state import ResearchDeps, ResearchState
+
+
+def get_model(provider: str, model: str) -> Any:
+    if provider == "ollama":
+        return OpenAIChatModel(
+            model_name=model,
+            provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
+        )
+    elif provider == "vllm":
+        return OpenAIChatModel(
+            model_name=model,
+            provider=OpenAIProvider(
+                base_url=f"{Config.VLLM_RESEARCH_BASE_URL or Config.VLLM_QA_BASE_URL}/v1",
+                api_key="none",
+            ),
+        )
+    else:
+        return f"{provider}:{model}"
+
+
+def log(deps: "ResearchDeps", state: "ResearchState", msg: str) -> None:
+    deps.emit_log(msg, state)
+
+
+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,  # pyright: ignore[reportAttributeAccessIssue]
+            }
+            for qa in context.qa_responses
+        ],
+        "insights": [
+            {
+                "id": insight.id,
+                "summary": insight.summary,
+                "status": insight.status.value,
+                "supporting_sources": insight.supporting_sources,
+                "originating_questions": insight.originating_questions,
+                "notes": insight.notes,
+            }
+            for insight in context.insights
+        ],
+        "gaps": [
+            {
+                "id": gap.id,
+                "description": gap.description,
+                "severity": gap.severity.value,
+                "blocking": gap.blocking,
+                "resolved": gap.resolved,
+                "resolved_by": gap.resolved_by,
+                "supporting_sources": gap.supporting_sources,
+                "notes": gap.notes,
+            }
+            for gap in context.gaps
+        ],
+    }
+    return format_as_xml(context_data, root_tag="research_context")
+
+
+def format_analysis_for_prompt(
+    analysis: InsightAnalysis | None,
+) -> str:
+    """Format the latest insight analysis as XML for prompts."""
+
+    if analysis is None:
+        return "<latest_analysis />"
+
+    data = {
+        "commentary": analysis.commentary,
+        "highlights": [
+            {
+                "id": insight.id,
+                "summary": insight.summary,
+                "status": insight.status.value,
+                "supporting_sources": insight.supporting_sources,
+                "originating_questions": insight.originating_questions,
+                "notes": insight.notes,
+            }
+            for insight in analysis.highlights
+        ],
+        "gap_assessments": [
+            {
+                "id": gap.id,
+                "description": gap.description,
+                "severity": gap.severity.value,
+                "blocking": gap.blocking,
+                "resolved": gap.resolved,
+                "resolved_by": gap.resolved_by,
+                "supporting_sources": gap.supporting_sources,
+                "notes": gap.notes,
+            }
+            for gap in analysis.gap_assessments
+        ],
+        "resolved_gaps": analysis.resolved_gaps,
+        "new_questions": analysis.new_questions,
+    }
+    return format_as_xml(data, root_tag="latest_analysis")

haiku_rag-0.11.0/src/haiku/rag/research/dependencies.py

@@ -0,0 +1,215 @@
+from collections.abc import Iterable
+
+from pydantic import BaseModel, Field
+from rich.console import Console
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.research.models import (
+    GapRecord,
+    InsightAnalysis,
+    InsightRecord,
+    SearchAnswer,
+)
+from haiku.rag.research.stream import ResearchStream
+
+
+class ResearchContext(BaseModel):
+    """Context shared across research agents."""
+
+    original_question: str = Field(description="The original research question")
+    sub_questions: list[str] = Field(
+        default_factory=list, description="Decomposed sub-questions"
+    )
+    qa_responses: list[SearchAnswer] = Field(
+        default_factory=list, description="Structured QA pairs used during research"
+    )
+    insights: list[InsightRecord] = Field(
+        default_factory=list, description="Key insights discovered"
+    )
+    gaps: list[GapRecord] = Field(
+        default_factory=list, description="Identified information gaps"
+    )
+
+    def add_qa_response(self, qa: SearchAnswer) -> None:
+        """Add a structured QA response (minimal context already included)."""
+        self.qa_responses.append(qa)
+
+    def upsert_insights(self, records: Iterable[InsightRecord]) -> list[InsightRecord]:
+        """Merge one or more insights into the shared context with deduplication."""
+
+        merged: list[InsightRecord] = []
+        for record in records:
+            candidate = InsightRecord.model_validate(record)
+            existing = next(
+                (ins for ins in self.insights if ins.id == candidate.id), None
+            )
+            if not existing:
+                existing = next(
+                    (ins for ins in self.insights if ins.summary == candidate.summary),
+                    None,
+                )
+
+            if existing:
+                existing.summary = candidate.summary
+                existing.status = candidate.status
+                if candidate.notes:
+                    existing.notes = candidate.notes
+                existing.supporting_sources = _merge_unique(
+                    existing.supporting_sources, candidate.supporting_sources
+                )
+                existing.originating_questions = _merge_unique(
+                    existing.originating_questions, candidate.originating_questions
+                )
+                merged.append(existing)
+            else:
+                candidate = candidate.model_copy(deep=True)
+                if candidate.id is None:  # pragma: no cover - defensive
+                    raise ValueError(
+                        "InsightRecord.id must be populated after validation"
+                    )
+                candidate_id: str = candidate.id
+                candidate.id = self._allocate_insight_id(candidate_id)
+                self.insights.append(candidate)
+                merged.append(candidate)
+
+        return merged
+
+    def upsert_gaps(self, records: Iterable[GapRecord]) -> list[GapRecord]:
+        """Merge one or more gap records into the shared context with deduplication."""
+
+        merged: list[GapRecord] = []
+        for record in records:
+            candidate = GapRecord.model_validate(record)
+            existing = next((gap for gap in self.gaps if gap.id == candidate.id), None)
+            if not existing:
+                existing = next(
+                    (
+                        gap
+                        for gap in self.gaps
+                        if gap.description == candidate.description
+                    ),
+                    None,
+                )
+
+            if existing:
+                existing.description = candidate.description
+                existing.severity = candidate.severity
+                existing.blocking = candidate.blocking
+                existing.resolved = candidate.resolved
+                if candidate.notes:
+                    existing.notes = candidate.notes
+                existing.supporting_sources = _merge_unique(
+                    existing.supporting_sources, candidate.supporting_sources
+                )
+                existing.resolved_by = _merge_unique(
+                    existing.resolved_by, candidate.resolved_by
+                )
+                merged.append(existing)
+            else:
+                candidate = candidate.model_copy(deep=True)
+                if candidate.id is None:  # pragma: no cover - defensive
+                    raise ValueError("GapRecord.id must be populated after validation")
+                candidate_id: str = candidate.id
+                candidate.id = self._allocate_gap_id(candidate_id)
+                self.gaps.append(candidate)
+                merged.append(candidate)
+
+        return merged
+
+    def mark_gap_resolved(
+        self, identifier: str, resolved_by: Iterable[str] | None = None
+    ) -> GapRecord | None:
+        """Mark a gap as resolved by identifier (id or description)."""
+
+        gap = self._find_gap(identifier)
+        if gap is None:
+            return None
+
+        gap.resolved = True
+        gap.blocking = False
+        if resolved_by:
+            gap.resolved_by = _merge_unique(gap.resolved_by, list(resolved_by))
+        return gap
+
+    def integrate_analysis(self, analysis: InsightAnalysis) -> None:
+        """Apply an analysis result to the shared context."""
+
+        merged_insights: list[InsightRecord] = []
+        if analysis.highlights:
+            merged_insights = self.upsert_insights(analysis.highlights)
+            analysis.highlights = merged_insights
+        if analysis.gap_assessments:
+            merged_gaps = self.upsert_gaps(analysis.gap_assessments)
+            analysis.gap_assessments = merged_gaps
+        if analysis.resolved_gaps:
+            resolved_by_list = (
+                [ins.id for ins in merged_insights if ins.id is not None]
+                if merged_insights
+                else None
+            )
+            for resolved in analysis.resolved_gaps:
+                self.mark_gap_resolved(resolved, resolved_by=resolved_by_list)
+        for question in analysis.new_questions:
+            if question not in self.sub_questions:
+                self.sub_questions.append(question)
+
+    def _allocate_insight_id(self, candidate_id: str) -> str:
+        taken: set[str] = set()
+        for ins in self.insights:
+            if ins.id is not None:
+                taken.add(ins.id)
+        return _allocate_sequential_id(candidate_id, taken)
+
+    def _allocate_gap_id(self, candidate_id: str) -> str:
+        taken: set[str] = set()
+        for gap in self.gaps:
+            if gap.id is not None:
+                taken.add(gap.id)
+        return _allocate_sequential_id(candidate_id, taken)
+
+    def _find_gap(self, identifier: str) -> GapRecord | None:
+        normalized = identifier.lower().strip()
+        for gap in self.gaps:
+            if gap.id is not None and gap.id == normalized:
+                return gap
+            if gap.description.lower().strip() == normalized:
+                return gap
+        return None
+
+
+class ResearchDependencies(BaseModel):
+    """Dependencies for research agents with multi-agent context."""
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    client: HaikuRAG = Field(description="RAG client for document operations")
+    context: ResearchContext = Field(description="Shared research context")
+    console: Console | None = None
+    stream: ResearchStream | None = Field(
+        default=None, description="Optional research event stream"
+    )
+
+
+def _merge_unique(existing: list[str], incoming: Iterable[str]) -> list[str]:
+    """Merge two iterables preserving order while removing duplicates."""
+
+    merged = list(existing)
+    seen = {item for item in existing if item}
+    for item in incoming:
+        if item and item not in seen:
+            merged.append(item)
+            seen.add(item)
+    return merged
+
+
+def _allocate_sequential_id(candidate: str, taken: set[str]) -> str:
+    slug = candidate
+    if slug not in taken:
+        return slug
+    base = slug
+    counter = 2
+    while True:
+        slug = f"{base}-{counter}"
+        if slug not in taken:
+            return slug
+        counter += 1