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

haiku/rag/mcp.py

@@ -7,12 +7,8 @@ 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
+from haiku.rag.store.models import SearchResult
+from haiku.rag.utils import format_citations
 
 
 class DocumentResult(BaseModel):
@@ -25,84 +21,92 @@ class DocumentResult(BaseModel):
     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")
+def create_mcp_server(
+    db_path: Path, config: AppConfig = Config, read_only: bool = False
+) -> FastMCP:
+    """Create an MCP server with the specified database path.
 
-    @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
+    Args:
+        db_path: Path to the database file.
+        config: Configuration to use.
+        read_only: If True, write tools (add_document_*, delete_document) are not registered.
+    """
+    mcp = FastMCP("haiku-rag")
 
+    # Write tools - only registered when not in read-only mode
+    if not read_only:
+
+        @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 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
+
+    # Read tools - always registered
     @mcp.tool()
-    async def search_documents(query: str, limit: int = 5) -> list[SearchResult]:
+    async def search_documents(
+        query: str, limit: int | None = None
+    ) -> 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
+            async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
+                return await rag.search(query, limit=limit)
         except Exception:
             return []
 
@@ -110,7 +114,7 @@ def create_mcp_server(db_path: Path, config: AppConfig = Config) -> FastMCP:
     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:
+            async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
                 document = await rag.get_document_by_id(document_id)
 
                 if document is None:
@@ -145,7 +149,7 @@ def create_mcp_server(db_path: Path, config: AppConfig = Config) -> FastMCP:
             List of DocumentResult instances matching the criteria.
         """
         try:
-            async with HaikuRAG(db_path, config=config) as rag:
+            async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
                 documents = await rag.list_documents(limit, offset, filter)
 
                 return [
@@ -163,15 +167,6 @@ def create_mcp_server(db_path: Path, config: AppConfig = Config) -> FastMCP:
         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,
@@ -189,23 +184,32 @@ def create_mcp_server(db_path: Path, config: AppConfig = Config) -> FastMCP:
             The answer as a string.
         """
         try:
-            async with HaikuRAG(db_path, config=config) as rag:
+            async with HaikuRAG(db_path, config=config, read_only=read_only) 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
+                    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,
+                    )
 
-                    graph = build_deep_qa_graph(config=config)
-                    context = DeepQAContext(
-                        original_question=question, use_citations=cite
+                    graph = build_research_graph(config=config)
+                    context = ResearchContext(original_question=question)
+                    state = ResearchState.from_config(
+                        context=context,
+                        config=config,
+                        max_iterations=2,
+                        confidence_threshold=0.0,
                     )
-                    state = DeepQAState.from_config(context=context, config=config)
-                    deps = DeepQADeps(client=rag)
+                    deps = ResearchDeps(client=rag)
 
                     result = await graph.run(state=state, deps=deps)
-                    answer = result.answer
+                    answer = result.executive_summary
+                    citations = []
                 else:
-                    answer = await rag.ask(question, cite=cite)
+                    answer, citations = await rag.ask(question)
+                if cite and citations:
+                    answer += "\n\n" + format_citations(citations)
                 return answer
         except Exception as e:
             return f"Error answering question: {e!s}"
@@ -230,7 +234,7 @@ def create_mcp_server(db_path: Path, config: AppConfig = Config) -> FastMCP:
             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:
+            async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
                 graph = build_research_graph(config=config)
                 context = ResearchContext(original_question=question)
                 state = ResearchState.from_config(context=context, config=config)

haiku/rag/monitor.py

@@ -23,11 +23,19 @@ class FileFilter(DefaultFilter):
         *,
         ignore_patterns: list[str] | None = None,
         include_patterns: list[str] | None = None,
+        supported_extensions: list[str] | None = None,
     ) -> None:
-        # Lazy import to avoid loading docling
-        from haiku.rag.reader import FileReader
+        if supported_extensions is None:
+            # Default to docling-local extensions if not provided
+            from haiku.rag.converters.docling_local import DoclingLocalConverter
+            from haiku.rag.converters.text_utils import TextFileHandler
+
+            supported_extensions = (
+                DoclingLocalConverter.docling_extensions
+                + TextFileHandler.text_extensions
+            )
 
