agno 1.8.1__py3-none-any.whl → 2.0.0a1__py3-none-any.whl

agno/vectordb/pineconedb/pineconedb.py

@@ -22,10 +22,10 @@ except ImportError:
     raise ImportError("The `pinecone` package is not installed, please install using `pip install pinecone`.")
 
 
-from agno.document import Document
-from agno.embedder import Embedder
-from agno.reranker.base import Reranker
-from agno.utils.log import log_debug, log_info, logger
+from agno.knowledge.document import Document
+from agno.knowledge.embedder import Embedder
+from agno.knowledge.reranker.base import Reranker
+from agno.utils.log import log_debug, log_info, log_warning, logger
 from agno.vectordb.base import VectorDb
 
 
@@ -114,7 +114,7 @@ class PineconeDb(VectorDb):
         # Embedder for embedding the document contents
         _embedder = embedder
         if _embedder is None:
-            from agno.embedder.openai import OpenAIEmbedder
+            from agno.knowledge.embedder.openai import OpenAIEmbedder
 
             _embedder = OpenAIEmbedder()
             log_info("Embedder not provided, using OpenAIEmbedder as default.")
@@ -177,6 +177,9 @@ class PineconeDb(VectorDb):
             if self.use_hybrid_search:
                 self.metric = "dotproduct"
 
