DataOps-platform/app/core/knowledge/sync.py

from __future__ import annotations

from collections.abc import Mapping, Sequence
from typing import Protocol

from app.core.knowledge.chunking import build_chunks, embedding_input_hash
from app.core.knowledge.contracts import (
    KnowledgeSnapshot,
    PointChange,
    PreparedPublication,
    SnapshotDiff,
    SyncResult,
)
from app.core.knowledge.diff import diff_snapshots


class KnowledgeRepository(Protocol):
    def load_active_snapshot(
        self, source_type: str, source_uid: str
    ) -> KnowledgeSnapshot | None: ...

    def load_embedding(self, embedding_hash: str) -> Sequence[float] | None: ...

    def activate(self, publication: PreparedPublication) -> None: ...


class EmbeddingProvider(Protocol):
    profile_key: str

    def embed(self, texts: list[str]) -> list[list[float]]: ...


def _initial_diff(snapshot: KnowledgeSnapshot) -> SnapshotDiff:
    return SnapshotDiff(
        added=tuple(
            PointChange(
                point_key=point.point_key,
                change_kind="added",
                new=point,
            )
            for point in snapshot.points
        )
    )


class KnowledgeSyncService:
    def __init__(
        self, *, repository: KnowledgeRepository, embedder: EmbeddingProvider
    ) -> None:
        self._repository = repository
        self._embedder = embedder

    def sync(self, snapshot: KnowledgeSnapshot) -> SyncResult:
        old = self._repository.load_active_snapshot(
            snapshot.source_type, snapshot.source_uid
        )
        if old and snapshot.source_revision < old.source_revision:
            raise ValueError("source revision cannot move backwards")
        if old and snapshot.source_revision == old.source_revision:
            if snapshot.source_snapshot_hash != old.source_snapshot_hash:
                raise ValueError(
                    "same source revision cannot have a different snapshot hash"
                )
            return SyncResult(
                status="canonical_active",
                source_type=snapshot.source_type,
                source_uid=snapshot.source_uid,
                source_revision=snapshot.source_revision,
                diff_counts={
                    "added": 0,
                    "modified": 0,
                    "deleted": 0,
                    "unchanged": len(snapshot.points),
                },
                embedded_chunk_count=0,
                reused_embedding_count=len(snapshot.points),
            )
        diff = diff_snapshots(old, snapshot) if old else _initial_diff(snapshot)
        chunks = build_chunks(snapshot)
        changed_point_keys = diff.changed_point_keys

        embeddings: dict[str, tuple[float, ...]] = {}
        missing_hashes: list[str] = []
        missing_texts: list[str] = []
        reused_count = 0
        for chunk in chunks:
            embedding_hash = embedding_input_hash(chunk, self._embedder.profile_key)
            existing = self._repository.load_embedding(embedding_hash)
            if existing is not None:
                reused_count += 1
                continue
            if chunk.primary_point_key not in changed_point_keys:
                raise RuntimeError(
                    "unchanged knowledge point is missing a reusable embedding"
                )
            missing_hashes.append(embedding_hash)
            missing_texts.append(chunk.content)

        if missing_texts:
            vectors = self._embedder.embed(missing_texts)
            if len(vectors) != len(missing_texts):
                raise RuntimeError(
                    "embedding provider returned an unexpected vector count"
                )
            embeddings = {
                embedding_hash: tuple(float(value) for value in vector)
                for embedding_hash, vector in zip(missing_hashes, vectors, strict=True)
            }

        publication = PreparedPublication(
            snapshot=snapshot,
            diff=diff,
            chunks=chunks,
            embeddings=embeddings,
            status="canonical_active",
        )
        self._repository.activate(publication)
        counts: Mapping[str, int] = {
            "added": len(diff.added),
            "modified": len(diff.modified),
            "deleted": len(diff.deleted),
            "unchanged": len(diff.unchanged),
        }
        return SyncResult(
            status=publication.status,
            source_type=snapshot.source_type,
            source_uid=snapshot.source_uid,
            source_revision=snapshot.source_revision,
            diff_counts=counts,
            embedded_chunk_count=len(embeddings),
            reused_embedding_count=reused_count,
        )