# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors # # SPDX-License-Identifier: Apache-2.0 # # Portions derived from https://github.com/microsoft/autogen are under the MIT License. # SPDX-License-Identifier: MIT import os import re import urllib.parse from typing import Any, Callable, Optional, Union from ....import_utils import optional_import_block, require_optional_import from .base import Document, ItemID, QueryResults, VectorDB from .utils import get_logger with optional_import_block(): import numpy as np import pgvector # noqa: F401 import psycopg from pgvector.psycopg import register_vector from sentence_transformers import SentenceTransformer PGVECTOR_MAX_BATCH_SIZE = os.environ.get("PGVECTOR_MAX_BATCH_SIZE", 40000) logger = get_logger(__name__) @require_optional_import(["psycopg", "sentence_transformers", "numpy"], "retrievechat-pgvector") class Collection: """A Collection object for PGVector. Attributes: client: The PGVector client. collection_name (str): The name of the collection. Default is "documents". embedding_function (Callable): The embedding function used to generate the vector representation. Default is None. SentenceTransformer("all-MiniLM-L6-v2").encode will be used when None. Models can be chosen from: https://huggingface.co/models?library=sentence-transformers metadata (Optional[dict[str, Any]]): The metadata of the collection. get_or_create (Optional): The flag indicating whether to get or create the collection. """ def __init__( self, client: Optional[Any] = None, collection_name: str = "ag2-docs", embedding_function: Optional[Callable[..., Any]] = None, metadata: Optional[Any] = None, get_or_create: Optional[Any] = None, ): """Initialize the Collection object. Args: client: The PostgreSQL client. collection_name: The name of the collection. Default is "documents". embedding_function: The embedding function used to generate the vector representation. metadata: The metadata of the collection. get_or_create: The flag indicating whether to get or create the collection. Returns: None """ self.client = client self.name = self.set_collection_name(collection_name) self.require_embeddings_or_documents = False self.ids = [] if embedding_function: self.embedding_function = embedding_function else: self.embedding_function = SentenceTransformer("all-MiniLM-L6-v2").encode self.metadata = metadata if metadata else {"hnsw:space": "ip", "hnsw:construction_ef": 32, "hnsw:M": 16} self.documents = "" self.get_or_create = get_or_create # This will get the model dimension size by computing the embeddings dimensions sentences = [ "The weather is lovely today in paradise.", ] embeddings = self.embedding_function(sentences) self.dimension = len(embeddings[0]) def set_collection_name(self, collection_name) -> str: name = re.sub("-", "_", collection_name) self.name = name return self.name def add( self, ids: list[ItemID], documents: Optional[list[Document]], embeddings: Optional[list[Any]] = None, metadatas: Optional[list[Any]] = None, ) -> None: """Add documents to the collection. Args: ids (List[ItemID]): A list of document IDs. embeddings (List): A list of document embeddings. Optional metadatas (List): A list of document metadatas. Optional documents (List): A list of documents. Returns: None """ cursor = self.client.cursor() sql_values = [] if embeddings is not None and metadatas is not None: for doc_id, embedding, metadata, document in zip(ids, embeddings, metadatas, documents): metadata = re.sub("'", '"', str(metadata)) sql_values.append((doc_id, embedding, metadata, document)) sql_string = f"INSERT INTO {self.name} (id, embedding, metadatas, documents)\nVALUES (%s, %s, %s, %s);\n" elif embeddings is not None: for doc_id, embedding, document in zip(ids, embeddings, documents): sql_values.append((doc_id, embedding, document)) sql_string = f"INSERT INTO {self.name} (id, embedding, documents) VALUES (%s, %s, %s);\n" elif metadatas is not None: for doc_id, metadata, document in zip(ids, metadatas, documents): metadata = re.sub("'", '"', str(metadata)) embedding = self.embedding_function(document) sql_values.append((doc_id, metadata, embedding, document)) sql_string = f"INSERT INTO {self.name} (id, metadatas, embedding, documents)\nVALUES (%s, %s, %s, %s);\n" else: for doc_id, document in zip(ids, documents): embedding = self.embedding_function(document) sql_values.append((doc_id, document, embedding)) sql_string = f"INSERT INTO {self.name} (id, documents, embedding)\nVALUES (%s, %s, %s);\n" logger.debug(f"Add SQL String:\n{sql_string}\n{sql_values}") cursor.executemany(sql_string, sql_values) cursor.close() def upsert( self, ids: list[ItemID], documents: list[Document], embeddings: Optional[list[Any]] = None, metadatas: Optional[list[Any]] = None, ) -> None: """Upsert documents into the collection. Args: ids (List[ItemID]): A list of document IDs. documents (List): A list of documents. embeddings (List): A list of document embeddings. metadatas (List): A list of document metadatas. Returns: None """ cursor = self.client.cursor() sql_values = [] if embeddings is not None and metadatas is not None: for doc_id, embedding, metadata, document in zip(ids, embeddings, metadatas, documents): metadata = re.sub("'", '"', str(metadata)) sql_values.append((doc_id, embedding, metadata, document, embedding, metadata, document)) sql_string = ( f"INSERT INTO {self.name} (id, embedding, metadatas, documents)\n" f"VALUES (%s, %s, %s, %s)\n" f"ON CONFLICT (id)\n" f"DO UPDATE SET embedding = %s,\n" f"metadatas = %s, documents = %s;\n" ) elif embeddings is not None: for doc_id, embedding, document in zip(ids, embeddings, documents): sql_values.append((doc_id, embedding, document, embedding, document)) sql_string = ( f"INSERT INTO {self.name} (id, embedding, documents) " f"VALUES (%s, %s, %s) ON CONFLICT (id)\n" f"DO UPDATE SET embedding = %s, documents = %s;\n" ) elif metadatas is not None: for doc_id, metadata, document in zip(ids, metadatas, documents): metadata = re.sub("'", '"', str(metadata)) embedding = self.embedding_function(document) sql_values.append((doc_id, metadata, embedding, document, metadata, document, embedding)) sql_string = ( f"INSERT INTO {self.name} (id, metadatas, embedding, documents)\n" f"VALUES (%s, %s, %s, %s)\n" f"ON CONFLICT (id)\n" f"DO UPDATE SET metadatas = %s, documents = %s, embedding = %s;\n" ) else: for doc_id, document in zip(ids, documents): embedding = self.embedding_function(document) sql_values.append((doc_id, document, embedding, document)) sql_string = ( f"INSERT INTO {self.name} (id, documents, embedding)\n" f"VALUES (%s, %s, %s)\n" f"ON CONFLICT (id)\n" f"DO UPDATE SET documents = %s;\n" ) logger.debug(f"Upsert SQL String:\n{sql_string}\n{sql_values}") cursor.executemany(sql_string, sql_values) cursor.close() def count(self) -> int: """Get the total number of documents in the collection. Returns: int: The total number of documents. """ cursor = self.client.cursor() query = f"SELECT COUNT(*) FROM {self.name}" cursor.execute(query) total = cursor.fetchone()[0] cursor.close() try: total = int(total) except (TypeError, ValueError): total = None return total def table_exists(self, table_name: str) -> bool: """Check if a table exists in the PostgreSQL database. Args: table_name (str): The name of the table to check. Returns: bool: True if the table exists, False otherwise. """ cursor = self.client.cursor() cursor.execute( """ SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name = %s ) """, (table_name,), ) exists = cursor.fetchone()[0] return exists def get( self, ids: Optional[str] = None, include: Optional[str] = None, where: Optional[str] = None, limit: Optional[Union[int, str]] = None, offset: Optional[Union[int, str]] = None, ) -> list[Document]: """Retrieve documents from the collection. Args: ids (Optional[List]): A list of document IDs. include (Optional): The fields to include. where (Optional): Additional filtering criteria. limit (Optional): The maximum number of documents to retrieve. offset (Optional): The offset for pagination. Returns: List: The retrieved documents. """ cursor = self.client.cursor() # Initialize variables for query components select_clause = "SELECT id, metadatas, documents, embedding" from_clause = f"FROM {self.name}" where_clause = "" limit_clause = "" offset_clause = "" # Handle include clause if include: select_clause = f"SELECT id, {', '.join(include)}, embedding" # Handle where clause if ids: where_clause = f"WHERE id IN ({', '.join(['%s' for _ in ids])})" elif where: where_clause = f"WHERE {where}" # Handle limit and offset clauses if limit: limit_clause = "LIMIT %s" if offset: offset_clause = "OFFSET %s" # Construct the full query query = f"{select_clause} {from_clause} {where_clause} {limit_clause} {offset_clause}" retrieved_documents = [] try: # Execute the query with the appropriate values if ids is not None: cursor.execute(query, ids) else: query_params = [] if limit: query_params.append(limit) if offset: query_params.append(offset) cursor.execute(query, query_params) retrieval = cursor.fetchall() for retrieved_document in retrieval: retrieved_documents.append( Document( id=retrieved_document[0].strip(), metadata=retrieved_document[1], content=retrieved_document[2], embedding=retrieved_document[3], ) ) except (psycopg.errors.UndefinedTable, psycopg.errors.UndefinedColumn) as e: logger.info(f"Error executing select on non-existent table: {self.name}. Creating it instead. Error: {e}") self.create_collection(collection_name=self.name, dimension=self.dimension) logger.info(f"Created table {self.name}") cursor.close() return retrieved_documents def update(self, ids: list[str], embeddings: list[Any], metadatas: list[Any], documents: list[Document]) -> None: """Update documents in the collection. Args: ids (List): A list of document IDs. embeddings (List): A list of document embeddings. metadatas (List): A list of document metadatas. documents (List): A list of documents. Returns: None """ cursor = self.client.cursor() sql_values = [] for doc_id, embedding, metadata, document in zip(ids, embeddings, metadatas, documents): sql_values.append((doc_id, embedding, metadata, document, doc_id, embedding, metadata, document)) sql_string = ( f"INSERT INTO {self.name} (id, embedding, metadata, document) " f"VALUES (%s, %s, %s, %s) " f"ON CONFLICT (id) " f"DO UPDATE SET id = %s, embedding = %s, " f"metadata = %s, document = %s;\n" ) logger.debug(f"Upsert SQL String:\n{sql_string}\n") cursor.executemany(sql_string, sql_values) cursor.close() @staticmethod def euclidean_distance(arr1: list[float], arr2: list[float]) -> float: """Calculate the Euclidean distance between two vectors. Parameters: - arr1 (List[float]): The first vector. - arr2 (List[float]): The second vector. Returns: - float: The Euclidean distance between arr1 and arr2. """ dist = np.linalg.norm(arr1 - arr2) return dist @staticmethod def cosine_distance(arr1: list[float], arr2: list[float]) -> float: """Calculate the cosine distance between two vectors. Parameters: - arr1 (List[float]): The first vector. - arr2 (List[float]): The second vector. Returns: - float: The cosine distance between arr1 and arr2. """ dist = np.dot(arr1, arr2) / (np.linalg.norm(arr1) * np.linalg.norm(arr2)) return dist @staticmethod def inner_product_distance(arr1: list[float], arr2: list[float]) -> float: """Calculate the Euclidean distance between two vectors. Parameters: - arr1 (List[float]): The first vector. - arr2 (List[float]): The second vector. Returns: - float: The Euclidean distance between arr1 and arr2. """ dist = np.linalg.norm(arr1 - arr2) return dist def query( self, query_texts: list[str], collection_name: Optional[str] = None, n_results: Optional[int] = 10, distance_type: Optional[str] = "euclidean", distance_threshold: Optional[float] = -1, include_embedding: Optional[bool] = False, ) -> QueryResults: """Query documents in the collection. Args: query_texts (List[str]): A list of query texts. collection_name (Optional[str]): The name of the collection. n_results (int): The maximum number of results to return. distance_type (Optional[str]): Distance search type - euclidean or cosine distance_threshold (Optional[float]): Distance threshold to limit searches include_embedding (Optional[bool]): Include embedding values in QueryResults Returns: QueryResults: The query results. """ if collection_name: self.name = collection_name clause = "ORDER BY" if distance_threshold == -1: distance_threshold = "" clause = "ORDER BY" elif distance_threshold > 0: distance_threshold = f"< {distance_threshold}" clause = "WHERE" cursor = self.client.cursor() results = [] for query_text in query_texts: vector = self.embedding_function(query_text, convert_to_tensor=False).tolist() if distance_type.lower() == "cosine": index_function = "<=>" elif distance_type.lower() == "euclidean": index_function = "<->" elif distance_type.lower() == "inner-product": index_function = "<#>" else: index_function = "<->" query = ( f"SELECT id, documents, embedding, metadatas " f"FROM {self.name} " f"{clause} embedding {index_function} '{vector!s}' {distance_threshold} " f"LIMIT {n_results}" ) cursor.execute(query) result = [] for row in cursor.fetchall(): fetched_document = Document(id=row[0].strip(), content=row[1], embedding=row[2], metadata=row[3]) fetched_document_array = self.convert_string_to_array(array_string=fetched_document.get("embedding")) if distance_type.lower() == "cosine": distance = self.cosine_distance(fetched_document_array, vector) elif distance_type.lower() == "euclidean": distance = self.euclidean_distance(fetched_document_array, vector) elif distance_type.lower() == "inner-product": distance = self.inner_product_distance(fetched_document_array, vector) else: distance = self.euclidean_distance(fetched_document_array, vector) if not include_embedding: fetched_document = Document(id=row[0].strip(), content=row[1], metadata=row[3]) result.append((fetched_document, distance)) results.append(result) cursor.close() logger.debug(f"Query Results: {results}") return results @staticmethod def convert_string_to_array(array_string: str) -> list[float]: """Convert a string representation of an array to a list of floats. Parameters: - array_string (str): The string representation of the array. Returns: - list: A list of floats parsed from the input string. If the input is not a string, it returns the input itself. """ if not isinstance(array_string, str): return array_string array_string = array_string.strip("[]") array = [float(num) for num in array_string.split()] return array def modify(self, metadata, collection_name: Optional[str] = None) -> None: """Modify metadata for the collection. Args: collection_name: The name of the collection. metadata: The new metadata. Returns: None """ if collection_name: self.name = collection_name cursor = self.client.cursor() cursor.execute("UPDATE collectionsSET metadata = '%s'WHERE collection_name = '%s';", (metadata, self.name)) cursor.close() def delete(self, ids: list[ItemID], collection_name: Optional[str] = None) -> None: """Delete documents from the collection. Args: ids (List[ItemID]): A list of document IDs to delete. collection_name (str): The name of the collection to delete. Returns: None """ if collection_name: self.name = collection_name cursor = self.client.cursor() id_placeholders = ", ".join(["%s" for _ in ids]) cursor.execute(f"DELETE FROM {self.name} WHERE id IN ({id_placeholders});", ids) cursor.close() def delete_collection(self, collection_name: Optional[str] = None) -> None: """Delete the entire collection. Args: collection_name (Optional[str]): The name of the collection to delete. Returns: None """ if collection_name: self.name = collection_name cursor = self.client.cursor() cursor.execute(f"DROP TABLE IF EXISTS {self.name}") cursor.close() def create_collection( self, collection_name: Optional[str] = None, dimension: Optional[Union[str, int]] = None ) -> None: """Create a new collection. Args: collection_name (Optional[str]): The name of the new collection. dimension (Optional[Union[str, int]]): The dimension size of the sentence embedding model Returns: None """ if collection_name: self.name = collection_name if dimension: self.dimension = dimension elif self.dimension is None: self.dimension = 384 cursor = self.client.cursor() cursor.execute( f"CREATE TABLE {self.name} (" f"documents text, id CHAR(8) PRIMARY KEY, metadatas JSONB, embedding vector({self.dimension}));" f"CREATE INDEX " f"ON {self.name} USING hnsw (embedding vector_l2_ops) WITH (m = {self.metadata['hnsw:M']}, " f"ef_construction = {self.metadata['hnsw:construction_ef']});" f"CREATE INDEX " f"ON {self.name} USING hnsw (embedding vector_cosine_ops) WITH (m = {self.metadata['hnsw:M']}, " f"ef_construction = {self.metadata['hnsw:construction_ef']});" f"CREATE INDEX " f"ON {self.name} USING hnsw (embedding vector_ip_ops) WITH (m = {self.metadata['hnsw:M']}, " f"ef_construction = {self.metadata['hnsw:construction_ef']});" ) cursor.close() @require_optional_import(["pgvector", "psycopg", "sentence_transformers"], "retrievechat-pgvector") class PGVectorDB(VectorDB): """A vector database that uses PGVector as the backend.""" def __init__( self, *, conn: Optional["psycopg.Connection"] = None, connection_string: Optional[str] = None, host: Optional[str] = None, port: Optional[Union[int, str]] = None, dbname: Optional[str] = None, username: Optional[str] = None, password: Optional[str] = None, connect_timeout: Optional[int] = 10, embedding_function: Callable = None, metadata: Optional[dict[str, Any]] = None, ) -> None: """Initialize the vector database. Note: connection_string or host + port + dbname must be specified Args: conn: psycopg.Connection | A customer connection object to connect to the database. A connection object may include additional key/values: https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING connection_string: "postgresql://username:password@hostname:port/database" | The PGVector connection string. Default is None. host: str | The host to connect to. Default is None. port: int | The port to connect to. Default is None. dbname: str | The database name to connect to. Default is None. username: str | The database username to use. Default is None. password: str | The database user password to use. Default is None. connect_timeout: int | The timeout to set for the connection. Default is 10. embedding_function: Callable | The embedding function used to generate the vector representation. Default is None. SentenceTransformer("all-MiniLM-L6-v2").encode will be used when None. Models can be chosen from: https://huggingface.co/models?library=sentence-transformers metadata: dict | The metadata of the vector database. Default is None. If None, it will use this setting: `{"hnsw:space": "ip", "hnsw:construction_ef": 30, "hnsw:M": 16}`. Creates Index on table using hnsw (embedding vector_l2_ops) WITH (m = hnsw:M) ef_construction = "hnsw:construction_ef". For more info: https://github.com/pgvector/pgvector?tab=readme-ov-file#hnsw Returns: None """ self.client = self.establish_connection( conn=conn, connection_string=connection_string, host=host, port=port, dbname=dbname, username=username, password=password, connect_timeout=connect_timeout, ) if embedding_function: self.embedding_function = embedding_function else: self.embedding_function = SentenceTransformer("all-MiniLM-L6-v2").encode self.metadata = metadata register_vector(self.client) self.active_collection = None def establish_connection( self, conn: Optional["psycopg.Connection"] = None, connection_string: Optional[str] = None, host: Optional[str] = None, port: Optional[Union[int, str]] = None, dbname: Optional[str] = None, username: Optional[str] = None, password: Optional[str] = None, connect_timeout: Optional[int] = 10, ) -> "psycopg.Connection": """Establishes a connection to a PostgreSQL database using psycopg. Args: conn: An existing psycopg connection object. If provided, this connection will be used. connection_string: A string containing the connection information. If provided, a new connection will be established using this string. host: The hostname of the PostgreSQL server. Used if connection_string is not provided. port: The port number to connect to at the server host. Used if connection_string is not provided. dbname: The database name. Used if connection_string is not provided. username: The username to connect as. Used if connection_string is not provided. password: The user's password. Used if connection_string is not provided. connect_timeout: Maximum wait for connection, in seconds. The default is 10 seconds. Returns: A psycopg.Connection object representing the established connection. Raises: PermissionError if no credentials are supplied psycopg.Error: If an error occurs while trying to connect to the database. """ try: if conn: self.client = conn elif connection_string: parsed_connection = urllib.parse.urlparse(connection_string) encoded_username = urllib.parse.quote(parsed_connection.username, safe="") encoded_password = urllib.parse.quote(parsed_connection.password, safe="") encoded_password = f":{encoded_password}@" encoded_host = urllib.parse.quote(parsed_connection.hostname, safe="") encoded_port = f":{parsed_connection.port}" encoded_database = urllib.parse.quote(parsed_connection.path[1:], safe="") connection_string_encoded = ( f"{parsed_connection.scheme}://{encoded_username}{encoded_password}" f"{encoded_host}{encoded_port}/{encoded_database}" ) self.client = psycopg.connect(conninfo=connection_string_encoded, autocommit=True) elif host: connection_string = "" if host: encoded_host = urllib.parse.quote(host, safe="") connection_string += f"host={encoded_host} " if port: connection_string += f"port={port} " if dbname: encoded_database = urllib.parse.quote(dbname, safe="") connection_string += f"dbname={encoded_database} " if username: encoded_username = urllib.parse.quote(username, safe="") connection_string += f"user={encoded_username} " if password: encoded_password = urllib.parse.quote(password, safe="") connection_string += f"password={encoded_password} " self.client = psycopg.connect( conninfo=connection_string, connect_timeout=connect_timeout, autocommit=True, ) else: logger.error("Credentials were not supplied...") raise PermissionError self.client.execute("CREATE EXTENSION IF NOT EXISTS vector") except psycopg.Error as e: logger.error("Error connecting to the database: ", e) raise e return self.client def create_collection( self, collection_name: str, overwrite: bool = False, get_or_create: bool = True ) -> Collection: """Create a collection in the vector database. Case 1. if the collection does not exist, create the collection. Case 2. the collection exists, if overwrite is True, it will overwrite the collection. Case 3. the collection exists and overwrite is False, if get_or_create is True, it will get the collection, otherwise it raise a ValueError. Args: collection_name: str | The name of the collection. overwrite: bool | Whether to overwrite the collection if it exists. Default is False. get_or_create: bool | Whether to get the collection if it exists. Default is True. Returns: Collection | The collection object. """ try: if self.active_collection and self.active_collection.name == collection_name: collection = self.active_collection else: collection = self.get_collection(collection_name) except ValueError: collection = None if collection is None: collection = Collection( client=self.client, collection_name=collection_name, embedding_function=self.embedding_function, get_or_create=get_or_create, metadata=self.metadata, ) collection.set_collection_name(collection_name=collection_name) collection.create_collection(collection_name=collection_name) return collection elif overwrite: self.delete_collection(collection_name) collection = Collection( client=self.client, collection_name=collection_name, embedding_function=self.embedding_function, get_or_create=get_or_create, metadata=self.metadata, ) collection.set_collection_name(collection_name=collection_name) collection.create_collection(collection_name=collection_name) return collection elif get_or_create: return collection elif not collection.table_exists(table_name=collection_name): collection = Collection( client=self.client, collection_name=collection_name, embedding_function=self.embedding_function, get_or_create=get_or_create, metadata=self.metadata, ) collection.set_collection_name(collection_name=collection_name) collection.create_collection(collection_name=collection_name) return collection else: raise ValueError(f"Collection {collection_name} already exists.") def get_collection(self, collection_name: str = None) -> Collection: """Get the collection from the vector database. Args: collection_name: str | The name of the collection. Default is None. If None, return the current active collection. Returns: Collection | The collection object. """ if collection_name is None: if self.active_collection is None: raise ValueError("No collection is specified.") else: logger.debug( f"No collection is specified. Using current active collection {self.active_collection.name}." ) else: if not (self.active_collection and self.active_collection.name == collection_name): self.active_collection = Collection( client=self.client, collection_name=collection_name, embedding_function=self.embedding_function, ) return self.active_collection def delete_collection(self, collection_name: str) -> None: """Delete the collection from the vector database. Args: collection_name: str | The name of the collection. Returns: None """ if self.active_collection: self.active_collection.delete_collection(collection_name) else: collection = self.get_collection(collection_name) collection.delete_collection(collection_name) if self.active_collection and self.active_collection.name == collection_name: self.active_collection = None def _batch_insert( self, collection: Collection, embeddings=None, ids=None, metadatas=None, documents=None, upsert=False ) -> None: batch_size = int(PGVECTOR_MAX_BATCH_SIZE) default_metadata = {"hnsw:space": "ip", "hnsw:construction_ef": 32, "hnsw:M": 16} default_metadatas = [default_metadata] * min(batch_size, len(documents)) for i in range(0, len(documents), min(batch_size, len(documents))): end_idx = i + min(batch_size, len(documents) - i) collection_kwargs = { "documents": documents[i:end_idx], "ids": ids[i:end_idx], "metadatas": metadatas[i:end_idx] if metadatas else default_metadatas, "embeddings": embeddings[i:end_idx] if embeddings else None, } if upsert: collection.upsert(**collection_kwargs) else: collection.add(**collection_kwargs) def insert_docs(self, docs: list[Document], collection_name: str = None, upsert: bool = False) -> None: """Insert documents into the collection of the vector database. Args: docs: List[Document] | A list of documents. Each document is a TypedDict `Document`. collection_name: str | The name of the collection. Default is None. upsert: bool | Whether to update the document if it exists. Default is False. kwargs: Dict | Additional keyword arguments. Returns: None """ if not docs: return if docs[0].get("content") is None: raise ValueError("The document content is required.") if docs[0].get("id") is None: raise ValueError("The document id is required.") documents = [doc.get("content") for doc in docs] ids = [doc.get("id") for doc in docs] collection = self.get_collection(collection_name) if docs[0].get("embedding") is None: logger.debug( "No content embedding is provided. " "Will use the VectorDB's embedding function to generate the content embedding." ) embeddings = None else: embeddings = [doc.get("embedding") for doc in docs] metadatas = None if docs[0].get("metadata") is None else [doc.get("metadata") for doc in docs] self._batch_insert(collection, embeddings, ids, metadatas, documents, upsert) def update_docs(self, docs: list[Document], collection_name: str = None) -> None: """Update documents in the collection of the vector database. Args: docs: List[Document] | A list of documents. collection_name: str | The name of the collection. Default is None. Returns: None """ self.insert_docs(docs, collection_name, upsert=True) def delete_docs(self, ids: list[ItemID], collection_name: str = None) -> None: """Delete documents from the collection of the vector database. Args: ids: List[ItemID] | A list of document ids. Each id is a typed `ItemID`. collection_name: str | The name of the collection. Default is None. kwargs: Dict | Additional keyword arguments. Returns: None """ collection = self.get_collection(collection_name) collection.delete(ids=ids, collection_name=collection_name) def retrieve_docs( self, queries: list[str], collection_name: str = None, n_results: int = 10, distance_threshold: float = -1, ) -> QueryResults: """Retrieve documents from the collection of the vector database based on the queries. Args: queries: List[str] | A list of queries. Each query is a string. collection_name: str | The name of the collection. Default is None. n_results: int | The number of relevant documents to return. Default is 10. distance_threshold: float | The threshold for the distance score, only distance smaller than it will be returned. Don't filter with it if `< 0`. Default is -1. kwargs: Dict | Additional keyword arguments. Returns: QueryResults | The query results. Each query result is a list of list of tuples containing the document and the distance. """ collection = self.get_collection(collection_name) if isinstance(queries, str): queries = [queries] results = collection.query( query_texts=queries, n_results=n_results, distance_threshold=distance_threshold, ) logger.debug(f"Retrieve Docs Results:\n{results}") return results def get_docs_by_ids( self, ids: list[ItemID] = None, collection_name: str = None, include=None, **kwargs ) -> list[Document]: """Retrieve documents from the collection of the vector database based on the ids. Args: ids: List[ItemID] | A list of document ids. If None, will return all the documents. Default is None. collection_name: str | The name of the collection. Default is None. include: List[str] | The fields to include. Default is None. If None, will include ["metadatas", "documents"], ids will always be included. kwargs: dict | Additional keyword arguments. Returns: List[Document] | The results. """ collection = self.get_collection(collection_name) include = include if include else ["metadatas", "documents"] results = collection.get(ids, include=include, **kwargs) logger.debug(f"Retrieve Documents by ID Results:\n{results}") return results