haiku.rag-slim 0.16.0__py3-none-any.whl → 0.24.0__py3-none-any.whl

haiku/rag/graph/agui/server.py

@@ -18,8 +18,7 @@ from starlette.responses import JSONResponse, StreamingResponse
 from starlette.routing import Route
 
 from haiku.rag.config.models import AGUIConfig
-from haiku.rag.graph.agui.emitter import AGUIEmitter
-from haiku.rag.graph.agui.events import AGUIEvent
+from haiku.rag.graph.agui.emitter import AGUIEmitter, AGUIEvent
 from haiku.rag.graph.agui.stream import stream_graph
 
 
@@ -151,23 +150,25 @@ def format_sse_event(event: AGUIEvent) -> str:
     return f"data: {event_json}\n\n"
 
 
-def create_agui_server(config: "AppConfig", db_path: Path | None = None) -> Starlette:
-    """Create AG-UI server with both research and deep ask endpoints.
+def create_agui_server(  # pragma: no cover
+    config: "AppConfig", db_path: Path | None = None
+) -> Starlette:
+    """Create AG-UI server with research endpoint.
 
     Args:
-        config: Application config with research and qa settings
+        config: Application config with research settings
         db_path: Optional database path override
 
     Returns:
-        Starlette app with research and deep ask endpoints
+        Starlette app with research endpoint
     """
     from haiku.rag.client import HaikuRAG
-    from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
-    from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
-    from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
     from haiku.rag.graph.research.dependencies import ResearchContext
     from haiku.rag.graph.research.graph import build_research_graph
-    from haiku.rag.graph.research.state import ResearchDeps, ResearchState
+    from haiku.rag.graph.research.state import (
+        ResearchDeps,
+        ResearchState,
+    )
 
     # Store client reference for proper lifecycle management
     _client_cache: dict[str, HaikuRAG] = {}
@@ -190,7 +191,14 @@ def create_agui_server(config: "AppConfig", db_path: Path | None = None) -> Star
             if messages:
                 question = messages[0].get("content", "")
         context = ResearchContext(original_question=question)
-        return ResearchState.from_config(context=context, config=config)
+        max_iterations = input_state.get("max_iterations")
+        confidence_threshold = input_state.get("confidence_threshold")
+        return ResearchState.from_config(
+            context=context,
+            config=config,
+            max_iterations=max_iterations,
+            confidence_threshold=confidence_threshold,
+        )
 
     def research_deps_factory(input_config: dict[str, Any]) -> ResearchDeps:
         effective_db_path = (
@@ -200,29 +208,7 @@ def create_agui_server(config: "AppConfig", db_path: Path | None = None) -> Star
         )
         return ResearchDeps(client=get_client(effective_db_path))
 
-    # Deep ask graph factories
-    def deep_ask_graph_factory() -> Graph:
-        return build_deep_qa_graph(config)
-
-    def deep_ask_state_factory(input_state: dict[str, Any]) -> DeepQAState:
-        question = input_state.get("question", "")
-        if not question:
-            messages = input_state.get("messages", [])
-            if messages:
-                question = messages[0].get("content", "")
-        use_citations = input_state.get("use_citations", False)
-        context = DeepQAContext(original_question=question, use_citations=use_citations)
-        return DeepQAState.from_config(context=context, config=config)
-
-    def deep_ask_deps_factory(input_config: dict[str, Any]) -> DeepQADeps:
-        effective_db_path = (
-            db_path
-            or input_config.get("db_path")
-            or config.storage.data_dir / "haiku.rag.lancedb"
-        )
-        return DeepQADeps(client=get_client(effective_db_path))
-
-    # Create event stream functions for each graph type
+    # Create event stream function
     async def research_event_stream(
         input_data: RunAgentInput,
     ) -> AsyncIterator[str]:
@@ -235,18 +221,6 @@ def create_agui_server(config: "AppConfig", db_path: Path | None = None) -> Star
             event_data = format_sse_event(event)
             yield event_data
 
-    async def deep_ask_event_stream(
-        input_data: RunAgentInput,
-    ) -> AsyncIterator[str]:
-        """Generate SSE event stream from deep ask graph execution."""
-        graph = deep_ask_graph_factory()
-        initial_state = deep_ask_state_factory(input_data.state)
-        deps = deep_ask_deps_factory(input_data.config)
-
-        async for event in stream_graph(graph, initial_state, deps):
-            event_data = format_sse_event(event)
-            yield event_data
-
     # Endpoint handlers
     async def stream_research(request: Request) -> StreamingResponse:
         """Research graph streaming endpoint."""
@@ -263,21 +237,6 @@ def create_agui_server(config: "AppConfig", db_path: Path | None = None) -> Star
             },
         )
 
-    async def stream_deep_ask(request: Request) -> StreamingResponse:
-        """Deep ask graph streaming endpoint."""
-        body = await request.json()
-        input_data = RunAgentInput(**body)
-
-        return StreamingResponse(
-            deep_ask_event_stream(input_data),
-            media_type="text/event-stream",
-            headers={
-                "Cache-Control": "no-cache",
-                "Connection": "keep-alive",
-                "X-Accel-Buffering": "no",
-            },
-        )
-
     async def health_check(_: Request) -> JSONResponse:
         """Health check endpoint."""
         return JSONResponse({"status": "healthy"})
@@ -285,7 +244,6 @@ def create_agui_server(config: "AppConfig", db_path: Path | None = None) -> Star
     # Define routes
     routes = [
         Route("/v1/research/stream", stream_research, methods=["POST"]),
-        Route("/v1/deep-ask/stream", stream_deep_ask, methods=["POST"]),
         Route("/health", health_check, methods=["GET"]),
     ]
 

haiku/rag/graph/agui/stream.py

@@ -8,8 +8,7 @@ from typing import Protocol, TypeVar
 from pydantic import BaseModel
 from pydantic_graph.beta import Graph
 
-from haiku.rag.graph.agui.emitter import AGUIEmitter
-from haiku.rag.graph.agui.events import AGUIEvent
+from haiku.rag.graph.agui.emitter import AGUIEmitter, AGUIEvent
 
 StateT = TypeVar("StateT", bound=BaseModel)
 ResultT = TypeVar("ResultT")

haiku/rag/graph/research/__init__.py

@@ -1,3 +1,6 @@
-from haiku.rag.graph.common.models import SearchAnswer
 from haiku.rag.graph.research.dependencies import ResearchContext, ResearchDependencies
-from haiku.rag.graph.research.models import EvaluationResult, ResearchReport
+from haiku.rag.graph.research.models import (
+    EvaluationResult,
+    ResearchReport,
+    SearchAnswer,
+)

haiku/rag/graph/research/dependencies.py

@@ -1,14 +1,12 @@
-from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any
 
-from pydantic import BaseModel, Field, PrivateAttr
+from pydantic import BaseModel, Field
 
 from haiku.rag.client import HaikuRAG
-from haiku.rag.graph.common.models import SearchAnswer
-from haiku.rag.graph.research.models import (
-    GapRecord,
-    InsightAnalysis,
-    InsightRecord,
-)
+from haiku.rag.store.models import SearchResult
+
+if TYPE_CHECKING:
+    from haiku.rag.graph.research.models import SearchAnswer
 
 
 class ResearchContext(BaseModel):
@@ -18,124 +16,14 @@ class ResearchContext(BaseModel):
     sub_questions: list[str] = Field(
         default_factory=list, description="Decomposed sub-questions"
     )
-    qa_responses: list[SearchAnswer] = Field(
+    qa_responses: list[Any] = 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"
-    )
-
-    # Private dict indexes for O(1) lookups
-    _insights_by_id: dict[str, InsightRecord] = PrivateAttr(default_factory=dict)
-    _gaps_by_id: dict[str, GapRecord] = PrivateAttr(default_factory=dict)
-
-    def model_post_init(self, __context: object) -> None:
-        """Build indexes after initialization."""
-        self._insights_by_id = {ins.id: ins for ins in self.insights}
-        self._gaps_by_id = {gap.id: gap for gap in self.gaps}
 
-    def add_qa_response(self, qa: SearchAnswer) -> None:
-        """Add a structured QA response (minimal context already included)."""
+    def add_qa_response(self, qa: "SearchAnswer") -> None:
+        """Add a structured QA response."""
         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 = self._insights_by_id.get(candidate.id)
-
-            if existing:
-                # Update existing insight
-                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:
-                # Add new insight
-                new_insight = candidate.model_copy(deep=True)
-                self.insights.append(new_insight)
-                self._insights_by_id[new_insight.id] = new_insight
-                merged.append(new_insight)
-
-        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 = self._gaps_by_id.get(candidate.id)
-
-            if existing:
-                # Update existing gap
-                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:
-                # Add new gap
-                new_gap = candidate.model_copy(deep=True)
-                self.gaps.append(new_gap)
-                self._gaps_by_id[new_gap.id] = new_gap
-                merged.append(new_gap)
-
-        return merged
-
-    def mark_gap_resolved(
-        self, identifier: str, resolved_by: Iterable[str] | None = None
-    ) -> GapRecord | None:
-        """Mark a gap as resolved by identifier."""
-        gap = self._gaps_by_id.get(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 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)
-
 
 class ResearchDependencies(BaseModel):
     """Dependencies for research agents with multi-agent context."""
@@ -144,8 +32,6 @@ class ResearchDependencies(BaseModel):
 
     client: HaikuRAG = Field(description="RAG client for document operations")
     context: ResearchContext = Field(description="Shared research context")
-
-
-def _merge_unique(existing: list[str], incoming: Iterable[str]) -> list[str]:
-    """Merge two iterables preserving order while removing duplicates."""
-    return [k for k in dict.fromkeys([*existing, *incoming]) if k]
+    search_results: list[SearchResult] = Field(
+        default_factory=list, description="Search results for citation resolution"
+    )