adapters

Bases: ABC

Backend-agnostic read/write interface for one vector collection.

Source code in vecparity/adapters/base.py

class VectorDBAdapter(ABC):
    """Backend-agnostic read/write interface for one vector collection."""

    @abstractmethod
    def get(self, id: str) -> VectorRecord | None:
        """Fetch a single record by id, or None if it doesn't exist."""

    @abstractmethod
    def upsert(self, records: list[VectorRecord]) -> None:
        """Insert or overwrite records, batched by the caller."""

    @abstractmethod
    def delete(self, ids: list[str]) -> None:
        """Remove records by id."""

    @abstractmethod
    def list_changed_since(self, cursor: float | None) -> Iterator[VectorRecord]:
        """Yield records created/updated after `cursor` (a unix timestamp).
        `cursor=None` means from the beginning (full backfill)."""

    @abstractmethod
    def search(self, vector: list[float], top_k: int) -> list[ScoredMatch]:
        """Run a similarity search, used only for parity verification."""

    @abstractmethod
    def count(self) -> int:
        """Total records in the collection, for sanity/pre-flight checks."""

count() abstractmethod

Total records in the collection, for sanity/pre-flight checks.

Source code in vecparity/adapters/base.py

@abstractmethod
def count(self) -> int:
    """Total records in the collection, for sanity/pre-flight checks."""

delete(ids) abstractmethod

Remove records by id.

Source code in vecparity/adapters/base.py

@abstractmethod
def delete(self, ids: list[str]) -> None:
    """Remove records by id."""

get(id) abstractmethod

Fetch a single record by id, or None if it doesn't exist.

Source code in vecparity/adapters/base.py

@abstractmethod
def get(self, id: str) -> VectorRecord | None:
    """Fetch a single record by id, or None if it doesn't exist."""

list_changed_since(cursor) abstractmethod

Yield records created/updated after cursor (a unix timestamp). cursor=None means from the beginning (full backfill).

Source code in vecparity/adapters/base.py

@abstractmethod
def list_changed_since(self, cursor: float | None) -> Iterator[VectorRecord]:
    """Yield records created/updated after `cursor` (a unix timestamp).
    `cursor=None` means from the beginning (full backfill)."""

search(vector, top_k) abstractmethod

Run a similarity search, used only for parity verification.

Source code in vecparity/adapters/base.py

@abstractmethod
def search(self, vector: list[float], top_k: int) -> list[ScoredMatch]:
    """Run a similarity search, used only for parity verification."""

upsert(records) abstractmethod

Insert or overwrite records, batched by the caller.

Source code in vecparity/adapters/base.py

@abstractmethod
def upsert(self, records: list[VectorRecord]) -> None:
    """Insert or overwrite records, batched by the caller."""

Bases: VectorDBAdapter

Source code in vecparity/adapters/memory.py

class MemoryAdapter(VectorDBAdapter):
    def __init__(self) -> None:
        self._store: dict[str, VectorRecord] = {}

    def get(self, id: str) -> VectorRecord | None:
        return self._store.get(id)

    def upsert(self, records: list[VectorRecord]) -> None:
        now = time.time()
        for r in records:
            if r.updated_at is None:
                r = r.model_copy(update={"updated_at": now})
            self._store[r.id] = r

    def delete(self, ids: list[str]) -> None:
        for id in ids:
            self._store.pop(id, None)

    def list_changed_since(self, cursor: float | None) -> Iterator[VectorRecord]:
        for r in self._store.values():
            if cursor is None or (r.updated_at or 0) > cursor:
                yield r

    def search(self, vector: list[float], top_k: int) -> list[ScoredMatch]:
        if not self._store:
            return []
        q = np.asarray(vector, dtype=np.float32)
        q_norm = q / (np.linalg.norm(q) + 1e-12)
        scored: list[ScoredMatch] = []
        for r in self._store.values():
            v = r.as_array()
            v_norm = v / (np.linalg.norm(v) + 1e-12)
            score = float(np.dot(q_norm, v_norm))
            scored.append(ScoredMatch(id=r.id, score=score, metadata=r.metadata))
        scored.sort(key=lambda m: m.score, reverse=True)
        return scored[:top_k]

    def count(self) -> int:
        return len(self._store)