-        self.extensions = tuple(FileReader.extensions)
+        self.extensions = tuple(supported_extensions)
         self.ignore_spec = (
             pathspec.PathSpec.from_lines(GitWildMatchPattern, ignore_patterns)
             if ignore_patterns
@@ -72,16 +80,33 @@ class FileWatcher:
         client: HaikuRAG,
         config: AppConfig = Config,
     ):
+        from haiku.rag.converters import get_converter
+
         self.paths = config.monitor.directories
         self.client = client
         self.ignore_patterns = config.monitor.ignore_patterns or None
         self.include_patterns = config.monitor.include_patterns or None
         self.delete_orphans = config.monitor.delete_orphans
+        self.supported_extensions = get_converter(config).supported_extensions
 
     async def observe(self):
+        if not self.paths:
+            logger.warning("No directories configured for monitoring")
+            return
+
+        # Validate all paths exist before attempting to watch
+        missing_paths = [p for p in self.paths if not Path(p).exists()]
+        if missing_paths:
+            raise FileNotFoundError(
+                f"Monitor directories do not exist: {missing_paths}. "
+                "Check your haiku.rag.yaml configuration."
+            )
+
         logger.info(f"Watching files in {self.paths}")
         filter = FileFilter(
-            ignore_patterns=self.ignore_patterns, include_patterns=self.include_patterns
+            ignore_patterns=self.ignore_patterns,
+            include_patterns=self.include_patterns,
+            supported_extensions=self.supported_extensions,
         )
         await self.refresh()
 
@@ -96,9 +121,6 @@ class FileWatcher:
                 await self._delete_document(Path(path))
 
     async def refresh(self):
-        # Lazy import to avoid loading docling
-        from haiku.rag.reader import FileReader
-
         # Delete orphaned documents in background if enabled
         if self.delete_orphans:
             logger.info("Starting orphan cleanup in background")
@@ -106,12 +128,14 @@ class FileWatcher:
 
         # Create filter to apply same logic as observe()
         filter = FileFilter(
-            ignore_patterns=self.ignore_patterns, include_patterns=self.include_patterns
+            ignore_patterns=self.ignore_patterns,
+            include_patterns=self.include_patterns,
+            supported_extensions=self.supported_extensions,
         )
 
         for path in self.paths:
             for f in Path(path).rglob("**/*"):
-                if f.is_file() and f.suffix in FileReader.extensions:
+                if f.is_file() and f.suffix in self.supported_extensions:
                     # Apply pattern filters
                     if filter(Change.added, str(f)):
                         await self._upsert_document(f)

haiku/rag/providers/__init__.py

@@ -0,0 +1,5 @@
+"""Provider clients for external services."""
+
+from haiku.rag.providers.docling_serve import DoclingServeClient
+
+__all__ = ["DoclingServeClient"]

haiku/rag/providers/docling_serve.py

