diff options
Diffstat (limited to '.venv/lib/python3.12/site-packages/pgvector/psycopg2')
5 files changed, 113 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/pgvector/psycopg2/__init__.py b/.venv/lib/python3.12/site-packages/pgvector/psycopg2/__init__.py new file mode 100644 index 00000000..7c952951 --- /dev/null +++ b/.venv/lib/python3.12/site-packages/pgvector/psycopg2/__init__.py @@ -0,0 +1,8 @@ +from .register import register_vector +from ..utils import HalfVector, SparseVector + +__all__ = [ + 'register_vector', + 'HalfVector', + 'SparseVector' +] diff --git a/.venv/lib/python3.12/site-packages/pgvector/psycopg2/halfvec.py b/.venv/lib/python3.12/site-packages/pgvector/psycopg2/halfvec.py new file mode 100644 index 00000000..b50e89b6 --- /dev/null +++ b/.venv/lib/python3.12/site-packages/pgvector/psycopg2/halfvec.py @@ -0,0 +1,25 @@ +from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type +from ..utils import HalfVector + + +class HalfvecAdapter: + def __init__(self, value): + self._value = value + + def getquoted(self): + return adapt(HalfVector._to_db(self._value)).getquoted() + + +def cast_halfvec(value, cur): + return HalfVector._from_db(value) + + +def register_halfvec_info(oid, array_oid, scope): + halfvec = new_type((oid,), 'HALFVEC', cast_halfvec) + register_type(halfvec, scope) + + if array_oid is not None: + halfvecarray = new_array_type((array_oid,), 'HALFVECARRAY', halfvec) + register_type(halfvecarray, scope) + + register_adapter(HalfVector, HalfvecAdapter) diff --git a/.venv/lib/python3.12/site-packages/pgvector/psycopg2/register.py b/.venv/lib/python3.12/site-packages/pgvector/psycopg2/register.py new file mode 100644 index 00000000..77528520 --- /dev/null +++ b/.venv/lib/python3.12/site-packages/pgvector/psycopg2/register.py @@ -0,0 +1,29 @@ +import psycopg2 +from psycopg2.extensions import cursor +from .halfvec import register_halfvec_info +from .sparsevec import register_sparsevec_info +from .vector import register_vector_info + + +# TODO make globally False by default in 0.4.0 +# note: register_adapter is always global +# TODO make arrays True by defalt in 0.4.0 +def register_vector(conn_or_curs=None, globally=True, arrays=False): + conn = conn_or_curs if hasattr(conn_or_curs, 'cursor') else conn_or_curs.connection + cur = conn.cursor(cursor_factory=cursor) + scope = None if globally else conn_or_curs + + # use to_regtype to get first matching type in search path + cur.execute("SELECT typname, oid FROM pg_type WHERE oid IN (to_regtype('vector'), to_regtype('_vector'), to_regtype('halfvec'), to_regtype('_halfvec'), to_regtype('sparsevec'), to_regtype('_sparsevec'))") + type_info = dict(cur.fetchall()) + + if 'vector' not in type_info: + raise psycopg2.ProgrammingError('vector type not found in the database') + + register_vector_info(type_info['vector'], type_info['_vector'] if arrays else None, scope) + + if 'halfvec' in type_info: + register_halfvec_info(type_info['halfvec'], type_info['_halfvec'] if arrays else None, scope) + + if 'sparsevec' in type_info: + register_sparsevec_info(type_info['sparsevec'], type_info['_sparsevec'] if arrays else None, scope) diff --git a/.venv/lib/python3.12/site-packages/pgvector/psycopg2/sparsevec.py b/.venv/lib/python3.12/site-packages/pgvector/psycopg2/sparsevec.py new file mode 100644 index 00000000..a542807b --- /dev/null +++ b/.venv/lib/python3.12/site-packages/pgvector/psycopg2/sparsevec.py @@ -0,0 +1,25 @@ +from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type +from ..utils import SparseVector + + +class SparsevecAdapter: + def __init__(self, value): + self._value = value + + def getquoted(self): + return adapt(SparseVector._to_db(self._value)).getquoted() + + +def cast_sparsevec(value, cur): + return SparseVector._from_db(value) + + +def register_sparsevec_info(oid, array_oid, scope): + sparsevec = new_type((oid,), 'SPARSEVEC', cast_sparsevec) + register_type(sparsevec, scope) + + if array_oid is not None: + sparsevecarray = new_array_type((array_oid,), 'SPARSEVECARRAY', sparsevec) + register_type(sparsevecarray, scope) + + register_adapter(SparseVector, SparsevecAdapter) diff --git a/.venv/lib/python3.12/site-packages/pgvector/psycopg2/vector.py b/.venv/lib/python3.12/site-packages/pgvector/psycopg2/vector.py new file mode 100644 index 00000000..9861f01f --- /dev/null +++ b/.venv/lib/python3.12/site-packages/pgvector/psycopg2/vector.py @@ -0,0 +1,26 @@ +import numpy as np +from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type +from ..utils import Vector + + +class VectorAdapter: + def __init__(self, value): + self._value = value + + def getquoted(self): + return adapt(Vector._to_db(self._value)).getquoted() + + +def cast_vector(value, cur): + return Vector._from_db(value) + + +def register_vector_info(oid, array_oid, scope): + vector = new_type((oid,), 'VECTOR', cast_vector) + register_type(vector, scope) + + if array_oid is not None: + vectorarray = new_array_type((array_oid,), 'VECTORARRAY', vector) + register_type(vectorarray, scope) + + register_adapter(np.ndarray, VectorAdapter) |