autogen.agentchat.contrib.vectordb.pgvectordb.Collection

Collection

Collection(
    client: Any | None = None,
    collection_name: str = 'ag2-docs',
    embedding_function: Callable[..., Any] | None = None,
    metadata: Any | None = None,
    get_or_create: Any | None = None
)

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.
Initialize the Collection object.

Parameters:

Name	Description
client	Type: Any | None Default: None
collection_name	Type: str Default: ‘ag2-docs’
embedding_function	Type: Callable[…, Any] | None Default: None
metadata	Type: Any | None Default: None
get_or_create	Type: Any | None Default: None

Instance Methods

add

add(
    self,
    ids: list[str | int],
    documents: list[Document] | None,
    embeddings: list[Any] | None = None,
    metadatas: list[Any] | None = None
) -> None

Add documents to the collection.

Parameters:

Name	Description
ids	A list of document IDs. Type: list[str | int]
documents	A list of documents. Type: list[Document] | None
embeddings	A list of document embeddings. Optional Type: list[typing.Any] | None Default: None
metadatas	A list of document metadatas. Optional Type: list[typing.Any] | None Default: None

Returns:

Type	Description
None	None

convert_string_to_array

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.

Parameters:

Name	Description
array_string	Type: str

Returns:

Type	Description
list[float]	- list: A list of floats parsed from the input string. If the input is not a string, it returns the input itself.

cosine_distance

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.

Parameters:

Name	Description
arr1	Type: list[float]
arr2	Type: list[float]

Returns:

Type	Description
float	- float: The cosine distance between arr1 and arr2.

count

count(self) -> int

Get the total number of documents in the collection.
Returns:
int: The total number of documents.

Returns:

Type	Description
int	int: The total number of documents.

create_collection

create_collection(
    self,
    collection_name: str | None = None,
    dimension: int | str | None = None
) -> None

Create a new collection.

Parameters:

Name	Description
collection_name	The name of the new collection. Type: str | None Default: None
dimension	The dimension size of the sentence embedding model Type: int | str | None Default: None

Returns:

Type	Description
None	None

delete

delete(
    self,
    ids: list[str | int],
    collection_name: str | None = None
) -> None

Delete documents from the collection.

Parameters:

Name	Description
ids	A list of document IDs to delete. Type: list[str | int]
collection_name	The name of the collection to delete. Type: str | None Default: None

Returns:

Type	Description
None	None

delete_collection

delete_collection(self, collection_name: str | None = None) -> None

Delete the entire collection.

Parameters:

Name	Description
collection_name	The name of the collection to delete. Type: str | None Default: None

Returns:

Type	Description
None	None

euclidean_distance

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.

Parameters:

Name	Description
arr1	Type: list[float]
arr2	Type: list[float]

Returns:

Type	Description
float	- float: The Euclidean distance between arr1 and arr2.

get

get(
    self,
    ids: str | None = None,
    include: str | None = None,
    where: str | None = None,
    limit: int | str | None = None,
    offset: int | str | None = None
) -> list[Document]

Retrieve documents from the collection.

Parameters:

Name	Description
ids	A list of document IDs. Type: str | None Default: None
include	The fields to include. Type: str | None Default: None
where	Additional filtering criteria. Type: str | None Default: None
limit	The maximum number of documents to retrieve. Type: int | str | None Default: None
offset	The offset for pagination. Type: int | str | None Default: None

Returns:

Type	Description
list[Document]	List: The retrieved documents.

inner_product_distance

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.

Parameters:

Name	Description
arr1	Type: list[float]
arr2	Type: list[float]

Returns:

Type	Description
float	- float: The Euclidean distance between arr1 and arr2.

modify

modify(
    self,
    metadata,
    collection_name: str | None = None
) -> None

Modify metadata for the collection.

Parameters:

Name	Description
metadata	The new metadata.
collection_name	The name of the collection. Type: str | None Default: None

Returns:

Type	Description
None	None

query

query(
    self,
    query_texts: list[str],
    collection_name: str | None = None,
    n_results: int | None = 10,
    distance_type: str | None = 'euclidean',
    distance_threshold: float | None = -1,
    include_embedding: bool | None = False
) -> list[list[tuple[Document, float]]]

Query documents in the collection.

Parameters:

Name	Description
query_texts	A list of query texts. Type: list[str]
collection_name	The name of the collection. Type: str | None Default: None
n_results	The maximum number of results to return. Type: int | None Default: 10
distance_type	Distance search type - euclidean or cosine Type: str | None Default: ‘euclidean’
distance_threshold	Distance threshold to limit searches Type: float | None Default: -1
include_embedding	Include embedding values in QueryResults Type: bool | None Default: False

Returns:

Type	Description
list[list[tuple[Document, float]]]	QueryResults: The query results.

set_collection_name

set_collection_name(self, collection_name) -> str

Parameters:

Name	Description
collection_name

table_exists

table_exists(self, table_name: str) -> bool

Check if a table exists in the PostgreSQL database.

Parameters:

Name	Description
table_name	The name of the table to check. Type: str

Returns:

Type	Description
bool	bool: True if the table exists, False otherwise.

update

update(
    self,
    ids: list[str],
    embeddings: list[Any],
    metadatas: list[Any],
    documents: list[Document]
) -> None

Update documents in the collection.

Parameters:

Name	Description
ids	A list of document IDs. Type: list[str]
embeddings	A list of document embeddings. Type: list[typing.Any]
metadatas	A list of document metadatas. Type: list[typing.Any]
documents	A list of documents. Type: list[Document]

Returns:

Type	Description
None	None

upsert

upsert(
    self,
    ids: list[str | int],
    documents: list[Document],
    embeddings: list[Any] | None = None,
    metadatas: list[Any] | None = None
) -> None

Upsert documents into the collection.

Parameters:

Name	Description
ids	A list of document IDs. Type: list[str | int]
documents	A list of documents. Type: list[Document]
embeddings	A list of document embeddings. Type: list[typing.Any] | None Default: None
metadatas	A list of document metadatas. Type: list[typing.Any] | None Default: None

Returns:

Type	Description
None	None