fischer-agentkit/src/agentkit/memory/embedder.py

"""Embedder 接口与实现 - 文本向量化"""

from __future__ import annotations

import hashlib
import logging
import os
import time
from abc import ABC, abstractmethod
from collections import OrderedDict

import httpx

logger = logging.getLogger(__name__)


class EmbeddingCache:
    """LRU cache for embedding vectors with TTL support.

    Key: SHA-256 hash of input text
    Value: (embedding vector, timestamp)
    """

    def __init__(self, max_size: int = 1000, ttl: int = 3600):
        """
        Args:
            max_size: Maximum number of entries in the cache.
            ttl: Time-to-live in seconds for cached entries.
        """
        self._max_size = max_size
        self._ttl = ttl
        self._cache: OrderedDict[str, tuple[list[float], float]] = OrderedDict()

    @staticmethod
    def _make_key(text: str) -> str:
        """Generate SHA-256 hash key from input text."""
        return hashlib.sha256(text.encode()).hexdigest()

    def get(self, text: str) -> list[float] | None:
        """Retrieve a cached embedding if present and not expired.

        Returns ``None`` on cache miss or if the entry has expired.
        """
        key = self._make_key(text)
        entry = self._cache.get(key)
        if entry is None:
            return None

        embedding, ts = entry
        if time.monotonic() - ts > self._ttl:
            # Expired — remove and report miss
            del self._cache[key]
            return None

        # Move to end (most recently used)
        self._cache.move_to_end(key)
        return embedding

    def put(self, text: str, embedding: list[float]) -> None:
        """Store an embedding in the cache, evicting the LRU entry if full."""
        key = self._make_key(text)
        if key in self._cache:
            self._cache.move_to_end(key)
        self._cache[key] = (embedding, time.monotonic())

        # Evict oldest entries if over capacity
        while len(self._cache) > self._max_size:
            self._cache.popitem(last=False)

    def clear(self) -> None:
        """Remove all entries from the cache."""
        self._cache.clear()


class Embedder(ABC):
    """文本嵌入抽象基类"""

    @abstractmethod
    async def embed(self, text: str) -> list[float]:
        """生成文本的嵌入向量"""
        ...

    @abstractmethod
    def get_dimension(self) -> int:
        """返回嵌入向量的维度"""
        ...


class OpenAIEmbedder(Embedder):
    """OpenAI Embeddings API 实现"""

    def __init__(
        self,
        api_key: str | None = None,
        model: str = "text-embedding-3-small",
        base_url: str | None = None,
        cache: EmbeddingCache | None = None,
    ):
        self._api_key = api_key
        self._model = model
        self._base_url = base_url
        self._dimension = 1536  # text-embedding-3-small 默认维度
        self._client: httpx.AsyncClient | None = None
        self._cache = cache

    def _get_client(self) -> httpx.AsyncClient:
        """Lazily create and reuse a single httpx.AsyncClient."""
        if self._client is None:
            self._client = httpx.AsyncClient(timeout=30.0)
        return self._client

    async def aclose(self) -> None:
        """Close the underlying httpx.AsyncClient."""
        if self._client is not None:
            await self._client.aclose()
            self._client = None

    async def __aenter__(self) -> "OpenAIEmbedder":
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        await self.aclose()

    async def embed(self, text: str) -> list[float]:
        """使用 OpenAI API 生成嵌入向量"""
        # Check cache first
        if self._cache is not None:
            cached = self._cache.get(text)
            if cached is not None:
                return cached

        try:
            api_key = self._api_key or os.environ.get("OPENAI_API_KEY", "")
            base_url = self._base_url or "https://api.openai.com/v1"

            client = self._get_client()
            response = await client.post(
                f"{base_url}/embeddings",
                headers={"Authorization": f"Bearer {api_key}"},
                json={"input": text, "model": self._model},
            )
            response.raise_for_status()
            data = response.json()
            embedding = data["data"][0]["embedding"]
            self._dimension = len(embedding)

            # Store in cache
            if self._cache is not None:
                self._cache.put(text, embedding)

            return embedding
        except (httpx.HTTPError, ValueError, KeyError, TypeError) as e:
            logger.error(f"OpenAI embedding failed: {e}")
            raise

    def get_dimension(self) -> int:
        return self._dimension


class MockEmbedder(Embedder):
    """Mock Embedder - 生成确定性伪嵌入向量，用于测试"""

    def __init__(self, dimension: int = 128):
        self._dimension = dimension

    async def embed(self, text: str) -> list[float]:
        """基于文本哈希生成确定性伪嵌入向量"""
        hash_bytes = hashlib.sha256(text.encode()).digest()
        vector = []
        for i in range(self._dimension):
            byte_idx = i % len(hash_bytes)
            vector.append(hash_bytes[byte_idx] / 255.0)
        # 归一化为单位向量
        magnitude = sum(x**2 for x in vector) ** 0.5
        if magnitude > 0:
            vector = [x / magnitude for x in vector]
        return vector

    def get_dimension(self) -> int:
        return self._dimension