pgvector 0.3.5__py3-none-any.whl

pgvector/asyncpg/__init__.py

@@ -0,0 +1,9 @@
+from .register import register_vector
+from ..utils import Vector, HalfVector, SparseVector
+
+__all__ = [
+    'register_vector',
+    'Vector',
+    'HalfVector',
+    'SparseVector'
+]

pgvector/asyncpg/register.py

@@ -0,0 +1,31 @@
+from ..utils import Vector, HalfVector, SparseVector
+
+
+async def register_vector(conn, schema='public'):
+    await conn.set_type_codec(
+        'vector',
+        schema=schema,
+        encoder=Vector._to_db_binary,
+        decoder=Vector._from_db_binary,
+        format='binary'
+    )
+
+    try:
+        await conn.set_type_codec(
+            'halfvec',
+            schema=schema,
+            encoder=HalfVector._to_db_binary,
+            decoder=HalfVector._from_db_binary,
+            format='binary'
+        )
+
+        await conn.set_type_codec(
+            'sparsevec',
+            schema=schema,
+            encoder=SparseVector._to_db_binary,
+            decoder=SparseVector._from_db_binary,
+            format='binary'
+        )
+    except ValueError as e:
+        if not str(e).startswith('unknown type:'):
+            raise e

pgvector/django/__init__.py

@@ -0,0 +1,26 @@
+from .bit import BitField
+from .extensions import VectorExtension
+from .functions import L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance
+from .halfvec import HalfVectorField
+from .indexes import IvfflatIndex, HnswIndex
+from .sparsevec import SparseVectorField
+from .vector import VectorField
+from ..utils import HalfVector, SparseVector
+
+__all__ = [
+    'VectorExtension',
+    'VectorField',
+    'HalfVectorField',
+    'BitField',
+    'SparseVectorField',
+    'IvfflatIndex',
+    'HnswIndex',
+    'L2Distance',
+    'MaxInnerProduct',
+    'CosineDistance',
+    'L1Distance',
+    'HammingDistance',
+    'JaccardDistance',
+    'HalfVector',
+    'SparseVector'
+]

pgvector/django/bit.py

@@ -0,0 +1,32 @@
+from django import forms
+from django.db.models import Field
+
+
+# https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/
+class BitField(Field):
+    description = 'Bit string'
+
+    def __init__(self, *args, length=None, **kwargs):
+        self.length = length
+        super().__init__(*args, **kwargs)
+
+    def deconstruct(self):
+        name, path, args, kwargs = super().deconstruct()
+        if self.length is not None:
+            kwargs['length'] = self.length
+        return name, path, args, kwargs
+
+    def db_type(self, connection):
+        if self.length is None:
+            return 'bit'
+        return 'bit(%d)' % self.length
+
+    def formfield(self, **kwargs):
+        return super().formfield(form_class=BitFormField, **kwargs)
+
+
+class BitFormField(forms.CharField):
+    def to_python(self, value):
+        if isinstance(value, str) and value == '':
+            return None
+        return super().to_python(value)

pgvector/django/extensions.py

@@ -0,0 +1,6 @@
+from django.contrib.postgres.operations import CreateExtension
+
+
+class VectorExtension(CreateExtension):
+    def __init__(self):
+        self.name = 'vector'

pgvector/django/functions.py

@@ -0,0 +1,55 @@
+from django.db.models import FloatField, Func, Value
+from ..utils import Vector, HalfVector, SparseVector
+
+
+class DistanceBase(Func):
+    output_field = FloatField()
+
+    def __init__(self, expression, vector, **extra):
+        if not hasattr(vector, 'resolve_expression'):
+            if isinstance(vector, HalfVector):
+                vector = Value(HalfVector._to_db(vector))
+            elif isinstance(vector, SparseVector):
+                vector = Value(SparseVector._to_db(vector))
+            else:
+                vector = Value(Vector._to_db(vector))
+        super().__init__(expression, vector, **extra)
+
+
+class BitDistanceBase(Func):
+    output_field = FloatField()
+
+    def __init__(self, expression, vector, **extra):
+        if not hasattr(vector, 'resolve_expression'):
+            vector = Value(vector)
+        super().__init__(expression, vector, **extra)
+
+
+class L2Distance(DistanceBase):
+    function = ''
+    arg_joiner = ' <-> '
+
+
+class MaxInnerProduct(DistanceBase):
+    function = ''
+    arg_joiner = ' <#> '
+
+
+class CosineDistance(DistanceBase):
+    function = ''
+    arg_joiner = ' <=> '
+
+
+class L1Distance(DistanceBase):
+    function = ''
+    arg_joiner = ' <+> '
+
+
+class HammingDistance(BitDistanceBase):
+    function = ''
+    arg_joiner = ' <~> '
+
+
+class JaccardDistance(BitDistanceBase):
+    function = ''
+    arg_joiner = ' <%%> '

