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

agno/vectordb/pgvector/pgvector.py

@@ -1,8 +1,10 @@
 import asyncio
+from hashlib import md5
 from math import sqrt
 from typing import Any, Dict, List, Optional, Union, cast
 
 try:
+    from sqlalchemy import update
     from sqlalchemy.dialects import postgresql
     from sqlalchemy.engine import Engine, create_engine
     from sqlalchemy.inspection import inspect
@@ -10,6 +12,7 @@ try:
     from sqlalchemy.schema import Column, Index, MetaData, Table
     from sqlalchemy.sql.expression import bindparam, desc, func, select, text
     from sqlalchemy.types import DateTime, String
+
 except ImportError:
     raise ImportError("`sqlalchemy` not installed. Please install using `pip install sqlalchemy psycopg`")
 
@@ -18,11 +21,10 @@ try:
 except ImportError:
     raise ImportError("`pgvector` not installed. Please install using `pip install pgvector`")
 
-from agno.document import Document
-from agno.embedder import Embedder
-from agno.reranker.base import Reranker
+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, logger
-from agno.utils.string import safe_content_hash
 from agno.vectordb.base import VectorDb
 from agno.vectordb.distance import Distance
 from agno.vectordb.pgvector.index import HNSW, Ivfflat
@@ -53,6 +55,7 @@ class PgVector(VectorDb):
         schema_version: int = 1,
         auto_upgrade_schema: bool = False,
         reranker: Optional[Reranker] = None,
+        use_batch: bool = False,
     ):
         """
         Initialize the PgVector instance.
@@ -93,10 +96,11 @@ class PgVector(VectorDb):
         self.db_url: Optional[str] = db_url
         self.db_engine: Engine = db_engine
         self.metadata: MetaData = MetaData(schema=self.schema)
+        self.use_batch: bool = use_batch
 
         # 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.")
@@ -155,6 +159,7 @@ class PgVector(VectorDb):
             Column("created_at", DateTime(timezone=True), server_default=func.now()),
             Column("updated_at", DateTime(timezone=True), onupdate=func.now()),
             Column("content_hash", String),
+            Column("content_id", String),
             extend_existing=True,
         )
 
@@ -162,7 +167,7 @@ class PgVector(VectorDb):
         Index(f"idx_{self.table_name}_id", table.c.id)
         Index(f"idx_{self.table_name}_name", table.c.name)
         Index(f"idx_{self.table_name}_content_hash", table.c.content_hash)
-
+        Index(f"idx_{self.table_name}_content_id", table.c.content_id)
         return table
 
     def get_table(self) -> Table:
@@ -229,23 +234,6 @@ class PgVector(VectorDb):
             logger.error(f"Error checking if record exists: {e}")
             return False
 
-    def doc_exists(self, document: Document) -> bool:
-        """
-        Check if a document with the same content hash exists in the table.
-
-        Args:
-            document (Document): The document to check.
-
-        Returns:
-            bool: True if the document exists, False otherwise.
-        """
-        content_hash = safe_content_hash(document.content)
-        return self._record_exists(self.table.c.content_hash, content_hash)
-
-    async def async_doc_exists(self, document: Document) -> bool:
-        """Check if document exists asynchronously by running in a thread."""
-        return await asyncio.to_thread(self.doc_exists, document)
-
     def name_exists(self, name: str) -> bool:
         """
         Check if a document with the given name exists in the table.
@@ -274,6 +262,12 @@ class PgVector(VectorDb):
         """
         return self._record_exists(self.table.c.id, id)
 
+    def content_hash_exists(self, content_hash: str) -> bool:
+        """
+        Check if a document with the given content hash exists in the table.
+        """
+        return self._record_exists(self.table.c.content_hash, content_hash)
+
     def _clean_content(self, content: str) -> str:
         """
         Clean the content by replacing null characters.
