haiku.rag 0.10.2__py3-none-any.whl → 0.14.0__py3-none-any.whl

haiku/rag/monitor.py

@@ -1,73 +0,0 @@
-from pathlib import Path
-
-from watchfiles import Change, DefaultFilter, awatch
-
-from haiku.rag.client import HaikuRAG
-from haiku.rag.logging import get_logger
-from haiku.rag.reader import FileReader
-from haiku.rag.store.models.document import Document
-
-logger = get_logger()
-
-
-class FileFilter(DefaultFilter):
-    def __init__(self, *, ignore_paths: list[Path] | None = None) -> None:
-        self.extensions = tuple(FileReader.extensions)
-        super().__init__(ignore_paths=ignore_paths)
-
-    def __call__(self, change: "Change", path: str) -> bool:
-        return path.endswith(self.extensions) and super().__call__(change, path)
-
-
-class FileWatcher:
-    def __init__(self, paths: list[Path], client: HaikuRAG):
-        self.paths = paths
-        self.client = client
-
-    async def observe(self):
-        logger.info(f"Watching files in {self.paths}")
-        filter = FileFilter()
-        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):
-        for path in self.paths:
-            for f in Path(path).rglob("**/*"):
-                if f.is_file() and f.suffix in FileReader.extensions:
-                    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)
-            if existing_doc:
-                doc = await self.client.create_document_from_source(str(file))
-                logger.info(f"Updated document {existing_doc.id} from {file}")
-                return doc
-            else:
-                doc = await self.client.create_document_from_source(str(file))
-                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_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

@@ -1,15 +0,0 @@
-from haiku.rag.client import HaikuRAG
-from haiku.rag.config import Config
-from haiku.rag.qa.agent import QuestionAnswerAgent
-
-
-def get_qa_agent(client: HaikuRAG, use_citations: bool = False) -> QuestionAnswerAgent:
-    provider = Config.QA_PROVIDER
-    model_name = Config.QA_MODEL
-
-    return QuestionAnswerAgent(
-        client=client,
-        provider=provider,
-        model=model_name,
-        use_citations=use_citations,
-    )

haiku/rag/qa/agent.py

