""" Copyright 2024-2026 ChatterMate Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. """ from typing import Any, Dict, List, Optional, Union from hashlib import md5 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.vectordb.pgvector.pgvector import PgVector from agno.vectordb.distance import Distance from agno.vectordb.pgvector.index import HNSW, Ivfflat from agno.vectordb.search import SearchType from sqlalchemy.engine import Engine class OptimizedPgVector(PgVector): """ OptimizedPgVector class that extends the standard PgVector implementation to avoid re-embedding documents that already have embeddings. This class optimizes the workflow for web crawling by allowing embeddings to be generated during the crawling phase, and then preserved during insertion into the database. """ def __init__( self, table_name: str, schema: str = "ai", db_url: Optional[str] = None, db_engine: Optional[Engine] = None, embedder: Optional[Embedder] = None, search_type: SearchType = SearchType.vector, vector_index: Union[Ivfflat, HNSW] = HNSW(), distance: Distance = Distance.cosine, prefix_match: bool = False, vector_score_weight: float = 0.5, content_language: str = "english", schema_version: int = 1, auto_upgrade_schema: bool = False, reranker: Optional[Reranker] = None, ): """Initialize the OptimizedPgVector with the same parameters as PgVector.""" super().__init__( table_name=table_name, schema=schema, db_url=db_url, db_engine=db_engine, embedder=embedder, search_type=search_type, vector_index=vector_index, distance=distance, prefix_match=prefix_match, vector_score_weight=vector_score_weight, content_language=content_language, schema_version=schema_version, auto_upgrade_schema=auto_upgrade_schema, reranker=reranker, ) def create(self) -> None: """Create the per-org table, tolerating a concurrent creator. agno's PgVector.create() checks table_exists() and then issues CREATE TABLE, which is not atomic. The knowledge processor runs two queue items at once, so an organisation whose first two sources are queued together has both workers find no table and both create it; the loser died with `duplicate key value violates unique constraint "pg_type_typname_nsp_index"` — Postgres rejecting the second table's implicit row type — and the user lost that source. Losing the race is not an error: the table the other worker created is the one we wanted. Anything else still raises. """ try: super().create() except Exception as e: if not self._is_duplicate_object_error(e): raise logger.debug( f"Table {self.schema}.{self.table_name} was created concurrently; continuing" ) @staticmethod def _is_duplicate_object_error(error: Exception) -> bool: """True when the failure is 'someone else already created this'. Matched on SQLSTATE rather than message text: 42P07 duplicate_table and 23505 unique_violation, the latter being how a clash on the pg_type catalog index surfaces. """ for exc in (error, getattr(error, 'orig', None)): sqlstate = getattr(exc, 'sqlstate', None) or getattr(exc, 'pgcode', None) if sqlstate in ('42P07', '23505'): return True return False def insert( self, documents: List[Document], filters: Optional[Dict[str, Any]] = None, batch_size: int = 100, ) -> None: """ Insert documents into the database, preserving existing embeddings. Args: 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. """ try: with self.Session() as sess: for i in range(0, len(documents), batch_size): batch_docs = documents[i : i + batch_size] try: # Prepare documents for insertion batch_records = [] for doc in batch_docs: try: # Embed document only if not already embedded (should be rare) if doc.embedding is None: if self.embedder is not None: doc.embed(embedder=self.embedder) logger.debug(f"Document embedding done for {doc.name}") else: logger.warning(f"Document '{doc.name}' has no embedding and no embedder available - skipping") continue # Merge filters into document metadata for searchability merged_meta_data = doc.meta_data.copy() if doc.meta_data else {} if filters: merged_meta_data.update(filters) logger.debug(f"Merged filters into meta_data: {filters}") cleaned_content = self._clean_content(doc.content) content_hash = md5(cleaned_content.encode()).hexdigest() _id = doc.id or content_hash record = { "id": _id, "name": doc.name, "meta_data": merged_meta_data, # Use merged metadata "filters": filters, "content": cleaned_content, "embedding": doc.embedding, "usage": doc.usage, "content_hash": content_hash, } batch_records.append(record) except Exception as e: logger.error(f"Error processing document '{doc.name}': {e}") # Insert the batch of records from sqlalchemy.dialects import postgresql insert_stmt = postgresql.insert(self.table) sess.execute(insert_stmt, batch_records) sess.commit() # Commit batch independently 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( self, documents: List[Document], filters: Optional[Dict[str, Any]] = None, batch_size: int = 100, ) -> None: """ Upsert (insert or update) documents in the database, preserving existing embeddings. 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] try: # Prepare documents for upserting batch_records = [] for doc in batch_docs: try: # Skip embedding if the document already has an embedding if doc.embedding is None: doc.embed(embedder=self.embedder) # Merge filters into document metadata for searchability merged_meta_data = doc.meta_data.copy() if doc.meta_data else {} if filters: merged_meta_data.update(filters) logger.debug(f"Merged filters into meta_data: {filters}") cleaned_content = self._clean_content(doc.content) content_hash = md5(cleaned_content.encode()).hexdigest() _id = doc.id or content_hash record = { "id": _id, "name": doc.name, "meta_data": merged_meta_data, # Use merged metadata "filters": filters, "content": cleaned_content, "embedding": doc.embedding, "usage": doc.usage, "content_hash": content_hash, } batch_records.append(record) except Exception as e: logger.error(f"Error processing document '{doc.name}': {e}") # Upsert the batch of records from sqlalchemy.dialects import postgresql insert_stmt = postgresql.insert(self.table).values(batch_records) upsert_stmt = insert_stmt.on_conflict_do_update( index_elements=["id"], set_=dict( 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, ), ) sess.execute(upsert_stmt) sess.commit() # Commit batch independently 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