Bases: VectorDBAdapter

Source code in vecparity/adapters/pgvector.py

class PgVectorAdapter(VectorDBAdapter):
    def __init__(
        self,
        conn: psycopg.Connection,
        table: str,
        id_col: str = "id",
        vector_col: str = "embedding",
        metadata_col: str = "metadata",
        updated_at_col: str = "updated_at",
        distance_op: str = "<=>",  # cosine; use "<->" for L2, "<#>" for inner product
    ) -> None:
        self.conn = conn
        self.table = table
        self.id_col = id_col
        self.vector_col = vector_col
        self.metadata_col = metadata_col
        self.updated_at_col = updated_at_col
        self.distance_op = distance_op

        # Registers the vector type adapter; without it inserts/searches fail.
        from pgvector.psycopg import register_vector

        register_vector(conn)

    def get(self, id: str) -> VectorRecord | None:
        with self.conn.cursor() as cur:
            cur.execute(
                f"SELECT {self.id_col}, {self.vector_col}, {self.metadata_col}, "
                f"{self.updated_at_col} FROM {self.table} WHERE {self.id_col} = %s",
                (id,),
            )
            row = cur.fetchone()
        return self._to_record(row) if row else None

    def upsert(self, records: list[VectorRecord]) -> None:
        now = time.time()
        with self.conn.cursor() as cur:
            for r in records:
                cur.execute(
                    f"""
                    INSERT INTO {self.table}
                        ({self.id_col}, {self.vector_col}, {self.metadata_col}, {self.updated_at_col})
                    VALUES (%s, %s, %s, %s)
                    ON CONFLICT ({self.id_col}) DO UPDATE SET
                        {self.vector_col} = EXCLUDED.{self.vector_col},
                        {self.metadata_col} = EXCLUDED.{self.metadata_col},
                        {self.updated_at_col} = EXCLUDED.{self.updated_at_col}
                    """,
                    (r.id, r.vector, json.dumps(r.metadata), r.updated_at or now),
                )
        self.conn.commit()

    def delete(self, ids: list[str]) -> None:
        with self.conn.cursor() as cur:
            cur.execute(f"DELETE FROM {self.table} WHERE {self.id_col} = ANY(%s)", (ids,))
        self.conn.commit()

    def list_changed_since(self, cursor: float | None) -> Iterator[VectorRecord]:
        with self.conn.cursor(name="vecparity_scroll") as cur:  # server-side cursor
            if cursor is None:
                cur.execute(
                    f"SELECT {self.id_col}, {self.vector_col}, {self.metadata_col}, "
                    f"{self.updated_at_col} FROM {self.table} ORDER BY {self.id_col}"
                )
            else:
                cur.execute(
                    f"SELECT {self.id_col}, {self.vector_col}, {self.metadata_col}, "
                    f"{self.updated_at_col} FROM {self.table} "
                    f"WHERE {self.updated_at_col} > %s ORDER BY {self.updated_at_col}",
                    (cursor,),
                )
            for row in cur:
                yield self._to_record(row)

    def search(self, vector: list[float], top_k: int) -> list[ScoredMatch]:
        # Needs an explicit ::vector cast; Postgres can't infer the type
        # for a bare query parameter the way it can for an INSERT column.
        with self.conn.cursor() as cur:
            cur.execute(
                f"""
                SELECT {self.id_col}, {self.metadata_col},
                       1 - ({self.vector_col} {self.distance_op} %s::vector) AS score
                FROM {self.table}
                ORDER BY {self.vector_col} {self.distance_op} %s::vector
                LIMIT %s
                """,
                (vector, vector, top_k),
            )
            rows = cur.fetchall()
        return [ScoredMatch(id=row[0], score=float(row[2]), metadata=row[1] or {}) for row in rows]

    def count(self) -> int:
        with self.conn.cursor() as cur:
            cur.execute(f"SELECT COUNT(*) FROM {self.table}")
            row = cur.fetchone()
            return int(row[0]) if row else 0

    def _to_record(self, row: tuple[Any, ...]) -> VectorRecord:
        id, vector, metadata, updated_at = row
        # pgvector's Vector wrapper isn't directly iterable, and pydantic
        # needs plain floats rather than numpy scalars.
        return VectorRecord(
            id=id,
            vector=[float(x) for x in vector.to_list()],
            metadata=metadata or {},
            updated_at=updated_at,
        )

