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

haiku/rag/graph/research/models.py

@@ -0,0 +1,166 @@
+import uuid
+from enum import Enum
+
+from pydantic import BaseModel, Field, field_validator
+
+
+def _deduplicate_list(items: list[str]) -> list[str]:
+    """Remove duplicates while preserving order."""
+    return list(dict.fromkeys(items))
+
+
+class InsightStatus(str, Enum):
+    OPEN = "open"
+    VALIDATED = "validated"
+    TENTATIVE = "tentative"
+
+
+class GapSeverity(str, Enum):
+    LOW = "low"
+    MEDIUM = "medium"
+    HIGH = "high"
+
+
+class TrackedRecord(BaseModel):
+    """Base model for tracked entities with sources and metadata."""
+
+    model_config = {"validate_assignment": True}
+
+    id: str = Field(
+        default_factory=lambda: str(uuid.uuid4())[:8],
+        description="Unique identifier for the record",
+    )
+    supporting_sources: list[str] = Field(
+        default_factory=list,
+        description="Source identifiers backing this record",
+    )
+    notes: str | None = Field(
+        default=None,
+        description="Optional elaboration or caveats",
+    )
+
+    @field_validator("supporting_sources", mode="before")
+    @classmethod
+    def deduplicate_sources(cls, v: list[str]) -> list[str]:
+        """Ensure supporting_sources has no duplicates."""
+        return _deduplicate_list(v) if v else []
+
+
+class InsightRecord(TrackedRecord):
+    """Structured insight with provenance and lifecycle metadata."""
+
+    summary: str = Field(description="Concise description of the insight")
+    status: InsightStatus = Field(
+        default=InsightStatus.OPEN,
+        description="Lifecycle status for the insight",
+    )
+    originating_questions: list[str] = Field(
+        default_factory=list,
+        description="Research sub-questions that produced this insight",
+    )
+
+    @field_validator("originating_questions", mode="before")
+    @classmethod
+    def deduplicate_questions(cls, v: list[str]) -> list[str]:
+        """Ensure originating_questions has no duplicates."""
+        return _deduplicate_list(v) if v else []
+
+
+class GapRecord(TrackedRecord):
+    """Structured representation of an identified research gap."""
+
+    description: str = Field(description="Concrete statement of what is missing")
+    severity: GapSeverity = Field(
+        default=GapSeverity.MEDIUM,
+        description="Severity of the gap for answering the main question",
+    )
+    blocking: bool = Field(
+        default=True,
+        description="Whether this gap blocks a confident answer",
+    )
+    resolved: bool = Field(
+        default=False,
+        description="Flag indicating if the gap has been resolved",
+    )
+    resolved_by: list[str] = Field(
+        default_factory=list,
+        description="Insight IDs or notes explaining how the gap was closed",
+    )
+
+    @field_validator("resolved_by", mode="before")
+    @classmethod
+    def deduplicate_resolved_by(cls, v: list[str]) -> list[str]:
+        """Ensure resolved_by has no duplicates."""
+        return _deduplicate_list(v) if v else []
+
+
+class InsightAnalysis(BaseModel):
+    """Output of the insight aggregation agent."""
+
+    highlights: list[InsightRecord] = Field(
+        default_factory=list,
+        description="New or updated insights discovered this iteration",
+    )
+    gap_assessments: list[GapRecord] = Field(
+        default_factory=list,
+        description="New or updated gap records based on current evidence",
+    )
+    resolved_gaps: list[str] = Field(
+        default_factory=list,
+        description="Gap identifiers or descriptions considered resolved",
+    )
+    new_questions: list[str] = Field(
+        default_factory=list,
+        max_length=3,
+        description="Up to three follow-up sub-questions to pursue next",
+    )
+    commentary: str = Field(
+        description="Short narrative summary of the incremental findings",
+    )
+
+
+class EvaluationResult(BaseModel):
+    """Result of analysis and evaluation."""
+
+    key_insights: list[str] = Field(
+        description="Main insights extracted from the research so far"
+    )
+    new_questions: list[str] = Field(
+        description="New sub-questions to add to the research (max 3)",
+        max_length=3,
+        default=[],
+    )
+    gaps: list[str] = Field(
+        description="Concrete information gaps that remain", default_factory=list
+    )
+    confidence_score: float = Field(
+        description="Confidence level in the completeness of research (0-1)",
+        ge=0.0,
+        le=1.0,
+    )
+    is_sufficient: bool = Field(
+        description="Whether the research is sufficient to answer the original question"
+    )
+    reasoning: str = Field(
+        description="Explanation of why the research is or isn't complete"
+    )
+
+
+class ResearchReport(BaseModel):
+    """Final research report structure."""
+
+    title: str = Field(description="Concise title for the research")
+    executive_summary: str = Field(description="Brief overview of key findings")
+    main_findings: list[str] = Field(
+        description="Primary research findings with supporting evidence"
+    )
+    conclusions: list[str] = Field(description="Evidence-based conclusions")
+    limitations: list[str] = Field(
+        description="Limitations of the current research", default=[]
+    )
+    recommendations: list[str] = Field(
+        description="Actionable recommendations based on findings", default=[]
+    )
+    sources_summary: str = Field(
+        description="Summary of sources used and their reliability"
+    )

