file_search.py

"""
Production-ready Google File Search Tool for RAG (Retrieval-Augmented Generation)

This module provides a clean interface for using Google's File Search API
with the Gemini models to perform document-based question answering.
"""

import os
import time
import mimetypes
from typing import List, Optional, Dict, Any, Union
from pathlib import Path
from dataclasses import dataclass
from enum import Enum

from google import genai
from google.genai import types
from dotenv import load_dotenv

# Load environment variables
load_dotenv(".env.local")

# Register markdown MIME type if not already registered
if not mimetypes.guess_type("file.md")[0]:
    mimetypes.add_type("text/markdown", ".md")


class Model(str, Enum):
    """Supported Gemini models for File Search"""

    GEMINI_2_5_FLASH = "gemini-2.5-flash"
    GEMINI_2_5_PRO = "gemini-2.5-pro"


@dataclass
class FileSearchConfig:
    """Configuration for File Search operations"""

    display_name: Optional[str] = None
    max_tokens_per_chunk: Optional[int] = None
    max_overlap_tokens: Optional[int] = None
    metadata: Optional[Dict[str, Any]] = None


@dataclass
class SearchResult:
    """Represents a search result from File Search"""

    text: str
    sources: Optional[List[str]] = None
    grounding_chunks: Optional[List[Any]] = None


