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

haiku/rag/monitor.py

@@ -0,0 +1,194 @@
+import asyncio
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import pathspec
+from pathspec.patterns.gitwildmatch import GitWildMatchPattern
+from watchfiles import Change, DefaultFilter, awatch
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.config import AppConfig, Config
+from haiku.rag.store.models.document import Document
+
+if TYPE_CHECKING:
+    pass
+
+logger = logging.getLogger(__name__)
+
+
+class FileFilter(DefaultFilter):
+    def __init__(
+        self,
+        *,
+        ignore_patterns: list[str] | None = None,
+        include_patterns: list[str] | None = None,
+    ) -> None:
+        # Lazy import to avoid loading docling
+        from haiku.rag.reader import FileReader
+
+        self.extensions = tuple(FileReader.extensions)
+        self.ignore_spec = (
+            pathspec.PathSpec.from_lines(GitWildMatchPattern, ignore_patterns)
+            if ignore_patterns
+            else None
+        )
+        self.include_spec = (
+            pathspec.PathSpec.from_lines(GitWildMatchPattern, include_patterns)
+            if include_patterns
+            else None
+        )
+        super().__init__()
+
+    def __call__(self, change: Change, path: str) -> bool:
+        if not self.include_file(path):
+            return False
+
+        # Apply default watchfiles filter
+        return super().__call__(change, path)
+
+    def include_file(self, path: str) -> bool:
+        """Check if a file should be included based on filters."""
+        # Check extension filter
+        if not path.endswith(self.extensions):
+            return False
+
+        # Apply include patterns if specified (whitelist mode)
+        if self.include_spec:
+            if not self.include_spec.match_file(path):
+                return False
+
+        # Apply ignore patterns (blacklist mode)
+        if self.ignore_spec:
+            if self.ignore_spec.match_file(path):
+                return False
+
+        return True
+
+
+class FileWatcher:
+    def __init__(
+        self,
+        client: HaikuRAG,
+        config: AppConfig = Config,
+    ):
+        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
+
+    async def observe(self):
+        logger.info(f"Watching files in {self.paths}")
+        filter = FileFilter(
+            ignore_patterns=self.ignore_patterns, include_patterns=self.include_patterns
+        )
+        await self.refresh()
+
+        async for changes in awatch(*self.paths, watch_filter=filter):
+            await self.handler(changes)
+
+    async def handler(self, changes: set[tuple[Change, str]]):
+        for change, path in changes:
+            if change == Change.added or change == Change.modified:
+                await self._upsert_document(Path(path))
+            elif change == Change.deleted:
+                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")
+            asyncio.create_task(self._delete_orphans())
+
+        # Create filter to apply same logic as observe()
+        filter = FileFilter(
+            ignore_patterns=self.ignore_patterns, include_patterns=self.include_patterns
+        )
+
+        for path in self.paths:
+            for f in Path(path).rglob("**/*"):
+                if f.is_file() and f.suffix in FileReader.extensions:
+                    # Apply pattern filters
+                    if filter(Change.added, str(f)):
+                        await self._upsert_document(f)
+
+    async def _upsert_document(self, file: Path) -> Document | None:
+        try:
+            uri = file.as_uri()
+            existing_doc = await self.client.get_document_by_uri(uri)
+
+            result = await self.client.create_document_from_source(str(file))
+            doc = result if isinstance(result, Document) else result[0]
+
+            if existing_doc:
+                # Check if document was actually updated by comparing updated_at timestamps
+                if doc.updated_at > existing_doc.updated_at:
+                    logger.info(f"Updated document {existing_doc.id} from {file}")
+                else:
+                    logger.info(
+                        f"Skipped unchanged document {existing_doc.id} from {file}"
+                    )
+            else:
+                logger.info(f"Created new document {doc.id} from {file}")
+
+            return doc
+        except Exception as e:
+            logger.error(f"Failed to upsert document from {file}: {e}")
+            return None
+
+    async def _delete_orphans(self):
+        """Delete documents whose source files no longer exist."""
+        try:
+            from urllib.parse import unquote, urlparse
+
+            # Create filter to apply same include/exclude logic
+            filter = FileFilter(
+                ignore_patterns=self.ignore_patterns,
+                include_patterns=self.include_patterns,
+            )
+
+            all_docs = await self.client.list_documents()
+
+            for doc in all_docs:
+                if not doc.uri or not doc.id:
+                    continue
+
+                # Only check file:// URIs
+                parsed = urlparse(doc.uri)
+                if parsed.scheme != "file":
+                    continue
+
+                # Convert URI to Path, decoding URL-encoded characters (like %20 for spaces)
+                file_path = Path(unquote(parsed.path))
+
+                # Check if file exists
+                if not file_path.exists():
+                    # Check if file is within monitored directories
+                    is_monitored = any(
+                        file_path.is_relative_to(monitored_path)
+                        for monitored_path in self.paths
+                    )
+
+                    # Check if file would have been included by filters
+                    if is_monitored and filter.include_file(str(file_path)):
+                        await self.client.delete_document(doc.id)
+                        logger.info(
+                            f"Deleted orphaned document {doc.id} for {file_path}"
+                        )
+        except Exception as e:
+            logger.error(f"Failed to delete orphaned documents: {e}")
+
+    async def _delete_document(self, file: Path):
+        try:
+            uri = file.as_uri()
+            existing_doc = await self.client.get_document_by_uri(uri)
+
+            if existing_doc and existing_doc.id:
+                await self.client.delete_document(existing_doc.id)
+                logger.info(f"Deleted document {existing_doc.id} for {file}")
+        except Exception as e:
+            logger.error(f"Failed to delete document for {file}: {e}")

