asyncpg-vector 0.1.0__py3-none-any.whl

asyncpg_vector/__init__.py

@@ -0,0 +1,193 @@
+"""
+PostgreSQL vector support for asyncpg.
+
+Registers data types `vector` and `halfvec` from the PostgreSQL extension `vector` to the asynchronous PostgreSQL
+client `asyncpg`, and marshals vector data to and from PostgreSQL database tables.
+
+:see: https://github.com/hunyadi/asyncpg_vector
+"""
+
+import base64
+import struct
+import sys
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from struct import pack
+from typing import ClassVar
+
+import asyncpg
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
+
+__version__ = "0.1.0"
+__author__ = "Levente Hunyadi"
+__copyright__ = "Copyright 2025, Levente Hunyadi"
+__license__ = "MIT"
+__maintainer__ = "Levente Hunyadi"
+__status__ = "Production"
+
+
+class BasicVector(ABC):
+    "Base class for PostgreSQL vector types."
+
+    __slots__ = ("_data",)
+
+    _data: bytes
+
+    def __init__(self, data: bytes) -> None:
+        self._data = data
+
+    def __eq__(self, value: object) -> bool:
+        if not isinstance(value, self.__class__):
+            return False
+
+        return self._data == value._data
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}({self._data!r})"
+
+    def size(self) -> int:
+        "Number of vector dimensions."
+
+        return len(self._data) // self.bytes_per_item()
+
+    @classmethod
+    @abstractmethod
+    def bytes_per_item(cls) -> int:
+        "Number of bytes per item in the vector."
+
+        ...
+
+    @abstractmethod
+    def to_float_list(self) -> list[float]:
+        "Converts the vector to a list of double-precision floating-point values."
+
+        ...
+
+    @classmethod
+    @abstractmethod
+    def from_float_list(cls, vec: Sequence[float]) -> Self:
+        "Creates a vector from a list of double-precision floating-point numbers."
+
+        ...
+
+    @classmethod
+    def from_float_base64(cls, base64_encoded: str) -> Self:
+        """
+        Creates a vector from a base64-encoded string storing a list of double-precision floating-point numbers.
+
+        This allows lossless transition of a floating-point value between APIs.
+        """
+
+        decoded_bytes = base64.b64decode(base64_encoded)
+        float_tuple = struct.unpack(f"{len(decoded_bytes) // 4}f", decoded_bytes)
+        return cls.from_float_list(float_tuple)
+
+    def to_database_binary(self) -> bytes:
+        "Writes the data into the PostgreSQL data transfer representation."
+
+        return pack(">HH", self.size(), 0) + self._data
+
+    @classmethod
+    def from_database_binary(cls, data: bytes) -> Self:
+        "Reads the data from the PostgreSQL data transfer representation."
+
+        size, _unused = struct.unpack(">HH", data[0:4])
+        if len(data) != 4 + cls.bytes_per_item() * size:
+            raise ValueError(f"expected size: {cls.bytes_per_item()} * {size}; got {len(data) - 4} bytes")
+        return cls(data[4:])
+
+    @classmethod
+    def _to_database_binary(cls, value: Self | list[float] | None) -> bytes | None:
+        "Writes a value or instance into the PostgreSQL data transfer representation."
+
+        if value is None:
+            return value  # asyncpg uses `None` for representing SQL `NULL`
+
+        match value:
+            case BasicVector():
+                obj = value
+            case list():
+                if value:
+                    list_item = value[0]
+                    if not isinstance(list_item, float):
+                        raise ValueError(f"unsupported list item type: {type(list_item).__name__}")
+
+                    obj = cls.from_float_list(value)
+                else:
+                    obj = cls(bytes())
+            case _:
+                raise ValueError(f"unsupported type: {type(value).__name__}")
+
+        return obj.to_database_binary()
+
+    @classmethod
+    def _from_database_binary(cls, value: bytes | None) -> Self | None:
+        "Creates an instance from the PostgreSQL data transfer representation."
+
+        if value is None:
+            return value
+
+        return cls.from_database_binary(value)
+
+
+class HalfVector(BasicVector):
+    "Implements the PostgreSQL extension type `halfvec`."
+
+    type_name: ClassVar[str] = "halfvec"
+    cosine_similarity: ClassVar[str] = "halfvec_cosine_ops"
+
+    @override
+    @classmethod
+    def bytes_per_item(cls) -> int:
+        return 2
+
+    @override
+    def to_float_list(self) -> list[float]:
+        return list(struct.unpack(f">{self.size()}e", self._data))
+
+    @override
+    @classmethod
+    def from_float_list(cls, vec: Sequence[float]) -> Self:
+        return cls(struct.pack(f">{len(vec)}e", *vec))
+
+
+class Vector(BasicVector):
+    "Implements the PostgreSQL extension type `vector`."
+
+    type_name: ClassVar[str] = "vector"
+    cosine_similarity: ClassVar[str] = "vector_cosine_ops"
+
+    @override
+    @classmethod
+    def bytes_per_item(cls) -> int:
+        return 4
+
+    @override
+    def to_float_list(self) -> list[float]:
+        return list(struct.unpack(f">{self.size()}f", self._data))
+
+    @override
+    @classmethod
+    def from_float_list(cls, vec: Sequence[float]) -> Self:
+        return cls(struct.pack(f">{len(vec)}f", *vec))
+
+    @override
+    @classmethod
+    def from_float_base64(cls, base64_encoded: str) -> Self:
+        return cls(base64.b64decode(base64_encoded))
+
+
+async def register_vector(conn: asyncpg.Connection, schema: str = "public") -> None:
+    "Registers `vector` extension types with Python module `asyncpg`."
+
+    await conn.set_type_codec("vector", schema=schema, encoder=Vector._to_database_binary, decoder=Vector._from_database_binary, format="binary")  # pyright: ignore[reportPrivateUsage]
+    await conn.set_type_codec("halfvec", schema=schema, encoder=HalfVector._to_database_binary, decoder=HalfVector._from_database_binary, format="binary")  # pyright: ignore[reportPrivateUsage]