+            if self.dimension is None:
+                raise ValueError("Dimension is not set for this Pinecone index")
+
             self.client.create_index(
                 name=self.name,
                 dimension=self.dimension,
@@ -234,6 +237,17 @@ class PineconeDb(VectorDb):
 
     def upsert(
         self,
+        content_hash: str,
+        documents: List[Document],
+        filters: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        if self.content_hash_exists(content_hash):
+            self._delete_by_content_hash(content_hash)
+        self._upsert(content_hash=content_hash, documents=documents, filters=filters)
+
+    def _upsert(
+        self,
+        content_hash: str,
         documents: List[Document],
         filters: Optional[Dict[str, Any]] = None,
         namespace: Optional[str] = None,
@@ -255,10 +269,22 @@ class PineconeDb(VectorDb):
         for document in documents:
             document.embed(embedder=self.embedder)
             document.meta_data["text"] = document.content
+            # Include name and content_id in metadata
+            metadata = document.meta_data.copy()
+            if filters:
+                metadata.update(filters)
+
+            if document.name:
+                metadata["name"] = document.name
+            if document.content_id:
+                metadata["content_id"] = document.content_id
+
+            metadata["content_hash"] = content_hash
+
             data_to_upsert = {
                 "id": document.id,
                 "values": document.embedding,
-                "metadata": document.meta_data,
+                "metadata": metadata,
             }
             if self.use_hybrid_search:
                 data_to_upsert["sparse_values"] = self.sparse_encoder.encode_documents(document.content)
@@ -273,6 +299,7 @@ class PineconeDb(VectorDb):
 
     async def async_upsert(
         self,
+        content_hash: str,
         documents: List[Document],
         filters: Optional[Dict[str, Any]] = None,
         namespace: Optional[str] = None,
@@ -293,7 +320,7 @@ class PineconeDb(VectorDb):
 
         # Process each batch in parallel
         async def process_batch(batch_docs):
-            return await asyncio.to_thread(self._prepare_vectors, batch_docs)
+            return await self._prepare_vectors(batch_docs)
 
         # Run all batches in parallel
         batch_vectors = await asyncio.gather(*[process_batch(batch) for batch in batches])
@@ -308,16 +335,25 @@ class PineconeDb(VectorDb):
 
         log_debug(f"Finished async upsert of {len(documents)} documents")
 
-    def _prepare_vectors(self, documents):
+    async def _prepare_vectors(self, documents: List[Document]) -> List[Dict[str, Any]]:
         """Prepare vectors for upsert."""
         vectors = []
+        embed_tasks = [document.async_embed(embedder=self.embedder) for document in documents]
+        await asyncio.gather(*embed_tasks, return_exceptions=True)
+
         for doc in documents:
-            doc.embed(embedder=self.embedder)
             doc.meta_data["text"] = doc.content
+            # Include name and content_id in metadata
+            metadata = doc.meta_data.copy()
+            if doc.name:
+                metadata["name"] = doc.name
+            if doc.content_id:
+                metadata["content_id"] = doc.content_id
+
             data_to_upsert = {
                 "id": doc.id,
                 "values": doc.embedding,
-                "metadata": doc.meta_data,
+                "metadata": metadata,
             }
             if self.use_hybrid_search:
                 data_to_upsert["sparse_values"] = self.sparse_encoder.encode_documents(doc.content)
@@ -333,9 +369,11 @@ class PineconeDb(VectorDb):
             show_progress=show_progress,
         )
 
-    async def async_insert(self, documents: List[Document], filters: Optional[Dict[str, Any]] = None) -> None:
-        """Pinecone doesn't support insert. Raise an error."""
-        raise NotImplementedError("Pinecone does not support insert operations. Use async_upsert instead.")
+    async def async_insert(
+        self, content_hash: str, documents: List[Document], filters: Optional[Dict[str, Any]] = None
+    ) -> None:
+        log_warning("Pinecone does not support insert operations. Redirecting to async_upsert instead.")
+        await self.async_upsert(content_hash=content_hash, documents=documents, filters=filters)
 
     def upsert_available(self) -> bool:
         """Check if upsert operation is available.
@@ -346,20 +384,9 @@ class PineconeDb(VectorDb):
         """
         return True
 
-    def insert(self, documents: List[Document], filters: Optional[Dict[str, Any]] = None) -> None:
-        """Insert documents into the index.
-
-        This method is not supported by Pinecone. Use `upsert` instead.
-
-        Args:
-            documents (List[Document]): The documents to insert.
-            filters (Optional[Dict[str, Any]], optional): The filters for the insert. Defaults to None.
-
-        Raises:
-            NotImplementedError: This method is not supported by Pinecone.
-
-        """
-        raise NotImplementedError("Pinecone does not support insert operations. Use upsert instead.")
+    def insert(self, content_hash: str, documents: List[Document], filters: Optional[Dict[str, Any]] = None) -> None:
+        log_warning("Pinecone does not support insert operations. Redirecting to upsert instead.")
+        self.upsert(content_hash=content_hash, documents=documents, filters=filters)
 
     def _hybrid_scale(self, dense: List[float], sparse: Dict[str, Any], alpha: float):
         """Hybrid vector scaling using a convex combination
@@ -479,3 +506,170 @@ class PineconeDb(VectorDb):
 
     async def async_drop(self) -> None:
         raise NotImplementedError(f"Async not supported on {self.__class__.__name__}.")
+
+    def delete_by_id(self, id: str) -> bool:
+        """Delete a document by ID."""
+        try:
+            self.index.delete(ids=[id])
+            return True
+        except Exception as e:
+            log_warning(f"Error deleting document with ID {id}: {e}")
+            return False
+
+    def delete_by_name(self, name: str) -> bool:
+        """Delete documents by name (stored in metadata)."""
+        try:
+            # Delete all documents where metadata.name equals the given name
+            self.index.delete(filter={"name": {"$eq": name}})
+            return True
+        except Exception as e:
+            log_warning(f"Error deleting documents with name {name}: {e}")
+            return False
+
+    def delete_by_metadata(self, metadata: Dict[str, Any]) -> bool:
+        """Delete documents by metadata."""
+        try:
+            # Build filter for metadata matching
+            filter_conditions = {}
+            for key, value in metadata.items():
+                filter_conditions[key] = {"$eq": value}
+
+            self.index.delete(filter=filter_conditions)
+            return True
+        except Exception as e:
+            log_warning(f"Error deleting documents with metadata {metadata}: {e}")
+            return False
+
+    def delete_by_content_id(self, content_id: str) -> bool:
+        """Delete documents by content ID (stored in metadata)."""
+        try:
+            # Delete all documents where metadata.content_id equals the given content_id
+            self.index.delete(filter={"content_id": {"$eq": content_id}})
+            return True
+        except Exception as e:
+            log_warning(f"Error deleting documents with content_id {content_id}: {e}")
+            return False
+
+    def get_count(self) -> int:
+        """Get the count of documents in the index."""
+        try:
+            # Pinecone doesn't have a direct count method, so we use describe_index_stats
+            stats = self.index.describe_index_stats()
+            # The stats include total_vector_count which gives us the count
+            return stats.total_vector_count
+        except Exception as e:
+            log_warning(f"Error getting document count: {e}")
+            return 0
+
+    def id_exists(self, id: str) -> bool:
+        """Check if a document with the given ID exists in the index.
+
+        Args:
+            id (str): The ID to check.
+
+        Returns:
+            bool: True if the document exists, False otherwise.
+        """
+        try:
+            response = self.index.fetch(ids=[id], namespace=self.namespace)
+            return len(response.vectors) > 0
+        except Exception as e:
+            log_warning(f"Error checking if ID {id} exists: {e}")
+            return False
+
+    def content_hash_exists(self, content_hash: str) -> bool:
+        """Check if documents with the given content hash exist in the index.
+
+        Args:
+            content_hash (str): The content hash to check.
+
+        Returns:
+            bool: True if documents with the content hash exist, False otherwise.
+        """
+        try:
+            # Use a dummy vector to perform a minimal query with filter
+            # We only need to check if any results exist
+            if self.dimension is None:
+                raise ValueError("Dimension is not set for this Pinecone index")
+            dummy_vector = [0.0] * self.dimension
+            response = self.index.query(
+                vector=dummy_vector,
+                top_k=1,
+                namespace=self.namespace,
+                filter={"content_hash": {"$eq": content_hash}},
+                include_metadata=False,
+                include_values=False,
+            )
+            return len(response.matches) > 0
+        except Exception as e:
+            log_warning(f"Error checking if content_hash {content_hash} exists: {e}")
+            return False
+
+    def _delete_by_content_hash(self, content_hash: str) -> bool:
+        """Delete documents by content hash (stored in metadata).
+
+        Args:
+            content_hash (str): The content hash to delete.
+
+        Returns:
+            bool: True if documents were deleted, False otherwise.
+        """
+        try:
+            # Delete all documents where metadata.content_hash equals the given content_hash
+            self.index.delete(filter={"content_hash": {"$eq": content_hash}}, namespace=self.namespace)
+            return True
+        except Exception as e:
+            log_warning(f"Error deleting documents with content_hash {content_hash}: {e}")
+            return False
+
+    def update_metadata(self, content_id: str, metadata: Dict[str, Any]) -> None:
+        """
+        Update the metadata for documents with the given content_id.
+
+        Args:
+            content_id (str): The content ID to update
+            metadata (Dict[str, Any]): The metadata to update
+        """
+        try:
+            # Query for vectors with the given content_id
+            query_response = self.index.query(
+                filter={"content_id": {"$eq": content_id}},
+                top_k=10000,  # Get all matching vectors
+                include_metadata=True,
+                namespace=self.namespace,
+            )
+
+            if not query_response.matches:
+                logger.debug(f"No documents found with content_id: {content_id}")
+                return
+
+            # Prepare updates for each matching vector
+            update_data = []
+            for match in query_response.matches:
+                vector_id = match.id
+                current_metadata = match.metadata or {}
+
+                # Merge existing metadata with new metadata
+                updated_metadata = current_metadata.copy()
+                updated_metadata.update(metadata)
+
+                if "filters" not in updated_metadata:
+                    updated_metadata["filters"] = {}
+                if isinstance(updated_metadata["filters"], dict):
+                    updated_metadata["filters"].update(metadata)
+                else:
+                    updated_metadata["filters"] = metadata
+
+                update_data.append({"id": vector_id, "metadata": updated_metadata})
+
+            # Update vectors in batches
+            batch_size = 100
+            for i in range(0, len(update_data), batch_size):
+                batch = update_data[i : i + batch_size]
+                self.index.update(vectors=batch, namespace=self.namespace)
+
+            logger.debug(f"Updated metadata for {len(update_data)} documents with content_id: {content_id}")
+
+        except Exception as e:
+            logger.error(f"Error updating metadata for content_id '{content_id}': {e}")
+            raise