Bases: VectorDBAdapter

Source code in vecparity/adapters/qdrant.py

class QdrantAdapter(VectorDBAdapter):
    def __init__(
        self,
        client: QdrantClient,
        collection: str,
        updated_at_field: str = "updated_at",
        scroll_batch_size: int = 256,
    ) -> None:
        self.client = client
        self.collection = collection
        self.updated_at_field = updated_at_field
        self.scroll_batch_size = scroll_batch_size

    def get(self, id: str) -> VectorRecord | None:
        points = self.client.retrieve(
            collection_name=self.collection, ids=[_point_id(id)], with_vectors=True
        )
        if not points:
            return None
        return self._to_record(points[0])

    def upsert(self, records: list[VectorRecord]) -> None:
        points = [
            qm.PointStruct(
                id=_point_id(r.id),
                vector=r.vector,
                payload={
                    **r.metadata,
                    self.updated_at_field: r.updated_at,
                    _ORIGINAL_ID_KEY: r.id,
                },
            )
            for r in records
        ]
        self.client.upsert(collection_name=self.collection, points=points)

    def delete(self, ids: list[str]) -> None:
        self.client.delete(
            collection_name=self.collection,
            points_selector=qm.PointIdsList(points=[_point_id(id) for id in ids]),
        )

    def list_changed_since(self, cursor: float | None) -> Iterator[VectorRecord]:
        query_filter = None
        if cursor is not None:
            query_filter = qm.Filter(
                must=[qm.FieldCondition(key=self.updated_at_field, range=qm.Range(gt=cursor))]
            )
        offset = None
        while True:
            points, offset = self.client.scroll(
                collection_name=self.collection,
                scroll_filter=query_filter,
                limit=self.scroll_batch_size,
                offset=offset,
                with_vectors=True,
            )
            for p in points:
                yield self._to_record(p)
            if offset is None:
                break

    def search(self, vector: list[float], top_k: int) -> list[ScoredMatch]:
        # qdrant-client >=1.10 dropped .search() in favor of .query_points().
        response = self.client.query_points(
            collection_name=self.collection,
            query=vector,
            limit=top_k,
            with_payload=True,
        )
        hits = response.points
        return [
            ScoredMatch(
                id=(h.payload or {}).get(_ORIGINAL_ID_KEY, str(h.id)),
                score=h.score,
                metadata={k: v for k, v in (h.payload or {}).items() if k != _ORIGINAL_ID_KEY},
            )
            for h in hits
        ]

    def count(self) -> int:
        return int(self.client.count(collection_name=self.collection).count)

    def _to_record(self, point: qm.Record) -> VectorRecord:
        payload = dict(point.payload or {})
        original_id = payload.pop(_ORIGINAL_ID_KEY, str(point.id))
        updated_at = payload.pop(self.updated_at_field, None)
        return VectorRecord(
            id=original_id,
            vector=list(point.vector or []),
            metadata=payload,
            updated_at=updated_at,
        )

Bases: VectorDBAdapter

Source code in vecparity/adapters/pinecone.py

