MemoryAgents/backend/app/embeddings.py

from __future__ import annotations

import hashlib
from typing import Any

import httpx

from .config import settings


class EmbeddingUnavailable(RuntimeError):
    pass


def embedding_fingerprint(content: str) -> str:
    """Invalidate stored vectors when either the text or model config changes."""
    payload = (
        f"{settings.embedding_base_url}\0{settings.embedding_model}\0"
        f"{settings.embedding_dimensions}\0{content}"
    )
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()


class OpenAICompatibleEmbeddingClient:
    """Embedding client for DashScope and other OpenAI-compatible endpoints."""

    @property
    def configured(self) -> bool:
        return bool(
            settings.embedding_api_key
            and settings.embedding_base_url
            and settings.embedding_model
        )

    async def embed(self, inputs: list[str]) -> list[list[float]]:
        if not self.configured:
            raise EmbeddingUnavailable(
                "未配置 EMBEDDING_API_KEY / EMBEDDING_BASE_URL / EMBEDDING_MODEL"
            )
        if not inputs:
            return []

        payload: dict[str, Any] = {
            "model": settings.embedding_model,
            "input": inputs,
            "encoding_format": "float",
        }
        if settings.embedding_dimensions:
            payload["dimensions"] = settings.embedding_dimensions

        async with httpx.AsyncClient(timeout=60) as client:
            response = await client.post(
                f"{settings.embedding_base_url}/embeddings",
                headers={"Authorization": f"Bearer {settings.embedding_api_key}"},
                json=payload,
            )
            response.raise_for_status()
            body = response.json()

        data = body.get("data")
        if not isinstance(data, list):
            raise EmbeddingUnavailable("Embedding API 返回中缺少 data 数组")
        ordered = sorted(data, key=lambda item: item.get("index", 0))
        vectors = [item.get("embedding") for item in ordered]
        if len(vectors) != len(inputs) or not all(isinstance(item, list) for item in vectors):
            raise EmbeddingUnavailable("Embedding API 返回的向量数量不匹配")
        if settings.embedding_dimensions and any(
            len(vector) != settings.embedding_dimensions for vector in vectors
        ):
            raise EmbeddingUnavailable(
                f"Embedding 向量维度与配置不一致，期望 {settings.embedding_dimensions}"
            )
        return [[float(value) for value in vector] for vector in vectors]


embedding_client = OpenAICompatibleEmbeddingClient()