pgvector/django/halfvec.py

@@ -0,0 +1,60 @@
+from django import forms
+from django.db.models import Field
+from ..utils import HalfVector
+
+
+# https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/
+class HalfVectorField(Field):
+    description = 'Half vector'
+    empty_strings_allowed = False
+
+    def __init__(self, *args, dimensions=None, **kwargs):
+        self.dimensions = dimensions
+        super().__init__(*args, **kwargs)
+
+    def deconstruct(self):
+        name, path, args, kwargs = super().deconstruct()
+        if self.dimensions is not None:
+            kwargs['dimensions'] = self.dimensions
+        return name, path, args, kwargs
+
+    def db_type(self, connection):
+        if self.dimensions is None:
+            return 'halfvec'
+        return 'halfvec(%d)' % self.dimensions
+
+    def from_db_value(self, value, expression, connection):
+        return HalfVector._from_db(value)
+
+    def to_python(self, value):
+        if value is None or isinstance(value, HalfVector):
+            return value
+        elif isinstance(value, str):
+            return HalfVector._from_db(value)
+        else:
+            return HalfVector(value)
+
+    def get_prep_value(self, value):
+        return HalfVector._to_db(value)
+
+    def value_to_string(self, obj):
+        return self.get_prep_value(self.value_from_object(obj))
+
+    def formfield(self, **kwargs):
+        return super().formfield(form_class=HalfVectorFormField, **kwargs)
+
+
+class HalfVectorWidget(forms.TextInput):
+    def format_value(self, value):
+        if isinstance(value, HalfVector):
+            value = value.to_list()
+        return super().format_value(value)
+
+
+class HalfVectorFormField(forms.CharField):
+    widget = HalfVectorWidget
+
+    def to_python(self, value):
+        if isinstance(value, str) and value == '':
+            return None
+        return super().to_python(value)

pgvector/django/indexes.py

@@ -0,0 +1,46 @@
+from django.contrib.postgres.indexes import PostgresIndex
+
+
+class IvfflatIndex(PostgresIndex):
+    suffix = 'ivfflat'
+
+    def __init__(self, *expressions, lists=None, **kwargs):
+        self.lists = lists
+        super().__init__(*expressions, **kwargs)
+
+    def deconstruct(self):
+        path, args, kwargs = super().deconstruct()
+        if self.lists is not None:
+            kwargs['lists'] = self.lists
+        return path, args, kwargs
+
+    def get_with_params(self):
+        with_params = []
+        if self.lists is not None:
+            with_params.append('lists = %d' % self.lists)
+        return with_params
+
+
+class HnswIndex(PostgresIndex):
+    suffix = 'hnsw'
+
+    def __init__(self, *expressions, m=None, ef_construction=None, **kwargs):
+        self.m = m
+        self.ef_construction = ef_construction
+        super().__init__(*expressions, **kwargs)
+
+    def deconstruct(self):
+        path, args, kwargs = super().deconstruct()
+        if self.m is not None:
+            kwargs['m'] = self.m
+        if self.ef_construction is not None:
+            kwargs['ef_construction'] = self.ef_construction
+        return path, args, kwargs
+
+    def get_with_params(self):
+        with_params = []
+        if self.m is not None:
+            with_params.append('m = %d' % self.m)
+        if self.ef_construction is not None:
+            with_params.append('ef_construction = %d' % self.ef_construction)
+        return with_params

pgvector/django/sparsevec.py

