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

agno/vectordb/surrealdb/surrealdb.py

@@ -11,8 +11,8 @@ except ImportError as e:
     msg = "The `surrealdb` package is not installed. Please install it via `pip install surrealdb`."
     raise ImportError(msg) from e
 
-from agno.document import Document
-from agno.embedder import Embedder
+from agno.knowledge.document import Document
+from agno.knowledge.embedder import Embedder
 from agno.utils.log import log_debug, log_error, log_info
 from agno.vectordb.base import VectorDb
 from agno.vectordb.distance import Distance
@@ -48,6 +48,32 @@ class SurrealDb(VectorDb):
         LIMIT 1
     """
 
+    CONTENT_HASH_EXISTS_QUERY: Final[str] = """
+        SELECT * FROM {collection}
+        WHERE meta_data.content_hash = $content_hash
+        LIMIT 1
+    """
+
+    DELETE_BY_ID_QUERY: Final[str] = """
+        DELETE FROM {collection}
+        WHERE id = $id
+    """
+
+    DELETE_BY_NAME_QUERY: Final[str] = """
+        DELETE FROM {collection}
+        WHERE meta_data.name = $name
+    """
+
+    DELETE_BY_METADATA_QUERY: Final[str] = """
+        DELETE FROM {collection}
+        WHERE {conditions}
+    """
+
+    DELETE_BY_CONTENT_ID_QUERY: Final[str] = """
+        DELETE FROM {collection}
+        WHERE content_id = $content_id
+    """
+
     UPSERT_QUERY: Final[str] = """
         UPSERT {thing}
         SET content = $content,
@@ -98,7 +124,7 @@ class SurrealDb(VectorDb):
         """
         # Embedder for embedding the document contents
         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.")
@@ -213,10 +239,41 @@ class SurrealDb(VectorDb):
         result = self.client.query(self.NAME_EXISTS_QUERY.format(collection=self.collection), {"name": name})
         return bool(self._extract_result(result))
 
-    def insert(self, documents: List[Document], filters: Optional[Dict[str, Any]] = None) -> None:
+    def id_exists(self, id: str) -> bool:
+        """Check if a document exists by its ID.
+
+        Args:
+            id: The ID of the document to check.
+
+        Returns:
+            True if the document exists, False otherwise.
+
+        """
+        log_debug(f"Checking if document exists by ID: {id}")
+        result = self.client.query(self.ID_EXISTS_QUERY.format(collection=self.collection), {"id": id})
+        return bool(self._extract_result(result))
+
+    def content_hash_exists(self, content_hash: str) -> bool:
+        """Check if a document exists by its content hash.
+
+        Args:
+            content_hash: The content hash of the document to check.
+
+        Returns:
+            True if the document exists, False otherwise.
+
+        """
+        log_debug(f"Checking if document exists by content hash: {content_hash}")
+        result = self.client.query(
+            self.CONTENT_HASH_EXISTS_QUERY.format(collection=self.collection), {"content_hash": content_hash}
+        )
+        return bool(self._extract_result(result))
+
+    def insert(self, content_hash: str, documents: List[Document], filters: Optional[Dict[str, Any]] = None) -> None:
         """Insert documents into the vector store.
 
         Args:
+            content_hash: The content hash for the documents.
             documents: A list of documents to insert.
             filters: A dictionary of filters to apply to the query.
 
@@ -224,15 +281,17 @@ class SurrealDb(VectorDb):
         for doc in documents:
             doc.embed(embedder=self.embedder)
             meta_data: Dict[str, Any] = doc.meta_data if isinstance(doc.meta_data, dict) else {}
+            meta_data["content_hash"] = content_hash
             data: Dict[str, Any] = {"content": doc.content, "embedding": doc.embedding, "meta_data": meta_data}
             if filters:
                 data["meta_data"].update(filters)
             self.client.create(self.collection, data)
 