@@ -288,6 +282,7 @@ class PgVector(VectorDb):
 
     def insert(
         self,
+        content_hash: str,
         documents: List[Document],
         filters: Optional[Dict[str, Any]] = None,
         batch_size: int = 100,
@@ -296,6 +291,7 @@ class PgVector(VectorDb):
         Insert documents into the database.
 
         Args:
+            content_hash (str): The content hash to insert.
             documents (List[Document]): List of documents to insert.
             filters (Optional[Dict[str, Any]]): Filters to apply to the documents.
             batch_size (int): Number of documents to insert in each batch.
@@ -310,7 +306,7 @@ class PgVector(VectorDb):
                         batch_records = []
                         for doc in batch_docs:
                             try:
-                                batch_records.append(self._get_document_record(doc, filters))
+                                batch_records.append(self._get_document_record(doc, filters, content_hash))
                             except Exception as e:
                                 logger.error(f"Error processing document '{doc.name}': {e}")
 
@@ -327,9 +323,62 @@ class PgVector(VectorDb):
             logger.error(f"Error inserting documents: {e}")
             raise
 
-    async def async_insert(self, documents: List[Document], filters: Optional[Dict[str, Any]] = None) -> None:
-        """Insert documents asynchronously by running in a thread."""
-        await asyncio.to_thread(self.insert, documents, filters)
+    async def async_insert(
+        self,
+        content_hash: str,
+        documents: List[Document],
+        filters: Optional[Dict[str, Any]] = None,
+        batch_size: int = 100,
+    ) -> None:
+        """Insert documents asynchronously with parallel embedding."""
+        try:
+            with self.Session() as sess:
+                for i in range(0, len(documents), batch_size):
+                    batch_docs = documents[i : i + batch_size]
+                    log_debug(f"Processing batch starting at index {i}, size: {len(batch_docs)}")
+                    try:
+                        embed_tasks = [doc.async_embed(embedder=self.embedder) for doc in batch_docs]
+                        await asyncio.gather(*embed_tasks, return_exceptions=True)
+
+                        # Prepare documents for insertion
+                        batch_records = []
+                        for doc in batch_docs:
+                            try:
+                                cleaned_content = self._clean_content(doc.content)
+                                record_id = doc.id or content_hash
+
+                                meta_data = doc.meta_data or {}
+                                if filters:
+                                    meta_data.update(filters)
+
+                                record = {
+                                    "id": record_id,
+                                    "name": doc.name,
+                                    "meta_data": doc.meta_data,
+                                    "filters": filters,
+                                    "content": cleaned_content,
+                                    "embedding": doc.embedding,
+                                    "usage": doc.usage,
+                                    "content_hash": content_hash,
+                                    "content_id": doc.content_id,
+                                }
+                                batch_records.append(record)
+                            except Exception as e:
+                                logger.error(f"Error processing document '{doc.name}': {e}")
+
+                        # Insert the batch of records
+                        if batch_records:
+                            insert_stmt = postgresql.insert(self.table)
+                            sess.execute(insert_stmt, batch_records)
+                            sess.commit()  # Commit batch independently
+                            log_info(f"Inserted batch of {len(batch_records)} documents.")
+                    except Exception as e:
+                        logger.error(f"Error with batch starting at index {i}: {e}")
+                        sess.rollback()  # Rollback the current batch if there's an error
+                        raise
+        except Exception as e:
+            logger.error(f"Error inserting documents: {e}")
+            raise
 
     def upsert_available(self) -> bool:
         """
@@ -342,6 +391,27 @@ class PgVector(VectorDb):
 
     def upsert(
         self,
+        content_hash: str,
+        documents: List[Document],
+        filters: Optional[Dict[str, Any]] = None,
+        batch_size: int = 100,
+    ) -> None:
+        """
+        Upsert documents by content hash.
+        First delete all documents with the same content hash.
+        Then upsert the new documents.
+        """
+        try:
+            if self.content_hash_exists(content_hash):
+                self._delete_by_content_hash(content_hash)
+            self._upsert(content_hash, documents, filters, batch_size)
+        except Exception as e:
+            logger.error(f"Error upserting documents by content hash: {e}")
+            raise
+
+    def _upsert(
+        self,
+        content_hash: str,
         documents: List[Document],
         filters: Optional[Dict[str, Any]] = None,
         batch_size: int = 100,
@@ -358,16 +428,22 @@ class PgVector(VectorDb):
             with self.Session() as sess:
                 for i in range(0, len(documents), batch_size):
                     batch_docs = documents[i : i + batch_size]
-                    log_debug(f"Processing batch starting at index {i}, size: {len(batch_docs)}")
+                    log_info(f"Processing batch starting at index {i}, size: {len(batch_docs)}")
                     try:
                         # Prepare documents for upserting
-                        batch_records = []
+                        batch_records_dict: Dict[str, Dict[str, Any]] = {}  # Use dict to deduplicate by ID
                         for doc in batch_docs:
                             try:
-                                batch_records.append(self._get_document_record(doc, filters))
+                                batch_records_dict[doc.id] = self._get_document_record(doc, filters, content_hash)  # type: ignore
                             except Exception as e:
                                 logger.error(f"Error processing document '{doc.name}': {e}")
 
+                        # Convert dict to list for upsert
+                        batch_records = list(batch_records_dict.values())
+                        if not batch_records:
+                            log_info("No valid records to upsert in this batch.")
+                            continue
+
                         # Upsert the batch of records
                         insert_stmt = postgresql.insert(self.table).values(batch_records)
                         upsert_stmt = insert_stmt.on_conflict_do_update(
@@ -380,6 +456,7 @@ class PgVector(VectorDb):
                                 "embedding": insert_stmt.excluded.embedding,
                                 "usage": insert_stmt.excluded.usage,
                                 "content_hash": insert_stmt.excluded.content_hash,
+                                "content_id": insert_stmt.excluded.content_id,
                             },
                         )
                         sess.execute(upsert_stmt)
@@ -393,30 +470,155 @@ class PgVector(VectorDb):
             logger.error(f"Error upserting documents: {e}")
             raise
 
-    def _get_document_record(self, doc: Document, filters: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+    def _get_document_record(
+        self, doc: Document, filters: Optional[Dict[str, Any]] = None, content_hash: str = ""
+    ) -> Dict[str, Any]:
         doc.embed(embedder=self.embedder)
         cleaned_content = self._clean_content(doc.content)
-        content_hash = safe_content_hash(doc.content)
-        _id = doc.id or content_hash
+        record_id = doc.id or content_hash
 
         meta_data = doc.meta_data or {}
         if filters:
             meta_data.update(filters)
 
         return {
-            "id": _id,
+            "id": record_id,
             "name": doc.name,
-            "meta_data": meta_data,
+            "meta_data": doc.meta_data,
             "filters": filters,
             "content": cleaned_content,
             "embedding": doc.embedding,
             "usage": doc.usage,
             "content_hash": content_hash,
+            "content_id": doc.content_id,
         }
 
-    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,
+        batch_size: int = 100,
+    ) -> None:
         """Upsert documents asynchronously by running in a thread."""
-        await asyncio.to_thread(self.upsert, documents, filters)
+        try:
+            if self.content_hash_exists(content_hash):
+                self._delete_by_content_hash(content_hash)
+            await self._async_upsert(content_hash, documents, filters, batch_size)
+        except Exception as e:
+            logger.error(f"Error upserting documents by content hash: {e}")
+            raise
+
+    async def _async_upsert(
+        self,
+        content_hash: str,
+        documents: List[Document],
+        filters: Optional[Dict[str, Any]] = None,
+        batch_size: int = 100,
+    ) -> None:
+        """
+        Upsert (insert or update) documents in the database.
+
+        Args:
+            documents (List[Document]): List of documents to upsert.
+            filters (Optional[Dict[str, Any]]): Filters to apply to the documents.
+            batch_size (int): Number of documents to upsert in each batch.
+        """
+        try:
+            with self.Session() as sess:
+                for i in range(0, len(documents), batch_size):
+                    batch_docs = documents[i : i + batch_size]
+                    log_info(f"Processing batch starting at index {i}, size: {len(batch_docs)}")
+                    try:
+                        embed_tasks = [doc.async_embed(embedder=self.embedder) for doc in batch_docs]
+                        await asyncio.gather(*embed_tasks, return_exceptions=True)
+
+                        # Prepare documents for upserting
+                        batch_records_dict = {}  # Use dict to deduplicate by ID
+                        for doc in batch_docs:
+                            try:
+                                cleaned_content = self._clean_content(doc.content)
+                                record_id = md5(cleaned_content.encode()).hexdigest()
+
+                                meta_data = doc.meta_data or {}
+                                if filters:
+                                    meta_data.update(filters)
+
+                                record = {
+                                    "id": record_id,  # use record_id as a reproducible id to avoid duplicates while upsert
+                                    "name": doc.name,
+                                    "meta_data": doc.meta_data,
+                                    "filters": filters,
+                                    "content": cleaned_content,
+                                    "embedding": doc.embedding,
+                                    "usage": doc.usage,
+                                    "content_hash": content_hash,
+                                    "content_id": doc.content_id,
+                                }
+                                batch_records_dict[record_id] = record  # This deduplicates by ID
+                            except Exception as e:
+                                logger.error(f"Error processing document '{doc.name}': {e}")
+
+                        # Convert dict to list for upsert
+                        batch_records = list(batch_records_dict.values())
+                        if not batch_records:
+                            log_info("No valid records to upsert in this batch.")
+                            continue
+
+                        # Upsert the batch of records
+                        insert_stmt = postgresql.insert(self.table).values(batch_records)
+                        upsert_stmt = insert_stmt.on_conflict_do_update(
+                            index_elements=["id"],
+                            set_={
+                                "name": insert_stmt.excluded.name,
+                                "meta_data": insert_stmt.excluded.meta_data,
+                                "filters": insert_stmt.excluded.filters,
+                                "content": insert_stmt.excluded.content,
+                                "embedding": insert_stmt.excluded.embedding,
+                                "usage": insert_stmt.excluded.usage,
+                                "content_hash": insert_stmt.excluded.content_hash,
+                                "content_id": insert_stmt.excluded.content_id,
+                            },
+                        )
+                        sess.execute(upsert_stmt)
+                        sess.commit()  # Commit batch independently
+                        log_info(f"Upserted batch of {len(batch_records)} documents.")
+                    except Exception as e:
+                        logger.error(f"Error with batch starting at index {i}: {e}")
+                        sess.rollback()  # Rollback the current batch if there's an error
+                        raise
+        except Exception as e:
+            logger.error(f"Error upserting documents: {e}")
+            raise
+
+    def update_metadata(self, content_id: str, metadata: Dict[str, Any]) -> None:
+        """
+        Update the metadata for a document.
+
+        Args:
+            id (str): The ID of the document.
+            metadata (Dict[str, Any]): The metadata to update.
+        """
+        try:
+            with self.Session() as sess:
+                # Merge JSONB instead of overwriting: coalesce(existing, '{}') || :new
+                stmt = (
+                    update(self.table)
+                    .where(self.table.c.content_id == content_id)
+                    .values(
+                        meta_data=func.coalesce(self.table.c.meta_data, text("'{}'::jsonb")).op("||")(
+                            bindparam("md", metadata, type_=postgresql.JSONB)
+                        ),
+                        filters=func.coalesce(self.table.c.filters, text("'{}'::jsonb")).op("||")(
+                            bindparam("ft", metadata, type_=postgresql.JSONB)
+                        ),
+                    )
+                )
+                sess.execute(stmt)
+                sess.commit()
+        except Exception as e:
+            logger.error(f"Error updating metadata for document {content_id}: {e}")
+            raise
 
     def search(self, query: str, limit: int = 5, filters: Optional[Dict[str, Any]] = None) -> List[Document]:
         """
@@ -1025,6 +1227,86 @@ class PgVector(VectorDb):
             sess.rollback()
             return False
 
+    def delete_by_id(self, id: str) -> bool:
+        """
+        Delete content by ID.
+        """
+        try:
+            with self.Session() as sess, sess.begin():
+                stmt = self.table.delete().where(self.table.c.id == id)
+                sess.execute(stmt)
+                sess.commit()
+                log_info(f"Deleted records with id '{id}' from table '{self.table.fullname}'.")
+                return True
+        except Exception as e:
+            logger.error(f"Error deleting rows from table '{self.table.fullname}': {e}")
+            sess.rollback()
+            return False
+
+    def delete_by_name(self, name: str) -> bool:
+        """
+        Delete content by name.
+        """
+        try:
+            with self.Session() as sess, sess.begin():
+                stmt = self.table.delete().where(self.table.c.name == name)
+                sess.execute(stmt)
+                sess.commit()
+                log_info(f"Deleted records with name '{name}' from table '{self.table.fullname}'.")
+                return True
+        except Exception as e:
+            logger.error(f"Error deleting rows from table '{self.table.fullname}': {e}")
+            sess.rollback()
+            return False
+
+    def delete_by_metadata(self, metadata: Dict[str, Any]) -> bool:
+        """
+        Delete content by metadata.
+        """
+        try:
+            with self.Session() as sess, sess.begin():
+                stmt = self.table.delete().where(self.table.c.meta_data.contains(metadata))
+                sess.execute(stmt)
+                sess.commit()
+                log_info(f"Deleted records with metadata '{metadata}' from table '{self.table.fullname}'.")
+                return True
+        except Exception as e:
+            logger.error(f"Error deleting rows from table '{self.table.fullname}': {e}")
+            sess.rollback()
+            return False
+
+    def delete_by_content_id(self, content_id: str) -> bool:
+        """
+        Delete content by content ID.
+        """
+        try:
+            with self.Session() as sess, sess.begin():
+                stmt = self.table.delete().where(self.table.c.content_id == content_id)
+                sess.execute(stmt)
+                sess.commit()
+                log_info(f"Deleted records with content ID '{content_id}' from table '{self.table.fullname}'.")
+                return True
+        except Exception as e:
+            logger.error(f"Error deleting rows from table '{self.table.fullname}': {e}")
+            sess.rollback()
+            return False
+
+    def _delete_by_content_hash(self, content_hash: str) -> bool:
+        """
+        Delete content by content hash.
+        """
+        try:
+            with self.Session() as sess, sess.begin():
+                stmt = self.table.delete().where(self.table.c.content_hash == content_hash)
+                sess.execute(stmt)
+                sess.commit()
+                log_info(f"Deleted records with content hash '{content_hash}' from table '{self.table.fullname}'.")
+                return True
+        except Exception as e:
+            logger.error(f"Error deleting rows from table '{self.table.fullname}': {e}")
+            sess.rollback()
+            return False
+
     def __deepcopy__(self, memo):
         """
         Create a deep copy of the PgVector instance, handling unpickleable attributes.