Repository: agno-agi/agno
Pull request: Update lancedb and gemini embedder
## Description
This makes a small update to Lance DB and update GeminiEmbedder to using the latest Gemini SDK

Fixes #1865
---

## Type of change

Please check the options that are relevant:

- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)
- [x] Model update (Addition or modification of models)
- [ ] Other (please describe):

---

## Checklist

- [ ] Adherence to standards: Code complies with Agno’s style guidelines and best practices.
- [ ] Formatting and validation: You have run `./scripts/format.sh` and `./scripts/validate.sh` to ensure code is formatted and linted.
- [ ] Self-review completed: A thorough review has been performed by the contributor(s).
- [ ] Documentation: Docstrings and comments have been added or updated for any complex logic.
- [ ] Examples and guides: Relevant cookbook examples have been included or updated (if applicable).
- [ ] Tested in a clean environment: Changes have been tested in a clean environment to confirm expected behavior.
- [ ] Tests (optional): Tests have been added or updated to cover any new or changed functionality.

---

## Additional Notes

Include any deployment notes, performance implications, security considerations, or other relevant information (e.g., screenshots or logs if applicable).


Problem statement:
ValueError: Vector column vector has variable length vectors Set on_bad_vectors='drop' to remove them, or set on_bad_vectors='fill' and fill_value=<value> to replace them.
Code:

```
txt_path = "/media/ahamad/ontent_large_pdf"
# # html_paths = get_sds_paths(pdf_dir, extn=".html")  
# chunking_strategy=SemanticChunking()

# vector_db=PgVector(
#     table_name="autozone_documents",
#     db_url="postgresql+psycopg://ai:ai@localhost:5432/ai",
#     embedder=OllamaEmbedder(),
# ),

vector_db = LanceDb(
    table_name="auto_zone",
    uri="/tmp/lancedb",
    embedder=OllamaEmbedder()
)

knowledge_base = PDFKnowledgeBase(
    path=txt_path,
    vector_db=vector_db,     
    # vector_db=ChromaDb(collection="auto_zone", embedder=OllamaEmbedder()),
    # chunking_strategy=SemanticChunking()
    chunking_strategy=FixedSizeChunking()
)

knowledge_base.load(recreate=False)
```
Error:

```
Exception has occurred: ValueError
Vector column vector has variable length vectors Set on_bad_vectors='drop' to remove them, or set on_bad_vectors='fill' and fill_value=<value> to replace them.
  File "/media/ahamad/work/projects/auto_zone/Autozone_experiments/knowledge_base.py", line 47, in <module>
    knowledge_base.load(recreate=False)
ValueError: Vector column vector has variable length vectors Set on_bad_vectors='drop' to remove them, or set on_bad_vectors='fill' and fill_value=<value> to replace them.
```

Hints from issue discussion:
@AhmadShaik can you please confirm with me if you are using Phidata or Agno? If Phidata then what version are you using right now? If Agno then can you please share the complete code with imports as well?
Got the same error using phidata 2.7.10 and Python 3.12.

Code:
```
from phi.agent import Agent
from phi.embedder.ollama import OllamaEmbedder
from phi.knowledge.pdf import PDFKnowledgeBase
from phi.vectordb.lancedb import LanceDb, SearchType
from phi.model.ollama import Ollama

# Create a knowledge base from a PDF

knowledge_base = PDFKnowledgeBase(
    path="tmp",
    # Use LanceDB as the vector database
    vector_db=LanceDb(
        table_name="recipes",
        uri="tmp/lancedb_agenticrag",
        search_type=SearchType.vector,
        embedder=OllamaEmbedder(model="all-minilm"),
    ),
)
# Comment out after first run as the knowledge base is loaded
knowledge_base.load()

agent = Agent(
    model=Ollama(id="llama3.2"),
    # Add the knowledge base to the agent
    knowledge=knowledge_base,
    show_tool_calls=True,
    markdown=True,
)
agent.print_response("Which ingredients are needed for beef wellington", stream=True)
```

