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

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

API Reference

​Collection

​Instance Methods

​add

​convert_string_to_array

​cosine_distance

​count

​create_collection

​delete

​delete_collection

​euclidean_distance

​get

​inner_product_distance

​modify

​query

​set_collection_name

​table_exists

​update

​upsert

Collection

Instance Methods

add

convert_string_to_array

cosine_distance

count

create_collection

delete

delete_collection

euclidean_distance

get

inner_product_distance

modify

query

set_collection_name

table_exists

update

upsert