-    def upsert(self, documents: List[Document], filters: Optional[Dict[str, Any]] = None) -> None:
+    def upsert(self, content_hash: str, documents: List[Document], filters: Optional[Dict[str, Any]] = None) -> None:
         """Upsert documents into the vector store.
 
         Args:
+            content_hash: The content hash for the documents.
             documents: A list of documents to upsert.
             filters: A dictionary of filters to apply to the query.
 
@@ -240,6 +299,7 @@ class SurrealDb(VectorDb):
         for doc in documents:
             doc.embed(embedder=self.embedder)
             meta_data: Dict[str, Any] = doc.meta_data if isinstance(doc.meta_data, dict) else {}
+            meta_data["content_hash"] = content_hash
             data: Dict[str, Any] = {"content": doc.content, "embedding": doc.embedding, "meta_data": meta_data}
             if filters:
                 data["meta_data"].update(filters)
@@ -321,6 +381,67 @@ class SurrealDb(VectorDb):
         self.client.query(self.DELETE_ALL_QUERY.format(collection=self.collection))
         return True
 
+    def delete_by_id(self, id: str) -> bool:
+        """Delete a document by its ID.
+
+        Args:
+            id: The ID of the document to delete.
+
+        Returns:
+            True if the document was deleted, False otherwise.
+
+        """
+        log_debug(f"Deleting document by ID: {id}")
+        result = self.client.query(self.DELETE_BY_ID_QUERY.format(collection=self.collection), {"id": id})
+        return bool(result)
+
+    def delete_by_name(self, name: str) -> bool:
+        """Delete documents by their name.
+
+        Args:
+            name: The name of the documents to delete.
+
+        Returns:
+            True if documents were deleted, False otherwise.
+
+        """
+        log_debug(f"Deleting documents by name: {name}")
+        result = self.client.query(self.DELETE_BY_NAME_QUERY.format(collection=self.collection), {"name": name})
+        return bool(result)
+
+    def delete_by_metadata(self, metadata: Dict[str, Any]) -> bool:
+        """Delete documents by their metadata.
+
+        Args:
+            metadata: The metadata to match for deletion.
+
+        Returns:
+            True if documents were deleted, False otherwise.
+
+        """
+        log_debug(f"Deleting documents by metadata: {metadata}")
+        conditions = [f"meta_data.{key} = ${key}" for key in metadata.keys()]
+        conditions_str = " AND ".join(conditions)
+        query = self.DELETE_BY_METADATA_QUERY.format(collection=self.collection, conditions=conditions_str)
+        result = self.client.query(query, metadata)
+        return bool(result)
+
+    def delete_by_content_id(self, content_id: str) -> bool:
+        """Delete documents by their content ID.
+
+        Args:
+            content_id: The content ID of the documents to delete.
+
+        Returns:
+            True if documents were deleted, False otherwise.
+
+        """
+        log_debug(f"Deleting documents by content ID: {content_id}")
+        result = self.client.query(
+            self.DELETE_BY_CONTENT_ID_QUERY.format(collection=self.collection), {"content_id": content_id}
+        )
+        return bool(result)
+
     @staticmethod
     def _extract_result(query_result: Union[List[Dict[str, Any]], Dict[str, Any]]) -> Union[List[Any], Dict[str, Any]]:
         """Extract the actual result from SurrealDB query response.
@@ -380,10 +501,13 @@ class SurrealDb(VectorDb):
         )
         return bool(self._extract_result(response))
 
-    async def async_insert(self, documents: List[Document], filters: Optional[Dict[str, Any]] = None) -> None:
+    async def async_insert(
+        self, content_hash: str, documents: List[Document], filters: Optional[Dict[str, Any]] = None
+    ) -> None:
         """Insert documents into the vector store asynchronously.
 
         Args:
+            content_hash: The content hash for the documents.
             documents: A list of documents to insert.
             filters: A dictionary of filters to apply to the query.
 
@@ -391,16 +515,20 @@ class SurrealDb(VectorDb):
         for doc in documents:
             doc.embed(embedder=self.embedder)
             meta_data: Dict[str, Any] = doc.meta_data if isinstance(doc.meta_data, dict) else {}
+            meta_data["content_hash"] = content_hash
             data: Dict[str, Any] = {"content": doc.content, "embedding": doc.embedding, "meta_data": meta_data}
             if filters:
                 data["meta_data"].update(filters)
             log_debug(f"Inserting document asynchronously: {doc.name} ({doc.meta_data})")
             await self.async_client.create(self.collection, data)
 
-    async def async_upsert(self, documents: List[Document], filters: Optional[Dict[str, Any]] = None) -> None:
+    async def async_upsert(
+        self, content_hash: str, documents: List[Document], filters: Optional[Dict[str, Any]] = None
+    ) -> None:
         """Upsert documents into the vector store asynchronously.
 
         Args:
+            content_hash: The content hash for the documents.
             documents: A list of documents to upsert.
             filters: A dictionary of filters to apply to the query.
 
@@ -408,6 +536,7 @@ class SurrealDb(VectorDb):
         for doc in documents:
             doc.embed(embedder=self.embedder)
             meta_data: Dict[str, Any] = doc.meta_data if isinstance(doc.meta_data, dict) else {}
+            meta_data["content_hash"] = content_hash
             data: Dict[str, Any] = {"content": doc.content, "embedding": doc.embedding, "meta_data": meta_data}
             if filters:
                 data["meta_data"].update(filters)
@@ -491,3 +620,54 @@ class SurrealDb(VectorDb):
 
         """
         return True
+
+    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 documents with the given content_id
+            query = f"SELECT * FROM {self.collection} WHERE content_id = $content_id"
+            result = self.client.query(query, {"content_id": content_id})
+
+            if not result or not result[0].get("result"):
+                log_debug(f"No documents found with content_id: {content_id}")
+                return
+
+            documents = result[0]["result"]
+            updated_count = 0
+
+            # Update each matching document
+            for doc in documents:
+                doc_id = doc["id"]
+                current_metadata = doc.get("meta_data", {})
+                current_filters = doc.get("filters", {})
+
+                # Merge existing metadata with new metadata
+                if isinstance(current_metadata, dict):
+                    updated_metadata = current_metadata.copy()
+                    updated_metadata.update(metadata)
+                else:
+                    updated_metadata = metadata
+
+                # Merge existing filters with new metadata
+                if isinstance(current_filters, dict):
+                    updated_filters = current_filters.copy()
+                    updated_filters.update(metadata)
+                else:
+                    updated_filters = metadata
+
+                # Update the document
+                update_query = f"UPDATE {doc_id} SET meta_data = $metadata, filters = $filters"
+                self.client.query(update_query, {"metadata": updated_metadata, "filters": updated_filters})
+                updated_count += 1
+
+            log_debug(f"Updated metadata for {updated_count} documents with content_id: {content_id}")
+
+        except Exception as e:
+            log_error(f"Error updating metadata for content_id '{content_id}': {e}")
+            raise