Error occurs on the line with `knowledge_base.load()`
```
ValueError: Vector column 'vector' has variable length vectors. Set on_bad_vectors='drop' to remove them, set on_bad_vectors='fill' and fill_value=<value> to replace them, or set on_bad_vectors='null' to replace them with null.
```

Thinking it is about recreation option.
@researchk8 I tried to replicate it locally and I think your issue is that your knowledge_base had a different embedder or structure before (i.e. there is existing files). You would have to recreate.

Please let me know if this resolves your issue.
We'll update the LanceDB implementation to allow for this.

Commit message: Merge

Patch to review:
diff --git a/cookbook/agent_concepts/knowledge/vector_dbs/lance_db.py b/cookbook/agent_concepts/knowledge/vector_dbs/lance_db.py
index ae0486db61..4291b00b1b 100644
--- a/cookbook/agent_concepts/knowledge/vector_dbs/lance_db.py
+++ b/cookbook/agent_concepts/knowledge/vector_dbs/lance_db.py
@@ -8,7 +8,7 @@
 # By default, it stores data in /tmp/lancedb
 vector_db = LanceDb(
     table_name="recipes",
-    uri="/tmp/lancedb",  # You can change this path to store data elsewhere
+    uri="tmp/lancedb",  # You can change this path to store data elsewhere
 )
 
 # Create knowledge base
@@ -17,7 +17,7 @@
     vector_db=vector_db,
 )
 
-knowledge_base.load(recreate=False)  # Comment out after first run
+knowledge_base.load(recreate=True)  # Comment out after first run
 
 # Create and use the agent
 agent = Agent(knowledge=knowledge_base, show_tool_calls=True)
diff --git a/libs/agno/agno/embedder/google.py b/libs/agno/agno/embedder/google.py
index b177d5c712..72bc73ebaa 100644
--- a/libs/agno/agno/embedder/google.py
+++ b/libs/agno/agno/embedder/google.py
@@ -1,34 +1,34 @@
 from dataclasses import dataclass
 from os import getenv
-from types import ModuleType
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple
 
 from agno.embedder.base import Embedder
 from agno.utils.log import logger
 
 try:
-    import google.generativeai as genai
-    from google.generativeai.types.text_types import BatchEmbeddingDict, EmbeddingDict
+    from google import genai
+    from google.genai import Client as GeminiClient
+    from google.genai.types import EmbedContentResponse
 except ImportError:
-    logger.error("`google-generativeai` not installed. Please install it using `pip install google-generativeai`")
-    raise
+    raise ImportError("`google-genai` not installed. Please install it using `pip install google-genai`")
 
 
 @dataclass
 class GeminiEmbedder(Embedder):
-    id: str = "models/text-embedding-004"
+    id: str = "text-embedding-004"
     task_type: str = "RETRIEVAL_QUERY"
     title: Optional[str] = None
     dimensions: Optional[int] = 768
     api_key: Optional[str] = None
     request_params: Optional[Dict[str, Any]] = None
     client_params: Optional[Dict[str, Any]] = None
-    gemini_client: Optional[ModuleType] = None
+    gemini_client: Optional[GeminiClient] = None
 
     @property
     def client(self):
         if self.gemini_client:
             return self.gemini_client
+
         _client_params: Dict[str, Any] = {}
 
         self.api_key = self.api_key or getenv("GOOGLE_API_KEY")
@@ -39,35 +39,39 @@ def client(self):
             _client_params["api_key"] = self.api_key
         if self.client_params:
             _client_params.update(self.client_params)
-        self.gemini_client = genai
-        self.gemini_client.configure(**_client_params)  # type: ignore
+
+        self.gemini_client = genai.Client(**_client_params)
+
         return self.gemini_client
 
