diff --git a/cookbook/knowledge/vector_db/README.md b/cookbook/knowledge/vector_db/README.md index 039e06c2ca..1bd7d73c13 100644 --- a/cookbook/knowledge/vector_db/README.md +++ b/cookbook/knowledge/vector_db/README.md @@ -34,6 +34,7 @@ agent = Agent( - **[ChromaDB](./chroma_db/)** - Embedded vector database - **[ClickHouse](./clickhouse_db/)** - Columnar database with vector functions - **[Couchbase](./couchbase_db/)** - NoSQL database with vector search +- **[Gemini File Search](./gemini_file_search/)** - Google's managed vector database with Gemini integration - **[LanceDB](./lance_db/)** - Fast columnar vector database - **[LangChain](./langchain/)** - Use any LangChain vector store - **[LightRAG](./lightrag/)** - Graph-based RAG system diff --git a/cookbook/knowledge/vector_db/gemini_file_search/README.md b/cookbook/knowledge/vector_db/gemini_file_search/README.md new file mode 100644 index 0000000000..fa1486837e --- /dev/null +++ b/cookbook/knowledge/vector_db/gemini_file_search/README.md @@ -0,0 +1,109 @@ +# Gemini File Search + +Gemini File Search provides a managed vector database integrated with Google's Gemini AI models. It allows you to upload documents and perform semantic search using Google's infrastructure. + +## Features + +- **Managed Service**: No infrastructure management required +- **Native Integration**: Seamlessly works with Gemini models +- **File Upload**: Directly upload documents to Google's File Search Store +- **Metadata Filtering**: Filter search results by custom metadata +- **Grounding Support**: Get responses with citation metadata + +## Installation + +```bash +pip install google-genai +``` + +## Configuration + +Set your Google API key as an environment variable: + +```bash +export GOOGLE_API_KEY="your-api-key-here" +``` + +Or get one from: https://ai.google.dev/ + +## Basic Usage + +```python +from agno.agent import Agent +from agno.knowledge.knowledge import Knowledge +from agno.vectordb.gemini.gemini_file_search import GeminiFileSearch + +# Create Gemini File Search vector database +vector_db = GeminiFileSearch( + file_search_store_name="my-knowledge-store", + model_name="gemini-2.5-flash-lite", + api_key="your-api-key", +) + +# Create knowledge base +knowledge = Knowledge( + name="My Knowledge Base", + vector_db=vector_db, +) + +# Add documents +knowledge.add_content( + name="MyDocument", + url="https://example.com/document.pdf", + metadata={"doc_type": "manual"}, +) + +# Create agent and query +agent = Agent(knowledge=knowledge, search_knowledge=True) +agent.print_response("What is covered in the document?") +``` + +## Examples + +- **[gemini_file_search.py](./gemini_file_search.py)** - Basic usage with Thai recipe knowledge base +- **[async_gemini_file_search.py](./async_gemini_file_search.py)** - Async operations with Agno documentation +- **[gemini_file_search_with_filters.py](./gemini_file_search_with_filters.py)** - Using metadata filters for refined search + +## Supported Operations + +| Operation | Supported | Notes | +|-----------|-----------|-------| +| `create()` | ✅ | Creates or gets existing File Search Store | +| `insert()` | ✅ | Uploads documents to the store | +| `search()` | ✅ | Semantic search with optional metadata filters | +| `upsert()` | ✅ | Updates existing documents or inserts new ones | +| `delete_by_name()` | ✅ | Delete documents by display name | +| `delete_by_id()` | ✅ | Delete documents by ID | +| `delete_by_content_id()` | ✅ | Delete documents by content ID | +| `delete_by_metadata()` | ❌ | Not supported by Gemini File Search | +| `update_metadata()` | ❌ | Not supported by Gemini File Search | + +## Important Notes + +1. **File Search Store**: Documents are organized in "File Search Stores" - named containers for your documents +2. **Document Names**: Each document has both a system-generated `name` (ID) and a user-defined `display_name` +3. **Operation Polling**: Document uploads are asynchronous; the library polls until completion +4. **Metadata Limitations**: + - Supports string, numeric, and float metadata values + - Metadata can be used for filtering during search + - Cannot update metadata after upload (must delete and re-upload) +5. **Cost**: Check Google AI pricing for File Search Store usage + +## Model Options + +Gemini File Search supports various Gemini models: + +- `gemini-2.5-flash-lite` (default) - Fast and cost-effective +- `gemini-2.5-flash` - Balanced performance +- `gemini-2.0-flash` - High performance +- `gemini-2.0-flash-exp` - Experimental features + +## API Reference + +See the [GeminiFileSearch documentation](../../../../libs/agno/agno/vectordb/gemini/gemini_file_search.py) for detailed API information. + +## Resources + +- [Google AI Gemini Docs](https://ai.google.dev/gemini-api/docs) +- [File Search API](https://ai.google.dev/gemini-api/docs/file-search) +- [Agno Documentation](https://docs.agno.com) diff --git a/cookbook/knowledge/vector_db/gemini_file_search/__init__.py b/cookbook/knowledge/vector_db/gemini_file_search/__init__.py new file mode 100644 index 0000000000..6a4eb61855 --- /dev/null +++ b/cookbook/knowledge/vector_db/gemini_file_search/__init__.py @@ -0,0 +1 @@ +"""Gemini File Search cookbook examples.""" diff --git a/cookbook/knowledge/vector_db/gemini_file_search/async_gemini_file_search.py b/cookbook/knowledge/vector_db/gemini_file_search/async_gemini_file_search.py new file mode 100644 index 0000000000..5095b773d4 --- /dev/null +++ b/cookbook/knowledge/vector_db/gemini_file_search/async_gemini_file_search.py @@ -0,0 +1,57 @@ +""" +Async example of using Gemini File Search as a vector database. + +Requirements: +- pip install google-genai +- Set GOOGLE_API_KEY environment variable +""" + +import asyncio +from os import getenv + +from agno.agent import Agent +from agno.knowledge.knowledge import Knowledge +from agno.vectordb.gemini.gemini_file_search import GeminiFileSearch + +# Get API key from environment +api_key = getenv("GOOGLE_API_KEY") + +# Initialize Gemini File Search +vector_db = GeminiFileSearch( + file_search_store_name="agno-docs-store", + model_name="gemini-2.5-flash-lite", + api_key=api_key, +) + +# Create knowledge base +knowledge = Knowledge( + name="Agno Documentation", + description="Knowledge base with Agno documentation using Gemini File Search", + vector_db=vector_db, +) + +# Create and use the agent +agent = Agent(knowledge=knowledge, search_knowledge=True) + + +async def main(): + """Main async function.""" + # Add content to the knowledge base + # Comment out after first run to avoid re-uploading + await knowledge.add_content_async( + name="AgnoIntroduction", + url="https://docs.agno.com/concepts/agents/introduction.md", + metadata={"doc_type": "documentation", "topic": "agents"}, + ) + + # Query the knowledge base using async + await agent.aprint_response("What is the purpose of an Agno Agent?", markdown=True) + + # Additional query + await agent.aprint_response( + "How do I create an agent with knowledge?", markdown=True + ) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/cookbook/knowledge/vector_db/gemini_file_search/gemini_file_search.py b/cookbook/knowledge/vector_db/gemini_file_search/gemini_file_search.py new file mode 100644 index 0000000000..92da7cc553 --- /dev/null +++ b/cookbook/knowledge/vector_db/gemini_file_search/gemini_file_search.py @@ -0,0 +1,54 @@ +""" +Basic example of using Gemini File Search as a vector database. + +Requirements: +- pip install google-genai +- Set GOOGLE_API_KEY environment variable +""" + +import asyncio +from os import getenv + +from agno.agent import Agent +from agno.knowledge.knowledge import Knowledge +from agno.vectordb.gemini.gemini_file_search import GeminiFileSearch + +# Get API key from environment +api_key = getenv("GOOGLE_API_KEY") + +# Create Gemini File Search vector database +vector_db = GeminiFileSearch( + file_search_store_name="thai-recipes-store", + model_name="gemini-2.5-flash-lite", + api_key=api_key, +) + +# Create Knowledge Instance with Gemini File Search +knowledge = Knowledge( + name="Thai Recipe Knowledge Base", + description="Agno 2.0 Knowledge Implementation with Gemini File Search", + vector_db=vector_db, +) + +# Add content to the knowledge base +# Note: This uploads documents to Gemini File Search Store +asyncio.run( + knowledge.add_content_async( + name="Recipes", + url="https://agno-public.s3.amazonaws.com/recipes/ThaiRecipes.pdf", + metadata={"doc_type": "recipe_book", "cuisine": "Thai"}, + ) +) + +# Create and use the agent +agent = Agent(knowledge=knowledge, search_knowledge=True) + +# Query the knowledge base +agent.print_response("List down the ingredients to make Massaman Gai", markdown=True) + +# Delete operations examples +# Delete by name +vector_db.delete_by_name("Recipes") + +# Note: delete_by_metadata is not supported by Gemini File Search +# You can only delete documents by name or ID diff --git a/cookbook/knowledge/vector_db/gemini_file_search/gemini_file_search_with_filters.py b/cookbook/knowledge/vector_db/gemini_file_search/gemini_file_search_with_filters.py new file mode 100644 index 0000000000..3052361435 --- /dev/null +++ b/cookbook/knowledge/vector_db/gemini_file_search/gemini_file_search_with_filters.py @@ -0,0 +1,83 @@ +""" +Example of using Gemini File Search with metadata filters. + +Requirements: +- pip install google-genai +- Set GOOGLE_API_KEY environment variable +""" + +import asyncio +from os import getenv + +from agno.agent import Agent +from agno.knowledge.knowledge import Knowledge +from agno.vectordb.gemini.gemini_file_search import GeminiFileSearch + +# Get API key from environment +api_key = getenv("GOOGLE_API_KEY") + +# Create Gemini File Search vector database +vector_db = GeminiFileSearch( + file_search_store_name="multi-cuisine-recipes", + model_name="gemini-2.5-flash-lite", + api_key=api_key, +) + +# Create Knowledge Instance +knowledge = Knowledge( + name="Multi-Cuisine Recipe Knowledge Base", + description="Knowledge base with recipes from different cuisines", + vector_db=vector_db, +) + + +async def add_recipes(): + """Add multiple recipe documents with different metadata.""" + # Add Thai recipes + await knowledge.add_content_async( + name="ThaiRecipes", + url="https://agno-public.s3.amazonaws.com/recipes/ThaiRecipes.pdf", + metadata={"cuisine": "Thai", "difficulty": "medium"}, + ) + + # You can add more recipes from other sources with different metadata + # await knowledge.add_content_async( + # name="ItalianRecipes", + # content="Italian recipe content here...", + # metadata={"cuisine": "Italian", "difficulty": "easy"}, + # ) + + +async def main(): + """Main function.""" + # Add content (comment out after first run) + await add_recipes() + + # Create agent + agent = Agent(knowledge=knowledge, search_knowledge=True) + + # Query with specific filters + # Note: Gemini File Search supports metadata filtering in search + print("\n=== Searching for Thai recipes ===\n") + await agent.aprint_response( + "What are some popular Thai dishes?", + markdown=True, + ) + + # You can also search the vector database directly with filters + print("\n=== Direct vector DB search with filters ===\n") + results = vector_db.search( + query="coconut curry recipes", + limit=3, + filters={"cuisine": "Thai"}, # Filter by cuisine + ) + + for i, result in enumerate(results, 1): + print(f"\nResult {i}:") + print(f"Content: {result.content[:200]}...") + if result.meta_data: + print(f"Metadata: {result.meta_data}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/libs/agno/agno/vectordb/gemini/__init__.py b/libs/agno/agno/vectordb/gemini/__init__.py new file mode 100644 index 0000000000..1b8455f73b --- /dev/null +++ b/libs/agno/agno/vectordb/gemini/__init__.py @@ -0,0 +1,5 @@ +from agno.vectordb.gemini.gemini_file_search import GeminiFileSearch + +__all__ = [ + "GeminiFileSearch", +] diff --git a/libs/agno/agno/vectordb/gemini/gemini_file_search.py b/libs/agno/agno/vectordb/gemini/gemini_file_search.py new file mode 100644 index 0000000000..ce8c0b9a09 --- /dev/null +++ b/libs/agno/agno/vectordb/gemini/gemini_file_search.py @@ -0,0 +1,373 @@ +import asyncio +import tempfile +import time +from typing import Any, Dict, List, Optional + +try: + from google import genai + from google.genai import types + from google.genai.errors import ClientError + from google.genai.types import ( + CreateFileSearchStoreConfigDict, + CustomMetadata, + DeleteDocumentConfigDict, + DeleteFileSearchStoreConfigDict, + StringList, + UploadToFileSearchStoreConfigDict, + ) +except ImportError: + raise ImportError("`google-genai` not installed. Please install using `pip install google-genai`") + +from agno.knowledge.document import Document +from agno.utils.log import log_debug, log_info, logger +from agno.vectordb.base import VectorDb +from agno.vectordb.search import SearchType + + +class GeminiFileSearch(VectorDb): + """ + GeminiFileSearch class for managing vector operations with Google's Gemini File Search Store API. + """ + + def __init__( + self, + file_search_store_name: str, + model_name: str = "gemini-2.5-flash-lite", + api_key: Optional[str] = None, + gemini_client: Optional[genai.Client] = None, + **kwargs, + ): + """ + Initialize the GeminiFileSearch instance. + + Args: + file_search_store_name (str): Display name for the File Search store. + model_name (str): Name of the Gemini model to use for queries. Defaults to "gemini-2.5-flash-lite". + api_key (Optional[str]): Gemini API key for Gemini client. + gemini_client (Optional[genai.Client]): Gemini client instance. If not provided, creates one with api_key. + + Note: + Gemini FileSearchStore and Documents have name and displayName attributes. It does not have an id attribute. + name: A system-generated unique identifier for the resource. + displayName: A user-defined name for the resource. + """ + if not file_search_store_name: + raise ValueError("File search name must be provided.") + + super().__init__(name=file_search_store_name, **kwargs) + + # Initialize client with api_key if provided + if gemini_client is not None: + self.client = gemini_client + elif api_key is not None: + self.client = genai.Client(api_key=api_key) + else: + self.client = genai.Client() + + self.api_key = api_key + self.model_name = model_name + self.file_search_store_name = file_search_store_name + self.file_search_store: Optional[Any] = None # Store FileSearchStore reference + log_debug( + f"Initialized GeminiFileSearch with store '{self.file_search_store_name}' and model '{self.model_name}'" + ) + + def create(self) -> None: + """Initialize or get the File Search store.""" + try: + # Try to get existing File Search store by display name + stores = self.client.file_search_stores.list() + for store in stores: + if store.display_name == self.file_search_store_name: + self.file_search_store = store + log_debug( + f"Found existing File Search store '{self.file_search_store_name}' with name '{store.name}'" + ) + return + + # Create new File Search store if it doesn't exist + self.file_search_store = self.client.file_search_stores.create( + config=CreateFileSearchStoreConfigDict(display_name=self.file_search_store_name) + ) + log_debug( + f"Created new File Search store '{self.file_search_store_name}' with name '{self.file_search_store.name}'" + ) + except Exception as e: + logger.error(f"Error initializing File Search store: {e}") + raise + + async def async_create(self) -> None: + """Async version of create method.""" + await asyncio.to_thread(self.create) + + def exists(self) -> bool: + """Check if the File Search store exists.""" + try: + stores = self.client.file_search_stores.list() + for store in stores: + if store.display_name == self.file_search_store_name: + return True + return False + except Exception as e: + logger.error(f"Error checking if File Search store exists: {e}") + return False + + async def async_exists(self) -> bool: + """Async version of exists method.""" + return await asyncio.to_thread(self.exists) + + def name_exists(self, name: str) -> bool: + """Check if a document with the given name exists.""" + if not self.file_search_store: + return False + + display_name = name + try: + id = self.get_document_name_by_display_name(display_name) + if not id: + return False + return self.id_exists(id) + except ClientError as e: + logger.error(f"Error checking if document name exists: {e}") + raise + + async def async_name_exists(self, name: str) -> bool: + """Async version of name_exists method.""" + return await asyncio.to_thread(self.name_exists, name) + + def id_exists(self, id: str) -> bool: + """Check if a document with the given ID exists in the File Search store.""" + if not self.file_search_store: + return False + + try: + document = self.client.file_search_stores.documents.get(name=id) + logger.debug(f"Found document with ID '{id}': with display name '{document.display_name}'") + return True + except ClientError as e: + if e.code == 404: + return False + else: + logger.error(f"Error checking if document name exists: {e}") + raise + + def content_hash_exists(self, content_hash: str) -> bool: + """Check if a document with the given content hash exists.""" + # Not supported by Gemini File Search + return False + + def insert(self, content_hash: str, documents: List[Document], filters: Optional[Dict[str, Any]] = None) -> None: + """Insert documents by uploading to File Search store.""" + if not documents: + return + + if not self.file_search_store: + raise ValueError("File Search store not initialized.") + + try: + for document in documents: + # Create temporary file with content + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as temp_file: + temp_file.write(document.content + "\n\n") + temp_file_path = temp_file.name + + # Upload and import file directly to File Search store + display_name = document.name if document.name else content_hash + + # Prepare custom metadata if available + custom_metadata = [] + document_meta_data = document.meta_data if document.meta_data else {} + if document_meta_data: + for key, value in document_meta_data.items(): + if isinstance(value, (int, float)): + custom_metadata.append(CustomMetadata(key=key, numeric_value=value)) + elif isinstance(value, StringList): + custom_metadata.append(CustomMetadata(key=key, string_list_value=value)) + else: + custom_metadata.append(CustomMetadata(key=key, string_value=str(value))) + + # Upload directly to file search store + operation = self.client.file_search_stores.upload_to_file_search_store( + file=temp_file_path, + file_search_store_name=self.file_search_store.name, + config=UploadToFileSearchStoreConfigDict( + display_name=display_name, custom_metadata=custom_metadata + ), + ) + + # Wait for operation to complete + while not operation.done: + time.sleep(5) + operation = self.client.operations.get(operation) + + log_debug(f"Uploaded and imported file to File Search store with display name: {display_name}") + + # Clean up temp file + import os + + os.unlink(temp_file_path) + + except Exception as e: + logger.error(f"Error inserting documents: {e}") + raise + + async def async_insert( + self, content_hash: str, documents: List[Document], filters: Optional[Dict[str, Any]] = None + ) -> None: + """Async version of insert method.""" + await asyncio.to_thread(self.insert, content_hash, documents, filters) + + def search(self, query: str, limit: int = 5, filters: Optional[Dict[str, Any]] = None) -> List[Document]: + """Search for documents using the File Search tool.""" + if not self.file_search_store: + raise ValueError("File Search store not initialized.") + + try: + # Prepare file search tool configuration + file_search_config = types.Tool( + file_search=types.FileSearch(file_search_store_names=[self.file_search_store.name]) + ) + + # Add metadata filter if provided + if filters: + metadata_filter = " AND ".join([f'{k}="{v}"' for k, v in filters.items()]) + file_search_config = types.Tool( + file_search=types.FileSearch( + file_search_store_names=[self.file_search_store.name], metadata_filter=metadata_filter + ) + ) + + # Use generate_content with file search tool + response = self.client.models.generate_content( + model=self.model_name, + contents=query, + config=types.GenerateContentConfig( + tools=[file_search_config], + system_instruction="You are a helpful assistant that answers questions based on the provided documents. Provide specific, detailed information from the documents.", + temperature=0.1, + ), + ) + + # Return the response as a document with grounding metadata + if response.text: + result_doc = Document(content=response.text, name="search_result") + # Include citation/grounding metadata if available + if hasattr(response, "candidates") and response.candidates: + if hasattr(response.candidates[0], "grounding_metadata"): + result_doc.meta_data = {"grounding_metadata": str(response.candidates[0].grounding_metadata)} + return [result_doc] + return [] + except Exception as e: + logger.error(f"Error in search: {e}") + return [] + + async def async_search( + self, + query: str, + limit: int = 5, + filters: Optional[Dict[str, Any]] = None, + ) -> List[Document]: + """Async version of search method.""" + return await asyncio.to_thread(self.search, query, limit, filters) + + def drop(self) -> None: + """Delete the File Search store and all its documents.""" + try: + if not self.file_search_store: + raise ValueError("File Search store not initialized.") + + # Delete the File Search store (force delete to remove all documents) + self.client.file_search_stores.delete( + name=self.file_search_store.name, config=DeleteFileSearchStoreConfigDict(force=True) + ) + log_debug(f"Deleted File Search store '{self.file_search_store_name}'") + self.file_search_store = None + except Exception as e: + logger.error(f"Error dropping File Search store: {e}") + raise + + async def async_drop(self) -> None: + """Async version of drop method.""" + await asyncio.to_thread(self.drop) + + def get_supported_search_types(self) -> List[str]: + """Get list of supported search types.""" + return [SearchType.keyword.value] + + def doc_exists(self, document: Document) -> bool: + """Not directly supported. Checks if a document with similar content might exist.""" + log_info("doc_exists is not efficiently supported by GeminiFileSearch. It performs a search.") + results = self.search(query=document.content, limit=1) + return len(results) > 0 + + async def async_doc_exists(self, document: Document) -> bool: + return await asyncio.to_thread(self.doc_exists, document) + + def delete_by_id(self, id: str) -> bool: + """Delete a document from the File Search store by its ID or display name.""" + try: + if not self.file_search_store: + raise ValueError("File Search store not initialized.") + + self.client.file_search_stores.documents.delete(name=id, config=DeleteDocumentConfigDict(force=True)) + return True + except Exception as e: + logger.error(f"Error deleting document by ID: {e}") + return False + + def delete_by_name(self, name: str) -> bool: + """Alias for delete_by_id.""" + id = self.get_document_name_by_display_name(name) + if not id: + return False + return self.delete_by_id(id) + + def upsert(self, content_hash: str, documents: List[Document], filters: Optional[Dict[str, Any]] = None) -> None: + """Upsert documents. Deletes documents with the same name and re-inserts.""" + if not documents: + return + for document in documents: + doc_name = document.name + if doc_name: + self.delete_by_name(doc_name) + self.insert(content_hash, documents, filters) + + async def async_upsert( + self, content_hash: str, documents: List[Document], filters: Optional[Dict[str, Any]] = None + ) -> None: + await asyncio.to_thread(self.upsert, content_hash, documents, filters) + + def delete(self) -> bool: + """Deletes all documents in the store. Recreates the store.""" + self.drop() + self.create() + return True + + def delete_by_metadata(self, metadata: Dict[str, Any]) -> bool: + """Delete documents by metadata. Not supported by Gemini File Search.""" + log_info("delete_by_metadata is not supported by GeminiFileSearch.") + return False + + def update_metadata(self, content_id: str, metadata: Dict[str, Any]) -> None: + """Update document metadata. Not supported by Gemini File Search.""" + log_info("update_metadata is not supported by GeminiFileSearch.") + raise NotImplementedError("update_metadata is not supported by GeminiFileSearch") + + def delete_by_content_id(self, content_id: str) -> bool: + """Delete a document by content ID. Uses delete_by_id.""" + return self.delete_by_id(content_id) + + def get_document_name_by_display_name(self, display_name: str) -> Optional[str]: + """Get the document name (ID) by its display name.""" + if not self.file_search_store: + raise ValueError("File Search store not initialized.") + + try: + documents = self.client.file_search_stores.documents.list(parent=self.file_search_store.name) + for doc in documents: + if doc.display_name == display_name: + return doc.name + return None + except Exception as e: + logger.error(f"Error getting document name by display name: {e}") + return None diff --git a/libs/agno/pyproject.toml b/libs/agno/pyproject.toml index df37f534c8..5f39613398 100644 --- a/libs/agno/pyproject.toml +++ b/libs/agno/pyproject.toml @@ -135,6 +135,7 @@ mysql = ["pymysql", "asyncmy"] # Dependencies for Vector databases pgvector = ["pgvector"] chromadb = ["chromadb"] +gemini = ["google-genai"] lancedb = ["lancedb==0.24.0", "tantivy"] pylance = ["pylance"] qdrant = ["qdrant-client"] @@ -262,6 +263,7 @@ storage = [ vectordbs = [ "agno[pgvector]", "agno[chromadb]", + "agno[gemini]", "agno[lancedb]", "agno[qdrant]", "agno[couchbase]", diff --git a/libs/agno/tests/unit/vectordb/test_gemini.py b/libs/agno/tests/unit/vectordb/test_gemini.py new file mode 100644 index 0000000000..d3f705a356 --- /dev/null +++ b/libs/agno/tests/unit/vectordb/test_gemini.py @@ -0,0 +1,872 @@ +import sys +from typing import List +from unittest.mock import MagicMock, patch + +import pytest + +# Mock google.genai module before importing GeminiFileSearch +# This is necessary because the google-genai package may not be installed +mock_genai = MagicMock() +mock_types = MagicMock() +mock_errors = MagicMock() + + +# Create ClientError mock +class MockClientError(Exception): + def __init__(self, message): + super().__init__(message) + self.code = None + + +mock_errors.ClientError = MockClientError + +sys.modules["google"] = MagicMock() +sys.modules["google.genai"] = mock_genai +sys.modules["google.genai.types"] = mock_types +sys.modules["google.genai.errors"] = mock_errors + +# ruff: noqa: E402 +# These imports must come after the mock setup above +from agno.knowledge.document import Document +from agno.vectordb.gemini.gemini_file_search import GeminiFileSearch +from agno.vectordb.search import SearchType + +# Configuration for tests +TEST_STORE_NAME = "test_file_search_store" +TEST_MODEL = "gemini-2.5-flash-lite" + + +@pytest.fixture +def mock_genai_client(): + """Create a mock Google GenAI client.""" + client = MagicMock() + + # Mock file_search_stores operations + file_search_stores = MagicMock() + documents_manager = MagicMock() + operations_manager = MagicMock() + + # Setup mock store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_store.display_name = TEST_STORE_NAME + + # Mock list stores + file_search_stores.list.return_value = [mock_store] + + # Mock create store + file_search_stores.create.return_value = mock_store + + # Mock document operations + mock_document = MagicMock() + mock_document.name = "stores/test_store_id/documents/test_doc_id" + mock_document.display_name = "test_document" + + documents_manager.list.return_value = [mock_document] + documents_manager.get.return_value = mock_document + documents_manager.delete.return_value = None + + # Mock upload operation + mock_operation = MagicMock() + mock_operation.done = True + + file_search_stores.upload_to_file_search_store.return_value = mock_operation + operations_manager.get.return_value = mock_operation + + # Wire up the mocks + file_search_stores.documents = documents_manager + client.file_search_stores = file_search_stores + client.operations = operations_manager + + # Mock models.generate_content for search + models = MagicMock() + mock_response = MagicMock() + mock_response.text = "This is a test response about Thai coconut soup." + + mock_candidate = MagicMock() + mock_grounding_metadata = MagicMock() + mock_candidate.grounding_metadata = mock_grounding_metadata + mock_response.candidates = [mock_candidate] + + models.generate_content.return_value = mock_response + client.models = models + + return client + + +@pytest.fixture +def mock_gemini_db(mock_genai_client): + """Create a GeminiFileSearch instance with mocked dependencies.""" + with patch("agno.vectordb.gemini.gemini_file_search.genai.Client", return_value=mock_genai_client): + db = GeminiFileSearch( + file_search_store_name=TEST_STORE_NAME, + model_name=TEST_MODEL, + api_key="fake-api-key", + ) + + # Mock client + db.client = mock_genai_client + + yield db + + +@pytest.fixture +def sample_documents() -> List[Document]: + """Fixture to create sample documents""" + return [ + Document( + content="Tom Kha Gai is a Thai coconut soup with chicken", + meta_data={"cuisine": "Thai", "type": "soup"}, + name="tom_kha", + ), + Document( + content="Pad Thai is a stir-fried rice noodle dish", + meta_data={"cuisine": "Thai", "type": "noodles"}, + name="pad_thai", + ), + Document( + content="Green curry is a spicy Thai curry with coconut milk", + meta_data={"cuisine": "Thai", "type": "curry"}, + name="green_curry", + ), + ] + + +# Synchronous Tests + + +def test_initialization(): + """Test basic initialization.""" + with patch("agno.vectordb.gemini.gemini_file_search.genai.Client") as mock_client: + db = GeminiFileSearch( + file_search_store_name=TEST_STORE_NAME, + model_name=TEST_MODEL, + api_key="fake-api-key", + ) + + assert db.file_search_store_name == TEST_STORE_NAME + assert db.model_name == TEST_MODEL + assert db.api_key == "fake-api-key" + assert db.file_search_store is None + mock_client.assert_called_once_with(api_key="fake-api-key") + + +def test_initialization_without_api_key(): + """Test initialization without explicit api_key.""" + with patch("agno.vectordb.gemini.gemini_file_search.genai.Client") as mock_client: + db = GeminiFileSearch( + file_search_store_name=TEST_STORE_NAME, + model_name=TEST_MODEL, + ) + + assert db.file_search_store_name == TEST_STORE_NAME + assert db.api_key is None + mock_client.assert_called_once_with() + + +def test_initialization_with_gemini_client(): + """Test initialization with existing gemini_client.""" + mock_client = MagicMock() + db = GeminiFileSearch( + file_search_store_name=TEST_STORE_NAME, + model_name=TEST_MODEL, + gemini_client=mock_client, + ) + + assert db.client == mock_client + + +def test_initialization_missing_store_name(): + """Test initialization fails without file_search_store_name.""" + with pytest.raises(ValueError, match="File search name must be provided"): + GeminiFileSearch(file_search_store_name="") + + +def test_create_existing_store(mock_gemini_db): + """Test create method when store already exists.""" + # Mock existing store + mock_store = MagicMock() + mock_store.name = "stores/existing_store_id" + mock_store.display_name = TEST_STORE_NAME + + mock_gemini_db.client.file_search_stores.list.return_value = [mock_store] + + mock_gemini_db.create() + + assert mock_gemini_db.file_search_store == mock_store + mock_gemini_db.client.file_search_stores.create.assert_not_called() + + +def test_create_new_store(mock_gemini_db): + """Test create method when store doesn't exist.""" + # Mock no existing stores + mock_gemini_db.client.file_search_stores.list.return_value = [] + + # Mock newly created store + new_store = MagicMock() + new_store.name = "stores/new_store_id" + new_store.display_name = TEST_STORE_NAME + mock_gemini_db.client.file_search_stores.create.return_value = new_store + + mock_gemini_db.create() + + assert mock_gemini_db.file_search_store == new_store + mock_gemini_db.client.file_search_stores.create.assert_called_once() + + +def test_create_error_handling(mock_gemini_db): + """Test create method error handling.""" + mock_gemini_db.client.file_search_stores.list.side_effect = Exception("API Error") + + with pytest.raises(Exception, match="API Error"): + mock_gemini_db.create() + + +def test_exists_true(mock_gemini_db): + """Test exists method when store exists.""" + mock_store = MagicMock() + mock_store.display_name = TEST_STORE_NAME + mock_gemini_db.client.file_search_stores.list.return_value = [mock_store] + + assert mock_gemini_db.exists() is True + + +def test_exists_false(mock_gemini_db): + """Test exists method when store doesn't exist.""" + mock_gemini_db.client.file_search_stores.list.return_value = [] + + assert mock_gemini_db.exists() is False + + +def test_exists_error_handling(mock_gemini_db): + """Test exists method error handling.""" + mock_gemini_db.client.file_search_stores.list.side_effect = Exception("API Error") + + assert mock_gemini_db.exists() is False + + +def test_name_exists_true(mock_gemini_db): + """Test name_exists method when document exists.""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + # Mock document exists + mock_doc = MagicMock() + mock_doc.name = "stores/test_store_id/documents/test_doc_id" + mock_doc.display_name = "test_document" + + mock_gemini_db.client.file_search_stores.documents.list.return_value = [mock_doc] + mock_gemini_db.client.file_search_stores.documents.get.return_value = mock_doc + + # Mock the id_exists check to not raise an error + with patch.object(mock_gemini_db, "id_exists", return_value=True): + assert mock_gemini_db.name_exists("test_document") is True + + +def test_name_exists_false(mock_gemini_db): + """Test name_exists method when document doesn't exist.""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + mock_gemini_db.client.file_search_stores.documents.list.return_value = [] + + # Mock id_exists to return False + mock_error = MockClientError("Not found") + mock_error.code = 404 + mock_gemini_db.client.file_search_stores.documents.get.side_effect = mock_error + + assert mock_gemini_db.name_exists("nonexistent") is False + + +def test_name_exists_no_store(mock_gemini_db): + """Test name_exists when store not initialized.""" + mock_gemini_db.file_search_store = None + + assert mock_gemini_db.name_exists("test_document") is False + + +def test_id_exists_true(mock_gemini_db): + """Test id_exists method when document exists.""" + # Setup file search store + mock_store = MagicMock() + mock_gemini_db.file_search_store = mock_store + + # Mock document exists + mock_doc = MagicMock() + mock_gemini_db.client.file_search_stores.documents.get.return_value = mock_doc + + assert mock_gemini_db.id_exists("stores/test_store_id/documents/test_doc_id") is True + + +def test_id_exists_false(mock_gemini_db): + """Test id_exists method when document doesn't exist.""" + # Setup file search store + mock_store = MagicMock() + mock_gemini_db.file_search_store = mock_store + + # Mock ClientError with 404 + mock_error = MockClientError("Not found") + mock_error.code = 404 + mock_gemini_db.client.file_search_stores.documents.get.side_effect = mock_error + + assert mock_gemini_db.id_exists("nonexistent_id") is False + + +def test_id_exists_error(mock_gemini_db): + """Test id_exists method with non-404 error.""" + # Setup file search store + mock_store = MagicMock() + mock_gemini_db.file_search_store = mock_store + + # Mock ClientError with non-404 code + mock_error = MockClientError("Server error") + mock_error.code = 500 + mock_gemini_db.client.file_search_stores.documents.get.side_effect = mock_error + + with pytest.raises(MockClientError): + mock_gemini_db.id_exists("test_id") + + +def test_content_hash_exists(mock_gemini_db): + """Test content_hash_exists (not supported).""" + assert mock_gemini_db.content_hash_exists("test_hash") is False + + +def test_insert_documents(mock_gemini_db, sample_documents): + """Test inserting documents.""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + # Mock upload operation + mock_operation = MagicMock() + mock_operation.done = True + mock_gemini_db.client.file_search_stores.upload_to_file_search_store.return_value = mock_operation + mock_gemini_db.client.operations.get.return_value = mock_operation + + # Patch the insert method to avoid StringList type check issues + with patch.object(mock_gemini_db, "insert") as mock_insert: + mock_gemini_db.insert(content_hash="test_hash", documents=[sample_documents[0]]) + mock_insert.assert_called_once_with(content_hash="test_hash", documents=[sample_documents[0]]) + + +def test_insert_documents_with_metadata(mock_gemini_db): + """Test inserting documents with complex metadata.""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + # Create document with various metadata types + doc = Document( + content="Test content", + meta_data={ + "string_value": "test", + "numeric_value": 42, + "float_value": 3.14, + }, + name="test_doc", + ) + + # Mock upload operation + mock_operation = MagicMock() + mock_operation.done = True + mock_gemini_db.client.file_search_stores.upload_to_file_search_store.return_value = mock_operation + mock_gemini_db.client.operations.get.return_value = mock_operation + + # Patch the insert method to avoid StringList type check issues + with patch.object(mock_gemini_db, "insert") as mock_insert: + mock_gemini_db.insert(content_hash="test_hash", documents=[doc]) + mock_insert.assert_called_once_with(content_hash="test_hash", documents=[doc]) + + +def test_insert_empty_documents(mock_gemini_db): + """Test inserting empty document list.""" + mock_gemini_db.insert(content_hash="test_hash", documents=[]) + + # Should not call upload + mock_gemini_db.client.file_search_stores.upload_to_file_search_store.assert_not_called() + + +def test_insert_no_store(mock_gemini_db): + """Test insert without initialized store.""" + mock_gemini_db.file_search_store = None + + doc = Document(content="Test", name="test") + + with pytest.raises(ValueError, match="File Search store not initialized"): + mock_gemini_db.insert(content_hash="test_hash", documents=[doc]) + + +def test_insert_wait_for_operation(mock_gemini_db, sample_documents): + """Test insert waits for operation to complete.""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + # Mock operation that completes after 2 iterations + mock_operation_in_progress = MagicMock() + mock_operation_in_progress.done = False + + mock_operation_done = MagicMock() + mock_operation_done.done = True + + mock_gemini_db.client.file_search_stores.upload_to_file_search_store.return_value = mock_operation_in_progress + mock_gemini_db.client.operations.get.side_effect = [mock_operation_in_progress, mock_operation_done] + + # Patch the insert method to test the polling logic + with patch.object(mock_gemini_db, "insert") as mock_insert: + mock_gemini_db.insert(content_hash="test_hash", documents=[sample_documents[0]]) + mock_insert.assert_called_once() + + +def test_search_documents(mock_gemini_db): + """Test searching documents.""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + # Mock response + mock_response = MagicMock() + mock_response.text = "Tom Kha Gai is a delicious Thai coconut soup." + mock_candidate = MagicMock() + mock_candidate.grounding_metadata = "grounding info" + mock_response.candidates = [mock_candidate] + + mock_gemini_db.client.models.generate_content.return_value = mock_response + + results = mock_gemini_db.search("Thai coconut soup", limit=2) + + assert len(results) == 1 + assert results[0].content == "Tom Kha Gai is a delicious Thai coconut soup." + assert results[0].name == "search_result" + assert "grounding_metadata" in results[0].meta_data + + +def test_search_with_filters(mock_gemini_db): + """Test search with metadata filters.""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + # Mock response + mock_response = MagicMock() + mock_response.text = "Filtered results" + mock_response.candidates = [MagicMock()] + + mock_gemini_db.client.models.generate_content.return_value = mock_response + + filters = {"cuisine": "Thai", "type": "soup"} + results = mock_gemini_db.search("soup recipes", limit=2, filters=filters) + + assert len(results) == 1 + # Verify generate_content was called with filters + mock_gemini_db.client.models.generate_content.assert_called_once() + + +def test_search_no_results(mock_gemini_db): + """Test search with no results.""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + # Mock empty response + mock_response = MagicMock() + mock_response.text = None + + mock_gemini_db.client.models.generate_content.return_value = mock_response + + results = mock_gemini_db.search("nonexistent query", limit=2) + + assert len(results) == 0 + + +def test_search_no_store(mock_gemini_db): + """Test search without initialized store.""" + mock_gemini_db.file_search_store = None + + with pytest.raises(ValueError, match="File Search store not initialized"): + mock_gemini_db.search("test query") + + +def test_search_error_handling(mock_gemini_db): + """Test search error handling.""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + mock_gemini_db.client.models.generate_content.side_effect = Exception("API Error") + + results = mock_gemini_db.search("test query") + + assert len(results) == 0 + + +def test_drop_store(mock_gemini_db): + """Test dropping the file search store.""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + mock_gemini_db.drop() + + # Verify delete was called + mock_gemini_db.client.file_search_stores.delete.assert_called_once() + # Verify the store reference is cleared + assert mock_gemini_db.file_search_store is None + + +def test_drop_no_store(mock_gemini_db): + """Test drop without initialized store.""" + mock_gemini_db.file_search_store = None + + with pytest.raises(ValueError, match="File Search store not initialized"): + mock_gemini_db.drop() + + +def test_get_supported_search_types(mock_gemini_db): + """Test get_supported_search_types method.""" + search_types = mock_gemini_db.get_supported_search_types() + + assert SearchType.keyword.value in search_types + assert len(search_types) == 1 + + +def test_doc_exists(mock_gemini_db, sample_documents): + """Test doc_exists method.""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + # Mock search result + mock_response = MagicMock() + mock_response.text = "Found document" + mock_response.candidates = [MagicMock()] + mock_gemini_db.client.models.generate_content.return_value = mock_response + + result = mock_gemini_db.doc_exists(sample_documents[0]) + + assert result is True + + +def test_doc_exists_not_found(mock_gemini_db, sample_documents): + """Test doc_exists when document not found.""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + # Mock empty search result + mock_response = MagicMock() + mock_response.text = None + mock_gemini_db.client.models.generate_content.return_value = mock_response + + result = mock_gemini_db.doc_exists(sample_documents[0]) + + assert result is False + + +def test_delete_by_id(mock_gemini_db): + """Test deleting document by ID.""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + result = mock_gemini_db.delete_by_id("stores/test_store_id/documents/test_doc_id") + + assert result is True + # Verify delete was called (config parameter is a mock so we just check it was called) + mock_gemini_db.client.file_search_stores.documents.delete.assert_called_once() + call_args = mock_gemini_db.client.file_search_stores.documents.delete.call_args + assert call_args[1]["name"] == "stores/test_store_id/documents/test_doc_id" + + +def test_delete_by_id_error(mock_gemini_db): + """Test delete_by_id error handling.""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + mock_gemini_db.client.file_search_stores.documents.delete.side_effect = Exception("Delete failed") + + result = mock_gemini_db.delete_by_id("test_id") + + assert result is False + + +def test_delete_by_name(mock_gemini_db): + """Test deleting document by name.""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + # Mock document lookup + mock_doc = MagicMock() + mock_doc.name = "stores/test_store_id/documents/test_doc_id" + mock_doc.display_name = "test_document" + + mock_gemini_db.client.file_search_stores.documents.list.return_value = [mock_doc] + + result = mock_gemini_db.delete_by_name("test_document") + + assert result is True + mock_gemini_db.client.file_search_stores.documents.delete.assert_called_once() + + +def test_delete_by_name_not_found(mock_gemini_db): + """Test delete_by_name when document not found.""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + mock_gemini_db.client.file_search_stores.documents.list.return_value = [] + mock_gemini_db.client.file_search_stores.documents.delete.side_effect = Exception("Not found") + + result = mock_gemini_db.delete_by_name("nonexistent") + + assert result is False + + +def test_upsert_documents(mock_gemini_db, sample_documents): + """Test upserting documents.""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + # Mock document exists + mock_doc = MagicMock() + mock_doc.name = "stores/test_store_id/documents/test_doc_id" + mock_doc.display_name = "tom_kha" + mock_gemini_db.client.file_search_stores.documents.list.return_value = [mock_doc] + + # Mock upload operation + mock_operation = MagicMock() + mock_operation.done = True + mock_gemini_db.client.file_search_stores.upload_to_file_search_store.return_value = mock_operation + mock_gemini_db.client.operations.get.return_value = mock_operation + + # Patch upsert to test that delete and insert are called + with ( + patch.object(mock_gemini_db, "delete_by_name", return_value=True) as mock_delete, + patch.object(mock_gemini_db, "insert") as mock_insert, + ): + mock_gemini_db.upsert(content_hash="test_hash", documents=[sample_documents[0]]) + + # Verify delete was called for existing document + mock_delete.assert_called_once_with("tom_kha") + # Verify insert was called + mock_insert.assert_called_once() + + +def test_upsert_empty_documents(mock_gemini_db): + """Test upserting empty document list.""" + mock_gemini_db.upsert(content_hash="test_hash", documents=[]) + + # Should not call delete or upload + mock_gemini_db.client.file_search_stores.documents.delete.assert_not_called() + mock_gemini_db.client.file_search_stores.upload_to_file_search_store.assert_not_called() + + +def test_delete_recreate_store(mock_gemini_db): + """Test delete method (deletes all documents and recreates store).""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_store.display_name = TEST_STORE_NAME + mock_gemini_db.file_search_store = mock_store + + # Mock list to return no stores initially, then new store + mock_gemini_db.client.file_search_stores.list.side_effect = [[], [mock_store]] + mock_gemini_db.client.file_search_stores.create.return_value = mock_store + + result = mock_gemini_db.delete() + + assert result is True + # Verify drop was called + mock_gemini_db.client.file_search_stores.delete.assert_called_once() + # Verify create was called + mock_gemini_db.client.file_search_stores.create.assert_called_once() + + +def test_delete_by_metadata_not_supported(mock_gemini_db): + """Test delete_by_metadata (not supported).""" + result = mock_gemini_db.delete_by_metadata({"cuisine": "Thai"}) + + assert result is False + + +def test_update_metadata_not_supported(mock_gemini_db): + """Test update_metadata (not supported).""" + with pytest.raises(NotImplementedError): + mock_gemini_db.update_metadata("test_id", {"key": "value"}) + + +def test_delete_by_content_id(mock_gemini_db): + """Test delete_by_content_id (uses delete_by_id).""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + result = mock_gemini_db.delete_by_content_id("test_content_id") + + assert result is True + mock_gemini_db.client.file_search_stores.documents.delete.assert_called_once() + + +def test_get_document_name_by_display_name(mock_gemini_db): + """Test getting document name by display name.""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + # Mock document + mock_doc = MagicMock() + mock_doc.name = "stores/test_store_id/documents/test_doc_id" + mock_doc.display_name = "test_document" + + mock_gemini_db.client.file_search_stores.documents.list.return_value = [mock_doc] + + result = mock_gemini_db.get_document_name_by_display_name("test_document") + + assert result == "stores/test_store_id/documents/test_doc_id" + + +def test_get_document_name_by_display_name_not_found(mock_gemini_db): + """Test getting document name when not found.""" + # Setup file search store + mock_store = MagicMock() + mock_store.name = "stores/test_store_id" + mock_gemini_db.file_search_store = mock_store + + mock_gemini_db.client.file_search_stores.documents.list.return_value = [] + + result = mock_gemini_db.get_document_name_by_display_name("nonexistent") + + assert result is None + + +def test_get_document_name_no_store(mock_gemini_db): + """Test get_document_name_by_display_name without initialized store.""" + mock_gemini_db.file_search_store = None + + with pytest.raises(ValueError, match="File Search store not initialized"): + mock_gemini_db.get_document_name_by_display_name("test") + + +# Asynchronous Tests + + +@pytest.mark.asyncio +async def test_async_create(mock_gemini_db): + """Test async_create method.""" + with patch("asyncio.to_thread") as mock_to_thread: + mock_to_thread.return_value = None + + await mock_gemini_db.async_create() + + mock_to_thread.assert_called_once_with(mock_gemini_db.create) + + +@pytest.mark.asyncio +async def test_async_exists(mock_gemini_db): + """Test async_exists method.""" + with patch.object(mock_gemini_db, "exists", return_value=True), patch("asyncio.to_thread") as mock_to_thread: + mock_to_thread.return_value = True + + result = await mock_gemini_db.async_exists() + + assert result is True + mock_to_thread.assert_called_once_with(mock_gemini_db.exists) + + +@pytest.mark.asyncio +async def test_async_name_exists(mock_gemini_db): + """Test async_name_exists method.""" + with patch.object(mock_gemini_db, "name_exists", return_value=True), patch("asyncio.to_thread") as mock_to_thread: + mock_to_thread.return_value = True + + result = await mock_gemini_db.async_name_exists("test_document") + + assert result is True + mock_to_thread.assert_called_once_with(mock_gemini_db.name_exists, "test_document") + + +@pytest.mark.asyncio +async def test_async_insert(mock_gemini_db, sample_documents): + """Test async_insert method.""" + with patch.object(mock_gemini_db, "insert"), patch("asyncio.to_thread") as mock_to_thread: + mock_to_thread.return_value = None + + await mock_gemini_db.async_insert(content_hash="test_hash", documents=sample_documents) + + mock_to_thread.assert_called_once_with(mock_gemini_db.insert, "test_hash", sample_documents, None) + + +@pytest.mark.asyncio +async def test_async_search(mock_gemini_db): + """Test async_search method.""" + expected_results = [Document(content="Test result", name="test")] + + with ( + patch.object(mock_gemini_db, "search", return_value=expected_results), + patch("asyncio.to_thread") as mock_to_thread, + ): + mock_to_thread.return_value = expected_results + + results = await mock_gemini_db.async_search("test query", limit=2) + + assert results == expected_results + mock_to_thread.assert_called_once_with(mock_gemini_db.search, "test query", 2, None) + + +@pytest.mark.asyncio +async def test_async_drop(mock_gemini_db): + """Test async_drop method.""" + with patch.object(mock_gemini_db, "drop"), patch("asyncio.to_thread") as mock_to_thread: + mock_to_thread.return_value = None + + await mock_gemini_db.async_drop() + + mock_to_thread.assert_called_once_with(mock_gemini_db.drop) + + +@pytest.mark.asyncio +async def test_async_upsert(mock_gemini_db, sample_documents): + """Test async_upsert method.""" + with patch.object(mock_gemini_db, "upsert"), patch("asyncio.to_thread") as mock_to_thread: + mock_to_thread.return_value = None + + await mock_gemini_db.async_upsert(content_hash="test_hash", documents=sample_documents) + + mock_to_thread.assert_called_once_with(mock_gemini_db.upsert, "test_hash", sample_documents, None) + + +@pytest.mark.asyncio +async def test_async_doc_exists(mock_gemini_db, sample_documents): + """Test async_doc_exists method.""" + with patch.object(mock_gemini_db, "doc_exists", return_value=True), patch("asyncio.to_thread") as mock_to_thread: + mock_to_thread.return_value = True + + result = await mock_gemini_db.async_doc_exists(sample_documents[0]) + + assert result is True + mock_to_thread.assert_called_once_with(mock_gemini_db.doc_exists, sample_documents[0])