haiku/rag/qa/__init__.py

@@ -0,0 +1,33 @@
+from haiku.rag.client import HaikuRAG
+from haiku.rag.config import AppConfig, Config
+from haiku.rag.qa.agent import QuestionAnswerAgent
+
+
+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.
+
+    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.
+
+    Returns:
+        A configured QuestionAnswerAgent instance.
+    """
+    provider = config.qa.provider
+    model_name = config.qa.model
+
+    return QuestionAnswerAgent(
+        client=client,
+        provider=provider,
+        model=model_name,
+        use_citations=use_citations,
+        system_prompt=system_prompt,
+    )

haiku/rag/qa/agent.py

@@ -0,0 +1,93 @@
+from pydantic import BaseModel, Field
+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 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"
+    )
+
+
+class Dependencies(BaseModel):
+    model_config = {"arbitrary_types_allowed": True}
+    client: HaikuRAG
+
+
+class QuestionAnswerAgent:
+    def __init__(
+        self,
+        client: HaikuRAG,
+        provider: str,
+        model: str,
+        use_citations: bool = False,
+        q: float = 0.0,
+        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)
+
+        self._agent = Agent(
+            model=model_obj,
+            deps_type=Dependencies,
+            system_prompt=system_prompt,
+            retries=3,
+        )
+
+        @self._agent.tool
+        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)
+
+            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"
+                ),
+            )
+        else:
+            # For all other providers, use the provider:model format
+            return f"{provider}:{model}"
+
+    async def answer(self, question: str) -> str:
+        """Answer a question using the RAG system."""
+        deps = Dependencies(client=self._client)
+        result = await self._agent.run(question, deps=deps)
+        return result.output

haiku/rag/qa/prompts.py

@@ -0,0 +1,60 @@
+QA_SYSTEM_PROMPT = """
+You are a knowledgeable assistant that helps users find information from a document knowledge base.
+
+Your process:
+1. When a user asks a question, use the search_documents tool to find relevant information
+2. Search with specific keywords and phrases from the user's question
+3. Review the search results and their relevance scores
+4. If you need additional context, perform follow-up searches with different keywords
+5. Provide a short and to the point comprehensive answer based only on the retrieved documents
+
+Guidelines:
+- Base your answers strictly on the provided document content
+- Quote or reference specific information when possible
+- If multiple documents contain relevant information, synthesize them coherently
+- Indicate when information is incomplete or when you need to search for additional context
+- If the retrieved documents don't contain sufficient information, clearly state: "I cannot find enough information in the knowledge base to answer this question."
+- For complex questions, consider breaking them down and performing multiple searches
+- Stick to the answer, do not ellaborate or provide context unless explicitly asked for it.
+
+Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents.
+/no_think
+"""
+
+QA_SYSTEM_PROMPT_WITH_CITATIONS = """
+You are a knowledgeable assistant that helps users find information from a document knowledge base.
+
+IMPORTANT: You MUST use the search_documents tool for every question. Do not answer any question without first searching the knowledge base.
+
+Your process:
+1. IMMEDIATELY call the search_documents tool with relevant keywords from the user's question
+2. Review the search results and their relevance scores
+3. If you need additional context, perform follow-up searches with different keywords
+4. Provide a short and to the point comprehensive answer based only on the retrieved documents
+5. Always include citations for the sources used in your answer
+
+Guidelines:
+- Base your answers strictly on the provided document content
+- If multiple documents contain relevant information, synthesize them coherently
+- Indicate when information is incomplete or when you need to search for additional context
+- If the retrieved documents don't contain sufficient information, clearly state: "I cannot find enough information in the knowledge base to answer this question."
+- For complex questions, consider breaking them down and performing multiple searches
+- Stick to the answer, do not ellaborate or provide context unless explicitly asked for it.
+- ALWAYS include citations at the end of your response using the format below
+
+Citation Format:
+After your answer, include a "Citations:" section that lists:
+- The document title (if available) or URI from each search result used
+- A brief excerpt (first 50-100 characters) of the content that supported your answer
+- Format: "Citations:\n- [document title or URI]: [content_excerpt]..."
+
+Example response format:
+[Your answer here]
+
+Citations:
+- /path/to/document1.pdf: "This document explains that AFMAN stands for Air Force Manual..."
+- /path/to/document2.pdf: "The manual provides guidance on military procedures and..."
+
+Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents.
+/no_think
+"""

haiku/rag/reader.py

@@ -0,0 +1,135 @@
+from pathlib import Path
+from typing import ClassVar
+
+from docling_core.types.doc.document import DoclingDocument
+
+from haiku.rag.utils import text_to_docling_document
+
+# Check if docling is available
+try:
+    import docling  # noqa: F401
+
+    DOCLING_AVAILABLE = True
+except ImportError:
+    DOCLING_AVAILABLE = False
+
+
+class FileReader:
+    # Extensions supported by docling
+    docling_extensions: ClassVar[list[str]] = [
+        ".adoc",
+        ".asc",
+        ".asciidoc",
+        ".bmp",
+        ".csv",
+        ".docx",
+        ".html",
+        ".xhtml",
+        ".jpeg",
+        ".jpg",
+        ".md",
+        ".pdf",
+        ".png",
+        ".pptx",
+        ".tiff",
+        ".xlsx",
+        ".xml",
+        ".webp",
+    ]
+
+    # Plain text extensions that we'll read directly
+    text_extensions: ClassVar[list[str]] = [
+        ".astro",
+        ".c",
+        ".cpp",
+        ".css",
+        ".go",
+        ".h",
+        ".hpp",
+        ".java",
+        ".js",
+        ".json",
+        ".kt",
+        ".mdx",
+        ".mjs",
+        ".php",
+        ".py",
+        ".rb",
+        ".rs",
+        ".svelte",
+        ".swift",
+        ".ts",
+        ".tsx",
+        ".txt",
+        ".vue",
+        ".yaml",
+        ".yml",
+    ]
+
+    # Code file extensions with their markdown language identifiers for syntax highlighting
+    code_markdown_identifier: ClassVar[dict[str, str]] = {
+        ".astro": "astro",
+        ".c": "c",
+        ".cpp": "cpp",
+        ".css": "css",
+        ".go": "go",
+        ".h": "c",
+        ".hpp": "cpp",
+        ".java": "java",
+        ".js": "javascript",
+        ".json": "json",
+        ".kt": "kotlin",
+        ".mjs": "javascript",
+        ".php": "php",
+        ".py": "python",
+        ".rb": "ruby",
+        ".rs": "rust",
+        ".svelte": "svelte",
+        ".swift": "swift",
+        ".ts": "typescript",
+        ".tsx": "tsx",
+        ".vue": "vue",
+        ".yaml": "yaml",
+        ".yml": "yaml",
+    }
+
+    extensions: ClassVar[list[str]] = docling_extensions + text_extensions
+
+    @staticmethod
+    def parse_file(path: Path) -> DoclingDocument:
+        try:
+            file_extension = path.suffix.lower()
+
+            if file_extension in FileReader.docling_extensions:
+                # Use docling for complex document formats
+                if not DOCLING_AVAILABLE:
+                    raise ImportError(
+                        "Docling is required for processing this file type. "
+                        "Install with: pip install haiku.rag-slim[docling]"
+                    )
+                from docling.document_converter import DocumentConverter
+
+                converter = DocumentConverter()
+                result = converter.convert(path)
+                return result.document
+            elif file_extension in FileReader.text_extensions:
+                # Read plain text files directly
+                content = path.read_text(encoding="utf-8")
+
+                # Wrap code files (but not plain txt) in markdown code blocks for better presentation
+                if file_extension in FileReader.code_markdown_identifier:
+                    language = FileReader.code_markdown_identifier[file_extension]
+                    content = f"```{language}\n{content}\n```"
+
+                # Convert text to DoclingDocument by wrapping as markdown
+                return text_to_docling_document(content, name=f"{path.stem}.md")
+            else:
+                # Fallback: try to read as text and convert to DoclingDocument
+                content = path.read_text(encoding="utf-8")
+                return text_to_docling_document(content, name=f"{path.stem}.md")
+        except ImportError:
+            raise
+        except Exception as e:
+            raise ValueError(
+                f"Failed to parse file: {path} - {type(e).__name__}: {e}"
+            ) from e

haiku/rag/reranking/__init__.py

@@ -0,0 +1,63 @@
+import os
+
+from haiku.rag.config import AppConfig, Config
+from haiku.rag.reranking.base import RerankerBase
+
+_reranker_cache: dict[int, RerankerBase | None] = {}
+
+
+def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
+    """
+    Factory function to get the appropriate reranker based on the configuration.
+    Returns None if reranking is disabled.
+
+    Args:
+        config: Configuration to use. Defaults to global Config.
+
+    Returns:
+        A reranker instance if configured, None otherwise.
+    """
+    # Use config id as cache key to support multiple configs
+    config_id = id(config)
+    if config_id in _reranker_cache:
+        return _reranker_cache[config_id]
+
+    reranker: RerankerBase | None = None
+
+    if config.reranking.provider == "mxbai":
+        try:
+            from haiku.rag.reranking.mxbai import MxBAIReranker
+
+            os.environ["TOKENIZERS_PARALLELISM"] = "true"
+            reranker = MxBAIReranker()
+        except ImportError:
+            reranker = None
+
+    elif config.reranking.provider == "cohere":
+        try:
+            from haiku.rag.reranking.cohere import CohereReranker
+
+            reranker = CohereReranker()
+        except ImportError:
+            reranker = None
+
+    elif config.reranking.provider == "vllm":
+        try:
+            from haiku.rag.reranking.vllm import VLLMReranker
+
+            reranker = VLLMReranker(config.reranking.model)
+        except ImportError:
+            reranker = None
+
+    elif config.reranking.provider == "zeroentropy":
+        try:
+            from haiku.rag.reranking.zeroentropy import ZeroEntropyReranker
+
+            # Use configured model or default to zerank-1
+            model = config.reranking.model or "zerank-1"
+            reranker = ZeroEntropyReranker(model)
+        except ImportError:
+            reranker = None
+
+    _reranker_cache[config_id] = reranker
+    return reranker

haiku/rag/reranking/base.py

@@ -0,0 +1,13 @@
+from haiku.rag.config import Config
+from haiku.rag.store.models.chunk import Chunk
+
+
+class RerankerBase:
+    _model: str = Config.reranking.model
+
+    async def rerank(
+        self, query: str, chunks: list[Chunk], top_n: int = 10
+    ) -> list[tuple[Chunk, float]]:
+        raise NotImplementedError(
+            "Reranker is an abstract class. Please implement the rerank method in a subclass."
+        )

haiku/rag/reranking/cohere.py

@@ -0,0 +1,34 @@
+from haiku.rag.reranking.base import RerankerBase
+from haiku.rag.store.models.chunk import Chunk
+
+try:
+    import cohere
+except ImportError as e:
+    raise ImportError(
+        "cohere is not installed. Please install it with `pip install cohere` or use the cohere optional dependency."
+    ) from e
+
+
+class CohereReranker(RerankerBase):
+    def __init__(self):
+        # Cohere SDK reads CO_API_KEY from environment by default
+        self._client = cohere.ClientV2()
+
+    async def rerank(
+        self, query: str, chunks: list[Chunk], top_n: int = 10
+    ) -> list[tuple[Chunk, float]]:
+        if not chunks:
+            return []
+
+        documents = [chunk.content for chunk in chunks]
+
+        response = self._client.rerank(
+            model=self._model, query=query, documents=documents, top_n=top_n
+        )
+
+        reranked_chunks = []
+        for result in response.results:
+            original_chunk = chunks[result.index]
+            reranked_chunks.append((original_chunk, result.relevance_score))
+
+        return reranked_chunks

haiku/rag/reranking/mxbai.py

@@ -0,0 +1,28 @@
+from mxbai_rerank import MxbaiRerankV2  # pyright: ignore[reportMissingImports]
+
+from haiku.rag.config import Config
+from haiku.rag.reranking.base import RerankerBase
+from haiku.rag.store.models.chunk import Chunk
+
+
+class MxBAIReranker(RerankerBase):
+    def __init__(self):
+        self._client = MxbaiRerankV2(
+            Config.reranking.model, disable_transformers_warnings=True
+        )
+
+    async def rerank(
+        self, query: str, chunks: list[Chunk], top_n: int = 10
+    ) -> list[tuple[Chunk, float]]:
+        if not chunks:
+            return []
+
+        documents = [chunk.content for chunk in chunks]
+
+        results = self._client.rank(query=query, documents=documents, top_k=top_n)
+        reranked_chunks = []
+        for result in results:
+            original_chunk = chunks[result.index]
+            reranked_chunks.append((original_chunk, result.score))
+
+        return reranked_chunks