class PineconeAdapter(VectorDBAdapter):
    def __init__(
        self,
        client: Pinecone,
        index_name: str,
        namespace: str = "",
        updated_at_field: str = "updated_at",
        fetch_batch_size: int = 100,
    ) -> None:
        self.index = client.Index(index_name)
        self.namespace = namespace
        self.updated_at_field = updated_at_field
        self.fetch_batch_size = fetch_batch_size

    def get(self, id: str) -> VectorRecord | None:
        result = self.index.fetch(ids=[id], namespace=self.namespace)
        vectors = result.vectors if hasattr(result, "vectors") else result.get("vectors", {})
        if id not in vectors:
            return None
        return self._to_record(id, vectors[id])

    def upsert(self, records: list[VectorRecord]) -> None:
        vectors = []
        for r in records:
            metadata: dict[str, Any] = {**r.metadata, self.updated_at_field: r.updated_at}
            vectors.append({"id": r.id, "values": r.vector, "metadata": metadata})
        self.index.upsert(vectors=vectors, namespace=self.namespace)

    def delete(self, ids: list[str]) -> None:
        self.index.delete(ids=ids, namespace=self.namespace)

    def list_changed_since(self, cursor: float | None) -> Iterator[VectorRecord]:
        # index.list() yields ListResponse pages, each holding a batch of
        # ListItem(id=...) entries, not bare id strings.
        id_batch: list[str] = []
        for page in self.index.list(namespace=self.namespace):
            for item in page.vectors:
                if item.id is None:
                    continue
                id_batch.append(item.id)
                if len(id_batch) >= self.fetch_batch_size:
                    yield from self._fetch_and_filter(id_batch, cursor)
                    id_batch = []
        if id_batch:
            yield from self._fetch_and_filter(id_batch, cursor)

    def search(self, vector: list[float], top_k: int) -> list[ScoredMatch]:
        result = self.index.query(
            vector=vector, top_k=top_k, namespace=self.namespace, include_metadata=True
        )
        matches = result.matches if hasattr(result, "matches") else result.get("matches", [])
        return [
            ScoredMatch(id=m.id, score=m.score, metadata=dict(m.metadata or {})) for m in matches
        ]

    def count(self) -> int:
        stats = self.index.describe_index_stats()
        namespaces = (
            stats.namespaces if hasattr(stats, "namespaces") else stats.get("namespaces", {})
        )
        ns = namespaces.get(self.namespace)
        if ns is None:
            return 0
        count = ns.vector_count if hasattr(ns, "vector_count") else ns.get("vector_count", 0)
        return int(count)

    def _fetch_and_filter(self, ids: list[str], cursor: float | None) -> Iterator[VectorRecord]:
        result = self.index.fetch(ids=ids, namespace=self.namespace)
        vectors = result.vectors if hasattr(result, "vectors") else result.get("vectors", {})
        for id, v in vectors.items():
            record = self._to_record(id, v)
            if cursor is None or (record.updated_at or 0) > cursor:
                yield record

    def _to_record(self, id: str, v: Any) -> VectorRecord:
        values = v.values if hasattr(v, "values") else v["values"]
        metadata = dict(v.metadata if hasattr(v, "metadata") else v.get("metadata", {}) or {})
        updated_at = metadata.pop(self.updated_at_field, None)
        return VectorRecord(id=id, vector=list(values), metadata=metadata, updated_at=updated_at)

Bases: VectorDBAdapter

Source code in vecparity/adapters/milvus.py