@@ -0,0 +1,55 @@
+from django import forms
+from django.db.models import Field
+from ..utils import SparseVector
+
+
+# https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/
+class SparseVectorField(Field):
+    description = 'Sparse vector'
+    empty_strings_allowed = False
+
+    def __init__(self, *args, dimensions=None, **kwargs):
+        self.dimensions = dimensions
+        super().__init__(*args, **kwargs)
+
+    def deconstruct(self):
+        name, path, args, kwargs = super().deconstruct()
+        if self.dimensions is not None:
+            kwargs['dimensions'] = self.dimensions
+        return name, path, args, kwargs
+
+    def db_type(self, connection):
+        if self.dimensions is None:
+            return 'sparsevec'
+        return 'sparsevec(%d)' % self.dimensions
+
+    def from_db_value(self, value, expression, connection):
+        return SparseVector._from_db(value)
+
+    def to_python(self, value):
+        return SparseVector._from_db(value)
+
+    def get_prep_value(self, value):
+        return SparseVector._to_db(value)
+
+    def value_to_string(self, obj):
+        return self.get_prep_value(self.value_from_object(obj))
+
+    def formfield(self, **kwargs):
+        return super().formfield(form_class=SparseVectorFormField, **kwargs)
+
+
+class SparseVectorWidget(forms.TextInput):
+    def format_value(self, value):
+        if isinstance(value, SparseVector):
+            value = value.to_text()
+        return super().format_value(value)
+
+
+class SparseVectorFormField(forms.CharField):
+    widget = SparseVectorWidget
+
+    def to_python(self, value):
+        if isinstance(value, str) and value == '':
+            return None
+        return super().to_python(value)

pgvector/django/vector.py

@@ -0,0 +1,73 @@
+from django import forms
+from django.db.models import Field
+import numpy as np
+from ..utils import Vector
+
+
+# https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/
+class VectorField(Field):
+    description = 'Vector'
+    empty_strings_allowed = False
+
+    def __init__(self, *args, dimensions=None, **kwargs):
+        self.dimensions = dimensions
+        super().__init__(*args, **kwargs)
+
+    def deconstruct(self):
+        name, path, args, kwargs = super().deconstruct()
+        if self.dimensions is not None:
+            kwargs['dimensions'] = self.dimensions
+        return name, path, args, kwargs
+
+    def db_type(self, connection):
+        if self.dimensions is None:
+            return 'vector'
+        return 'vector(%d)' % self.dimensions
+
+    def from_db_value(self, value, expression, connection):
+        return Vector._from_db(value)
+
+    def to_python(self, value):
+        if isinstance(value, list):
+            return np.array(value, dtype=np.float32)
+        return Vector._from_db(value)
+
+    def get_prep_value(self, value):
+        return Vector._to_db(value)
+
+    def value_to_string(self, obj):
+        return self.get_prep_value(self.value_from_object(obj))
+
+    def validate(self, value, model_instance):
+        if isinstance(value, np.ndarray):
+            value = value.tolist()
+        super().validate(value, model_instance)
+
+    def run_validators(self, value):
+        if isinstance(value, np.ndarray):
+            value = value.tolist()
+        super().run_validators(value)
+
+    def formfield(self, **kwargs):
+        return super().formfield(form_class=VectorFormField, **kwargs)
+
+
+class VectorWidget(forms.TextInput):
+    def format_value(self, value):
+        if isinstance(value, np.ndarray):
+            value = value.tolist()
+        return super().format_value(value)
+
+
+class VectorFormField(forms.CharField):
+    widget = VectorWidget
+
+    def has_changed(self, initial, data):
+        if isinstance(initial, np.ndarray):
+            initial = initial.tolist()
+        return super().has_changed(initial, data)
+
+    def to_python(self, value):
+        if isinstance(value, str) and value == '':
+            return None
+        return super().to_python(value)

pgvector/peewee/__init__.py

@@ -0,0 +1,14 @@
+from .bit import FixedBitField
+from .halfvec import HalfVectorField
+from .sparsevec import SparseVectorField
+from .vector import VectorField
+from ..utils import HalfVector, SparseVector
+
+__all__ = [
+    'VectorField',
+    'HalfVectorField',
+    'FixedBitField',
+    'SparseVectorField',
+    'HalfVector',
+    'SparseVector'
+]