class FileSearchStore:
    """
    A production-ready wrapper for Google's File Search Store API.

    This class handles the creation, management, and querying of File Search stores
    with proper error handling and resource management.
    """

    def __init__(self, api_key: Optional[str] = None, store_name: Optional[str] = None):
        """
        Initialize the FileSearchStore.

        Args:
            api_key: Optional API key. If not provided, will use GEMINI_API_KEY from environment.
            store_name: Optional existing store name to connect to.
        """
        self.api_key = api_key or os.getenv("GEMINI_API_KEY")
        if not self.api_key:
            raise ValueError(
                "GEMINI_API_KEY not found. Please set it in .env.local or pass it directly."
            )

        # Set the API key in environment for google-genai
        os.environ["GEMINI_API_KEY"] = self.api_key

        self.client = genai.Client()
        self.store = None
        self.store_name = store_name

        if store_name:
            self._connect_to_store(store_name)

    def _connect_to_store(self, store_name: str):
        """Connect to an existing File Search store."""
        try:
            self.store = self.client.file_search_stores.get(name=store_name)
            self.store_name = store_name
            print(f"Connected to existing store: {store_name}")
        except Exception as e:
            print(f"Could not connect to store {store_name}: {e}")
            self.store = None

    def create_store(self, display_name: Optional[str] = None) -> str:
        """
        Create a new File Search store.

        Args:
            display_name: Optional human-readable name for the store.

        Returns:
            The store name (ID) of the created store.
        """
        config: Dict[str, Any] = {}
        if display_name:
            config["display_name"] = display_name

        try:
            self.store = self.client.file_search_stores.create(
                config=config if config else None  # type: ignore
            )
            self.store_name = str(self.store.name)  # Ensure it's a string
            print(f"Created new store: {self.store_name}")
            if display_name:
                print(f"Display name: {display_name}")
            return self.store_name
        except Exception as e:
            raise RuntimeError(f"Failed to create File Search store: {e}")

    def upload_file(
        self,
        file_path: Union[str, Path],
        display_name: Optional[str] = None,
        chunking_config: Optional[Dict[str, Any]] = None,
        metadata: Optional[List[Dict[str, Any]]] = None,
        wait: bool = True,
        timeout: int = 300,
    ) -> bool:
        """
        Upload a file to the File Search store.

        Args:
            file_path: Path to the file to upload.
            display_name: Optional display name for the file (shown in citations).
            chunking_config: Optional chunking configuration.
            metadata: Optional custom metadata for the file.
            wait: Whether to wait for the upload to complete.
            timeout: Maximum time to wait in seconds.

        Returns:
            True if upload was successful, False otherwise.
        """
        if not self.store:
            raise RuntimeError(
                "No store connected. Please create or connect to a store first."
            )

        file_path = Path(file_path)
        if not file_path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")

        # Detect MIME type
        mime_type, _ = mimetypes.guess_type(str(file_path))
        if not mime_type:
            # Fallback to application/octet-stream if detection fails
            mime_type = "application/octet-stream"

        # Build configuration
        config: Dict[str, Any] = {}
        if display_name:
            config["display_name"] = display_name

        # Always set MIME type to avoid detection issues
        config["mime_type"] = mime_type

        if chunking_config:
            config["chunking_config"] = chunking_config

        if metadata:
            config["custom_metadata"] = metadata

        try:
            print(f"Uploading file: {file_path} (MIME type: {mime_type})")
            if self.store_name is None:
                raise RuntimeError("Store name is not set")

            operation = self.client.file_search_stores.upload_to_file_search_store(
                file=str(file_path),
                file_search_store_name=self.store_name,
                config=config if config else None,  # type: ignore
            )

            if wait:
                return self._wait_for_operation(operation, timeout)

            return True

        except Exception as e:
            print(f"Failed to upload file: {e}")
            return False

    def _wait_for_operation(self, operation: Any, timeout: int = 300) -> bool:
        """
        Wait for an operation to complete.

        Args:
            operation: The operation to wait for.
            timeout: Maximum time to wait in seconds.

        Returns:
            True if operation completed successfully, False otherwise.
        """
        start_time = time.time()

        while not operation.done:
            if time.time() - start_time > timeout:
                print(f"Operation timed out after {timeout} seconds")
                return False

            time.sleep(2)
            try:
                operation = self.client.operations.get(operation)
            except Exception as e:
                print(f"Error checking operation status: {e}")
                return False

        print("Operation completed successfully")
        return True

    def query(
        self,
        prompt: str,
        model: Model = Model.GEMINI_2_5_FLASH,
        metadata_filter: Optional[str] = None,
        response_format: str = "markdown",
        temperature: Optional[float] = None,
        max_output_tokens: Optional[int] = None,
    ) -> SearchResult:
        """
        Query the File Search store with a prompt.

        Args:
            prompt: The question or prompt to send to the model.
            model: The Gemini model to use.
            metadata_filter: Optional metadata filter for document selection.
            response_format: Format instruction for the response (e.g., "markdown", "json").
            temperature: Optional temperature for response generation.
            max_output_tokens: Optional maximum number of output tokens.

        Returns:
            SearchResult containing the response text and metadata.
        """
        if not self.store:
            raise RuntimeError(
                "No store connected. Please create or connect to a store first."
            )

        # Format the prompt with response format instruction
        formatted_prompt = prompt
        if response_format == "markdown":
            formatted_prompt = f"{prompt}\n\n(Please format your response in markdown with clear sections and bullet points where appropriate)"
        elif response_format == "json":
            formatted_prompt = (
                f"{prompt}\n\n(Please format your response as valid JSON)"
            )

        # Build File Search configuration
        if self.store_name is None:
            raise RuntimeError("Store name is not set")

        file_search_config = types.FileSearch(file_search_store_names=[self.store_name])

        if metadata_filter:
            file_search_config.metadata_filter = metadata_filter

        # Build generation configuration
        gen_config = types.GenerateContentConfig(
            tools=[types.Tool(file_search=file_search_config)]
        )

        if temperature is not None:
            gen_config.temperature = temperature

        if max_output_tokens is not None:
            gen_config.max_output_tokens = max_output_tokens

        try:
            response = self.client.models.generate_content(
                model=model.value, contents=formatted_prompt, config=gen_config
            )

            # Extract grounding chunks and sources if available
            grounding_chunks = None
            sources: List[str] = []
            if response.candidates and response.candidates[0].grounding_metadata:
                grounding_metadata = response.candidates[0].grounding_metadata
                # Extract grounding chunks
                if hasattr(grounding_metadata, "grounding_chunks"):
                    grounding_chunks = grounding_metadata.grounding_chunks
                    # Extract unique source titles
                    if grounding_chunks:
                        sources = list(
                            {
                                str(chunk.retrieved_context.title)
                                for chunk in grounding_chunks
                                if hasattr(chunk, "retrieved_context")
                                and chunk.retrieved_context
                                and hasattr(chunk.retrieved_context, "title")
                            }
                        )

            return SearchResult(
                text=response.text or "",
                sources=sources,
                grounding_chunks=grounding_chunks,
            )

        except Exception as e:
            raise RuntimeError(f"Query failed: {e}")

    def list_stores(self) -> List[Dict[str, str]]:
        """
        List all available File Search stores.

        Returns:
            List of store information dictionaries.
        """
        stores = []
        try:
            for store in self.client.file_search_stores.list():
                store_info = {
                    "name": store.name,
                    "display_name": getattr(store, "display_name", "N/A"),
                    "create_time": str(getattr(store, "create_time", "N/A")),
                }
                stores.append(store_info)
        except Exception as e:
            print(f"Error listing stores: {e}")

        return stores

    def get_document_counts(self) -> Dict[str, int]:
        """
        Get document counts from the store metadata.
        This provides accurate counts without pagination issues.

        Returns:
            Dictionary with active, pending, failed, total counts and size_bytes
        """
        if not self.store_name:
            raise RuntimeError("No store connected. Please connect to a store first.")

        try:
            store = self.client.file_search_stores.get(name=self.store_name)

            active = int(getattr(store, "active_documents_count", 0) or 0)
            pending = int(getattr(store, "pending_documents_count", 0) or 0)
            failed = int(getattr(store, "failed_documents_count", 0) or 0)

            return {
                "active": active,
                "pending": pending,
                "failed": failed,
                "total": active + pending + failed,
                "size_bytes": int(getattr(store, "size_bytes", 0) or 0),
            }
        except Exception as e:
            print(f"Error getting document counts: {e}")
            return {
                "active": 0,
                "pending": 0,
                "failed": 0,
                "total": 0,
                "size_bytes": 0,
            }

    def list_documents(self) -> List[Dict[str, Any]]:
        """
        List all documents in the current File Search store.
        Note: Due to SDK pagination bug, this only returns first page (typically 10 docs).
        Use get_document_counts() for accurate total count.

        Returns:
            List of document information dictionaries.
        """
        if not self.store or not self.store_name:
            raise RuntimeError("No store connected. Please connect to a store first.")

        documents = []
        try:
            # The iterator automatically handles pagination
            # Just iterate through all documents
            for doc in self.client.file_search_stores.documents.list(
                parent=self.store_name
            ):
                doc_info = {
                    "name": getattr(doc, "name", ""),
                    "display_name": getattr(doc, "display_name", "N/A"),
                    "create_time": str(getattr(doc, "create_time", "N/A")),
                    "update_time": str(getattr(doc, "update_time", "N/A")),
                    "mime_type": getattr(doc, "mime_type", "N/A"),
                    "size_bytes": getattr(doc, "size_bytes", 0),
                    "state": str(getattr(doc, "state", "N/A")),
                }

                # Extract custom metadata if available
                if hasattr(doc, "custom_metadata") and doc.custom_metadata:
                    metadata = {}
                    for item in doc.custom_metadata:
                        key = getattr(item, "key", "")
                        # Check which value type is set (not just which attribute exists)
                        if (
                            hasattr(item, "numeric_value")
                            and item.numeric_value is not None
                        ):
                            metadata[key] = item.numeric_value
                        elif hasattr(item, "string_value") and item.string_value:
                            metadata[key] = item.string_value
                    doc_info["custom_metadata"] = metadata

                documents.append(doc_info)
        except Exception as e:
            print(f"Error listing documents: {e}")

        return documents

    def delete_store(self, force: bool = True) -> bool:
        """
        Delete the current File Search store.

        Args:
            force: Whether to force deletion even if store contains files.

        Returns:
            True if deletion was successful, False otherwise.
        """
        if not self.store:
            print("No store to delete")
            return False

        try:
            if self.store_name is None:
                raise RuntimeError("Store name is not set")

            self.client.file_search_stores.delete(
                name=self.store_name, config={"force": force}
            )
            print(f"Deleted store: {self.store_name}")
            self.store = None
            self.store_name = None
            return True
        except Exception as e:
            print(f"Failed to delete store: {e}")
            return False

    def __enter__(self):
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - does not delete the store by default."""
        pass


class FileSearchManager:
    """
    High-level manager for File Search operations.
    Provides convenient methods for common RAG workflows.
    """

    def __init__(self, api_key: Optional[str] = None):
        """Initialize the FileSearchManager."""
        self.api_key = api_key or os.getenv("GEMINI_API_KEY")
        if not self.api_key:
            raise ValueError(
                "GEMINI_API_KEY not found. Please set it in .env.local or pass it directly."
            )

        self.stores: Dict[str, FileSearchStore] = {}

    def create_rag_pipeline(
        self,
        name: str,
        files: List[Union[str, Path]],
        display_names: Optional[List[str]] = None,
    ) -> FileSearchStore:
        """
        Create a complete RAG pipeline with multiple files.

        Args:
            name: Name for the pipeline (store).
            files: List of file paths to include in the pipeline.
            display_names: Optional list of display names for the files.

        Returns:
            Configured FileSearchStore ready for queries.
        """
        store = FileSearchStore(api_key=self.api_key)
        store.create_store(display_name=name)

        if display_names and len(display_names) != len(files):
            raise ValueError("display_names must have the same length as files")

        for i, file_path in enumerate(files):
            display_name = display_names[i] if display_names else Path(file_path).name
            success = store.upload_file(file_path, display_name=display_name)
            if not success:
                print(f"Warning: Failed to upload {file_path}")

        self.stores[name] = store
        return store

    def quick_query(
        self, file_path: Union[str, Path], prompt: str, cleanup: bool = True
    ) -> str:
        """
        Quick one-shot query on a single file.

        Args:
            file_path: Path to the file to query.
            prompt: The question to ask about the file.
            cleanup: Whether to delete the store after querying.

        Returns:
            The response text.
        """
        store = FileSearchStore(api_key=self.api_key)
        try:
            store.create_store(display_name="quick-query")
            store.upload_file(file_path, display_name=Path(file_path).name)
            result = store.query(prompt)
            return result.text
        finally:
            if cleanup:
                store.delete_store()