class MilvusAdapter(VectorDBAdapter):
    def __init__(
        self,
        client: MilvusClient,
        collection_name: str,
        id_field: str = "id",
        vector_field: str = "vector",
        metadata_field: str = "metadata",
        updated_at_field: str = "updated_at",
        page_size: int = 256,
    ) -> None:
        self.client = client
        self.collection_name = collection_name
        self.id_field = id_field
        self.vector_field = vector_field
        self.metadata_field = metadata_field
        self.updated_at_field = updated_at_field
        self.page_size = page_size
        self._output_fields = [id_field, vector_field, metadata_field, updated_at_field]

    def get(self, id: str) -> VectorRecord | None:
        rows = self.client.get(
            self.collection_name,
            ids=[id],
            output_fields=self._output_fields,
            consistency_level="Strong",
        )
        if not rows:
            return None
        return self._to_record(rows[0])

    def upsert(self, records: list[VectorRecord]) -> None:
        data = [
            {
                self.id_field: r.id,
                self.vector_field: r.vector,
                self.metadata_field: r.metadata,
                self.updated_at_field: r.updated_at,
            }
            for r in records
        ]
        self.client.upsert(self.collection_name, data=data)

    def delete(self, ids: list[str]) -> None:
        self.client.delete(self.collection_name, ids=ids)

    def list_changed_since(self, cursor: float | None) -> Iterator[VectorRecord]:
        filter_expr = f"{self.updated_at_field} > {cursor}" if cursor is not None else ""
        offset = 0
        while True:
            rows = self.client.query(
                self.collection_name,
                filter=filter_expr,
                output_fields=self._output_fields,
                limit=self.page_size,
                offset=offset,
                consistency_level="Strong",
            )
            if not rows:
                break
            for row in rows:
                yield self._to_record(row)
            if len(rows) < self.page_size:
                break
            offset += len(rows)

    def search(self, vector: list[float], top_k: int) -> list[ScoredMatch]:
        results = self.client.search(
            self.collection_name,
            data=[vector],
            limit=top_k,
            output_fields=[self.id_field, self.metadata_field],
            consistency_level="Strong",
        )
        matches = []
        for hit in results[0]:
            entity = hit.get("entity", {})
            matches.append(
                ScoredMatch(
                    id=str(entity.get(self.id_field, hit.get("id"))),
                    score=float(hit["distance"]),
                    metadata=dict(entity.get(self.metadata_field, {}) or {}),
                )
            )
        return matches

    def count(self) -> int:
        result = self.client.query(
            self.collection_name,
            filter="",
            output_fields=["count(*)"],
            consistency_level="Strong",
        )
        if result and "count(*)" in result[0]:
            return int(result[0]["count(*)"])
        # Fallback: some pymilvus versions need the expression form instead.
        stats = self.client.get_collection_stats(self.collection_name)
        return int(stats.get("row_count", 0))

    def _to_record(self, row: dict[str, Any]) -> VectorRecord:
        metadata = dict(row.get(self.metadata_field) or {})
        return VectorRecord(
            id=str(row[self.id_field]),
            vector=[float(x) for x in row[self.vector_field]],
            metadata=metadata,
            updated_at=row.get(self.updated_at_field),
        )

Bases: VectorDBAdapter

Source code in vecparity/adapters/weaviate.py

class WeaviateAdapter(VectorDBAdapter):
    def __init__(
        self,
        collection: Collection,
        updated_at_field: str = "updated_at",
        page_size: int = 256,
    ) -> None:
        self.collection = collection
        self.updated_at_field = updated_at_field
        self.page_size = page_size

    def get(self, id: str) -> VectorRecord | None:
        obj = self.collection.query.fetch_object_by_id(_object_id(id), include_vector=True)
        if obj is None:
            return None
        return self._to_record(obj)

    def upsert(self, records: list[VectorRecord]) -> None:
        for r in records:
            oid = _object_id(r.id)
            properties = {
                **r.metadata,
                self.updated_at_field: r.updated_at,
                _ORIGINAL_ID_KEY: r.id,
            }
            if self.collection.data.exists(oid):
                self.collection.data.replace(uuid=oid, properties=properties, vector=r.vector)
            else:
                self.collection.data.insert(uuid=oid, properties=properties, vector=r.vector)

    def delete(self, ids: list[str]) -> None:
        for id in ids:
            self.collection.data.delete_by_id(_object_id(id))

    def list_changed_since(self, cursor: float | None) -> Iterator[VectorRecord]:
        filters = None
        if cursor is not None:
            filters = wvc.query.Filter.by_property(self.updated_at_field).greater_than(cursor)
        after = None
        while True:
            result = self.collection.query.fetch_objects(
                filters=filters,
                limit=self.page_size,
                after=after,
                include_vector=True,
            )
            objects = result.objects
            if not objects:
                break
            for obj in objects:
                yield self._to_record(obj)
            if len(objects) < self.page_size:
                break
            after = objects[-1].uuid

    def search(self, vector: list[float], top_k: int) -> list[ScoredMatch]:
        result = self.collection.query.near_vector(
            near_vector=vector,
            limit=top_k,
            return_metadata=wvc.query.MetadataQuery(distance=True),
        )
        matches = []
        for obj in result.objects:
            props = dict(obj.properties)
            original_id = props.pop(_ORIGINAL_ID_KEY, str(obj.uuid))
            props.pop(self.updated_at_field, None)
            distance = (
                obj.metadata.distance if obj.metadata and obj.metadata.distance is not None else 0.0
            )
            matches.append(ScoredMatch(id=original_id, score=1.0 - distance, metadata=props))
        return matches

    def count(self) -> int:
        result = self.collection.aggregate.over_all(total_count=True)
        return int(result.total_count or 0)

    def _to_record(self, obj: object) -> VectorRecord:
        properties = dict(obj.properties)  # type: ignore[attr-defined]
        original_id = properties.pop(_ORIGINAL_ID_KEY, str(obj.uuid))  # type: ignore[attr-defined]
        updated_at = properties.pop(self.updated_at_field, None)
        vector = obj.vector.get("default", []) if obj.vector else []  # type: ignore[attr-defined]
        return VectorRecord(
            id=original_id,
            vector=[float(x) for x in vector],
            metadata=properties,
            updated_at=updated_at,
        )