haiku/rag/graph/research/prompts.py

@@ -0,0 +1,107 @@
+INSIGHT_AGENT_PROMPT = """You are the insight aggregation specialist for the
+research workflow.
+
+Inputs available:
+- Original research question and sub-questions
+- Question–answer pairs with supporting snippets and sources
+- Existing insights and gaps (with status metadata)
+
+Tasks:
+1. Extract new or refined insights that advance understanding of the question.
+2. Update gap status, creating new gap entries when necessary and marking
+   resolved ones explicitly.
+3. Suggest up to 3 high-impact follow-up sub_questions that would close the
+   most important remaining gaps.
+
+Output format (map directly to fields):
+- highlights: list of insights with fields {summary, status, supporting_sources,
+  originating_questions, notes}. Use status one of {validated, open, tentative}.
+- gap_assessments: list of gaps with fields {description, severity, blocking,
+  resolved, resolved_by, supporting_sources, notes}. Severity must be one of
+  {low, medium, high}. resolved_by may reference related insight summaries if no
+  stable identifier yet.
+- resolved_gaps: list of identifiers or descriptions for gaps now closed.
+- new_questions: up to 3 standalone, specific sub-questions (no duplicates with
+  existing ones).
+- commentary: 1–3 sentences summarizing what changed this round.
+
+Guidance:
+- Be concise and avoid repeating previously recorded information unless it
+  changed materially.
+- Tie supporting_sources to the evidence used; omit if unavailable.
+- Only propose new sub_questions that directly address remaining gaps.
+- When marking a gap as resolved, ensure the rationale is clear via
+  resolved_by or notes."""
+
+DECISION_AGENT_PROMPT = """You are the research governor responsible for making
+stop/go decisions.
+
+Inputs available:
+- Original research question and current plan
+- Full insight ledger with status metadata
+- Up-to-date gap tracker, including resolved indicators
+- Latest insight analysis summary (highlights, gap changes, new questions)
+- Previous evaluation decision (if any)
+
+Tasks:
+1. Determine whether the collected evidence now answers the original question.
+2. Provide a confidence_score in [0,1] that reflects coverage, evidence quality,
+   and agreement across sources.
+3. List the highest-priority gaps that still block a confident answer. Reference
+   existing gap descriptions rather than inventing new ones.
+4. Optionally propose up to 3 new sub_questions only if they are not already in
+   the current backlog.
+
+Strictness:
+- Only mark research as sufficient when every critical aspect of the main
+  question is addressed with reliable, corroborated evidence.
+- Treat unresolved high-severity or blocking gaps as a hard stop.
+
+Output fields must line up with EvaluationResult:
+- key_insights: concise bullet-ready statements of the most decision-relevant
+  insights (cite status if helpful).
+- new_questions: follow-up sub-questions (max 3) meeting the specificity rules.
+- gaps: list remaining blockers; reuse wording from the tracked gaps when
+  possible to aid downstream reconciliation.
+- confidence_score: numeric in [0,1].
+- is_sufficient: true only when no blocking gaps remain.
+- reasoning: short narrative tying the decision to evidence coverage.
+
+Remember: prefer maintaining continuity with the structured context over
+introducing new terminology."""
+
+SYNTHESIS_AGENT_PROMPT = """You are a synthesis specialist producing the final
+research report.
+
+Goals:
+1. Synthesize all gathered information into a coherent narrative.
+2. Present findings clearly and concisely.
+3. Draw evidence‑based conclusions and recommendations.
+4. State limitations and uncertainties transparently.
+
+Report guidelines (map to output fields):
+- title: concise (5–12 words), informative.
+- executive_summary: 3–5 sentences summarizing the overall answer.
+- main_findings: 4–8 one‑sentence bullets; each reflects evidence from the
+  research (do not include inline citations or snippet text).
+- conclusions: 2–4 bullets that follow logically from findings.
+- recommendations: 2–5 actionable bullets tied to findings.
+- limitations: 1–3 bullets describing key constraints or uncertainties.
+- sources_summary: 2–4 sentences summarizing sources used and their reliability.
+
+Style:
+- Base all content solely on the collected evidence.
+- Be professional, objective, and specific.
+- Avoid meta commentary and refrain from speculation beyond the evidence."""
+
+PRESEARCH_AGENT_PROMPT = """You are a rapid research surveyor.
+
+Task:
+- Call gather_context once on the main question to obtain relevant text from
+  the knowledge base (KB).
+- Read that context and produce a short natural‑language summary of 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/graph/research/state.py

@@ -0,0 +1,85 @@
+import asyncio
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, Field
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.graph.research.dependencies import ResearchContext
+from haiku.rag.graph.research.models import (
+    EvaluationResult,
+    InsightAnalysis,
+    ResearchReport,
+)
+
+if TYPE_CHECKING:
+    from haiku.rag.config.models import AppConfig
+    from haiku.rag.graph.agui.emitter import AGUIEmitter
+
+
+@dataclass
+class ResearchDeps:
+    """Dependencies for research graph execution."""
+
+    client: HaikuRAG
+    agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None
+    semaphore: asyncio.Semaphore | None = None
+
+    def emit_log(self, message: str, state: "ResearchState | None" = None) -> None:
+        """Emit a log message through AG-UI events.
+
+        Args:
+            message: The message to log
+            state: Optional state to include in state update
+        """
+        if self.agui_emitter:
+            self.agui_emitter.log(message)
+            if state:
+                self.agui_emitter.update_state(state)
+
+
+class ResearchState(BaseModel):
+    """Research graph state model.
+
+    Fully JSON-serializable Pydantic model suitable for AG-UI state synchronization.
+    """
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    context: ResearchContext = Field(
+        description="Shared research context with questions, insights, and gaps"
+    )
+    iterations: int = Field(default=0, description="Current iteration number")
+    max_iterations: int = Field(default=3, description="Maximum allowed iterations")
+    confidence_threshold: float = Field(
+        default=0.8, description="Confidence threshold for completion", ge=0.0, le=1.0
+    )
+    max_concurrency: int = Field(
+        default=1, description="Maximum concurrent search operations", ge=1
+    )
+    last_eval: EvaluationResult | None = Field(
+        default=None, description="Last evaluation result"
+    )
+    last_analysis: InsightAnalysis | None = Field(
+        default=None, description="Last insight analysis"
+    )
+
+    @classmethod
+    def from_config(
+        cls, context: ResearchContext, config: "AppConfig"
+    ) -> "ResearchState":
+        """Create a ResearchState from an AppConfig.
+
+        Args:
+            context: The ResearchContext containing the question and settings
+            config: The AppConfig object (uses config.research for state parameters)
+
+        Returns:
+            A configured ResearchState instance
+        """
+        return cls(
+            context=context,
+            max_iterations=config.research.max_iterations,
+            confidence_threshold=config.research.confidence_threshold,
+            max_concurrency=config.research.max_concurrency,
+        )

haiku/rag/logging.py

@@ -0,0 +1,56 @@
+import logging
+import warnings
+
+from rich.console import Console
+from rich.logging import RichHandler
+
+
+def get_logger() -> logging.Logger:
+    """Return the library logger configured with a Rich handler."""
+    logger = logging.getLogger("haiku.rag")
+
+    handler = RichHandler(
+        console=Console(stderr=True),
+        rich_tracebacks=True,
+    )
+    formatter = logging.Formatter("%(message)s")
+    handler.setFormatter(formatter)
+
+    logger.setLevel(logging.INFO)
+
+    # Remove any existing handlers to avoid duplicates on reconfiguration
+    for hdlr in logger.handlers[:]:
+        logger.removeHandler(hdlr)
+
+    logger.addHandler(handler)
+    # Do not let messages propagate to the root logger
+    logger.propagate = False
+    return logger
+
+
+def configure_cli_logging(level: int = logging.INFO) -> logging.Logger:
+    """Configure logging for CLI runs.
+
+    - Silence ALL non-haiku.rag loggers by detaching root handlers and setting
+      their level to ERROR.
+    - Attach a Rich handler only to the "haiku.rag" logger.
+    - Prevent propagation so only our logger prints in the CLI.
+    """
+    # Silence root logger completely
+    root = logging.getLogger()
+    for hdlr in root.handlers[:]:
+        root.removeHandler(hdlr)
+    root.setLevel(logging.ERROR)
+
+    # Optionally silence some commonly noisy libraries explicitly as a safeguard
+    for noisy in ("httpx", "httpcore", "docling", "urllib3", "asyncio"):
+        logging.getLogger(noisy).setLevel(logging.ERROR)
+        logging.getLogger(noisy).propagate = False
+
+    # Configure and return our app logger
+    logger = get_logger()
+    logger.setLevel(level)
+    logger.propagate = False
+
+    warnings.filterwarnings("ignore")
+    return logger

haiku/rag/mcp.py

@@ -0,0 +1,245 @@
+from pathlib import Path
+from typing import Any
+
+from fastmcp import FastMCP
+from pydantic import BaseModel
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.config import AppConfig, Config
+from haiku.rag.graph.research.models import ResearchReport
+
+
+class SearchResult(BaseModel):
+    document_id: str
+    content: str
+    score: float
+
+
+class DocumentResult(BaseModel):
+    id: str | None
+    content: str
+    uri: str | None = None
+    title: str | None = None
+    metadata: dict[str, Any] = {}
+    created_at: str
+    updated_at: str
+
+
+def create_mcp_server(db_path: Path, config: AppConfig = Config) -> FastMCP:
+    """Create an MCP server with the specified database path."""
+    mcp = FastMCP("haiku-rag")
+
+    @mcp.tool()
+    async def add_document_from_file(
+        file_path: str,
+        metadata: dict[str, Any] | None = None,
+        title: str | None = None,
+    ) -> str | None:
+        """Add a document to the RAG system from a file path."""
+        try:
+            async with HaikuRAG(db_path, config=config) as rag:
+                result = await rag.create_document_from_source(
+                    Path(file_path), title=title, metadata=metadata or {}
+                )
+                # Handle both single document and list of documents (directories)
+                if isinstance(result, list):
+                    return result[0].id if result else None
+                return result.id
+        except Exception:
+            return None
+
+    @mcp.tool()
+    async def add_document_from_url(
+        url: str, metadata: dict[str, Any] | None = None, title: str | None = None
+    ) -> str | None:
+        """Add a document to the RAG system from a URL."""
+        try:
+            async with HaikuRAG(db_path, config=config) as rag:
+                result = await rag.create_document_from_source(
+                    url, title=title, metadata=metadata or {}
+                )
+                # Handle both single document and list of documents
+                if isinstance(result, list):
+                    return result[0].id if result else None
+                return result.id
+        except Exception:
+            return None
+
+    @mcp.tool()
+    async def add_document_from_text(
+        content: str,
+        uri: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        title: str | None = None,
+    ) -> str | None:
+        """Add a document to the RAG system from text content."""
+        try:
+            async with HaikuRAG(db_path, config=config) as rag:
+                document = await rag.create_document(
+                    content, uri, title=title, metadata=metadata or {}
+                )
+                return document.id
+        except Exception:
+            return None
+
+    @mcp.tool()
+    async def search_documents(query: str, limit: int = 5) -> list[SearchResult]:
+        """Search the RAG system for documents using hybrid search (vector similarity + full-text search)."""
+        try:
+            async with HaikuRAG(db_path, config=config) as rag:
+                results = await rag.search(query, limit)
+
+                search_results = []
+                for chunk, score in results:
+                    assert chunk.document_id is not None, (
+                        "Chunk document_id should not be None in search results"
+                    )
+                    search_results.append(
+                        SearchResult(
+                            document_id=chunk.document_id,
+                            content=chunk.content,
+                            score=score,
+                        )
+                    )
+
+                return search_results
+        except Exception:
+            return []
+
+    @mcp.tool()
+    async def get_document(document_id: str) -> DocumentResult | None:
+        """Get a document by its ID."""
+        try:
+            async with HaikuRAG(db_path, config=config) as rag:
+                document = await rag.get_document_by_id(document_id)
+
+                if document is None:
+                    return None
+
+                return DocumentResult(
+                    id=document.id,
+                    content=document.content,
+                    uri=document.uri,
+                    title=document.title,
+                    metadata=document.metadata,
+                    created_at=str(document.created_at),
+                    updated_at=str(document.updated_at),
+                )
+        except Exception:
+            return None
+
+    @mcp.tool()
+    async def list_documents(
+        limit: int | None = None,
+        offset: int | None = None,
+        filter: str | None = None,
+    ) -> list[DocumentResult]:
+        """List all documents with optional pagination and filtering.
+
+        Args:
+            limit: Maximum number of documents to return.
+            offset: Number of documents to skip.
+            filter: Optional SQL WHERE clause to filter documents.
+
+        Returns:
+            List of DocumentResult instances matching the criteria.
+        """
+        try:
+            async with HaikuRAG(db_path, config=config) as rag:
+                documents = await rag.list_documents(limit, offset, filter)
+
+                return [
+                    DocumentResult(
+                        id=doc.id,
+                        content=doc.content,
+                        uri=doc.uri,
+                        title=doc.title,
+                        metadata=doc.metadata,
+                        created_at=str(doc.created_at),
+                        updated_at=str(doc.updated_at),
+                    )
+                    for doc in documents
+                ]
+        except Exception:
+            return []
+
+    @mcp.tool()
+    async def delete_document(document_id: str) -> bool:
+        """Delete a document by its ID."""
+        try:
+            async with HaikuRAG(db_path, config=config) as rag:
+                return await rag.delete_document(document_id)
+        except Exception:
+            return False
+
+    @mcp.tool()
+    async def ask_question(
+        question: str,
+        cite: bool = False,
+        deep: bool = False,
+    ) -> str:
+        """Ask a question using the QA agent.
+
+        Args:
+            question: The question to ask.
+            cite: Whether to include citations in the response.
+            deep: Use deep multi-agent QA for complex questions that require decomposition.
+
+        Returns:
+            The answer as a string.
+        """
+        try:
+            async with HaikuRAG(db_path, config=config) as rag:
+                if deep:
+                    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
+
+                    graph = build_deep_qa_graph(config=config)
+                    context = DeepQAContext(
+                        original_question=question, use_citations=cite
+                    )
+                    state = DeepQAState.from_config(context=context, config=config)
+                    deps = DeepQADeps(client=rag)
+
+                    result = await graph.run(state=state, deps=deps)
+                    answer = result.answer
+                else:
+                    answer = await rag.ask(question, cite=cite)
+                return answer
+        except Exception as e:
+            return f"Error answering question: {e!s}"
+
+    @mcp.tool()
+    async def research_question(
+        question: str,
+    ) -> ResearchReport | None:
+        """Run multi-agent research to investigate a complex question.
+
+        The research process uses multiple agents to plan, search, evaluate, and synthesize
+        information iteratively until confidence threshold is met or max iterations reached.
+
+        Args:
+            question: The research question to investigate.
+
+        Returns:
+            A research report with findings, or None if an error occurred.
+        """
+        try:
+            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
+
+            async with HaikuRAG(db_path, config=config) as rag:
+                graph = build_research_graph(config=config)
+                context = ResearchContext(original_question=question)
+                state = ResearchState.from_config(context=context, config=config)
+                deps = ResearchDeps(client=rag)
+
+                result = await graph.run(state=state, deps=deps)
+
+                return result
+        except Exception:
+            return None
+
+    return mcp