about summary refs log tree commit diff
path: root/.venv/lib/python3.12/site-packages/vecs/client.py
diff options
context:
space:
mode:
authorS. Solomon Darnell2025-03-28 21:52:21 -0500
committerS. Solomon Darnell2025-03-28 21:52:21 -0500
commit4a52a71956a8d46fcb7294ac71734504bb09bcc2 (patch)
treeee3dc5af3b6313e921cd920906356f5d4febc4ed /.venv/lib/python3.12/site-packages/vecs/client.py
parentcc961e04ba734dd72309fb548a2f97d67d578813 (diff)
downloadgn-ai-master.tar.gz
two version of R2R are here HEAD master
Diffstat (limited to '.venv/lib/python3.12/site-packages/vecs/client.py')
-rw-r--r--.venv/lib/python3.12/site-packages/vecs/client.py246
1 files changed, 246 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/vecs/client.py b/.venv/lib/python3.12/site-packages/vecs/client.py
new file mode 100644
index 00000000..89bb3e3d
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/vecs/client.py
@@ -0,0 +1,246 @@
+"""
+Defines the 'Client' class
+
+Importing from the `vecs.client` directly is not supported.
+All public classes, enums, and functions are re-exported by the top level `vecs` module.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Optional
+
+from deprecated import deprecated
+from sqlalchemy import MetaData, create_engine, text
+from sqlalchemy.orm import sessionmaker
+
+from vecs.adapter import Adapter
+from vecs.exc import CollectionNotFound
+
+if TYPE_CHECKING:
+    from vecs.collection import Collection
+
+
+class Client:
+    """
+    The `vecs.Client` class serves as an interface to a PostgreSQL database with pgvector support. It facilitates
+    the creation, retrieval, listing and deletion of vector collections, while managing connections to the
+    database.
+
+    A `Client` instance represents a connection to a PostgreSQL database. This connection can be used to create
+    and manipulate vector collections, where each collection is a group of vector records in a PostgreSQL table.
+
+    The `vecs.Client` class can be also supports usage as a context manager to ensure the connection to the database
+    is properly closed after operations, or used directly.
+
+    Example usage:
+
+        DB_CONNECTION = "postgresql://<user>:<password>@<host>:<port>/<db_name>"
+
+        with vecs.create_client(DB_CONNECTION) as vx:
+            # do some work
+            pass
+
+        # OR
+
+        vx = vecs.create_client(DB_CONNECTION)
+        # do some work
+        vx.disconnect()
+    """
+
+    def __init__(self, connection_string: str):
+        """
+        Initialize a Client instance.
+
+        Args:
+            connection_string (str): A string representing the database connection information.
+
+        Returns:
+            None
+        """
+        self.engine = create_engine(connection_string)
+        self.meta = MetaData(schema="vecs")
+        self.Session = sessionmaker(self.engine)
+
+        with self.Session() as sess:
+            with sess.begin():
+                sess.execute(text("create schema if not exists vecs;"))
+                sess.execute(text("create extension if not exists vector;"))
+                self.vector_version: str = sess.execute(
+                    text(
+                        "select installed_version from pg_available_extensions where name = 'vector' limit 1;"
+                    )
+                ).scalar_one()
+
+    def _supports_hnsw(self):
+        return (
+            not self.vector_version.startswith("0.4")
+            and not self.vector_version.startswith("0.3")
+            and not self.vector_version.startswith("0.2")
+            and not self.vector_version.startswith("0.1")
+            and not self.vector_version.startswith("0.0")
+        )
+
+    def get_or_create_collection(
+        self,
+        name: str,
+        *,
+        dimension: Optional[int] = None,
+        adapter: Optional[Adapter] = None,
+    ) -> Collection:
+        """
+        Get a vector collection by name, or create it if no collection with
+        *name* exists.
+
+        Args:
+            name (str): The name of the collection.
+
+        Keyword Args:
+            dimension (int): The dimensionality of the vectors in the collection.
+            pipeline (int): The dimensionality of the vectors in the collection.
+
+        Returns:
+            Collection: The created collection.
+
+        Raises:
+            CollectionAlreadyExists: If a collection with the same name already exists
+        """
+        from vecs.collection import Collection
+
+        adapter_dimension = adapter.exported_dimension if adapter else None
+
+        collection = Collection(
+            name=name,
+            dimension=dimension or adapter_dimension,  # type: ignore
+            client=self,
+            adapter=adapter,
+        )
+
+        return collection._create_if_not_exists()
+
+    @deprecated("use Client.get_or_create_collection")
+    def create_collection(self, name: str, dimension: int) -> Collection:
+        """
+        Create a new vector collection.
+
+        Args:
+            name (str): The name of the collection.
+            dimension (int): The dimensionality of the vectors in the collection.
+
+        Returns:
+            Collection: The created collection.
+
+        Raises:
+            CollectionAlreadyExists: If a collection with the same name already exists
+        """
+        from vecs.collection import Collection
+
+        return Collection(name, dimension, self)._create()
+
+    @deprecated("use Client.get_or_create_collection")
+    def get_collection(self, name: str) -> Collection:
+        """
+        Retrieve an existing vector collection.
+
+        Args:
+            name (str): The name of the collection.
+
+        Returns:
+            Collection: The retrieved collection.
+
+        Raises:
+            CollectionNotFound: If no collection with the given name exists.
+        """
+        from vecs.collection import Collection
+
+        query = text(
+            f"""
+        select
+            relname as table_name,
+            atttypmod as embedding_dim
+        from
+            pg_class pc
+            join pg_attribute pa
+                on pc.oid = pa.attrelid
+        where
+            pc.relnamespace = 'vecs'::regnamespace
+            and pc.relkind = 'r'
+            and pa.attname = 'vec'
+            and not pc.relname ^@ '_'
+            and pc.relname = :name
+        """
+        ).bindparams(name=name)
+        with self.Session() as sess:
+            query_result = sess.execute(query).fetchone()
+
+            if query_result is None:
+                raise CollectionNotFound("No collection found with requested name")
+
+            name, dimension = query_result
+            return Collection(
+                name,
+                dimension,
+                self,
+            )
+
+    def list_collections(self) -> List["Collection"]:
+        """
+        List all vector collections.
+
+        Returns:
+            list[Collection]: A list of all collections.
+        """
+        from vecs.collection import Collection
+
+        return Collection._list_collections(self)
+
+    def delete_collection(self, name: str) -> None:
+        """
+        Delete a vector collection.
+
+        If no collection with requested name exists, does nothing.
+
+        Args:
+            name (str): The name of the collection.
+
+        Returns:
+            None
+        """
+        from vecs.collection import Collection
+
+        Collection(name, -1, self)._drop()
+        return
+
+    def disconnect(self) -> None:
+        """
+        Disconnect the client from the database.
+
+        Returns:
+            None
+        """
+        self.engine.dispose()
+        return
+
+    def __enter__(self) -> "Client":
+        """
+        Enable use of the 'with' statement.
+
+        Returns:
+            Client: The current instance of the Client.
+        """
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """
+        Disconnect the client on exiting the 'with' statement context.
+
+        Args:
+            exc_type: The exception type, if any.
+            exc_val: The exception value, if any.
+            exc_tb: The traceback, if any.
+
+        Returns:
+            None
+        """
+        self.disconnect()
+        return