pgvector/peewee/bit.py

@@ -0,0 +1,21 @@
+from peewee import Expression, Field
+
+
+class FixedBitField(Field):
+    field_type = 'bit'
+
+    def __init__(self, max_length=None, *args, **kwargs):
+        self.max_length = max_length
+        super(FixedBitField, self).__init__(*args, **kwargs)
+
+    def get_modifiers(self):
+        return self.max_length and [self.max_length] or None
+
+    def _distance(self, op, vector):
+        return Expression(lhs=self, op=op, rhs=self.to_value(vector))
+
+    def hamming_distance(self, vector):
+        return self._distance('<~>', vector)
+
+    def jaccard_distance(self, vector):
+        return self._distance('<%%>', vector)

pgvector/peewee/halfvec.py

@@ -0,0 +1,34 @@
+from peewee import Expression, Field
+from ..utils import HalfVector
+
+
+class HalfVectorField(Field):
+    field_type = 'halfvec'
+
+    def __init__(self, dimensions=None, *args, **kwargs):
+        self.dimensions = dimensions
+        super(HalfVectorField, self).__init__(*args, **kwargs)
+
+    def get_modifiers(self):
+        return self.dimensions and [self.dimensions] or None
+
+    def db_value(self, value):
+        return HalfVector._to_db(value)
+
+    def python_value(self, value):
+        return HalfVector._from_db(value)
+
+    def _distance(self, op, vector):
+        return Expression(lhs=self, op=op, rhs=self.to_value(vector))
+
+    def l2_distance(self, vector):
+        return self._distance('<->', vector)
+
+    def max_inner_product(self, vector):
+        return self._distance('<#>', vector)
+
+    def cosine_distance(self, vector):
+        return self._distance('<=>', vector)
+
+    def l1_distance(self, vector):
+        return self._distance('<+>', vector)

pgvector/peewee/sparsevec.py

@@ -0,0 +1,34 @@
+from peewee import Expression, Field
+from ..utils import SparseVector
+
+
+class SparseVectorField(Field):
+    field_type = 'sparsevec'
+
+    def __init__(self, dimensions=None, *args, **kwargs):
+        self.dimensions = dimensions
+        super(SparseVectorField, self).__init__(*args, **kwargs)
+
+    def get_modifiers(self):
+        return self.dimensions and [self.dimensions] or None
+
+    def db_value(self, value):
+        return SparseVector._to_db(value)
+
+    def python_value(self, value):
+        return SparseVector._from_db(value)
+
+    def _distance(self, op, vector):
+        return Expression(lhs=self, op=op, rhs=self.to_value(vector))
+
+    def l2_distance(self, vector):
+        return self._distance('<->', vector)
+
+    def max_inner_product(self, vector):
+        return self._distance('<#>', vector)
+
+    def cosine_distance(self, vector):
+        return self._distance('<=>', vector)
+
+    def l1_distance(self, vector):
+        return self._distance('<+>', vector)

pgvector/peewee/vector.py

@@ -0,0 +1,34 @@
+from peewee import Expression, Field
+from ..utils import Vector
+
+
+class VectorField(Field):
+    field_type = 'vector'
+
+    def __init__(self, dimensions=None, *args, **kwargs):
+        self.dimensions = dimensions
+        super(VectorField, self).__init__(*args, **kwargs)
+
+    def get_modifiers(self):
+        return self.dimensions and [self.dimensions] or None
+
+    def db_value(self, value):
+        return Vector._to_db(value)
+
+    def python_value(self, value):
+        return Vector._from_db(value)
+
+    def _distance(self, op, vector):
+        return Expression(lhs=self, op=op, rhs=self.to_value(vector))
+
+    def l2_distance(self, vector):
+        return self._distance('<->', vector)
+
+    def max_inner_product(self, vector):
+        return self._distance('<#>', vector)
+
+    def cosine_distance(self, vector):
+        return self._distance('<=>', vector)
+
+    def l1_distance(self, vector):
+        return self._distance('<+>', vector)

pgvector/psycopg/__init__.py