@@ -1,91 +0,0 @@
-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,
-    ):
-        self._client = client
-
-        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,
-        )
-
-        @self._agent.tool
-        async def search_documents(
-            ctx: RunContext[Dependencies],
-            query: str,
-            limit: int = 3,
-        ) -> list[SearchResult]:
-            """Search the knowledge base for relevant documents."""
-
-            # Remove quotes from queries as this requires positional indexing in lancedb
-            query = query.replace('"', "")
-            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.OLLAMA_BASE_URL}/v1"),
-            )
-        elif provider == "vllm":
-            return OpenAIChatModel(
-                model_name=model,
-                provider=OpenAIProvider(
-                    base_url=f"{Config.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

@@ -1,60 +0,0 @@
-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

@@ -1,115 +0,0 @@
-from pathlib import Path
-from typing import ClassVar
-
-from docling.document_converter import DocumentConverter
-from docling_core.types.doc.document import DoclingDocument
-
-from haiku.rag.utils import text_to_docling_document
-
-
-class FileReader:
-    # Extensions supported by docling
-    docling_extensions: ClassVar[list[str]] = [
-        ".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
-                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 Exception:
-            raise ValueError(f"Failed to parse file: {path}")

haiku/rag/reranking/__init__.py

@@ -1,34 +0,0 @@
-from haiku.rag.config import Config
-from haiku.rag.reranking.base import RerankerBase
-
-_reranker: RerankerBase | None = None
-
-
-def get_reranker() -> RerankerBase | None:
-    """
-    Factory function to get the appropriate reranker based on the configuration.
-    Returns None if if reranking is disabled.
-    """
-    global _reranker
-    if _reranker is not None:
-        return _reranker
-
-    if Config.RERANK_PROVIDER == "mxbai":
-        try:
-            from haiku.rag.reranking.mxbai import MxBAIReranker
-
-            _reranker = MxBAIReranker()
-            return _reranker
-        except ImportError:
-            return None
-
-    if Config.RERANK_PROVIDER == "cohere":
-        try:
-            from haiku.rag.reranking.cohere import CohereReranker
-
-            _reranker = CohereReranker()
-            return _reranker
-        except ImportError:
-            return None
-
-    return None

haiku/rag/reranking/base.py

@@ -1,13 +0,0 @@
-from haiku.rag.config import Config
-from haiku.rag.store.models.chunk import Chunk
-
-
-class RerankerBase:
-    _model: str = Config.RERANK_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

@@ -1,34 +0,0 @@
-from haiku.rag.config import Config
-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):
-        self._client = cohere.ClientV2(api_key=Config.COHERE_API_KEY)
-
-    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

@@ -1,28 +0,0 @@
-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.RERANK_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

haiku/rag/reranking/vllm.py

@@ -1,44 +0,0 @@
-import httpx
-
-from haiku.rag.config import Config
-from haiku.rag.reranking.base import RerankerBase
-from haiku.rag.store.models.chunk import Chunk
-
-
-class VLLMReranker(RerankerBase):
-    def __init__(self, model: str):
-        self._model = model
-        self._base_url = Config.VLLM_RERANK_BASE_URL
-
-    async def rerank(
-        self, query: str, chunks: list[Chunk], top_n: int = 10
-    ) -> list[tuple[Chunk, float]]:
-        if not chunks:
-            return []
-
-        # Prepare documents for reranking
-        documents = [chunk.content for chunk in chunks]
-
-        async with httpx.AsyncClient() as client:
-            response = await client.post(
-                f"{self._base_url}/v1/rerank",
-                json={"model": self._model, "query": query, "documents": documents},
-                headers={
-                    "accept": "application/json",
-                    "Content-Type": "application/json",
-                },
-            )
-            response.raise_for_status()
-
-            result = response.json()
-
-            # Extract scores and pair with chunks
-            scored_chunks = []
-            for item in result.get("results", []):
-                index = item["index"]
-                score = item["relevance_score"]
-                scored_chunks.append((chunks[index], score))
-
-            # Sort by score (descending) and return top_n
-            scored_chunks.sort(key=lambda x: x[1], reverse=True)
-            return scored_chunks[:top_n]

haiku/rag/research/__init__.py

@@ -1,20 +0,0 @@
-from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies
-from haiku.rag.research.graph import (
-    PlanNode,
-    ResearchDeps,
-    ResearchState,
-    build_research_graph,
-)
-from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer
-
-__all__ = [
-    "ResearchDependencies",
-    "ResearchContext",
-    "SearchAnswer",
-    "EvaluationResult",
-    "ResearchReport",
-    "ResearchDeps",
-    "ResearchState",
-    "PlanNode",
-    "build_research_graph",
-]

haiku/rag/research/common.py

@@ -1,53 +0,0 @@
-from typing import 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
-
-
-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(console, msg: str) -> None:
-    if console:
-        console.print(msg)
-
-
-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": context.insights,
-        "gaps": context.gaps,
-    }
-    return format_as_xml(context_data, root_tag="research_context")

haiku/rag/research/dependencies.py

@@ -1,47 +0,0 @@
-from pydantic import BaseModel, Field
-from rich.console import Console
-
-from haiku.rag.client import HaikuRAG
-from haiku.rag.research.models import SearchAnswer
-
-
-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[str] = Field(
-        default_factory=list, description="Key insights discovered"
-    )
-    gaps: list[str] = 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 add_insight(self, insight: str) -> None:
-        """Add a key insight."""
-        if insight not in self.insights:
-            self.insights.append(insight)
-
-    def add_gap(self, gap: str) -> None:
-        """Identify an information gap."""
-        if gap not in self.gaps:
-            self.gaps.append(gap)
-
-
-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

haiku/rag/research/graph.py

@@ -1,29 +0,0 @@
-from pydantic_graph import Graph
-
-from haiku.rag.research.models import ResearchReport
-from haiku.rag.research.nodes.evaluate import EvaluateNode
-from haiku.rag.research.nodes.plan import PlanNode
-from haiku.rag.research.nodes.search import SearchDispatchNode
-from haiku.rag.research.nodes.synthesize import SynthesizeNode
-from haiku.rag.research.state import ResearchDeps, ResearchState
-
-__all__ = [
-    "PlanNode",
-    "SearchDispatchNode",
-    "EvaluateNode",
-    "SynthesizeNode",
-    "ResearchState",
-    "ResearchDeps",
-    "build_research_graph",
-]
-
-
-def build_research_graph() -> Graph[ResearchState, ResearchDeps, ResearchReport]:
-    return Graph(
-        nodes=[
-            PlanNode,
-            SearchDispatchNode,
-            EvaluateNode,
-            SynthesizeNode,
-        ]
-    )