"""Embedder 接口与实现 - 文本向量化""" import hashlib import logging import os import time from abc import ABC, abstractmethod from collections import OrderedDict from typing import Any 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: Any = None self._cache = cache def _get_client(self): """Lazily create and reuse a single httpx.AsyncClient.""" if self._client is None: import httpx 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 Exception 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