asyncpg_vector-0.1.0.dist-info/METADATA

@@ -0,0 +1,132 @@
+Metadata-Version: 2.4
+Name: asyncpg_vector
+Version: 0.1.0
+Summary: PostgreSQL vector support for asyncpg
+Author-email: Levente Hunyadi <hunyadi@gmail.com>
+Maintainer-email: Levente Hunyadi <hunyadi@gmail.com>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/hunyadi/asyncpg_vector
+Project-URL: Source, https://github.com/hunyadi/asyncpg_vector
+Keywords: asyncpg,postgresql-extension,vector,vector-search
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Framework :: AsyncIO
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Topic :: Database :: Front-Ends
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Typing :: Typed
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: asyncpg>=0.30
+Requires-Dist: typing-extensions>=4.15; python_version < "3.11"
+Provides-Extra: dev
+Requires-Dist: asyncpg-stubs>=0.30; extra == "dev"
+Requires-Dist: build>=1.3; extra == "dev"
+Requires-Dist: mypy>=1.18; extra == "dev"
+Requires-Dist: ruff>=0.14; extra == "dev"
+Dynamic: license-file
+
+# PostgreSQL vector support for asyncpg
+
+Adds PostgreSQL [vector](https://github.com/pgvector/pgvector) support for Python.
+
+Registers data types `vector` and `halfvec` from the PostgreSQL extension `vector` to the asynchronous PostgreSQL client `asyncpg`, and marshals vector data to and from PostgreSQL database tables.
+
+Internally, the data is packed into a Python `bytes` object, with single-precision float vectors stored on 4 bytes per item (for class `Vector`) and half-precision float vectors stored on 2 bytes per item (for class `HalfVector`). Data is (un)packed using `struct` from the standard library as necessary.
+
+This module provides functionality similar to [pgvector-python](https://github.com/pgvector/pgvector-python) but imports minimum dependencies (e.g. no dependency on `numpy`).
+
+## Setup
+
+#### Install the package
+
+```sh
+pip install asyncpg_vector
+```
+
+#### Initialize
+
+Register vector types with your database connection or connection pool:
+
+**Connection**:
+
+```python
+from asyncpg_vector import register_vector
+
+async def main() -> None:
+    ...
+
+    await conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
+    await register_vector(conn)
+```
+
+**Pool**:
+
+```python
+from asyncpg_vector import register_vector
+
+async def init_connection(conn: asyncpg.Connection) -> None:
+    await conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
+    await register_vector(conn)
+
+async def main() -> None:
+    ...
+
+    pool = await asyncpg.create_pool(..., init=init_connection)
+```
+
+#### Perform similarity search
+
+First, create a table and an index:
+
+```python
+async def create(conn: asyncpg.Connection) -> None:
+    await conn.execute("""
+        CREATE TABLE IF NOT EXISTS items
+        (
+            id bigint NOT NULL GENERATED ALWAYS AS IDENTITY,
+            content text NOT NULL,
+            embedding halfvec(1536) NOT NULL,
+            CONSTRAINT pk_items PRIMARY KEY (id)
+        );
+
+        CREATE INDEX IF NOT EXISTS embedding_index ON items
+        USING hnsw (embedding halfvec_cosine_ops);
+    """)
+```
+
+Next, find documents in a knowledge base that match a search phrase using vector similarity with approximate nearest neighbor semantics:
+
+```python
+from asyncpg_vector import HalfVector
+
+async def search(conn: asyncpg.Connection, phrase: str) -> list[str]:
+    ...
+
+    embedding_response = await ai_client.embeddings.create(
+        input=phrase,
+        model="text-embedding-3-small",
+        encoding_format="base64"
+    )
+    embedding = HalfVector.from_float_base64(embedding_response.data[0].embedding)
+    query = """
+        SELECT
+            id,
+            content,
+            embedding <=> $1 AS distance
+        FROM items
+        ORDER BY distance
+        LIMIT 5
+    """
+    rows = await conn.fetch(query, embedding)
+
+    ...
+```

asyncpg_vector-0.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+asyncpg_vector/__init__.py,sha256=0Tn14_Kza21A7VerWo8PJxZQYf7RC1KX9tLBAxaEbSA,6117
+asyncpg_vector/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+asyncpg_vector-0.1.0.dist-info/licenses/LICENSE,sha256=rx4jD36wX8TyLZaR2HEOJ6TphFPjKUqoCSSYWzwWNRk,1093
+asyncpg_vector-0.1.0.dist-info/METADATA,sha256=e4ALWMZu08sXLT0X4qIayq0LcFJ3LJNsTMYAEUC_L40,4405
+asyncpg_vector-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+asyncpg_vector-0.1.0.dist-info/top_level.txt,sha256=HlvoexcqaYDtjDfX9A6Nibq4_cEfyhj03nAqouJ4Tpg,15
+asyncpg_vector-0.1.0.dist-info/zip-safe,sha256=frcCV1k9oG9oKj3dpUqdJg1PxRT2RSN_XKdLCPjaYaY,2
+asyncpg_vector-0.1.0.dist-info/RECORD,,

asyncpg_vector-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

asyncpg_vector-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Levente Hunyadi
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

asyncpg_vector-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+asyncpg_vector

asyncpg_vector-0.1.0.dist-info/zip-safe

@@ -0,0 +1 @@
+