Bases: VectorDBAdapter

Source code in vecparity/adapters/chroma.py

class ChromaAdapter(VectorDBAdapter):
    def __init__(
        self,
        collection: Collection,
        updated_at_field: str = "updated_at",
        page_size: int = 256,
    ) -> None:
        self.collection = collection
        self.updated_at_field = updated_at_field
        self.page_size = page_size

    def get(self, id: str) -> VectorRecord | None:
        result = self.collection.get(ids=[id], include=["embeddings", "metadatas"])
        ids = result["ids"]
        if not ids:
            return None
        embeddings = result["embeddings"]
        metadatas = result["metadatas"]
        assert embeddings is not None and metadatas is not None  # guaranteed by include=
        return self._to_record(ids[0], embeddings[0], dict(metadatas[0] or {}))

    def upsert(self, records: list[VectorRecord]) -> None:
        self.collection.upsert(
            ids=[r.id for r in records],
            embeddings=[r.vector for r in records],  # type: ignore[arg-type]
            metadatas=[{**r.metadata, self.updated_at_field: r.updated_at} for r in records],
        )

    def delete(self, ids: list[str]) -> None:
        self.collection.delete(ids=ids)

    def list_changed_since(self, cursor: float | None) -> Iterator[VectorRecord]:
        where = {self.updated_at_field: {"$gt": cursor}} if cursor is not None else None
        offset = 0
        while True:
            result = self.collection.get(
                where=where,  # type: ignore[arg-type]
                limit=self.page_size,
                offset=offset,
                include=["embeddings", "metadatas"],
            )
            ids = result["ids"]
            if not ids:
                break
            embeddings = result["embeddings"]
            metadatas = result["metadatas"]
            assert embeddings is not None and metadatas is not None
            for id, vector, metadata in zip(ids, embeddings, metadatas, strict=True):
                yield self._to_record(id, vector, dict(metadata or {}))
            if len(ids) < self.page_size:
                break
            offset += len(ids)

    def search(self, vector: list[float], top_k: int) -> list[ScoredMatch]:
        result = self.collection.query(
            query_embeddings=[vector],  # type: ignore[arg-type]
            n_results=top_k,
            include=["distances", "metadatas"],
        )
        raw_ids = result["ids"][0]
        distances = result["distances"]
        metadatas = result["metadatas"]
        assert distances is not None and metadatas is not None
        matches = []
        for id, distance, metadata in zip(raw_ids, distances[0], metadatas[0], strict=True):
            meta = dict(metadata or {})
            meta.pop(self.updated_at_field, None)
            matches.append(ScoredMatch(id=id, score=1.0 - distance, metadata=meta))
        return matches

    def count(self) -> int:
        return self.collection.count()

    def _to_record(self, id: str, vector: Any, metadata: dict[str, Any]) -> VectorRecord:
        metadata = dict(metadata)
        updated_at = metadata.pop(self.updated_at_field, None)
        return VectorRecord(
            id=id,
            vector=[float(x) for x in vector],
            metadata=metadata,
            updated_at=updated_at,
        )