mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-17 23:51:19 +00:00
fa249a3e38
- Integrate AnthropicChatModel with new AnthropicAsLLM component - Add component formatters for OpenAI and Anthropic chat models - Implement token counter component with estimated token counting - Create base client component for ReMe service communication - Refactor BaseComponent to remove app_context parameter from _start/_close - Update embedding model base class to remove retry logic and use npz cache - Add job component for sequential step execution with BaseJob - Implement step component base class for LLM workflow execution - Enhance application context with proper type annotations - Update component initialization to pass app_context automatically - Remove asyncio dependency from embedding model cache operations
318 lines
13 KiB
Python
318 lines
13 KiB
Python
"""Base embedding model with caching and batching support."""
|
|
|
|
import hashlib
|
|
import os
|
|
import time
|
|
|
|
import numpy as np
|
|
from abc import abstractmethod
|
|
from collections import OrderedDict
|
|
from pathlib import Path
|
|
|
|
from ..base_component import BaseComponent
|
|
from ...enumeration import ComponentEnum
|
|
from ...schema import BaseNode
|
|
|
|
|
|
class BaseEmbeddingModel(BaseComponent):
|
|
"""Abstract base class for embedding models with LRU cache.
|
|
|
|
Provides:
|
|
- LRU in-memory cache with disk persistence (npz)
|
|
- Automatic text truncation to max_input_length
|
|
- Batch embedding support
|
|
"""
|
|
|
|
component_type = ComponentEnum.EMBEDDING_MODEL
|
|
|
|
def __init__(
|
|
self,
|
|
api_key: str | None = None,
|
|
base_url: str | None = None,
|
|
model_name: str = "",
|
|
dimensions: int = 1024,
|
|
pass_dimensions: bool = False,
|
|
max_batch_size: int = 10,
|
|
max_input_length: int = 8192,
|
|
cache_dir: str | Path = ".reme",
|
|
max_cache_size: int = 2000,
|
|
enable_cache: bool = True,
|
|
encoding: str = "utf-8",
|
|
**kwargs,
|
|
):
|
|
"""Initialize embedding model configuration.
|
|
|
|
Args:
|
|
api_key: API key for the embedding service.
|
|
base_url: Base URL for the embedding service.
|
|
model_name: Name of the embedding model.
|
|
dimensions: Vector dimensions.
|
|
pass_dimensions: Whether to pass dimensions parameter to API.
|
|
max_batch_size: Maximum batch size for embedding requests.
|
|
max_input_length: Maximum input text length.
|
|
cache_dir: Directory for cache storage.
|
|
max_cache_size: Maximum LRU cache size.
|
|
enable_cache: Whether to enable caching.
|
|
encoding: Text encoding for cache file operations.
|
|
"""
|
|
super().__init__(**kwargs)
|
|
self.api_key: str = api_key or os.environ.get("EMBEDDING_API_KEY", "")
|
|
self.base_url: str = base_url or os.environ.get("EMBEDDING_BASE_URL", "")
|
|
self.model_name = model_name
|
|
self.dimensions = dimensions
|
|
self.pass_dimensions = pass_dimensions
|
|
self.max_batch_size = max_batch_size
|
|
self.max_input_length = max_input_length
|
|
self.cache_dir = cache_dir
|
|
self.max_cache_size = max_cache_size
|
|
self.enable_cache = enable_cache
|
|
self.encoding = encoding
|
|
|
|
self._embedding_cache: OrderedDict[str, list[float]] = OrderedDict()
|
|
self._cache_hits = 0
|
|
self._cache_misses = 0
|
|
self.cache_path: Path = Path(self.cache_dir)
|
|
|
|
async def _start(self) -> None:
|
|
"""Load cache on start."""
|
|
assert self.app_context is not None, "app_context must be provided"
|
|
working_path = Path(self.app_context.app_config.working_dir)
|
|
working_path.embedding_cache = self.cache_path
|
|
self._load_cache()
|
|
|
|
async def _close(self) -> None:
|
|
"""Save cache on close."""
|
|
self._save_cache()
|
|
|
|
def _truncate_text(self, text: str) -> str:
|
|
"""Truncate text to max_input_length."""
|
|
return text[: self.max_input_length] if len(text) > self.max_input_length else text
|
|
|
|
def _validate_and_adjust_embedding(self, embedding: list[float]) -> list[float]:
|
|
"""Adjust embedding dimensions to match expected dimensions."""
|
|
actual_len = len(embedding)
|
|
if actual_len == self.dimensions:
|
|
return embedding
|
|
|
|
if actual_len < self.dimensions:
|
|
self.logger.warning(
|
|
f"[ACTUAL_EMB_LENGTH] Embedding {actual_len} < expected {self.dimensions}, padding with zeros",
|
|
)
|
|
return embedding + [0.0] * (self.dimensions - actual_len)
|
|
|
|
self.logger.warning(
|
|
f"[ACTUAL_EMB_LENGTH] Embedding {actual_len} > expected {self.dimensions}, truncating",
|
|
)
|
|
return embedding[: self.dimensions]
|
|
|
|
def _get_cache_key(self, text: str) -> str:
|
|
"""Generate cache key from text + model_name + dimensions."""
|
|
cache_string = f"{text}|{self.model_name}|{self.dimensions}"
|
|
return hashlib.sha256(cache_string.encode(self.encoding)).hexdigest()
|
|
|
|
def _get_cache_file_path(self) -> Path:
|
|
"""Return path to the cache npz file."""
|
|
return self.cache_path / "embedding_cache.npz"
|
|
|
|
def _load_cache(self) -> None:
|
|
"""Load embedding cache from disk (npz format)."""
|
|
if not self.enable_cache:
|
|
return
|
|
|
|
self.cache_path.mkdir(parents=True, exist_ok=True)
|
|
cache_file = self._get_cache_file_path()
|
|
if not cache_file.exists():
|
|
self.logger.info(f"No cache file at {cache_file}, starting empty")
|
|
return
|
|
|
|
try:
|
|
load_start = time.time()
|
|
data = np.load(cache_file)
|
|
keys = data["keys"]
|
|
embeddings = data["embeddings"]
|
|
|
|
loaded_count = 0
|
|
for key, emb in zip(keys, embeddings):
|
|
key_str = str(key)
|
|
emb_list = emb.tolist()
|
|
if len(emb_list) != self.dimensions:
|
|
self.logger.warning(
|
|
f"Cache dimension mismatch for {key_str}: "
|
|
f"expected {self.dimensions}, got {len(emb_list)}",
|
|
)
|
|
continue
|
|
if len(self._embedding_cache) >= self.max_cache_size:
|
|
self.logger.info(f"Cache limit reached ({self.max_cache_size}), loaded {loaded_count}")
|
|
break
|
|
self._embedding_cache[key_str] = emb_list
|
|
loaded_count += 1
|
|
|
|
self.logger.info(f"Loaded {loaded_count} embeddings from {cache_file} in {time.time() - load_start:.2f}s")
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to load cache from {cache_file}: {e}, deleting file")
|
|
try:
|
|
cache_file.unlink()
|
|
except Exception as del_e:
|
|
self.logger.error(f"Failed to delete cache file: {del_e}")
|
|
|
|
def _save_cache(self) -> None:
|
|
"""Save embedding cache to disk (npz format)."""
|
|
if not self.enable_cache or not self._embedding_cache:
|
|
return
|
|
|
|
cache_file = self._get_cache_file_path()
|
|
try:
|
|
keys = []
|
|
embeddings = []
|
|
for cache_key, embedding in self._embedding_cache.items():
|
|
if len(embedding) != self.dimensions:
|
|
self.logger.warning(f"Cache dimension mismatch for {cache_key}")
|
|
continue
|
|
keys.append(cache_key)
|
|
embeddings.append(embedding)
|
|
|
|
np.savez(
|
|
cache_file,
|
|
keys=np.array(keys, dtype=str),
|
|
embeddings=np.array(embeddings, dtype=np.float32),
|
|
)
|
|
self.logger.info(f"Saved {len(keys)} embeddings to {cache_file}")
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to save cache to {cache_file}: {e}")
|
|
|
|
def _get_from_cache(self, text: str) -> list[float] | None:
|
|
"""Retrieve embedding from cache if available."""
|
|
if not self.enable_cache:
|
|
return None
|
|
|
|
cache_key = self._get_cache_key(text)
|
|
if cache_key not in self._embedding_cache:
|
|
self._cache_misses += 1
|
|
return None
|
|
|
|
embeddings = self._embedding_cache[cache_key]
|
|
if len(embeddings) != self.dimensions:
|
|
self.logger.warning("Cached embedding dimension mismatch, removing entry")
|
|
del self._embedding_cache[cache_key]
|
|
self._cache_misses += 1
|
|
return None
|
|
|
|
self._embedding_cache.move_to_end(cache_key)
|
|
self._cache_hits += 1
|
|
preview = text[:50] + "..." if len(text) > 50 else text
|
|
self.logger.info(f"Cache hit: {preview} (hits: {self._cache_hits}, misses: {self._cache_misses})")
|
|
return embeddings
|
|
|
|
def _put_to_cache(self, text: str, embedding: list[float]) -> None:
|
|
"""Store embedding in cache with LRU eviction."""
|
|
if not self.enable_cache or self.max_cache_size <= 0:
|
|
return
|
|
|
|
cache_key = self._get_cache_key(text)
|
|
if len(embedding) != self.dimensions:
|
|
self.logger.warning(f"[PUT_TO_CACHE] Dimension mismatch for {cache_key}")
|
|
return
|
|
|
|
if len(self._embedding_cache) >= self.max_cache_size and cache_key not in self._embedding_cache:
|
|
self._embedding_cache.popitem(last=False)
|
|
|
|
self._embedding_cache[cache_key] = embedding
|
|
self._embedding_cache.move_to_end(cache_key)
|
|
|
|
def get_cache_stats(self) -> dict[str, int | float]:
|
|
"""Return cache statistics: size, hits, misses, hit_rate."""
|
|
total = self._cache_hits + self._cache_misses
|
|
hit_rate = self._cache_hits / total if total > 0 else 0.0
|
|
return {
|
|
"cache_size": len(self._embedding_cache),
|
|
"max_cache_size": self.max_cache_size,
|
|
"cache_hits": self._cache_hits,
|
|
"cache_misses": self._cache_misses,
|
|
"hit_rate": hit_rate,
|
|
}
|
|
|
|
def clear_cache(self) -> None:
|
|
"""Clear in-memory cache and reset statistics."""
|
|
self._embedding_cache.clear()
|
|
self._cache_hits = 0
|
|
self._cache_misses = 0
|
|
|
|
@abstractmethod
|
|
async def _get_embeddings(self, input_text: list[str], **kwargs) -> list[list[float]]:
|
|
"""Fetch embeddings for a batch of texts. Override in subclasses."""
|
|
|
|
async def get_embedding(self, input_text: str, **kwargs) -> list[float]:
|
|
"""Get embedding for a single text with cache."""
|
|
truncated_text = self._truncate_text(input_text)
|
|
cached = self._get_from_cache(truncated_text)
|
|
if cached is not None:
|
|
return cached
|
|
|
|
try:
|
|
result = await self._get_embeddings([truncated_text], **kwargs)
|
|
if result and len(result) == 1:
|
|
embedding = self._validate_and_adjust_embedding(result[0])
|
|
self._put_to_cache(truncated_text, embedding)
|
|
return embedding
|
|
self.logger.warning(
|
|
f"Model {self.model_name} returned {len(result) if result else 0} results, expected 1",
|
|
)
|
|
except Exception as e:
|
|
self.logger.error(f"Model {self.model_name} failed: {e}")
|
|
return []
|
|
|
|
async def get_embeddings(self, input_text: list[str], **kwargs) -> list[list[float]]:
|
|
"""Get embeddings for multiple texts with cache and batching."""
|
|
truncated_texts = [self._truncate_text(t) for t in input_text]
|
|
results: list[list[float] | None] = [None] * len(truncated_texts)
|
|
texts_to_compute: list[tuple[int, str]] = []
|
|
|
|
for idx, text in enumerate(truncated_texts):
|
|
cached = self._get_from_cache(text)
|
|
if cached is not None:
|
|
results[idx] = cached
|
|
else:
|
|
texts_to_compute.append((idx, text))
|
|
|
|
if texts_to_compute:
|
|
uncached_texts = [text for _, text in texts_to_compute]
|
|
for i in range(0, len(uncached_texts), self.max_batch_size):
|
|
batch_texts = uncached_texts[i: i + self.max_batch_size]
|
|
batch_indices = [idx for idx, _ in texts_to_compute[i: i + self.max_batch_size]]
|
|
|
|
try:
|
|
batch_embeddings = await self._get_embeddings(batch_texts, **kwargs)
|
|
if batch_embeddings and len(batch_embeddings) == len(batch_texts):
|
|
for orig_idx, text, embedding in zip(batch_indices, batch_texts, batch_embeddings):
|
|
adjusted = self._validate_and_adjust_embedding(embedding)
|
|
results[orig_idx] = adjusted
|
|
self._put_to_cache(text, adjusted)
|
|
else:
|
|
self.logger.warning(
|
|
f"Batch returned {len(batch_embeddings) if batch_embeddings else 0} "
|
|
f"results for {len(batch_texts)} inputs",
|
|
)
|
|
for orig_idx in batch_indices:
|
|
if results[orig_idx] is None:
|
|
results[orig_idx] = []
|
|
except Exception as e:
|
|
self.logger.error(f"Model {self.model_name} batch failed: {e}")
|
|
for orig_idx in batch_indices:
|
|
if results[orig_idx] is None:
|
|
results[orig_idx] = []
|
|
|
|
return [r if r is not None else [] for r in results]
|
|
|
|
async def get_node_embeddings(self, nodes: list[BaseNode], **kwargs) -> list[BaseNode]:
|
|
"""Get embeddings for a list of nodes and assign to node.embedding."""
|
|
texts = [node.text for node in nodes]
|
|
embeddings = await self.get_embeddings(texts, **kwargs)
|
|
|
|
if len(embeddings) == len(nodes):
|
|
for node, vec in zip(nodes, embeddings):
|
|
node.embedding = vec
|
|
else:
|
|
self.logger.warning(f"Mismatch: {len(embeddings)} vectors for {len(nodes)} nodes, skipping assignment")
|
|
return nodes
|
|
|
|
|