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

haiku/rag/qa/agent.py

@@ -1,89 +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 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_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 URI from each search result used
-- A brief excerpt (first 50-100 characters) of the content that supported your answer
-- Format: "Citations:\n- [document_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
-
-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,37 +0,0 @@
-"""Multi-agent research workflow for advanced RAG queries."""
-
-from haiku.rag.research.base import (
-    BaseResearchAgent,
-    ResearchOutput,
-    SearchAnswer,
-    SearchResult,
-)
-from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies
-from haiku.rag.research.evaluation_agent import (
-    AnalysisEvaluationAgent,
-    EvaluationResult,
-)
-from haiku.rag.research.orchestrator import ResearchOrchestrator, ResearchPlan
-from haiku.rag.research.presearch_agent import PresearchSurveyAgent
-from haiku.rag.research.search_agent import SearchSpecialistAgent
-from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent
-
-__all__ = [
-    # Base classes
-    "BaseResearchAgent",
-    "ResearchDependencies",
-    "ResearchContext",
-    "SearchResult",
-    "ResearchOutput",
-    # Specialized agents
-    "SearchAnswer",
-    "SearchSpecialistAgent",
-    "PresearchSurveyAgent",
-    "AnalysisEvaluationAgent",
-    "EvaluationResult",
-    "SynthesisAgent",
-    "ResearchReport",
-    # Orchestrator
-    "ResearchOrchestrator",
-    "ResearchPlan",
-]

haiku/rag/research/base.py

@@ -1,130 +0,0 @@
-from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any
-
-from pydantic import BaseModel, Field
-from pydantic_ai import Agent
-from pydantic_ai.models.openai import OpenAIChatModel
-from pydantic_ai.output import ToolOutput
-from pydantic_ai.providers.ollama import OllamaProvider
-from pydantic_ai.providers.openai import OpenAIProvider
-from pydantic_ai.run import AgentRunResult
-
-from haiku.rag.config import Config
-
-if TYPE_CHECKING:
-    from haiku.rag.research.dependencies import ResearchDependencies
-
-
-class BaseResearchAgent[T](ABC):
-    """Base class for all research agents."""
-
-    def __init__(
-        self,
-        provider: str,
-        model: str,
-        output_type: type[T],
-    ):
-        self.provider = provider
-        self.model = model
-        self.output_type = output_type
-
-        model_obj = self._get_model(provider, model)
-
-        # Import deps type lazily to avoid circular import during module load
-        from haiku.rag.research.dependencies import ResearchDependencies
-
-        # If the agent is expected to return plain text, pass `str` directly.
-        # Otherwise, wrap the model with ToolOutput for robust tool-handling retries.
-        agent_output_type: Any
-        if self.output_type is str:  # plain text output
-            agent_output_type = str
-        else:
-            agent_output_type = ToolOutput(self.output_type, max_retries=3)
-
-        self._agent = Agent(
-            model=model_obj,
-            deps_type=ResearchDependencies,
-            output_type=agent_output_type,
-            system_prompt=self.get_system_prompt(),
-        )
-
-        # Register tools
-        self.register_tools()
-
-    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_RESEARCH_BASE_URL or Config.VLLM_QA_BASE_URL}/v1",
-                    api_key="none",
-                ),
-            )
-        else:
-            # For all other providers, use the provider:model format
-            return f"{provider}:{model}"
-
-    @abstractmethod
-    def get_system_prompt(self) -> str:
-        """Return the system prompt for this agent."""
-        pass
-
-    @abstractmethod
-    def register_tools(self) -> None:
-        """Register agent-specific tools."""
-        pass
-
-    async def run(
-        self, prompt: str, deps: "ResearchDependencies", **kwargs
-    ) -> AgentRunResult[T]:
-        """Execute the agent."""
-        return await self._agent.run(prompt, deps=deps, **kwargs)
-
-    @property
-    def agent(self) -> Agent[Any, T]:
-        """Access the underlying Pydantic AI agent."""
-        return self._agent
-
-
-class SearchResult(BaseModel):
-    """Standard search result format."""
-
-    content: str
-    score: float
-    document_uri: str
-    metadata: dict[str, Any] = Field(default_factory=dict)
-
-
-class ResearchOutput(BaseModel):
-    """Standard research output format."""
-
-    summary: str
-    detailed_findings: list[str]
-    sources: list[str]
-    confidence: float
-
-
-class SearchAnswer(BaseModel):
-    """Structured output for the SearchSpecialist agent."""
-
-    query: str = Field(description="The search query that was performed")
-    answer: str = Field(description="The answer generated based on the context")
-    context: list[str] = Field(
-        description=(
-            "Only the minimal set of relevant snippets (verbatim) that directly "
-            "support the answer"
-        )
-    )
-    sources: list[str] = Field(
-        description=(
-            "Document URIs corresponding to the snippets actually used in the"
-            " answer (one URI per snippet; omit if none)"
-        ),
-        default_factory=list,
-    )

haiku/rag/research/dependencies.py

@@ -1,45 +0,0 @@
-from pydantic import BaseModel, Field
-
-from haiku.rag.client import HaikuRAG
-from haiku.rag.research.base 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")

haiku/rag/research/evaluation_agent.py

@@ -1,42 +0,0 @@
-from pydantic import BaseModel, Field
-
-from haiku.rag.research.base import BaseResearchAgent
-from haiku.rag.research.prompts import EVALUATION_AGENT_PROMPT
-
-
-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=[],
-    )
-    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 AnalysisEvaluationAgent(BaseResearchAgent[EvaluationResult]):
-    """Agent that analyzes findings and evaluates research completeness."""
-
-    def __init__(self, provider: str, model: str) -> None:
-        super().__init__(provider, model, output_type=EvaluationResult)
-
-    def get_system_prompt(self) -> str:
-        return EVALUATION_AGENT_PROMPT
-
-    def register_tools(self) -> None:
-        """No additional tools needed - uses LLM capabilities directly."""
-        pass