-    def _response(self, text: str) -> Union[EmbeddingDict, BatchEmbeddingDict]:
-        _request_params: Dict[str, Any] = {
-            "content": text,
-            "model": self.id,
-            "output_dimensionality": self.dimensions,
-            "task_type": self.task_type,
-            "title": self.title,
-        }
+    def _response(self, text: str) -> EmbedContentResponse:
+        _request_params: Dict[str, Any] = {"contents": text, "model": self.id, "config": {}}
+        if self.dimensions:
+            _request_params["config"]["output_dimensionality"] = self.dimensions
+        if self.task_type:
+            _request_params["config"]["task_type"] = self.task_type
+        if self.title:
+            _request_params["config"]["title"] = self.title
+        if not _request_params["config"]:
+            del _request_params["config"]
+
         if self.request_params:
             _request_params.update(self.request_params)
-        return self.client.embed_content(**_request_params)
+        return self.client.models.embed_content(**_request_params)
 
     def get_embedding(self, text: str) -> List[float]:
         response = self._response(text=text)
         try:
-            return response.get("embedding", [])  # type: ignore
+            return response.embeddings[0].values
         except Exception as e:
             logger.warning(e)
             return []
 
     def get_embedding_and_usage(self, text: str) -> Tuple[List[float], Optional[Dict]]:
         response = self._response(text=text)
-        usage = None
+        usage = response.metadata.billable_character_count if response.metadata else None
         try:
-            return response.get("embedding", []), usage  # type: ignore
+            return response.embeddings[0].values, usage
         except Exception as e:
             logger.warning(e)
             return [], usage
diff --git a/libs/agno/agno/knowledge/csv.py b/libs/agno/agno/knowledge/csv.py
index d43c11783f..3cf95fefbf 100644
--- a/libs/agno/agno/knowledge/csv.py
+++ b/libs/agno/agno/knowledge/csv.py
@@ -1,6 +1,8 @@
 from pathlib import Path
 from typing import Iterator, List, Union
 
+from pydantic import Field
+
 from agno.document import Document
 from agno.document.reader.csv_reader import CSVReader
 from agno.knowledge.agent import AgentKnowledge
@@ -8,6 +10,7 @@
 
 class CSVKnowledgeBase(AgentKnowledge):
     path: Union[str, Path]
+    exclude_files: List[str] = Field(default_factory=list)
     reader: CSVReader = CSVReader()
 
     @property
@@ -23,6 +26,10 @@ def document_lists(self) -> Iterator[List[Document]]:
 
         if _csv_path.exists() and _csv_path.is_dir():
             for _csv in _csv_path.glob("**/*.csv"):
+                if _csv.name in self.exclude_files:
+                    continue
                 yield self.reader.read(file=_csv)
         elif _csv_path.exists() and _csv_path.is_file() and _csv_path.suffix == ".csv":
+            if _csv_path.name in self.exclude_files:
+                return
             yield self.reader.read(file=_csv_path)
diff --git a/libs/agno/agno/knowledge/pdf.py b/libs/agno/agno/knowledge/pdf.py
index ba6f40e778..dcdae886a5 100644
--- a/libs/agno/agno/knowledge/pdf.py
+++ b/libs/agno/agno/knowledge/pdf.py
@@ -1,6 +1,8 @@
 from pathlib import Path
 from typing import Iterator, List, Union
 
+from pydantic import Field
+
 from agno.document import Document
 from agno.document.reader.pdf_reader import PDFImageReader, PDFReader
 from agno.knowledge.agent import AgentKnowledge
@@ -8,6 +10,9 @@
 
 class PDFKnowledgeBase(AgentKnowledge):
     path: Union[str, Path]
+
+    exclude_files: List[str] = Field(default_factory=list)
+
     reader: Union[PDFReader, PDFImageReader] = PDFReader()
 
     @property
@@ -23,6 +28,10 @@ def document_lists(self) -> Iterator[List[Document]]:
 
         if _pdf_path.exists() and _pdf_path.is_dir():
             for _pdf in _pdf_path.glob("**/*.pdf"):
