import hashlib import logging import mimetypes import tempfile from collections.abc import AsyncGenerator from pathlib import Path from urllib.parse import urlparse import httpx from haiku.rag.config import Config from haiku.rag.reranking import get_reranker from haiku.rag.store.engine import Store from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.document import Document from haiku.rag.store.repositories.chunk import ChunkRepository from haiku.rag.store.repositories.document import DocumentRepository from haiku.rag.store.repositories.settings import SettingsRepository logger = logging.getLogger(__name__) class HaikuRAG: """High-level haiku-rag client.""" def __init__( self, db_path: Path = Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb", skip_validation: bool = False, ): """Initialize the RAG client with a database path. Args: db_path: Path to the database file. skip_validation: Whether to skip configuration validation on database load. """ self.store = Store(db_path, skip_validation=skip_validation) self.document_repository = DocumentRepository(self.store) self.chunk_repository = ChunkRepository(self.store) async def __aenter__(self): """Async context manager entry.""" return self async def __aexit__(self, exc_type, exc_val, exc_tb): # noqa: ARG002 """Async context manager exit.""" # Wait for any pending vacuum to complete before closing async with self.store._vacuum_lock: pass self.close() return False async def _create_document_with_docling( self, docling_document, uri: str | None = None, title: str | None = None, metadata: dict | None = None, chunks: list[Chunk] | None = None, ) -> Document: """Create a new document from DoclingDocument.""" content = docling_document.export_to_markdown() document = Document( content=content, uri=uri, title=title, metadata=metadata or {}, ) return await self.document_repository._create_with_docling( document, docling_document, chunks ) async def create_document( self, content: str, uri: str | None = None, title: str | None = None, metadata: dict | None = None, chunks: list[Chunk] | None = None, ) -> Document: """Create a new document with optional URI and metadata. Args: content: The text content of the document. uri: Optional URI identifier for the document. metadata: Optional metadata dictionary. chunks: Optional list of pre-created chunks to use instead of generating new ones. Returns: The created Document instance. """ # Lazy import to avoid loading docling from haiku.rag.utils import text_to_docling_document # Convert content to DoclingDocument for processing docling_document = text_to_docling_document(content) document = Document( content=content, uri=uri, title=title, metadata=metadata or {}, ) return await self.document_repository._create_with_docling( document, docling_document, chunks ) async def create_document_from_source( self, source: str | Path, title: str | None = None, metadata: dict | None = None ) -> Document | list[Document]: """Create or update document(s) from a file path, directory, or URL. Checks if a document with the same URI already exists: - If MD5 is unchanged, returns existing document - If MD5 changed, updates the document - If no document exists, creates a new one Args: source: File path, directory (as string or Path), or URL to parse title: Optional title (only used for single files, not directories) metadata: Optional metadata dictionary Returns: Document instance (created, updated, or existing) for single files/URLs List of Document instances for directories Raises: ValueError: If the file/URL cannot be parsed or doesn't exist httpx.RequestError: If URL request fails """ # Lazy import to avoid loading docling from haiku.rag.reader import FileReader # Normalize metadata metadata = metadata or {} # Check if it's a URL source_str = str(source) parsed_url = urlparse(source_str) if parsed_url.scheme in ("http", "https"): return await self._create_or_update_document_from_url( source_str, title=title, metadata=metadata ) elif parsed_url.scheme == "file": # Handle file:// URI by converting to path source_path = Path(parsed_url.path) else: # Handle as regular file path source_path = Path(source) if isinstance(source, str) else source # Handle directories if source_path.is_dir(): documents = [] supported_extensions = set(FileReader.extensions) for file_path in source_path.rglob("*"): if ( file_path.is_file() and file_path.suffix.lower() in supported_extensions ): doc = await self._create_document_from_file( file_path, title=None, metadata=metadata ) documents.append(doc) return documents # Handle single file return await self._create_document_from_file( source_path, title=title, metadata=metadata ) async def _create_document_from_file( self, source_path: Path, title: str | None = None, metadata: dict | None = None ) -> Document: """Create or update a document from a single file path. Args: source_path: Path to the file title: Optional title metadata: Optional metadata dictionary Returns: Document instance (created, updated, or existing) Raises: ValueError: If the file cannot be parsed or doesn't exist """ # Lazy import to avoid loading docling from haiku.rag.reader import FileReader metadata = metadata or {} if source_path.suffix.lower() not in FileReader.extensions: raise ValueError(f"Unsupported file extension: {source_path.suffix}") if not source_path.exists(): raise ValueError(f"File does not exist: {source_path}") uri = source_path.absolute().as_uri() md5_hash = hashlib.md5(source_path.read_bytes()).hexdigest() # Get content type from file extension (do before early return) content_type, _ = mimetypes.guess_type(str(source_path)) if not content_type: content_type = "application/octet-stream" # Merge metadata with contentType and md5 metadata.update({"contentType": content_type, "md5": md5_hash}) # Check if document already exists existing_doc = await self.get_document_by_uri(uri) if existing_doc and existing_doc.metadata.get("md5") == md5_hash: # MD5 unchanged; update title/metadata if provided updated = False if title is not None and title != existing_doc.title: existing_doc.title = title updated = True if metadata: existing_doc.metadata = {**(existing_doc.metadata or {}), **metadata} updated = True if updated: return await self.document_repository.update(existing_doc) return existing_doc # Parse file only when content changed or new document docling_document = FileReader.parse_file(source_path) if existing_doc: # Update existing document existing_doc.content = docling_document.export_to_markdown() existing_doc.metadata = metadata if title is not None: existing_doc.title = title return await self.document_repository._update_with_docling( existing_doc, docling_document ) else: # Create new document using DoclingDocument return await self._create_document_with_docling( docling_document=docling_document, uri=uri, title=title, metadata=metadata, ) async def _create_or_update_document_from_url( self, url: str, title: str | None = None, metadata: dict | None = None ) -> Document: """Create or update a document from a URL by downloading and parsing the content. Checks if a document with the same URI already exists: - If MD5 is unchanged, returns existing document - If MD5 changed, updates the document - If no document exists, creates a new one Args: url: URL to download and parse metadata: Optional metadata dictionary Returns: Document instance (created, updated, or existing) Raises: ValueError: If the content cannot be parsed httpx.RequestError: If URL request fails """ # Lazy import to avoid loading docling from haiku.rag.reader import FileReader metadata = metadata or {} async with httpx.AsyncClient() as client: response = await client.get(url) response.raise_for_status() md5_hash = hashlib.md5(response.content).hexdigest() # Get content type early (used for potential no-op update) content_type = response.headers.get("content-type", "").lower() # Check if document already exists existing_doc = await self.get_document_by_uri(url) if existing_doc and existing_doc.metadata.get("md5") == md5_hash: # MD5 unchanged; update title/metadata if provided updated = False if title is not None and title != existing_doc.title: existing_doc.title = title updated = True metadata.update({"contentType": content_type, "md5": md5_hash}) if metadata: existing_doc.metadata = { **(existing_doc.metadata or {}), **metadata, } updated = True if updated: return await self.document_repository.update(existing_doc) return existing_doc file_extension = self._get_extension_from_content_type_or_url( url, content_type ) if file_extension not in FileReader.extensions: raise ValueError( f"Unsupported content type/extension: {content_type}/{file_extension}" ) # Create a temporary file with the appropriate extension with tempfile.NamedTemporaryFile( mode="wb", suffix=file_extension ) as temp_file: temp_file.write(response.content) temp_file.flush() # Ensure content is written to disk temp_path = Path(temp_file.name) # Parse the content using FileReader docling_document = FileReader.parse_file(temp_path) # Merge metadata with contentType and md5 metadata.update({"contentType": content_type, "md5": md5_hash}) if existing_doc: existing_doc.content = docling_document.export_to_markdown() existing_doc.metadata = metadata if title is not None: existing_doc.title = title return await self.document_repository._update_with_docling( existing_doc, docling_document ) else: return await self._create_document_with_docling( docling_document=docling_document, uri=url, title=title, metadata=metadata, ) def _get_extension_from_content_type_or_url( self, url: str, content_type: str ) -> str: """Determine file extension from content type or URL.""" # Common content type mappings content_type_map = { "text/html": ".html", "text/plain": ".txt", "text/markdown": ".md", "application/pdf": ".pdf", "application/json": ".json", "text/csv": ".csv", "application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx", "application/vnd.openxmlformats-officedocument.presentationml.presentation": ".pptx", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": ".xlsx", } # Try content type first for ct, ext in content_type_map.items(): if ct in content_type: return ext # Try URL extension parsed_url = urlparse(url) path = Path(parsed_url.path) if path.suffix: return path.suffix.lower() # Default to .html for web content return ".html" async def get_document_by_id(self, document_id: str) -> Document | None: """Get a document by its ID. Args: document_id: The unique identifier of the document. Returns: The Document instance if found, None otherwise. """ return await self.document_repository.get_by_id(document_id) async def get_document_by_uri(self, uri: str) -> Document | None: """Get a document by its URI. Args: uri: The URI identifier of the document. Returns: The Document instance if found, None otherwise. """ return await self.document_repository.get_by_uri(uri) async def update_document(self, document: Document) -> Document: """Update an existing document.""" # Lazy import to avoid loading docling from haiku.rag.utils import text_to_docling_document # Convert content to DoclingDocument docling_document = text_to_docling_document(document.content) return await self.document_repository._update_with_docling( document, docling_document ) async def delete_document(self, document_id: str) -> bool: """Delete a document by its ID.""" return await self.document_repository.delete(document_id) async def list_documents( self, limit: int | None = None, offset: int | None = None ) -> list[Document]: """List all documents with optional pagination. Args: limit: Maximum number of documents to return. offset: Number of documents to skip. Returns: List of Document instances. """ return await self.document_repository.list_all(limit=limit, offset=offset) async def search( self, query: str, limit: int = 5, search_type: str = "hybrid" ) -> list[tuple[Chunk, float]]: """Search for relevant chunks using the specified search method with optional reranking. Args: query: The search query string. limit: Maximum number of results to return. search_type: Type of search - "vector", "fts", or "hybrid" (default). Returns: List of (chunk, score) tuples ordered by relevance. """ # Get reranker if available reranker = get_reranker() if reranker is None: # No reranking - return direct search results return await self.chunk_repository.search(query, limit, search_type) # Get more initial results (3X) for reranking search_limit = limit * 3 search_results = await self.chunk_repository.search( query, search_limit, search_type ) # Apply reranking chunks = [chunk for chunk, _ in search_results] reranked_results = await reranker.rerank(query, chunks, top_n=limit) # Return reranked results with scores from reranker return reranked_results async def expand_context( self, search_results: list[tuple[Chunk, float]], radius: int = Config.CONTEXT_CHUNK_RADIUS, ) -> list[tuple[Chunk, float]]: """Expand search results with adjacent chunks, merging overlapping chunks. Args: search_results: List of (chunk, score) tuples from search. radius: Number of adjacent chunks to include before/after each chunk. Defaults to CONTEXT_CHUNK_RADIUS config setting. Returns: List of (chunk, score) tuples with expanded and merged context chunks. """ if radius == 0: return search_results # Group chunks by document_id to handle merging within documents document_groups = {} for chunk, score in search_results: doc_id = chunk.document_id if doc_id not in document_groups: document_groups[doc_id] = [] document_groups[doc_id].append((chunk, score)) results = [] for doc_id, doc_chunks in document_groups.items(): # Get all expanded ranges for this document expanded_ranges = [] for chunk, score in doc_chunks: adjacent_chunks = await self.chunk_repository.get_adjacent_chunks( chunk, radius ) all_chunks = adjacent_chunks + [chunk] # Get the range of orders for this expanded chunk orders = [c.order for c in all_chunks] min_order = min(orders) max_order = max(orders) expanded_ranges.append( { "original_chunk": chunk, "score": score, "min_order": min_order, "max_order": max_order, "all_chunks": sorted(all_chunks, key=lambda c: c.order), } ) # Merge overlapping/adjacent ranges merged_ranges = self._merge_overlapping_ranges(expanded_ranges) # Create merged chunks for merged_range in merged_ranges: combined_content_parts = [c.content for c in merged_range["all_chunks"]] # Use the first original chunk for metadata original_chunk = merged_range["original_chunks"][0] merged_chunk = Chunk( id=original_chunk.id, document_id=original_chunk.document_id, content="".join(combined_content_parts), metadata=original_chunk.metadata, document_uri=original_chunk.document_uri, document_title=original_chunk.document_title, document_meta=original_chunk.document_meta, ) # Use the highest score from merged chunks best_score = max(merged_range["scores"]) results.append((merged_chunk, best_score)) return results def _merge_overlapping_ranges(self, expanded_ranges): """Merge overlapping or adjacent expanded ranges.""" if not expanded_ranges: return [] # Sort by min_order sorted_ranges = sorted(expanded_ranges, key=lambda x: x["min_order"]) merged = [] current = { "min_order": sorted_ranges[0]["min_order"], "max_order": sorted_ranges[0]["max_order"], "original_chunks": [sorted_ranges[0]["original_chunk"]], "scores": [sorted_ranges[0]["score"]], "all_chunks": sorted_ranges[0]["all_chunks"], } for range_info in sorted_ranges[1:]: # Check if ranges overlap or are adjacent (max_order + 1 >= min_order) if current["max_order"] >= range_info["min_order"] - 1: # Merge ranges current["max_order"] = max( current["max_order"], range_info["max_order"] ) current["original_chunks"].append(range_info["original_chunk"]) current["scores"].append(range_info["score"]) # Merge all_chunks and deduplicate by order all_chunks_dict = {} for chunk in current["all_chunks"] + range_info["all_chunks"]: order = chunk.order all_chunks_dict[order] = chunk current["all_chunks"] = [ all_chunks_dict[order] for order in sorted(all_chunks_dict.keys()) ] else: # No overlap, add current to merged and start new merged.append(current) current = { "min_order": range_info["min_order"], "max_order": range_info["max_order"], "original_chunks": [range_info["original_chunk"]], "scores": [range_info["score"]], "all_chunks": range_info["all_chunks"], } # Add the last range merged.append(current) return merged async def ask( self, question: str, cite: bool = False, system_prompt: str | None = None ) -> str: """Ask a question using the configured QA agent. Args: question: The question to ask. cite: Whether to include citations in the response. system_prompt: Optional custom system prompt for the QA agent. Returns: The generated answer as a string. """ from haiku.rag.qa import get_qa_agent qa_agent = get_qa_agent(self, use_citations=cite, system_prompt=system_prompt) return await qa_agent.answer(question) async def rebuild_database(self) -> AsyncGenerator[str, None]: """Rebuild the database by deleting all chunks and re-indexing all documents. For documents with URIs: - Re-adds from source if source exists - Re-embeds from existing content if source is missing For documents without URIs: - Re-creates chunks from existing content Yields: int: The ID of the document currently being processed """ # Lazy import to avoid loading docling from haiku.rag.utils import text_to_docling_document await self.chunk_repository.delete_all() self.store.recreate_embeddings_table() # Update settings to current config settings_repo = SettingsRepository(self.store) settings_repo.save_current_settings() documents = await self.list_documents() for doc in documents: assert doc.id is not None, "Document ID should not be None" if doc.uri: # Document has a URI - check if source is accessible source_accessible = False parsed_url = urlparse(doc.uri) try: if parsed_url.scheme == "file": # Check if file exists source_path = Path(parsed_url.path) source_accessible = source_path.exists() elif parsed_url.scheme in ("http", "https"): # For URLs, we'll try to create and catch errors source_accessible = True else: source_accessible = False except Exception: source_accessible = False if source_accessible: # Source exists - delete and recreate from source try: await self.delete_document(doc.id) new_doc = await self.create_document_from_source( source=doc.uri, metadata=doc.metadata or {} ) # URIs always point to single files/URLs, never directories assert isinstance(new_doc, Document) assert new_doc.id is not None, ( "New document ID should not be None" ) yield new_doc.id except Exception as e: logger.error( "Error recreating document from source %s: %s", doc.uri, e, ) continue else: # Source missing - re-embed from existing content logger.warning( "Source missing for %s, re-embedding from content", doc.uri ) docling_document = text_to_docling_document(doc.content) await self.chunk_repository.create_chunks_for_document( doc.id, docling_document ) yield doc.id else: # Document without URI - re-create chunks from existing content docling_document = text_to_docling_document(doc.content) await self.chunk_repository.create_chunks_for_document( doc.id, docling_document ) yield doc.id # Final maintenance: centralized vacuum to curb disk usage try: await self.store.vacuum() except Exception: pass async def vacuum(self) -> None: """Optimize and clean up old versions across all tables.""" await self.store.vacuum() def close(self): """Close the underlying store connection.""" self.store.close()