@@ -0,0 +1,11 @@
+from .register import register_vector, register_vector_async
+from ..utils import Bit, HalfVector, SparseVector, Vector
+
+__all__ = [
+    'register_vector',
+    'register_vector_async',
+    'Vector',
+    'HalfVector',
+    'Bit',
+    'SparseVector'
+]

pgvector/psycopg/bit.py

@@ -0,0 +1,31 @@
+from psycopg.adapt import Dumper
+from psycopg.pq import Format
+from ..utils import Bit
+
+
+class BitDumper(Dumper):
+
+    format = Format.TEXT
+
+    def dump(self, obj):
+        return Bit._to_db(obj).encode('utf8')
+
+
+class BitBinaryDumper(BitDumper):
+
+    format = Format.BINARY
+
+    def dump(self, obj):
+        return Bit._to_db_binary(obj)
+
+
+def register_bit_info(context, info):
+    info.register(context)
+
+    # add oid to anonymous class for set_types
+    text_dumper = type('', (BitDumper,), {'oid': info.oid})
+    binary_dumper = type('', (BitBinaryDumper,), {'oid': info.oid})
+
+    adapters = context.adapters
+    adapters.register_dumper(Bit, text_dumper)
+    adapters.register_dumper(Bit, binary_dumper)

pgvector/psycopg/halfvec.py

@@ -0,0 +1,53 @@
+from psycopg.adapt import Loader, Dumper
+from psycopg.pq import Format
+from ..utils import HalfVector
+
+
+class HalfVectorDumper(Dumper):
+
+    format = Format.TEXT
+
+    def dump(self, obj):
+        return HalfVector._to_db(obj).encode('utf8')
+
+
+class HalfVectorBinaryDumper(HalfVectorDumper):
+
+    format = Format.BINARY
+
+    def dump(self, obj):
+        return HalfVector._to_db_binary(obj)
+
+
+class HalfVectorLoader(Loader):
+
+    format = Format.TEXT
+
+    def load(self, data):
+        if isinstance(data, memoryview):
+            data = bytes(data)
+        return HalfVector._from_db(data.decode('utf8'))
+
+
+class HalfVectorBinaryLoader(HalfVectorLoader):
+
+    format = Format.BINARY
+
+    def load(self, data):
+        if isinstance(data, memoryview):
+            data = bytes(data)
+        return HalfVector._from_db_binary(data)
+
+
+def register_halfvec_info(context, info):
+    info.register(context)
+
+    # add oid to anonymous class for set_types
+    text_dumper = type('', (HalfVectorDumper,), {'oid': info.oid})
+    binary_dumper = type('', (HalfVectorBinaryDumper,), {'oid': info.oid})
+
+    adapters = context.adapters
+    adapters.register_dumper(HalfVector, text_dumper)
+    adapters.register_dumper(HalfVector, binary_dumper)
+    adapters.register_loader(info.oid, HalfVectorLoader)
+    adapters.register_loader(info.oid, HalfVectorBinaryLoader)

pgvector/psycopg/register.py

@@ -0,0 +1,37 @@
+from psycopg.types import TypeInfo
+from .bit import register_bit_info
+from .halfvec import register_halfvec_info
+from .sparsevec import register_sparsevec_info
+from .vector import register_vector_info
+
+
+def register_vector(context):
+    info = TypeInfo.fetch(context, 'vector')
+    register_vector_info(context, info)
+
+    info = TypeInfo.fetch(context, 'bit')
+    register_bit_info(context, info)
+
+    info = TypeInfo.fetch(context, 'halfvec')
+    if info is not None:
+        register_halfvec_info(context, info)
+
+    info = TypeInfo.fetch(context, 'sparsevec')
+    if info is not None:
+        register_sparsevec_info(context, info)
+
+
+async def register_vector_async(context):
+    info = await TypeInfo.fetch(context, 'vector')
+    register_vector_info(context, info)
+
+    info = await TypeInfo.fetch(context, 'bit')
+    register_bit_info(context, info)
+
+    info = await TypeInfo.fetch(context, 'halfvec')
+    if info is not None:
+        register_halfvec_info(context, info)
+
+    info = await TypeInfo.fetch(context, 'sparsevec')
+    if info is not None:
+        register_sparsevec_info(context, info)