+                if _pdf.name in self.exclude_files:
+                    continue
                 yield self.reader.read(pdf=_pdf)
         elif _pdf_path.exists() and _pdf_path.is_file() and _pdf_path.suffix == ".pdf":
+            if _pdf_path.name in self.exclude_files:
+                return
             yield self.reader.read(pdf=_pdf_path)
diff --git a/libs/agno/agno/vectordb/lancedb/lance_db.py b/libs/agno/agno/vectordb/lancedb/lance_db.py
index 13d99f2580..951d5d7843 100644
--- a/libs/agno/agno/vectordb/lancedb/lance_db.py
+++ b/libs/agno/agno/vectordb/lancedb/lance_db.py
@@ -18,6 +18,25 @@
 
 
 class LanceDb(VectorDb):
+    """
+    LanceDB is a vector database that uses LanceDB as the underlying database.
+
+    Args:
+        uri: The URI of the LanceDB database.
+        table: The LanceDB table instance to use.
+        table_name: The name of the LanceDB table to use.
+        connection: The LanceDB connection to use.
+        api_key: The API key to use for the LanceDB connection.
+        embedder: The embedder to use when embedding the document contents.
+        search_type: The search type to use when searching for documents.
+        distance: The distance metric to use when searching for documents.
+        nprobes: The number of probes to use when searching for documents.
+        reranker: The reranker to use when reranking documents.
+        use_tantivy: Whether to use Tantivy for full text search.
+        on_bad_vectors: What to do if the vector is bad. One of "error", "drop", "fill", "null".
+        fill_value: The value to fill the vector with if on_bad_vectors is "fill".
+    """
+
     def __init__(
         self,
         uri: lancedb.URI = "/tmp/lancedb",
@@ -31,6 +50,8 @@ def __init__(
         nprobes: Optional[int] = None,
         reranker: Optional[Reranker] = None,
         use_tantivy: bool = True,
+        on_bad_vectors: Optional[str] = None,  # One of "error", "drop", "fill", "null".
+        fill_value: Optional[float] = None,  # Only used if on_bad_vectors is "fill"
     ):
         # Embedder for embedding the document contents
         if embedder is None:
@@ -84,6 +105,8 @@ def __init__(
 
         self.reranker: Optional[Reranker] = reranker
         self.nprobes: Optional[int] = nprobes
+        self.on_bad_vectors: Optional[str] = on_bad_vectors
+        self.fill_value: Optional[float] = fill_value
         self.fts_index_exists = False
         self.use_tantivy = use_tantivy
 
@@ -128,11 +151,16 @@ def doc_exists(self, document: Document) -> bool:
         Args:
             document (Document): Document to validate
         """
-        if self.table is not None:
-            cleaned_content = document.content.replace("\x00", "\ufffd")
-            doc_id = md5(cleaned_content.encode()).hexdigest()
-            result = self.table.search().where(f"{self._id}='{doc_id}'").to_arrow()
-            return len(result) > 0
+        try:
+            if self.table is not None:
+                cleaned_content = document.content.replace("\x00", "\ufffd")
+                doc_id = md5(cleaned_content.encode()).hexdigest()
+                result = self.table.search().where(f"{self._id}='{doc_id}'").to_arrow()
+                return len(result) > 0
+        except Exception:
+            # Search sometimes fails with stale cache data, it means the doc doesn't exist
+            return False
+
         return False
 
     def insert(self, documents: List[Document], filters: Optional[Dict[str, Any]] = None) -> None:
@@ -176,7 +204,11 @@ def insert(self, documents: List[Document], filters: Optional[Dict[str, Any]] =
             logger.debug("No new data to insert")
             return
 
-        self.table.add(data)
+        if self.on_bad_vectors is not None:
+            self.table.add(data, on_bad_vectors=self.on_bad_vectors, fill_value=self.fill_value)
+        else:
+            self.table.add(data)
+
         logger.debug(f"Inserted {len(data)} documents")
 
     def upsert(self, documents: List[Document], filters: Optional[Dict[str, Any]] = None) -> None:
@@ -326,7 +358,17 @@ def delete(self) -> bool:
         return False
 
     def name_exists(self, name: str) -> bool:
+        """Check if a document with the given name exists in the database"""
         if self.table is None:
             return False
-        result = self.table.search().select(["payload"]).to_pandas()
-        return any(row["payload"].get("name") == name for _, row in result.iterrows())
+
+        try:
+            result = self.table.search().select(["payload"]).to_pandas()
+            # Convert the JSON strings in payload column to dictionaries
+            payloads = result["payload"].apply(json.loads)
+
+            # Check if the name exists in any of the payloads
+            return any(payload.get("name") == name for payload in payloads)
+        except Exception as e:
+            logger.error(f"Error checking name existence: {e}")
+            return False
diff --git a/libs/agno/tests/unit/vectordb/test_chromadb.py b/libs/agno/tests/unit/vectordb/test_chromadb.py
index b13690aa74..80caf3d851 100644
--- a/libs/agno/tests/unit/vectordb/test_chromadb.py
+++ b/libs/agno/tests/unit/vectordb/test_chromadb.py
@@ -138,8 +138,7 @@ def test_get_count(chroma_db, sample_documents):
     assert chroma_db.get_count() == 3
 
 
-@pytest.mark.asyncio
-async def test_error_handling(chroma_db):
+def test_error_handling(chroma_db):
     """Test error handling scenarios"""
     # Test search with invalid query
     results = chroma_db.search("")
diff --git a/libs/agno/tests/unit/vectordb/test_lancedb.py b/libs/agno/tests/unit/vectordb/test_lancedb.py
new file mode 100644
index 0000000000..3a30db1e30
--- /dev/null
+++ b/libs/agno/tests/unit/vectordb/test_lancedb.py
@@ -0,0 +1,185 @@
+import os
+import shutil
+from typing import List
+
+import pytest
+
+from agno.document import Document
+from agno.embedder.openai import OpenAIEmbedder
+from agno.vectordb.lancedb import LanceDb
+from agno.vectordb.search import SearchType
+
+TEST_TABLE = "test_table"
+TEST_PATH = "tmp/test_lancedb"
+
+
+@pytest.fixture
+def lance_db():
+    """Fixture to create and clean up a LanceDb instance"""
+    # Ensure the test directory exists with proper permissions
+    os.makedirs(TEST_PATH, exist_ok=True)
+
+    # Clean up any existing data before the test
+    if os.path.exists(TEST_PATH):
+        shutil.rmtree(TEST_PATH)
+        os.makedirs(TEST_PATH)
+
+    db = LanceDb(uri=TEST_PATH, table_name=TEST_TABLE, embedder=OpenAIEmbedder())
+    db.create()
+    yield db
+
+    # Cleanup after test
+    try:
+        db.drop()
+    except Exception:
+        pass
+
+    if os.path.exists(TEST_PATH):
+        shutil.rmtree(TEST_PATH)
+
+
+@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",
+        ),
+    ]
+
+
+def test_create_table(lance_db):
+    """Test creating a table"""
+    assert lance_db.exists() is True
+    assert lance_db.get_count() == 0
+
+
+def test_insert_documents(lance_db, sample_documents):
+    """Test inserting documents"""
+    lance_db.insert(sample_documents)
+    assert lance_db.get_count() == 3
+
+
+def test_vector_search(lance_db, sample_documents):
+    """Test vector search"""
+    lance_db.insert(sample_documents)
+
+    # Search for coconut-related dishes
+    results = lance_db.vector_search("coconut dishes", limit=2)
+    assert len(results) == 2
+    assert any("coconut" in doc.content.lower() for doc in results)
+
+
+def test_keyword_search(lance_db, sample_documents):
+    """Test keyword search"""
+    lance_db.search_type = SearchType.keyword
+    lance_db.insert(sample_documents)
+
+    results = lance_db.search("spicy curry", limit=1)
+    assert len(results) == 1
+    assert "curry" in results[0].content.lower()
+
+
+def test_hybrid_search(lance_db, sample_documents):
+    """Test hybrid search"""
+    lance_db.search_type = SearchType.hybrid
+    lance_db.insert(sample_documents)
+
+    results = lance_db.search("Thai soup", limit=2)
+    assert len(results) == 2
+    assert any("thai" in doc.content.lower() for doc in results)
+
+
+def test_upsert_documents(lance_db, sample_documents):
+    """Test upserting documents"""
+    # Initial insert
+    lance_db.insert([sample_documents[0]])
+    assert lance_db.get_count() == 1
+
+    # Upsert same document with different content
+    modified_doc = Document(
+        content="Tom Kha Gai is a spicy and sour Thai coconut soup",
+        meta_data={"cuisine": "Thai", "type": "soup"},
+        name="tom_kha",
+    )
+    lance_db.upsert([modified_doc])
+
+    # Search to verify the update
+    results = lance_db.search("spicy and sour", limit=1)
+    assert len(results) == 1
+    assert "spicy and sour" in results[0].content
+
+
+def test_doc_exists(lance_db, sample_documents):
+    """Test document existence check"""
+    lance_db.insert([sample_documents[0]])
+    assert lance_db.doc_exists(sample_documents[0]) is True
+
+
+def test_name_exists(lance_db, sample_documents):
+    """Test name existence check"""
+    lance_db.insert([sample_documents[0]])
+    assert lance_db.name_exists("tom_kha") is True
+    assert lance_db.name_exists("nonexistent") is False
+
+
+def test_get_count(lance_db, sample_documents):
+    """Test document count"""
+    assert lance_db.get_count() == 0
+    lance_db.insert(sample_documents)
+    assert lance_db.get_count() == 3
+
+
+def test_error_handling(lance_db):
+    """Test error handling scenarios"""
+    # Test search with empty query
+    results = lance_db.search("")
+    assert len(results) == 0
+
+    # Test inserting empty document list
+    lance_db.insert([])
+    assert lance_db.get_count() == 0
+
+
+def test_bad_vectors_handling():
+    """Test handling of bad vectors"""
+    db = LanceDb(uri=TEST_PATH, table_name="test_bad_vectors", on_bad_vectors="fill", fill_value=0.0)
+    db.create()
+
+    try:
+        # Test with a document that might have bad vectors
+        doc = Document(content="Test document", meta_data={}, name="test")
+        db.insert([doc])
+        assert db.get_count() == 1
+    finally:
+        db.drop()
+        if os.path.exists(TEST_PATH):
+            shutil.rmtree(TEST_PATH)
+
+
+def test_custom_embedder():
+    """Test using a custom embedder"""
+    os.makedirs(TEST_PATH, exist_ok=True)
+
+    custom_embedder = OpenAIEmbedder()
+    db = LanceDb(uri=TEST_PATH, table_name=TEST_TABLE, embedder=custom_embedder)
+    db.create()
+    assert db.embedder == custom_embedder
+
+    try:
+        db.drop()
+    finally:
+        if os.path.exists(TEST_PATH):
+            shutil.rmtree(TEST_PATH)


Code review comments:
diff_hunk: @@ -17,7 +17,7 @@
     vector_db=vector_db,
 )
 
-knowledge_base.load(recreate=False)  # Comment out after first run
+knowledge_base.load(recreate=True)  # Comment out after first run
path: cookbook/agent_concepts/knowledge/vector_dbs/lance_db.py
text: @user1:
```suggestion
knowledge_base.load(recreate=False)  # Comment out after first run
```