@@ -0,0 +1,108 @@
+"""Shared client for docling-serve async API."""
+
+import asyncio
+from typing import Any
+
+import httpx
+
+
+class DoclingServeClient:
+    """Client for docling-serve async workflow.
+
+    Handles the submit → poll → fetch pattern used by both conversion and chunking.
+    """
+
+    def __init__(self, base_url: str, api_key: str | None = None, timeout: float = 300):
+        self.base_url = base_url.rstrip("/")
+        self.api_key = api_key
+        self.timeout = timeout
+
+    def _get_headers(self) -> dict[str, str]:
+        """Get headers for API requests."""
+        headers: dict[str, str] = {}
+        if self.api_key:
+            headers["X-Api-Key"] = self.api_key
+        return headers
+
+    async def submit_and_poll(
+        self,
+        endpoint: str,
+        files: dict[str, Any],
+        data: dict[str, Any],
+        name: str = "document",
+    ) -> dict[str, Any]:
+        """Submit a task and poll until completion.
+
+        Args:
+            endpoint: The async endpoint path (e.g., "/v1/convert/file/async")
+            files: Files to upload
+            data: Form data parameters
+            name: Name for error messages
+
+        Returns:
+            The result dictionary from the completed task
+
+        Raises:
+            ValueError: If the task fails or service is unavailable
+        """
+        headers = self._get_headers()
+
+        try:
+            async with httpx.AsyncClient(timeout=self.timeout) as client:
+                # Submit async task
+                submit_url = f"{self.base_url}{endpoint}"
+                response = await client.post(
+                    submit_url,
+                    files=files,
+                    data=data,
+                    headers=headers,
+                )
+                response.raise_for_status()
+                submit_result = response.json()
+                task_id = submit_result.get("task_id")
+
+                if not task_id:
+                    raise ValueError("docling-serve did not return a task_id")
+
+                # Poll for completion
+                poll_url = f"{self.base_url}/v1/status/poll/{task_id}"
+                while True:
+                    poll_response = await client.get(poll_url, headers=headers)
+                    poll_response.raise_for_status()
+                    poll_result = poll_response.json()
+                    status = poll_result.get("task_status")
+
+                    if status == "success":
+                        break
+                    elif status in ("failure", "error"):
+                        raise ValueError(
+                            f"docling-serve task failed for {name}: {poll_result}"
+                        )
+
+                    await asyncio.sleep(1)
+
+                # Fetch result
+                result_url = f"{self.base_url}/v1/result/{task_id}"
+                result_response = await client.get(result_url, headers=headers)
+                result_response.raise_for_status()
+                return result_response.json()
+
+        except httpx.ConnectError as e:
+            raise ValueError(
+                f"Could not connect to docling-serve at {self.base_url}. "
+                f"Ensure the service is running and accessible. Error: {e}"
+            )
+        except httpx.TimeoutException as e:
+            raise ValueError(
+                f"Request to docling-serve timed out after {self.timeout}s. Error: {e}"
+            )
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == 401:
+                raise ValueError(
+                    "Authentication failed. Check your API key configuration."
+                )
+            raise ValueError(f"HTTP error from docling-serve: {e}")
+        except ValueError:
+            raise
+        except Exception as e:
+            raise ValueError(f"Failed to process via docling-serve: {e}")

haiku/rag/qa/__init__.py

@@ -1,33 +1,35 @@
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import AppConfig, Config
 from haiku.rag.qa.agent import QuestionAnswerAgent
+from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT
+from haiku.rag.utils import build_prompt
 
 
 def get_qa_agent(
     client: HaikuRAG,
     config: AppConfig = Config,
-    use_citations: bool = False,
     system_prompt: str | None = None,
 ) -> QuestionAnswerAgent:
-    """
-    Factory function to get a QA agent based on the configuration.
+    """Factory function to get a QA agent based on the configuration.
 
     Args:
         client: HaikuRAG client instance.
         config: Configuration to use. Defaults to global Config.
