haiku.rag/src/haiku/rag/client.py

import tempfile
from pathlib import Path
from typing import Literal
from urllib.parse import urlparse

import httpx

from haiku.rag.reader import FileReader
from haiku.rag.store.engine import Store
from haiku.rag.store.models.document import Document
from haiku.rag.store.repositories.document import DocumentRepository


class RAGClient:
    """High-level haiku-rag client."""

    def __init__(self, db_path: Path | Literal[":memory:"]):
        """Initialize the RAG client with a database path."""
        self.store = Store(db_path)
        self.document_repository = DocumentRepository(self.store)

    async def create_document(
        self, content: str, uri: str | None = None, metadata: dict | None = None
    ) -> Document:
        """Create a new document with optional URI and metadata."""
        document = Document(
            content=content,
            uri=uri,
            metadata=metadata or {},
        )
        return await self.document_repository.create(document)

    async def create_document_from_source(
        self, source: str | Path, metadata: dict | None = None
    ) -> Document:
        """Create a document from a file path or URL.

        Args:
            source: File path (as string or Path) or URL to parse
            metadata: Optional metadata dictionary

        Returns:
            Created Document instance

        Raises:
            ValueError: If the file/URL cannot be parsed or doesn't exist
            httpx.RequestError: If URL request fails
        """

        # 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_document_from_url(source_str, metadata)

        # Handle as file path
        source_path = Path(source) if isinstance(source, str) else source
        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}")

        content = FileReader.parse_file(source_path)

        # Create the document
        return await self.create_document(
            content=content, uri=str(source_path.resolve()), metadata=metadata
        )

    async def _create_document_from_url(
        self, url: str, metadata: dict | None = None
    ) -> Document:
        """Create a document from a URL by downloading and parsing the content.

        Args:
            url: URL to download and parse
            metadata: Optional metadata dictionary

        Returns:
            Created Document instance

        Raises:
            ValueError: If the content cannot be parsed
            httpx.RequestError: If URL request fails
        """
        async with httpx.AsyncClient() as client:
            response = await client.get(url)
            response.raise_for_status()

            # Get content type to determine file extension
            content_type = response.headers.get("content-type", "").lower()

            # Try to determine file extension from content type or URL
            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, delete=False
            ) as temp_file:
                temp_file.write(response.content)
                temp_path = Path(temp_file.name)

            try:
                # Parse the content using FileReader
                content = FileReader.parse_file(temp_path)

                # Create the document with the original URL as URI
                return await self.create_document(
                    content=content, uri=url, metadata=metadata
                )
            finally:
                # Clean up temporary file
                temp_path.unlink(missing_ok=True)

    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: int) -> Document | None:
        """Get a document by its ID."""
        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."""
        return await self.document_repository.get_by_uri(uri)

    async def update_document(self, document: Document) -> Document:
        """Update an existing document."""
        return await self.document_repository.update(document)

    async def delete_document(self, document_id: int) -> 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."""
        return await self.document_repository.list_all(limit=limit, offset=offset)

    def close(self):
        """Close the underlying store connection."""
        self.store.close()