app/services/vector_store_contract.py

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Protocol, runtime_checkable


@dataclass(frozen=True)
class EmbeddingMetadata:
    embedding_model: str
    embedding_dim: int


@dataclass(frozen=True)
class VectorDocument:
    tenant_id: str
    entry_id: str
    text: str
    embedding: list[float]
    embedding_metadata: EmbeddingMetadata
    metadata: dict[str, str | int | float | bool | None] = field(default_factory=dict)
    vector_id: str | None = None


@dataclass(frozen=True)
class SearchResult:
    entry_id: str
    text: str
    similarity: float
    metadata: dict[str, str | int | float | bool | None] = field(default_factory=dict)


def build_vector_document(
    *,
    tenant_id: str,
    entry_id: str,
    text: str,
    embedding: list[float],
    embedding_model: str,
    metadata: dict[str, str | int | float | bool | None] | None = None,
) -> VectorDocument:
    return VectorDocument(
        tenant_id=tenant_id,
        entry_id=entry_id,
        text=text,
        embedding=embedding,
        embedding_metadata=EmbeddingMetadata(
            embedding_model=embedding_model,
            embedding_dim=len(embedding),
        ),
        metadata=metadata or {},
        vector_id=entry_id,
    )


@runtime_checkable
class VectorStoreContract(Protocol):
    def upsert(self, document: VectorDocument) -> None:
        raise NotImplementedError

    def query(
        self, tenant_id: str, query_embedding: list[float], n_results: int
    ) -> list[SearchResult]:
        raise NotImplementedError

    def delete_docs(self, tenant_id: str, entry_ids: list[str]) -> None:
        raise NotImplementedError

    def delete_tenant(self, tenant_id: str) -> None:
        raise NotImplementedError

    def health_check(self) -> None:
        raise NotImplementedError