-        use_citations: Whether to include citations in responses.
-        system_prompt: Optional custom system prompt.
+        system_prompt: Optional custom system prompt (overrides config).
 
     Returns:
         A configured QuestionAnswerAgent instance.
     """
-    provider = config.qa.provider
-    model_name = config.qa.model
+    # Determine the base prompt: explicit > config > default
+    if system_prompt is None:
+        system_prompt = config.prompts.qa or QA_SYSTEM_PROMPT
+
+    # Prepend system_context if configured
+    system_prompt = build_prompt(system_prompt, config)
 
     return QuestionAnswerAgent(
         client=client,
-        provider=provider,
-        model=model_name,
-        use_citations=use_citations,
+        model_config=config.qa.model,
+        config=config,
         system_prompt=system_prompt,
     )

haiku/rag/qa/agent.py

@@ -1,49 +1,38 @@
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
 from pydantic_ai import Agent, RunContext
-from pydantic_ai.models.openai import OpenAIChatModel
-from pydantic_ai.providers.ollama import OllamaProvider
-from pydantic_ai.providers.openai import OpenAIProvider
+from pydantic_ai.output import ToolOutput
 
 from haiku.rag.client import HaikuRAG
-from haiku.rag.config import Config
-from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT, QA_SYSTEM_PROMPT_WITH_CITATIONS
-
-
-class SearchResult(BaseModel):
-    content: str = Field(description="The document text content")
-    score: float = Field(description="Relevance score (higher is more relevant)")
-    document_uri: str = Field(
-        description="Source title (if available) or URI/path of the document"
-    )
+from haiku.rag.config.models import AppConfig, ModelConfig
+from haiku.rag.graph.research.models import Citation, RawSearchAnswer, resolve_citations
+from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT
+from haiku.rag.store.models import SearchResult
+from haiku.rag.utils import get_model
 
 
 class Dependencies(BaseModel):
     model_config = {"arbitrary_types_allowed": True}
     client: HaikuRAG
+    search_results: list[SearchResult] = []
+    search_filter: str | None = None
 
 
 class QuestionAnswerAgent:
     def __init__(
         self,
         client: HaikuRAG,
-        provider: str,
-        model: str,
-        use_citations: bool = False,
-        q: float = 0.0,
+        model_config: ModelConfig,
+        config: AppConfig | None = None,
         system_prompt: str | None = None,
     ):
         self._client = client
-
-        if system_prompt is None:
-            system_prompt = (
-                QA_SYSTEM_PROMPT_WITH_CITATIONS if use_citations else QA_SYSTEM_PROMPT
-            )
-        model_obj = self._get_model(provider, model)
+        model_obj = get_model(model_config, config)
 
         self._agent = Agent(
             model=model_obj,
             deps_type=Dependencies,
-            system_prompt=system_prompt,
+            output_type=ToolOutput(RawSearchAnswer, max_retries=3),
+            instructions=system_prompt or QA_SYSTEM_PROMPT,
             retries=3,
         )
 
@@ -51,43 +40,36 @@ class QuestionAnswerAgent:
         async def search_documents(
             ctx: RunContext[Dependencies],
             query: str,
-            limit: int = 3,
-        ) -> list[SearchResult]:
-            """Search the knowledge base for relevant documents."""
-            search_results = await ctx.deps.client.search(query, limit=limit)
-            expanded_results = await ctx.deps.client.expand_context(search_results)
+            limit: int | None = None,
+        ) -> str:
+            """Search the knowledge base for relevant documents.
 
-            return [
-                SearchResult(
-                    content=chunk.content,
-                    score=score,
-                    document_uri=(chunk.document_title or chunk.document_uri or ""),
-                )
-                for chunk, score in expanded_results
-            ]
-
-    def _get_model(self, provider: str, model: str):
-        """Get the appropriate model object for the provider."""
-        if provider == "ollama":
-            return OpenAIChatModel(
-                model_name=model,
-                provider=OllamaProvider(
-                    base_url=f"{Config.providers.ollama.base_url}/v1"
-                ),
-            )
-        elif provider == "vllm":
-            return OpenAIChatModel(
-                model_name=model,
-                provider=OpenAIProvider(
-                    base_url=f"{Config.providers.vllm.qa_base_url}/v1", api_key="none"
-                ),
+            Returns results with chunk IDs and relevance scores.
+            Reference results by their chunk_id in cited_chunks.
+            """
+            results = await ctx.deps.client.search(
+                query, limit=limit, filter=ctx.deps.search_filter
             )
-        else:
-            # For all other providers, use the provider:model format
-            return f"{provider}:{model}"
+            results = await ctx.deps.client.expand_context(results)
+            # Store results for citation resolution
+            ctx.deps.search_results = results
+            # Format with metadata for agent context
+            parts = [r.format_for_agent() for r in results]
+            return "\n\n".join(parts) if parts else "No results found."
+
+    async def answer(
+        self, question: str, filter: str | None = None
+    ) -> tuple[str, list[Citation]]:
+        """Answer a question using the RAG system.
+
+        Args:
+            question: The question to answer
+            filter: SQL WHERE clause to filter documents
 
-    async def answer(self, question: str) -> str:
-        """Answer a question using the RAG system."""
-        deps = Dependencies(client=self._client)
+        Returns:
+            Tuple of (answer text, list of resolved citations)
+        """
+        deps = Dependencies(client=self._client, search_filter=filter)
         result = await self._agent.run(question, deps=deps)
-        return result.output
+        citations = resolve_citations(result.output.cited_chunks, deps.search_results)
+        return result.output.answer, citations