From 5c28dd8be22b6a6b9bffcfa0182641c9045502d9 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 13 May 2026 12:22:59 +0800 Subject: [PATCH] up --- .../embedding/base_embedding_model.py | 175 ++----- .../embedding/openai_embedding_model.py | 8 +- reme2/component/file_store/base_file_store.py | 75 +-- reme2/component/file_store/bm25_lite.py | 6 +- .../component/file_store/local_file_store.py | 130 +++--- .../file_watcher/base_file_watcher.py | 433 +----------------- .../file_watcher/base_file_watcher_old.py | 427 +++++++++++++++++ .../file_watcher/light_file_watcher.py | 11 +- reme2/mcp/steps/memory_retriever.py | 2 +- reme2/memory/memory_io.py | 4 +- reme2/schema/file_edge.py | 5 +- reme2/schema/file_node.py | 7 + reme2/utils/similarity_utils.py | 8 - test/reme2/test_bm25_lite.py | 30 +- test/test_cache_memory_usage.py | 11 +- test/test_embedding_cache.py | 427 ----------------- 16 files changed, 629 insertions(+), 1130 deletions(-) create mode 100644 reme2/component/file_watcher/base_file_watcher_old.py delete mode 100644 test/test_embedding_cache.py diff --git a/reme2/component/embedding/base_embedding_model.py b/reme2/component/embedding/base_embedding_model.py index 179a28cb..f9564c3e 100644 --- a/reme2/component/embedding/base_embedding_model.py +++ b/reme2/component/embedding/base_embedding_model.py @@ -1,8 +1,5 @@ -"""Base embedding model with caching and batching support.""" - import hashlib import os -import time from abc import abstractmethod from collections import OrderedDict from pathlib import Path @@ -15,13 +12,7 @@ from ...schema import EmbNode 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 - """ + """Embedding model with LRU cache and disk persistence.""" component_type = ComponentEnum.EMBEDDING_MODEL @@ -34,14 +25,13 @@ class BaseEmbeddingModel(BaseComponent): pass_dimensions: bool = False, max_batch_size: int = 10, max_input_length: int = 8192, - max_cache_size: int = 2000, + max_cache_size: int = 5000, enable_cache: bool = True, - cache_name: str = "", **kwargs, ): 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.api_key = api_key or os.environ.get("EMBEDDING_API_KEY", "") + self.base_url = base_url or os.environ.get("EMBEDDING_BASE_URL", "") self.model_name = model_name self.dimensions = dimensions self.pass_dimensions = pass_dimensions @@ -49,193 +39,128 @@ class BaseEmbeddingModel(BaseComponent): self.max_input_length = max_input_length self.max_cache_size = max_cache_size self.enable_cache = enable_cache - self._embedding_cache: OrderedDict[str, list[float]] = OrderedDict() - self._cache_hits = 0 - self._cache_misses = 0 - self.working_dir = self.app_context.app_config.working_dir if self.app_context is not None else "" - self.cache_name: str = cache_name or self.name - self.cache_path: Path = Path(self.working_dir) / "embedding_cache" / f"{self.cache_name}.npz" - - def clear_cache(self) -> None: - """Clear in-memory cache and reset statistics.""" - self._embedding_cache.clear() - self._cache_hits = 0 - self._cache_misses = 0 + @property + def cache_path(self) -> Path: + working_dir = self.app_context.app_config.working_dir if self.app_context else "" + return Path(working_dir) / "embedding_cache" / f"{self.name}.npz" async def _start(self) -> None: - """Load cache on start.""" - self.clear_cache() + self._embedding_cache.clear() self._load_cache() async def _close(self) -> None: - """Save cache on close.""" self._save_cache() - 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"Embedding dim {actual_len} < expected {self.dimensions}, padding") - return embedding + [0.0] * (self.dimensions - actual_len) - self.logger.warning(f"Embedding dim {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.""" return hashlib.sha256(f"{text}|{self.model_name}|{self.dimensions}".encode()).hexdigest() def _load_cache(self) -> None: - """Load embedding cache from disk (npz format).""" if not self.enable_cache: return self.cache_path.parent.mkdir(parents=True, exist_ok=True) if not self.cache_path.exists(): - self.logger.info(f"No cache file at {self.cache_path}, starting empty") return - load_start = time.time() try: data = np.load(self.cache_path) except Exception: - self.logger.exception(f"Failed to load cache from {self.cache_path}, deleting file") self.cache_path.unlink(missing_ok=True) return - loaded_count = 0 for key, emb in zip(data["keys"], data["embeddings"]): emb_list = emb.tolist() if len(emb_list) != self.dimensions: - self.logger.warning( - f"Cache dimension mismatch for {key}: 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[str(key)] = emb_list - loaded_count += 1 - - self.logger.info(f"Loaded {loaded_count} embeddings from {self.cache_path} in {time.time() - load_start:.2f}s") def _save_cache(self) -> None: - """Save embedding cache to disk (npz format).""" if not self.enable_cache or not self._embedding_cache: return keys, embeddings = [], [] - for cache_key, embedding in self._embedding_cache.items(): - keys.append(cache_key) - embeddings.append(embedding) + for k, v in self._embedding_cache.items(): + keys.append(k) + embeddings.append(v) try: np.savez(self.cache_path, keys=np.array(keys, dtype=str), embeddings=np.array(embeddings, dtype=np.float32)) - except Exception as e: - self.logger.error(f"Failed to save cache to {self.cache_path}: {e}") - return - - self.logger.info(f"Saved {len(keys)} embeddings to {self.cache_path}") + except Exception: + pass 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 + key = self._get_cache_key(text) + if key not in self._embedding_cache: return None - self._embedding_cache.move_to_end(cache_key) - self._cache_hits += 1 - return self._embedding_cache[cache_key] + self._embedding_cache.move_to_end(key) + return self._embedding_cache[key] 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 - if len(embedding) != self.dimensions: + if not self.enable_cache or self.max_cache_size <= 0 or len(embedding) != self.dimensions: return - cache_key = self._get_cache_key(text) - if len(self._embedding_cache) >= self.max_cache_size and cache_key not in self._embedding_cache: + key = self._get_cache_key(text) + if len(self._embedding_cache) >= self.max_cache_size and 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 - 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": self._cache_hits / total if total > 0 else 0.0, - } + self._embedding_cache[key] = embedding + self._embedding_cache.move_to_end(key) @abstractmethod - async def _get_embeddings(self, input_text: list[str], **kwargs) -> list[list[float] | None]: - """Fetch embeddings for a batch of texts. Override in subclasses.""" + async def _get_embeddings(self, input_text: list[str], **kwargs) -> list[list[float] | None]: ... async def get_embedding(self, input_text: str, **kwargs) -> list[float] | None: - """Get embedding for a single text with cache. Returns None on failure.""" results = await self.get_embeddings([input_text], **kwargs) return results[0] if results else None async def get_embeddings(self, input_text: list[str], **kwargs) -> list[list[float] | None]: - """Get embeddings for multiple texts with cache and batching.""" - # TODO change to bytes instead of str - truncated_texts = [t[: self.max_input_length] for t in input_text] - results: list[list[float] | None] = [None] * len(truncated_texts) - texts_to_compute: list[tuple[int, str]] = [] + truncated = [t[: self.max_input_length] for t in input_text] + results: list[list[float] | None] = [None] * len(truncated) + to_compute: list[tuple[int, str]] = [] - for idx, text in enumerate(truncated_texts): + for idx, text in enumerate(truncated): cached = self._get_from_cache(text) if cached is not None: results[idx] = cached else: - texts_to_compute.append((idx, text)) + 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_to_compute[i: i + self.max_batch_size] - batch_indices = [idx for idx, _ in batch] - batch_texts = [text for _, text in batch] + if to_compute: + for i in range(0, len(to_compute), self.max_batch_size): + batch = to_compute[i: i + self.max_batch_size] + indices = [idx for idx, _ in batch] + texts = [text for _, text in batch] 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", - ) - except Exception as e: - self.logger.error(f"Model {self.model_name} batch failed: {e}") + embeddings = await self._get_embeddings(texts, **kwargs) + if embeddings and len(embeddings) == len(texts): + for orig_idx, text, emb in zip(indices, texts, embeddings): + if emb is None: + continue + if len(emb) != self.dimensions: + if len(emb) < self.dimensions: + emb = emb + [0.0] * (self.dimensions - len(emb)) + else: + emb = emb[: self.dimensions] + results[orig_idx] = emb + self._put_to_cache(text, emb) + except Exception: + pass return results async def get_node_embeddings(self, nodes: list[EmbNode], **kwargs) -> list[EmbNode]: - """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) + embeddings = await self.get_embeddings([n.text for n in nodes], **kwargs) if len(embeddings) == len(nodes): for node, vec in zip(nodes, embeddings): if vec is not None: node.embedding = vec - else: - self.logger.warning(f"Embedding failed for node, skipping assignment") - else: - self.logger.warning(f"Mismatch: {len(embeddings)} vectors for {len(nodes)} nodes, skipping assignment") - return nodes + return nodes \ No newline at end of file diff --git a/reme2/component/embedding/openai_embedding_model.py b/reme2/component/embedding/openai_embedding_model.py index 6de89b33..a6fabc8f 100644 --- a/reme2/component/embedding/openai_embedding_model.py +++ b/reme2/component/embedding/openai_embedding_model.py @@ -27,7 +27,11 @@ class OpenAIEmbeddingModel(BaseEmbeddingModel): if self._client is None: raise RuntimeError("Client not initialized. Call _start() first.") - create_kwargs: dict = {"model": self.model_name, "input": input_text, **kwargs} + create_kwargs: dict = { + "model": self.model_name, + "input": input_text, + **kwargs, + } if self.pass_dimensions: create_kwargs["dimensions"] = self.dimensions @@ -35,8 +39,8 @@ class OpenAIEmbeddingModel(BaseEmbeddingModel): result: list[list[float] | None] = [None] * len(input_text) for emb in completion.data: - vec = getattr(emb, "embedding", None) or getattr(emb, "dense_embedding", None) if 0 <= emb.index < len(input_text): + vec = emb.embedding or getattr(emb, "dense_embedding", None) if vec is not None: result[emb.index] = list(vec) else: diff --git a/reme2/component/file_store/base_file_store.py b/reme2/component/file_store/base_file_store.py index 276de1fa..68d98d9d 100644 --- a/reme2/component/file_store/base_file_store.py +++ b/reme2/component/file_store/base_file_store.py @@ -1,9 +1,10 @@ -"""Abstract base class for file stores — persistence + search for nodes & chunks.""" +"""Abstract base class for file stores.""" import re from abc import abstractmethod from pathlib import Path +from .bm25_lite import BM25Lite from ..base_component import BaseComponent from ..embedding import BaseEmbeddingModel from ...enumeration import ComponentEnum @@ -11,7 +12,7 @@ from ...schema import FileChunk, FileNode class BaseFileStore(BaseComponent): - """Abstract file-store engine: persistence + search for nodes & chunks.""" + """Abstract file-store engine.""" component_type = ComponentEnum.FILE_STORE @@ -19,14 +20,17 @@ class BaseFileStore(BaseComponent): self, store_name: str, embedding_model: str = "default", + tokenizer: str = "default", fts_enabled: bool = True, **kwargs, ): + """Initialize the file store.""" super().__init__(**kwargs) if not re.match(r"^[a-zA-Z0-9_]+$", store_name): raise ValueError(f"Invalid store name '{store_name}'. Only alphanumeric and underscores allowed.") self.store_name = store_name or self.name self._embedding_model_name = embedding_model + self.tokenizer = tokenizer self.fts_enabled = fts_enabled self.embedding_model: BaseEmbeddingModel | None = None @@ -37,9 +41,10 @@ class BaseFileStore(BaseComponent): if not self.vector_enabled and not self.fts_enabled: raise ValueError("At least one of embedding_model or fts_enabled must be set.") - # Lifecycle + self.bm25: BM25Lite | None = None async def _start(self) -> None: + """Start the file store.""" if self.vector_enabled and self.app_context is not None: model_dict = self.app_context.components.get(ComponentEnum.EMBEDDING_MODEL, {}) if self._embedding_model_name not in model_dict: @@ -49,66 +54,32 @@ class BaseFileStore(BaseComponent): raise TypeError(f"Expected BaseEmbeddingModel, got {type(model).__name__}") self.embedding_model = model + if self.fts_enabled: + self.bm25 = BM25Lite(index_dir=self.store_path, tokenizer=self.tokenizer, app_context=self.app_context) + if self.bm25 is not None: + await self.bm25.start() + async def _close(self) -> None: + """Close the file store and release resources.""" self.embedding_model = None + if self.fts_enabled and self.bm25 is not None: + await self.bm25.close() # Composite operations - async def upsert(self, node: FileNode, chunks: list[FileChunk]) -> None: - await self.upsert_node(node) - await self.upsert_chunks(node.path, chunks) + async def upsert_file(self, node: FileNode, chunks: list[FileChunk]) -> None: + """Upsert a node and its associated chunks.""" async def delete(self, path: str) -> None: - await self.delete_chunks(path) - await self.delete_node(path) + """Delete a node and all its associated chunks by path.""" - # Abstract: node operations - - @abstractmethod - async def upsert_node(self, node: FileNode) -> None: - """Persist a node, replacing any prior entry for `node.path`.""" - - @abstractmethod - async def get_node(self, path: str) -> FileNode | None: - """Fetch a node by path, or None if absent.""" - - @abstractmethod - async def delete_node(self, path: str) -> None: - """Delete the node entry for `path`.""" - - # Abstract: chunk operations - - @abstractmethod - async def upsert_chunks(self, path: str, chunks: list[FileChunk]) -> None: - """Insert or replace all chunks for `path`. Handles embedding internally.""" - - @abstractmethod - async def get_chunks(self, path: str) -> list[FileChunk]: - """All chunks for `path`.""" - - @abstractmethod - async def delete_chunks(self, path: str) -> None: - """Delete all chunks for `path`.""" - - # Abstract: search + async def reindex(self) -> None: + """Re-index all nodes and chunks in the store.""" @abstractmethod async def vector_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]: - """Vector similarity search.""" + """Perform vector similarity search.""" @abstractmethod async def keyword_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]: - """Full-text / keyword search.""" - - # Internal helpers - - async def get_embeddings(self, texts: list[str]) -> list[list[float] | None] | None: - """Embed texts. Returns None if vector search disabled or on API error.""" - if not self.vector_enabled or not texts or not self.embedding_model: - return None - try: - return await self.embedding_model.get_embeddings(texts) - except Exception as e: - self.logger.warning(f"[{self.store_name}] Disabling vector search: {e}") - self.vector_enabled = False - return None + """Perform full-text keyword search.""" diff --git a/reme2/component/file_store/bm25_lite.py b/reme2/component/file_store/bm25_lite.py index f35ceed4..1265a6a2 100644 --- a/reme2/component/file_store/bm25_lite.py +++ b/reme2/component/file_store/bm25_lite.py @@ -212,12 +212,12 @@ class BM25Lite(BaseComponent): self._idf_cache[token_id] = math.log(1 + (self.n_docs - df + 0.5) / (df + 0.5)) if df else 0.0 return self._idf_cache[token_id] - def retrieve(self, query: str, k: int = 3) -> dict[str, float]: + def retrieve(self, query: str, limit: int = 3) -> dict[str, float]: """Search for documents matching the query. Args: query: Search query string. - k: Maximum number of results to return. + limit: Maximum number of results to return. Returns: Dict mapping document IDs to BM25 scores, sorted by score descending. @@ -238,7 +238,7 @@ class BM25Lite(BaseComponent): tf_score = tf * (self.k1 + 1) / (tf + self.k1 * (1 - self.b + self.b * doc_len / avg_len)) scores[doc_id] = scores.get(doc_id, 0.0) + idf * tf_score - return dict(sorted(scores.items(), key=lambda x: x[1], reverse=True)[:k]) if scores else {} + return dict(sorted(scores.items(), key=lambda x: x[1], reverse=True)[:limit]) if scores else {} def dump(self): """Persist index to disk via pickle.""" diff --git a/reme2/component/file_store/local_file_store.py b/reme2/component/file_store/local_file_store.py index 907d72d6..6ca09f8f 100644 --- a/reme2/component/file_store/local_file_store.py +++ b/reme2/component/file_store/local_file_store.py @@ -2,13 +2,16 @@ from pathlib import Path +import aiofiles import numpy as np from pydantic import BaseModel from .base_file_store import BaseFileStore from ..component_registry import R from ...schema import FileChunk, FileNode -from ...utils import batch_cosine_similarity +from ...utils import batch_cosine_similarity, get_logger + +logger = get_logger() @R.register("local") @@ -26,60 +29,62 @@ class LocalFileStore(BaseFileStore): # Lifecycle async def _start(self) -> None: - self._load(self._nodes_file, self._nodes, FileNode, "path") - self._load(self._chunks_file, self._chunks, FileChunk, "id") await super()._start() - self.logger.info( - f"LocalFileStore '{self.store_name}' ready: " - f"{len(self._nodes)} nodes, {len(self._chunks)} chunks" - ) + await self._load(self._nodes_file, self._nodes, FileNode, "path") + await self._load(self._chunks_file, self._chunks, FileChunk, "id") + self.logger.info(f"LocalFileStore '{self.store_name}' ready: " + f"{len(self._nodes)} nodes, {len(self._chunks)} chunks") async def _close(self) -> None: - self._dump(self._nodes_file, self._nodes.values()) - self._dump(self._chunks_file, self._chunks.values()) + await self._dump(self._nodes_file, list(self._nodes.values())) + await self._dump(self._chunks_file, list(self._chunks.values())) self._nodes.clear() self._chunks.clear() await super()._close() - def _load(self, file: Path, target: dict, model: type[BaseModel], key: str) -> None: + async def _load(self, file: Path, target: dict, model: type[BaseModel], key: str) -> None: if not file.exists(): return + target.clear() try: - for line in file.read_text(encoding=self._encoding).splitlines(): - if line.strip(): - obj = model.model_validate_json(line) - target[getattr(obj, key)] = obj + async with aiofiles.open(file, encoding=self._encoding) as f: + async for line in f: + line = line.strip() + if line: + obj = model.model_validate_json(line) + target[getattr(obj, key)] = obj except Exception as e: - self.logger.warning(f"Failed to load {file}: {e}") + self.logger.exception(f"Failed to load {file}: {e}") - def _dump(self, file: Path, items) -> None: + async def _dump(self, file: Path, items: list[BaseModel]) -> None: try: content = "\n".join(o.model_dump_json() for o in items) tmp = file.with_suffix(".tmp") - tmp.write_text(content, encoding=self._encoding) + async with aiofiles.open(tmp, "w", encoding=self._encoding) as f: + await f.write(content) tmp.replace(file) except Exception as e: - self.logger.error(f"Failed to write {file}: {e}") + self.logger.exception(f"Failed to write {file}: {e}") # Node operations async def upsert_node(self, node: FileNode) -> None: self._nodes[node.path] = node - async def get_node(self, path: str) -> FileNode | None: + async def get_node_by_path(self, path: str) -> FileNode | None: return self._nodes.get(path) - async def delete_node(self, path: str) -> None: - self._nodes.pop(path, None) + async def delete_node_by_path(self, path: str) -> FileNode | None: + return self._nodes.pop(path, None) # Chunk operations - async def upsert_chunks(self, path: str, chunks: list[FileChunk]) -> None: - existing = await self.get_chunks(path) - cached = {c.hash: c.embedding for c in existing if c.embedding} + async def upsert_chunks_with_path(self, path: str, chunks: list[FileChunk]) -> None: + existing = await self.get_chunks_by_path(path) + cached = {c.id: c.embedding for c in existing if c.embedding} - await self.delete_chunks(path) + await self.delete_chunks_by_path(path) if not chunks: return @@ -87,8 +92,8 @@ class LocalFileStore(BaseFileStore): for c in chunks: if c.embedding: continue - if c.hash in cached: - c.embedding = cached[c.hash] + if c.id in cached: + c.embedding = cached[c.id] elif c.text: needs_embed.append(c) @@ -102,33 +107,35 @@ class LocalFileStore(BaseFileStore): for c in chunks: self._chunks[c.id] = c - async def get_chunks(self, path: str) -> list[FileChunk]: + async def get_chunks_by_path(self, path: str) -> list[FileChunk]: chunks = [c for c in self._chunks.values() if c.path == path] chunks.sort(key=lambda c: c.start_line) return chunks - async def delete_chunks(self, path: str) -> None: + async def delete_chunks_by_path(self, path: str) -> None: stale = [cid for cid, c in self._chunks.items() if c.path == path] for cid in stale: del self._chunks[cid] # Search - async def vector_search( - self, query: str, limit: int, search_filter: dict, - ) -> list[FileChunk]: - if not self.vector_enabled or not query: + async def vector_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]: + if self.embedding_model is None or not query: return [] - embeddings = await self.get_embeddings([query]) - if not embeddings or not embeddings[0]: + + query_embedding = await self.embedding_model.get_embedding(query) + if not query_embedding: return [] candidates = [c for c in self._chunks.values() if c.embedding] + emb_missing = [c for c in self._chunks.values() if not c.embedding] + if emb_missing: + logger.warning(f"Embedding missing for {len(emb_missing)} chunks") if not candidates: return [] chunk_embs = np.array([c.embedding for c in candidates]) - similarities = batch_cosine_similarity(np.array([embeddings[0]]), chunk_embs)[0] + similarities = batch_cosine_similarity(np.array([query_embedding]), chunk_embs)[0] results = [ c.model_copy(update={"scores": {"vector": float(s), "score": float(s)}}) @@ -137,33 +144,32 @@ class LocalFileStore(BaseFileStore): results.sort(key=lambda r: r.score, reverse=True) return results[:limit] - async def keyword_search( - self, query: str, limit: int, search_filter: dict, - ) -> list[FileChunk]: - if not self.fts_enabled or not query.split(): + async def keyword_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]: + if not self.fts_enabled or self.bm25 is None: return [] - results = [ - c.model_copy(update={"scores": {"keyword": s, "score": s}}) - for c in self._chunks.values() - if (s := self._keyword_score(query, c.text)) > 0 - ] - results.sort(key=lambda r: r.score, reverse=True) - return results[:limit] + query = query.strip() + if not query: + return [] - # Helpers + doc_id_score_dict = self.bm25.retrieve(query, limit=limit) + results = [] + for doc_id, score in doc_id_score_dict.items(): + chunk = self._chunks.get(doc_id) + if chunk: + results.append(chunk.model_copy(update={"scores": {"keyword": score, "score": score}})) + + return results + + # Reindex + + async def reindex_vector(self) -> None: + if not self.vector_enabled or self.embedding_model is None: + return + + await self.embedding_model.get_node_embeddings(list(self._chunks.values())) + + async def reindex_keyword(self) -> None: + if not self.fts_enabled or self.bm25 is None: + return - @staticmethod - def _keyword_score(query: str, text: str) -> float: - """Word-overlap score with phrase bonus. Range [0, 1].""" - words = query.split() - if not words or not text: - return 0.0 - text_lower = text.lower() - matches = sum(1 for w in words if w.lower() in text_lower) - if matches == 0: - return 0.0 - base = matches / len(words) - if len(words) > 1 and query.lower() in text_lower: - base = min(1.0, base + 0.2) - return base diff --git a/reme2/component/file_watcher/base_file_watcher.py b/reme2/component/file_watcher/base_file_watcher.py index bf2d84c1..08beb5cc 100644 --- a/reme2/component/file_watcher/base_file_watcher.py +++ b/reme2/component/file_watcher/base_file_watcher.py @@ -1,427 +1,38 @@ -"""Base file watcher with watchfiles integration. - -Per change → single-step pipeline: - added/modified → node, chunks = parser.parse(path) - file_store.upsert(node, chunks) - deleted → file_store.delete_chunks + delete_node - -The parser owns: chunking and edge extraction (text-only, no embedding). -The file_store owns: persistence of node + chunks AND the embedding -pipeline — `upsert` hash-diffs incoming chunks against its persisted -ones, reuses cached embeddings for unchanged blocks, and calls the -embedding API only for new-hash blocks. -The watcher owns: scheduling, cancel-on-modify, retry/timeout, and -startup recovery (re-parsing files whose mtime drifted while offline). -""" - -import asyncio -import os -import time -from pathlib import Path - -from watchfiles import Change, awatch - from ..base_component import BaseComponent -from ..file_parser import BaseFileParser -from ..file_store import BaseFileStore from ...enumeration import ComponentEnum +from ..file_parser import BaseFileParser class BaseFileWatcher(BaseComponent): - """Watches a directory and feeds the file_store via single-step parse+upsert.""" - component_type = ComponentEnum.FILE_WATCHER - _META_DIR = ".reme" - def __init__( - self, - recursive: bool = False, - debounce: int = 2000, - file_store: str = "default", - default_parser: str | None = None, - rebuild_index_on_start: bool = False, - poll_delay_ms: int = 2000, - parse_max_attempts: int = 3, - parse_retry_backoff: float = 2.0, - parse_task_timeout: float = 300.0, - **kwargs, + self, + watch_path: list[str], + recursive: bool = False, + file_parser: str = "default", + file_store: str = "default", + debounce: int = 2000, + poll_delay_ms: int = 2000, + **kwargs, ): super().__init__(**kwargs) - self._file_store_name: str = file_store - self._default_parser_name: str | None = default_parser - self.file_store: BaseFileStore | None = None - self._suffix_to_parser: dict[str, BaseFileParser] = {} - self._default_parser: BaseFileParser | None = None + self.watch_path: list[str] = watch_path self.recursive: bool = recursive + self.file_parser_name: str = file_parser + self.file_store: str = file_store self.debounce: int = debounce - self.rebuild_index_on_start: bool = rebuild_index_on_start self.poll_delay_ms: int = poll_delay_ms - self._stop_event = asyncio.Event() - self._watch_task: asyncio.Task | None = None + self.file_parser: BaseFileParser | None = None - # Parse pipeline configuration. - self._parse_max_attempts: int = max(1, int(parse_max_attempts)) - self._parse_retry_backoff: float = float(parse_retry_backoff) - self._parse_task_timeout: float = float(parse_task_timeout) - self._tasks: dict[str, asyncio.Task] = {} - self._mgmt_lock: asyncio.Lock | None = None - self._run_lock: asyncio.Lock | None = None - self._last_failure: dict[str, float] = {} + async def _start(self) -> None: + await super()._start() + if self.app_context is not None: + file_parser_dict = self.app_context.components.get(ComponentEnum.FILE_PARSER, {}) + if self.file_parser not in file_parser_dict: + raise ValueError(f"File parser {self.file_parser} not found") + self.file_parser = file_parser_dict[self.file_parser] - @property - def watch_path(self) -> str: - """Vault root the watcher is bound to (== file_store.working_dir). - - Available after _start. Reads the live value from the bound - store so config has a single source of truth. - """ - if self.file_store is None or self.file_store.working_dir is None: - raise RuntimeError( - "watch_path requires a started file_store with working_dir set", - ) - return str(self.file_store.working_dir) - - @property - def _meta_path(self) -> Path: - return (Path(self.watch_path) / self._META_DIR).resolve() - - # -- Lifecycle ---------------------------------------------------------- - - async def _start(self): - """Resolve file_store + parsers, sync any disk drift, start watch loop.""" - if self._file_store_name: - assert self.app_context is not None, "app_context must be provided" - - stores = self.app_context.components.get(ComponentEnum.FILE_STORE, {}) - if self._file_store_name not in stores: - raise ValueError(f"File store '{self._file_store_name}' not found.") - store = stores[self._file_store_name] - if not isinstance(store, BaseFileStore): - raise TypeError(f"Expected BaseFileStore, got {type(store).__name__}") - if store.working_dir is None: - raise ValueError( - f"File store '{self._file_store_name}' has no working_dir; " - f"set components.file_store.{self._file_store_name}.working_dir in config", - ) - self.file_store = store - - parsers = self.app_context.components.get(ComponentEnum.FILE_PARSER, {}) - for parser in parsers.values(): - if isinstance(parser, BaseFileParser): - for suffix in parser.suffixes: - self._suffix_to_parser[suffix] = parser - - if self._default_parser_name and self._default_parser_name in parsers: - parser = parsers[self._default_parser_name] - if isinstance(parser, BaseFileParser): - self._default_parser = parser - - if not self._suffix_to_parser and not self._default_parser: - self.logger.warning("No file parsers registered") - - async def _initialize_and_watch(): - await self._initial_sync_and_recovery() - await self._watch_loop() - - self._mgmt_lock = asyncio.Lock() - self._run_lock = asyncio.Lock() - self._last_failure.clear() - self._stop_event.clear() - self._watch_task = asyncio.create_task(_initialize_and_watch()) - self.logger.info(f"Started watching: {self.watch_path}") - - async def _close(self): - """Stop watching and release resources.""" - self._stop_event.set() - if self._watch_task and not self._watch_task.done(): - self._watch_task.cancel() - try: - await self._watch_task - except asyncio.CancelledError: - pass - - # Cancel all in-flight parse tasks before tearing down state. - tasks = list(self._tasks.values()) - for t in tasks: - if not t.done(): - t.cancel() - if tasks: - await asyncio.gather(*tasks, return_exceptions=True) - self._tasks.clear() - self._last_failure.clear() - - self._watch_task = None - self._stop_event.clear() - self._mgmt_lock = None - self._run_lock = None - self.file_store = None - self._suffix_to_parser.clear() - self._default_parser = None - self.logger.info("Stopped watching") - - # -- Startup sync ------------------------------------------------------- - - async def _initial_sync_and_recovery(self) -> None: - """At startup: file_store has its persisted state loaded already; we - diff against current disk for any changes that happened while we - were offline, then re-enqueue them through the parse pipeline.""" - await self._sync_with_disk() - - async def _sync_with_disk(self) -> bool: - """Diff cached graph (in file_store) vs disk state; reindex changes.""" - if self.file_store is None: - return False - - watch_path = Path(self.watch_path) - if not watch_path.exists(): - self.logger.warning(f"Watch path does not exist: {watch_path}") - return False - - on_disk: dict[str, float] = {} - if watch_path.is_file(): - on_disk[str(watch_path.resolve())] = watch_path.stat().st_mtime - elif watch_path.is_dir(): - iterator = watch_path.rglob("*") if self.recursive else watch_path.iterdir() - for fp in iterator: - if not fp.is_file(): - continue - abs_path = str(fp.resolve()) - if not self._watch_filter(abs_path): - continue - if not self._get_parser(abs_path): - continue - try: - on_disk[abs_path] = fp.stat().st_mtime - except OSError: - continue - - cached_paths = set(self.file_store.nodes.keys()) - disk_paths = set(on_disk.keys()) - - added = disk_paths - cached_paths - deleted = cached_paths - disk_paths - modified: set[str] = set() - for p in disk_paths & cached_paths: - cached_node = self.file_store.get_node(p) - if cached_node is None or cached_node.st_mtime != on_disk[p]: - modified.add(p) - - unchanged = len(disk_paths & cached_paths) - len(modified) - if not (added or deleted or modified): - self.logger.info(f"Index up-to-date ({unchanged} files cached, 0 changes)") - return False - - self.logger.info( - f"Incremental sync: +{len(added)} ~{len(modified)} -{len(deleted)} (unchanged {unchanged})", - ) - - changes: set[tuple[Change, str]] = set() - for p in deleted: - changes.add((Change.deleted, p)) - for p in added: - changes.add((Change.added, p)) - for p in modified: - changes.add((Change.modified, p)) - await self.on_changes(changes) - return True - - # -- Watch loop --------------------------------------------------------- - - async def _interruptible_sleep(self, seconds: float): - try: - await asyncio.wait_for(self._stop_event.wait(), timeout=seconds) - except asyncio.TimeoutError: - pass - - async def _watch_loop(self): - while not self._stop_event.is_set(): - if not Path(self.watch_path).exists(): - self.logger.warning(f"Watch path does not exist: {self.watch_path}, waiting 10 seconds...") - await self._interruptible_sleep(10) - continue - - try: - self.logger.info(f"Starting watch on: {self.watch_path}") - async for changes in awatch( - self.watch_path, - watch_filter=lambda _, p: self._watch_filter(p), - recursive=self.recursive, - debounce=self.debounce, - poll_delay_ms=self.poll_delay_ms, - stop_event=self._stop_event, - ): - if self._stop_event.is_set(): - break - await self.on_changes(changes) - - except FileNotFoundError as e: - self.logger.error(f"Watch path no longer exists: {e}, restarting in 10 seconds...") - if not self._stop_event.is_set(): - await self._interruptible_sleep(10) - - except Exception as e: - self.logger.error(f"Error in watch loop: {e}, restarting in 10 seconds...", exc_info=True) - if not self._stop_event.is_set(): - await self._interruptible_sleep(10) - - def _watch_filter(self, path: str) -> bool: - resolved = Path(path).resolve() - if resolved == self._meta_path: - return False - return self._meta_path not in resolved.parents - - def _get_parser(self, path: str) -> BaseFileParser | None: - suffix = Path(path).suffix.lower() - return self._suffix_to_parser.get(suffix, self._default_parser) - - # -- Change dispatch ---------------------------------------------------- - - async def on_changes(self, changes: set[tuple[Change, str]]) -> None: - if self.file_store is None: - self.logger.warning("File store not initialized, skipping changes") - return - - for change_type, path in changes: - try: - if change_type in (Change.added, Change.modified): - await self._on_modified(path) - elif change_type == Change.deleted: - await self._on_deleted(path) - except FileNotFoundError: - self.logger.warning(f"File not found: {path}, skipping") - except PermissionError: - self.logger.warning(f"Permission denied: {path}, skipping") - except Exception as e: - self.logger.opt(exception=True).error("Error processing {p}: {err}", p=path, err=str(e)) - - async def _on_modified(self, path: str) -> None: - if not Path(path).is_file(): - return - parser = self._get_parser(path) - if parser is None: - self.logger.debug(f"No parser for {path}, skipping") - return - await self._submit_parse_task(path, parser) - - async def _on_deleted(self, path: str) -> None: - # Cancel any in-flight parse task FIRST so a delayed cancellation - # can't race the cleanup writes below. - assert self.file_store is not None, "_on_deleted requires file_store" - targets = [path, *self._descendant_indexed_paths(path)] - for p in targets: - await self._cancel_parse_task(p) - await self.file_store.delete_chunks(p) - await self.file_store.delete_node(p) - if len(targets) == 1: - self.logger.info(f"Deleted {path}") - else: - self.logger.info(f"Deleted directory {path} ({len(targets) - 1} indexed children)") - - def _descendant_indexed_paths(self, path: str) -> list[str]: - """Indexed file paths that live beneath `path` as a directory. - - watchfiles emits a single Change.deleted for a removed directory - (no per-file events), so we have to find the orphans ourselves. - Resolves both sides to handle symlinks and trailing-separator - drift between the watcher and the file_store keys. - """ - assert self.file_store is not None - try: - prefix = str(Path(path).resolve()) + os.sep - except OSError: - prefix = path.rstrip(os.sep) + os.sep - return [p for p in self.file_store.nodes if p.startswith(prefix)] - - # -- Parse task pipeline ------------------------------------------------ - - async def _submit_parse_task(self, path: str, parser: BaseFileParser) -> None: - """Cancel any prior task + spawn a new one. Atomic under _mgmt_lock.""" - assert self._mgmt_lock is not None, "_submit_parse_task before _start()" - async with self._mgmt_lock: - await self._cancel_locked(path) - self._tasks[path] = asyncio.create_task( - self._run_parse_task(path, parser), - ) - - async def _cancel_parse_task(self, path: str) -> None: - """Cancel + await any in-flight parse task for `path`.""" - assert self._mgmt_lock is not None, "_cancel_parse_task before _start()" - async with self._mgmt_lock: - await self._cancel_locked(path) - - async def _cancel_locked(self, path: str) -> None: - """Cancel + await the task for `path`. Caller holds `_mgmt_lock`.""" - task = self._tasks.pop(path, None) - if task is not None and not task.done(): - task.cancel() - try: - await task - except (asyncio.CancelledError, Exception): - pass - - async def flush_parse_tasks(self) -> None: - """Wait until all parse tasks finish (loops because retries spawn new ones).""" - while self._tasks: - await asyncio.gather(*list(self._tasks.values()), return_exceptions=True) - - def pending_parse_count(self) -> int: - return sum(1 for t in self._tasks.values() if not t.done()) - - def failed_parse_paths(self) -> dict[str, float]: - return dict(self._last_failure) - - async def _run_parse_task(self, path: str, parser: BaseFileParser) -> None: - """One per-path parse task: parse → upsert. - - The store owns hash-diff + embedding: it inspects its own - existing chunks for `path`, reuses cached embeddings for - unchanged blocks, and only calls the embedding API for new-hash - blocks. The whole upsert is one call from this side. - """ - try: - for attempt in range(1, self._parse_max_attempts + 1): - try: - assert self.file_store is not None - node, chunks = await asyncio.wait_for( - parser.parse(path), - timeout=self._parse_task_timeout, - ) - - assert self._run_lock is not None - async with self._run_lock: - await asyncio.wait_for( - self.file_store.upsert(node, chunks), - timeout=self._parse_task_timeout, - ) - self._last_failure.pop(path, None) - - self.logger.info( - f"indexed {path}: chunks={len(chunks)} edges={len(node.edges)}", - ) - return - except asyncio.CancelledError: - raise - except FileNotFoundError: - self.logger.debug(f"parse task {path}: file gone, skipping") - self._last_failure.pop(path, None) - return - except Exception as e: - self._last_failure[path] = time.time() - if attempt < self._parse_max_attempts: - backoff = self._parse_retry_backoff * attempt - self.logger.warning( - f"parse task {path} attempt {attempt} failed " - f"({type(e).__name__}: {e}); retry in {backoff:.1f}s", - ) - await asyncio.sleep(backoff) - continue - self.logger.error( - f"parse task {path} giving up after {attempt} attempts: " - f"{type(e).__name__}: {e}", - ) - return - finally: - if self._tasks.get(path) is asyncio.current_task(): - self._tasks.pop(path, None) + async def _close(self) -> None: + await super()._close() \ No newline at end of file diff --git a/reme2/component/file_watcher/base_file_watcher_old.py b/reme2/component/file_watcher/base_file_watcher_old.py new file mode 100644 index 00000000..babf7320 --- /dev/null +++ b/reme2/component/file_watcher/base_file_watcher_old.py @@ -0,0 +1,427 @@ +"""Base file watcher with watchfiles integration. + +Per change → single-step pipeline: + added/modified → node, chunks = parser.parse(path) + file_store.upsert(node, chunks) + deleted → file_store.delete_chunks + delete_node + +The parser owns: chunking and edge extraction (text-only, no embedding). +The file_store owns: persistence of node + chunks AND the embedding +pipeline — `upsert` hash-diffs incoming chunks against its persisted +ones, reuses cached embeddings for unchanged blocks, and calls the +embedding API only for new-hash blocks. +The watcher owns: scheduling, cancel-on-modify, retry/timeout, and +startup recovery (re-parsing files whose mtime drifted while offline). +""" + +import asyncio +import os +import time +from pathlib import Path + +from watchfiles import Change, awatch + +from ..base_component import BaseComponent +from ..file_parser import BaseFileParser +from ..file_store import BaseFileStore +from ...enumeration import ComponentEnum + + +class BaseFileWatcher(BaseComponent): + """Watches a directory and feeds the file_store via single-step parse+upsert.""" + + component_type = ComponentEnum.FILE_WATCHER + + _META_DIR = ".reme" + + def __init__( + self, + recursive: bool = False, + debounce: int = 2000, + file_store: str = "default", + default_parser: str | None = None, + rebuild_index_on_start: bool = False, + poll_delay_ms: int = 2000, + parse_max_attempts: int = 3, + parse_retry_backoff: float = 2.0, + parse_task_timeout: float = 300.0, + **kwargs, + ): + super().__init__(**kwargs) + self._file_store_name: str = file_store + self._default_parser_name: str | None = default_parser + self.file_store: BaseFileStore | None = None + self._suffix_to_parser: dict[str, BaseFileParser] = {} + self._default_parser: BaseFileParser | None = None + self.recursive: bool = recursive + self.debounce: int = debounce + self.rebuild_index_on_start: bool = rebuild_index_on_start + self.poll_delay_ms: int = poll_delay_ms + + self._stop_event = asyncio.Event() + self._watch_task: asyncio.Task | None = None + + # Parse pipeline configuration. + self._parse_max_attempts: int = max(1, int(parse_max_attempts)) + self._parse_retry_backoff: float = float(parse_retry_backoff) + self._parse_task_timeout: float = float(parse_task_timeout) + self._tasks: dict[str, asyncio.Task] = {} + self._mgmt_lock: asyncio.Lock | None = None + self._run_lock: asyncio.Lock | None = None + self._last_failure: dict[str, float] = {} + + @property + def watch_path(self) -> str: + """Vault root the watcher is bound to (== file_store.working_dir). + + Available after _start. Reads the live value from the bound + store so config has a single source of truth. + """ + if self.file_store is None or self.file_store.working_dir is None: + raise RuntimeError( + "watch_path requires a started file_store with working_dir set", + ) + return str(self.file_store.working_dir) + + @property + def _meta_path(self) -> Path: + return (Path(self.watch_path) / self._META_DIR).resolve() + + # -- Lifecycle ---------------------------------------------------------- + + async def _start(self): + """Resolve file_store + parsers, sync any disk drift, start watch loop.""" + if self._file_store_name: + assert self.app_context is not None, "app_context must be provided" + + stores = self.app_context.components.get(ComponentEnum.FILE_STORE, {}) + if self._file_store_name not in stores: + raise ValueError(f"File store '{self._file_store_name}' not found.") + store = stores[self._file_store_name] + if not isinstance(store, BaseFileStore): + raise TypeError(f"Expected BaseFileStore, got {type(store).__name__}") + if store.working_dir is None: + raise ValueError( + f"File store '{self._file_store_name}' has no working_dir; " + f"set components.file_store.{self._file_store_name}.working_dir in config", + ) + self.file_store = store + + parsers = self.app_context.components.get(ComponentEnum.FILE_PARSER, {}) + for parser in parsers.values(): + if isinstance(parser, BaseFileParser): + for suffix in parser.suffixes: + self._suffix_to_parser[suffix] = parser + + if self._default_parser_name and self._default_parser_name in parsers: + parser = parsers[self._default_parser_name] + if isinstance(parser, BaseFileParser): + self._default_parser = parser + + if not self._suffix_to_parser and not self._default_parser: + self.logger.warning("No file parsers registered") + + async def _initialize_and_watch(): + await self._initial_sync_and_recovery() + await self._watch_loop() + + self._mgmt_lock = asyncio.Lock() + self._run_lock = asyncio.Lock() + self._last_failure.clear() + self._stop_event.clear() + self._watch_task = asyncio.create_task(_initialize_and_watch()) + self.logger.info(f"Started watching: {self.watch_path}") + + async def _close(self): + """Stop watching and release resources.""" + self._stop_event.set() + if self._watch_task and not self._watch_task.done(): + self._watch_task.cancel() + try: + await self._watch_task + except asyncio.CancelledError: + pass + + # Cancel all in-flight parse tasks before tearing down state. + tasks = list(self._tasks.values()) + for t in tasks: + if not t.done(): + t.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + self._tasks.clear() + self._last_failure.clear() + + self._watch_task = None + self._stop_event.clear() + self._mgmt_lock = None + self._run_lock = None + self.file_store = None + self._suffix_to_parser.clear() + self._default_parser = None + self.logger.info("Stopped watching") + + # -- Startup sync ------------------------------------------------------- + + async def _initial_sync_and_recovery(self) -> None: + """At startup: file_store has its persisted state loaded already; we + diff against current disk for any changes that happened while we + were offline, then re-enqueue them through the parse pipeline.""" + await self._sync_with_disk() + + async def _sync_with_disk(self) -> bool: + """Diff cached graph (in file_store) vs disk state; reindex changes.""" + if self.file_store is None: + return False + + watch_path = Path(self.watch_path) + if not watch_path.exists(): + self.logger.warning(f"Watch path does not exist: {watch_path}") + return False + + on_disk: dict[str, float] = {} + if watch_path.is_file(): + on_disk[str(watch_path.resolve())] = watch_path.stat().st_mtime + elif watch_path.is_dir(): + iterator = watch_path.rglob("*") if self.recursive else watch_path.iterdir() + for fp in iterator: + if not fp.is_file(): + continue + abs_path = str(fp.resolve()) + if not self._watch_filter(abs_path): + continue + if not self._get_parser(abs_path): + continue + try: + on_disk[abs_path] = fp.stat().st_mtime + except OSError: + continue + + cached_paths = set(self.file_store.nodes.keys()) + disk_paths = set(on_disk.keys()) + + added = disk_paths - cached_paths + deleted = cached_paths - disk_paths + modified: set[str] = set() + for p in disk_paths & cached_paths: + cached_node = self.file_store.get_node_by_path(p) + if cached_node is None or cached_node.st_mtime != on_disk[p]: + modified.add(p) + + unchanged = len(disk_paths & cached_paths) - len(modified) + if not (added or deleted or modified): + self.logger.info(f"Index up-to-date ({unchanged} files cached, 0 changes)") + return False + + self.logger.info( + f"Incremental sync: +{len(added)} ~{len(modified)} -{len(deleted)} (unchanged {unchanged})", + ) + + changes: set[tuple[Change, str]] = set() + for p in deleted: + changes.add((Change.deleted, p)) + for p in added: + changes.add((Change.added, p)) + for p in modified: + changes.add((Change.modified, p)) + await self.on_changes(changes) + return True + + # -- Watch loop --------------------------------------------------------- + + async def _interruptible_sleep(self, seconds: float): + try: + await asyncio.wait_for(self._stop_event.wait(), timeout=seconds) + except asyncio.TimeoutError: + pass + + async def _watch_loop(self): + while not self._stop_event.is_set(): + if not Path(self.watch_path).exists(): + self.logger.warning(f"Watch path does not exist: {self.watch_path}, waiting 10 seconds...") + await self._interruptible_sleep(10) + continue + + try: + self.logger.info(f"Starting watch on: {self.watch_path}") + async for changes in awatch( + self.watch_path, + watch_filter=lambda _, p: self._watch_filter(p), + recursive=self.recursive, + debounce=self.debounce, + poll_delay_ms=self.poll_delay_ms, + stop_event=self._stop_event, + ): + if self._stop_event.is_set(): + break + await self.on_changes(changes) + + except FileNotFoundError as e: + self.logger.error(f"Watch path no longer exists: {e}, restarting in 10 seconds...") + if not self._stop_event.is_set(): + await self._interruptible_sleep(10) + + except Exception as e: + self.logger.error(f"Error in watch loop: {e}, restarting in 10 seconds...", exc_info=True) + if not self._stop_event.is_set(): + await self._interruptible_sleep(10) + + def _watch_filter(self, path: str) -> bool: + resolved = Path(path).resolve() + if resolved == self._meta_path: + return False + return self._meta_path not in resolved.parents + + def _get_parser(self, path: str) -> BaseFileParser | None: + suffix = Path(path).suffix.lower() + return self._suffix_to_parser.get(suffix, self._default_parser) + + # -- Change dispatch ---------------------------------------------------- + + async def on_changes(self, changes: set[tuple[Change, str]]) -> None: + if self.file_store is None: + self.logger.warning("File store not initialized, skipping changes") + return + + for change_type, path in changes: + try: + if change_type in (Change.added, Change.modified): + await self._on_modified(path) + elif change_type == Change.deleted: + await self._on_deleted(path) + except FileNotFoundError: + self.logger.warning(f"File not found: {path}, skipping") + except PermissionError: + self.logger.warning(f"Permission denied: {path}, skipping") + except Exception as e: + self.logger.opt(exception=True).error("Error processing {p}: {err}", p=path, err=str(e)) + + async def _on_modified(self, path: str) -> None: + if not Path(path).is_file(): + return + parser = self._get_parser(path) + if parser is None: + self.logger.debug(f"No parser for {path}, skipping") + return + await self._submit_parse_task(path, parser) + + async def _on_deleted(self, path: str) -> None: + # Cancel any in-flight parse task FIRST so a delayed cancellation + # can't race the cleanup writes below. + assert self.file_store is not None, "_on_deleted requires file_store" + targets = [path, *self._descendant_indexed_paths(path)] + for p in targets: + await self._cancel_parse_task(p) + await self.file_store.delete_chunks_by_path(p) + await self.file_store.delete_node_by_path(p) + if len(targets) == 1: + self.logger.info(f"Deleted {path}") + else: + self.logger.info(f"Deleted directory {path} ({len(targets) - 1} indexed children)") + + def _descendant_indexed_paths(self, path: str) -> list[str]: + """Indexed file paths that live beneath `path` as a directory. + + watchfiles emits a single Change.deleted for a removed directory + (no per-file events), so we have to find the orphans ourselves. + Resolves both sides to handle symlinks and trailing-separator + drift between the watcher and the file_store keys. + """ + assert self.file_store is not None + try: + prefix = str(Path(path).resolve()) + os.sep + except OSError: + prefix = path.rstrip(os.sep) + os.sep + return [p for p in self.file_store.nodes if p.startswith(prefix)] + + # -- Parse task pipeline ------------------------------------------------ + + async def _submit_parse_task(self, path: str, parser: BaseFileParser) -> None: + """Cancel any prior task + spawn a new one. Atomic under _mgmt_lock.""" + assert self._mgmt_lock is not None, "_submit_parse_task before _start()" + async with self._mgmt_lock: + await self._cancel_locked(path) + self._tasks[path] = asyncio.create_task( + self._run_parse_task(path, parser), + ) + + async def _cancel_parse_task(self, path: str) -> None: + """Cancel + await any in-flight parse task for `path`.""" + assert self._mgmt_lock is not None, "_cancel_parse_task before _start()" + async with self._mgmt_lock: + await self._cancel_locked(path) + + async def _cancel_locked(self, path: str) -> None: + """Cancel + await the task for `path`. Caller holds `_mgmt_lock`.""" + task = self._tasks.pop(path, None) + if task is not None and not task.done(): + task.cancel() + try: + await task + except (asyncio.CancelledError, Exception): + pass + + async def flush_parse_tasks(self) -> None: + """Wait until all parse tasks finish (loops because retries spawn new ones).""" + while self._tasks: + await asyncio.gather(*list(self._tasks.values()), return_exceptions=True) + + def pending_parse_count(self) -> int: + return sum(1 for t in self._tasks.values() if not t.done()) + + def failed_parse_paths(self) -> dict[str, float]: + return dict(self._last_failure) + + async def _run_parse_task(self, path: str, parser: BaseFileParser) -> None: + """One per-path parse task: parse → upsert. + + The store owns hash-diff + embedding: it inspects its own + existing chunks for `path`, reuses cached embeddings for + unchanged blocks, and only calls the embedding API for new-hash + blocks. The whole upsert is one call from this side. + """ + try: + for attempt in range(1, self._parse_max_attempts + 1): + try: + assert self.file_store is not None + node, chunks = await asyncio.wait_for( + parser.parse(path), + timeout=self._parse_task_timeout, + ) + + assert self._run_lock is not None + async with self._run_lock: + await asyncio.wait_for( + self.file_store.upsert(node, chunks), + timeout=self._parse_task_timeout, + ) + self._last_failure.pop(path, None) + + self.logger.info( + f"indexed {path}: chunks={len(chunks)} edges={len(node.edges)}", + ) + return + except asyncio.CancelledError: + raise + except FileNotFoundError: + self.logger.debug(f"parse task {path}: file gone, skipping") + self._last_failure.pop(path, None) + return + except Exception as e: + self._last_failure[path] = time.time() + if attempt < self._parse_max_attempts: + backoff = self._parse_retry_backoff * attempt + self.logger.warning( + f"parse task {path} attempt {attempt} failed " + f"({type(e).__name__}: {e}); retry in {backoff:.1f}s", + ) + await asyncio.sleep(backoff) + continue + self.logger.error( + f"parse task {path} giving up after {attempt} attempts: " + f"{type(e).__name__}: {e}", + ) + return + finally: + if self._tasks.get(path) is asyncio.current_task(): + self._tasks.pop(path, None) diff --git a/reme2/component/file_watcher/light_file_watcher.py b/reme2/component/file_watcher/light_file_watcher.py index bebcc498..d0115b89 100644 --- a/reme2/component/file_watcher/light_file_watcher.py +++ b/reme2/component/file_watcher/light_file_watcher.py @@ -1,16 +1,7 @@ -"""Lightweight file watcher for Markdown files only.""" - -from pathlib import Path - from .base_file_watcher import BaseFileWatcher from ..component_registry import R -_MD_SUFFIXES = {".md", ".markdown"} - @R.register("light") class LightFileWatcher(BaseFileWatcher): - """Watches only Markdown files and delegates parsing to registered parsers.""" - - def _watch_filter(self, path: str) -> bool: - return super()._watch_filter(path) and Path(path).suffix.lower() in _MD_SUFFIXES + ... diff --git a/reme2/mcp/steps/memory_retriever.py b/reme2/mcp/steps/memory_retriever.py index a5c5f5e7..961d73ce 100644 --- a/reme2/mcp/steps/memory_retriever.py +++ b/reme2/mcp/steps/memory_retriever.py @@ -84,7 +84,7 @@ def _serialize_chunk(chunk, file_store, extras: dict | None = None) -> dict: `extras` lets the caller attach step-specific fields (e.g. `graph_hop`). """ item = chunk.model_dump(exclude_none=True, exclude={"embedding"}) - node = file_store.get_node(chunk.path) + node = file_store.get_node_by_path(chunk.path) if node is not None: item["file_metadata"] = node.metadata item["file_st_mtime"] = node.st_mtime diff --git a/reme2/memory/memory_io.py b/reme2/memory/memory_io.py index 632b70ef..44781393 100644 --- a/reme2/memory/memory_io.py +++ b/reme2/memory/memory_io.py @@ -62,7 +62,7 @@ async def get_file( On-disk frontmatter is the source of truth — the file_store cache may lag a write that hasn't been picked up by the watcher yet. """ - node = file_store.get_node(path) + node = file_store.get_node_by_path(path) result: dict = {"path": path, "exists": False} if node is not None: result.update({ @@ -80,7 +80,7 @@ async def get_file( result["metadata"] = dict(post.metadata) if include_chunks: - chunks = await file_store.get_chunks(path) + chunks = await file_store.get_chunks_by_path(path) result["chunks"] = [c.model_dump(exclude_none=True) for c in chunks] return result diff --git a/reme2/schema/file_edge.py b/reme2/schema/file_edge.py index 16ca2d90..f8866572 100644 --- a/reme2/schema/file_edge.py +++ b/reme2/schema/file_edge.py @@ -10,10 +10,7 @@ maintainer's job, not the schema's. ## Inline forms recognised by `parse_wikilinks` - [[X]] bare wikilink → predicate=None - extends:: [[X]] line-level Dataview → predicate="extends" - [extends:: [[X]]] inline-bracketed → predicate="extends" - extends:: [[A]], [[B]] multi-target → 2 edges + """ from __future__ import annotations diff --git a/reme2/schema/file_node.py b/reme2/schema/file_node.py index 73d34dc4..1a1eab91 100644 --- a/reme2/schema/file_node.py +++ b/reme2/schema/file_node.py @@ -2,6 +2,12 @@ from pydantic import BaseModel, Field, ConfigDict class FileEdge(BaseModel): + """ Format: + [[X]] bare wikilink → predicate=None + extends:: [[X]] line-level Dataview → predicate="extends" + [extends:: [[X]]] inline-bracketed → predicate="extends" + extends:: [[A]], [[B]] multi-target → 2 edges + """ link: str = Field(default=...) predicate: str | None = Field(default=None) @@ -31,4 +37,5 @@ class FileNode(BaseModel): path: str = Field(default=...) st_mtime: float = Field(default=...) edges: list[FileEdge] = Field(default_factory=list) + chunk_ids: list[str] = Field(default_factory=list) front_matter: FileFrontMatter = Field(default_factory=FileFrontMatter) diff --git a/reme2/utils/similarity_utils.py b/reme2/utils/similarity_utils.py index 8c7031a0..65741c45 100644 --- a/reme2/utils/similarity_utils.py +++ b/reme2/utils/similarity_utils.py @@ -63,14 +63,6 @@ def batch_cosine_similarity(nd_array1: np.ndarray, nd_array2: np.ndarray) -> np. Raises: ValueError: If embedding dimensions don't match between arrays. - - Examples: - >>> import numpy as np - >>> arr1 = np.array([[1.0, 0.0], [0.0, 1.0]]) - >>> arr2 = np.array([[1.0, 0.0], [1.0, 1.0]]) - >>> batch_cosine_similarity(arr1, arr2) - array([[1. , 0.70710678], - [0. , 0.70710678]]) """ if nd_array1.shape[1] != nd_array2.shape[1]: raise ValueError( diff --git a/test/reme2/test_bm25_lite.py b/test/reme2/test_bm25_lite.py index c242358b..a04a950a 100644 --- a/test/reme2/test_bm25_lite.py +++ b/test/reme2/test_bm25_lite.py @@ -114,7 +114,7 @@ def test_retrieve_basic(): } bm25.add_docs(docs) - results = bm25.retrieve("python", k=3) + results = bm25.retrieve("python", limit=3) assert len(results) <= 3 assert "doc1" in results or "doc3" in results @@ -136,10 +136,10 @@ def test_retrieve_with_limit(): } bm25.add_docs(docs) - results = bm25.retrieve("python", k=3) + results = bm25.retrieve("python", limit=3) assert len(results) == 3 - results = bm25.retrieve("python", k=5) + results = bm25.retrieve("python", limit=5) assert len(results) == 5 await bm25.close() @@ -158,10 +158,10 @@ def test_retrieve_empty_query(): docs = {"doc1": "hello world"} bm25.add_docs(docs) - results = bm25.retrieve("", k=3) + results = bm25.retrieve("", limit=3) assert results == {} - results = bm25.retrieve("unknownxyz", k=3) + results = bm25.retrieve("unknownxyz", limit=3) assert results == {} await bm25.close() @@ -177,7 +177,7 @@ def test_retrieve_empty_index(): with tempfile.TemporaryDirectory() as tmpdir: bm25 = await create_bm25(Path(tmpdir)) - results = bm25.retrieve("python", k=3) + results = bm25.retrieve("python", limit=3) assert results == {} await bm25.close() @@ -200,7 +200,7 @@ def test_update_doc(): assert bm25.n_docs == 1 assert bm25.total_len != old_len - results = bm25.retrieve("java", k=1) + results = bm25.retrieve("java", limit=1) assert "doc1" in results await bm25.close() @@ -227,7 +227,7 @@ def test_remove_doc(): assert bm25.n_docs == 1 assert "doc1" not in bm25.doc_meta - results = bm25.retrieve("hello", k=2) + results = bm25.retrieve("hello", limit=2) assert "doc1" not in results assert "doc2" in results @@ -325,7 +325,7 @@ def test_reindex_with_docs(): assert "doc2" in bm25.doc_meta assert len(bm25.vocab) < len(old_vocab) - results = bm25.retrieve("hello", k=1) + results = bm25.retrieve("hello", limit=1) assert "doc2" in results await bm25.close() @@ -362,7 +362,7 @@ def test_persistence(): for doc_id, meta in old_doc_meta.items(): assert doc_id in bm25_new.doc_meta - results = bm25_new.retrieve("hello", k=2) + results = bm25_new.retrieve("hello", limit=2) assert "doc1" in results or "doc2" in results await bm25_new.close() @@ -382,7 +382,7 @@ def test_custom_params(): assert bm25.b == 0.5 bm25.add_docs({"doc1": "test document"}) - results = bm25.retrieve("test", k=1) + results = bm25.retrieve("test", limit=1) assert "doc1" in results await bm25.close() @@ -405,7 +405,7 @@ def test_chinese_text(): } bm25.add_docs(docs) - results = bm25.retrieve("北京", k=2) + results = bm25.retrieve("北京", limit=2) assert len(results) <= 2 assert "doc1" in results or "doc2" in results @@ -429,10 +429,10 @@ def test_mixed_chinese_english(): } bm25.add_docs(docs) - results = bm25.retrieve("Python", k=3) + results = bm25.retrieve("Python", limit=3) assert len(results) > 0 - results = bm25.retrieve("编程", k=2) + results = bm25.retrieve("编程", limit=2) assert len(results) > 0 await bm25.close() @@ -503,7 +503,7 @@ def test_score_ordering(): } bm25.add_docs(docs) - results = bm25.retrieve("python", k=3) + results = bm25.retrieve("python", limit=3) scores = list(results.values()) for i in range(len(scores) - 1): diff --git a/test/test_cache_memory_usage.py b/test/test_cache_memory_usage.py index ed0020f9..27369dfe 100644 --- a/test/test_cache_memory_usage.py +++ b/test/test_cache_memory_usage.py @@ -158,10 +158,7 @@ def print_usage_guidelines(): print(" • ~420 MB per 50,000 entries") print(" • Scale linearly: ~8.4 MB per 1,000 entries\n") - print("3️⃣ Monitor cache statistics:") - print(" • Use model.get_cache_stats() to check hit rate") - print(" • If hit rate < 50%, consider reducing cache size") - print(" • If hit rate > 90%, you might benefit from larger cache\n") + print("3️⃣ Monitor cache size:\n") print("4️⃣ Configuration examples:\n") @@ -193,10 +190,8 @@ def print_usage_guidelines(): print(" max_cache_size=1000, # ~8.4 MB") print(" )") print() - print(" # Monitor cache performance") - print(" stats = model.get_cache_stats()") - print(" print(f'Hit rate: {stats[\"hit_rate\"]:.1%}')") - print(" print(f'Memory used: ~{stats[\"cache_size\"] * 8.4:.1f} KB')") + print(" # Monitor cache") + print(" print(f'Cache entries: {len(model._embedding_cache)}')") if __name__ == "__main__": diff --git a/test/test_embedding_cache.py b/test/test_embedding_cache.py deleted file mode 100644 index 5cc27e53..00000000 --- a/test/test_embedding_cache.py +++ /dev/null @@ -1,427 +0,0 @@ -""" -Async unit tests for embedding cache functionality. - -Tests cover: -- Cache hit/miss tracking -- LRU eviction policy -- Cache statistics -- Performance improvements with repeated queries -- Cache clearing - -Usage: - python test_embedding_cache.py -""" - -# flake8: noqa: E402 -# pylint: disable=C0413 - -import asyncio -import shutil -import tempfile -from typing import List - -from reme.core.utils import load_env - -load_env() - -from reme.core.embedding import OpenAIEmbeddingModel - - -def get_test_texts() -> List[str]: - """Create test texts for embedding cache testing.""" - return [ - "What is machine learning?", - "How does neural network work?", - "Explain artificial intelligence", - "Define deep learning", - "What is data science?", - ] - - -async def test_cache_basic_functionality(): - """Test basic cache hit/miss functionality.""" - print(f"\n{'='*60}") - print("Test 1: Basic Cache Functionality") - print(f"{'='*60}") - - temp_dir = tempfile.mkdtemp() - try: - model = OpenAIEmbeddingModel( - model_name="text-embedding-v4", - dimensions=1024, - max_cache_size=100, - max_retries=2, - raise_exception=True, - cache_dir=temp_dir, - ) - - test_text = "Hello, this is a test sentence for embedding cache." - - print(f"Input text: {test_text}") - print(f"Cache directory: {temp_dir}") - - # First call - should be a cache miss - print("\n1️⃣ First embedding call (cold cache):") - embedding1 = await model.get_embedding(test_text) - stats1 = model.get_cache_stats() - - print(f" Embedding dimension: {len(embedding1)}") - print(f" Cache size: {stats1['cache_size']}") - print(f" Cache hits: {stats1['cache_hits']}") - print(f" Cache misses: {stats1['cache_misses']}") - print(f" Hit rate: {stats1['hit_rate']:.2%}") - - assert len(embedding1) == 1024, "Embedding dimension mismatch" - assert stats1["cache_misses"] == 1, "Should have 1 cache miss" - assert stats1["cache_hits"] == 0, "Should have 0 cache hits" - assert stats1["cache_size"] == 1, "Cache should have 1 entry" - - # Second call with same text - should be a cache hit - print("\n2️⃣ Second embedding call (same text):") - embedding2 = await model.get_embedding(test_text) - stats2 = model.get_cache_stats() - - print(f" Cache hits: {stats2['cache_hits']}") - print(f" Cache misses: {stats2['cache_misses']}") - print(f" Hit rate: {stats2['hit_rate']:.2%}") - - assert embedding1 == embedding2, "Cached embedding should be identical" - assert stats2["cache_hits"] == 1, "Should have 1 cache hit" - assert stats2["cache_misses"] == 1, "Should still have 1 cache miss" - assert stats2["hit_rate"] == 0.5, "Hit rate should be 50%" - - await model.close() - print("\n✓ PASSED: Basic cache functionality works correctly") - finally: - shutil.rmtree(temp_dir, ignore_errors=True) - - -async def test_batch_cache_efficiency(): - """Test cache efficiency with batch embeddings including duplicates.""" - print(f"\n{'='*60}") - print("Test 2: Batch Cache Efficiency") - print(f"{'='*60}") - - temp_dir = tempfile.mkdtemp() - try: - model = OpenAIEmbeddingModel( - model_name="text-embedding-v4", - dimensions=1024, - max_cache_size=1000, - max_retries=2, - raise_exception=True, - cache_dir=temp_dir, - ) - - texts = get_test_texts() - - # Create a list with duplicates - texts_with_duplicates = texts + texts[:3] # 5 unique + 3 duplicates = 8 total - - print(f"Processing {len(texts_with_duplicates)} texts (5 unique + 3 duplicates)") - - # First batch - print("\n1️⃣ First batch (cold cache):") - embeddings1 = await model.get_embeddings(texts) - stats1 = model.get_cache_stats() - - print(f" Embeddings generated: {len(embeddings1)}") - print(f" Cache size: {stats1['cache_size']}") - print(f" Cache misses: {stats1['cache_misses']}") - print(f" Cache hits: {stats1['cache_hits']}") - - assert len(embeddings1) == len(texts), "Embeddings count mismatch" - assert stats1["cache_size"] == len(texts), f"Cache should have {len(texts)} entries" - assert stats1["cache_misses"] == len(texts), "All should be cache misses" - - # Second batch with duplicates - print("\n2️⃣ Second batch (with duplicates):") - embeddings2 = await model.get_embeddings(texts_with_duplicates) - stats2 = model.get_cache_stats() - - print(f" Embeddings generated: {len(embeddings2)}") - print(f" Cache hits: {stats2['cache_hits']}") - print(f" Cache misses: {stats2['cache_misses']}") - print(f" Hit rate: {stats2['hit_rate']:.2%}") - - assert len(embeddings2) == len(texts_with_duplicates), "Embeddings count mismatch" - assert stats2["cache_hits"] >= 3, "Should have at least 3 cache hits from duplicates" - - # Verify embeddings are identical for duplicated texts - for i in range(3): - assert embeddings2[i] == embeddings2[len(texts) + i], f"Duplicate {i} should have identical embedding" - - await model.close() - print("\n✓ PASSED: Batch cache efficiently handles duplicates") - finally: - shutil.rmtree(temp_dir, ignore_errors=True) - - -async def test_cache_lru_eviction(): - """Test LRU cache eviction policy.""" - print(f"\n{'='*60}") - print("Test 3: LRU Cache Eviction") - print(f"{'='*60}") - - temp_dir = tempfile.mkdtemp() - try: - # Create model with small cache size - model = OpenAIEmbeddingModel( - model_name="text-embedding-v4", - dimensions=1024, - max_cache_size=3, # Small cache for testing eviction - max_retries=2, - raise_exception=True, - cache_dir=temp_dir, - ) - - texts = get_test_texts()[:5] # Use 5 texts, cache size is 3 - - print(f"Cache size limit: {model.max_cache_size}") - print(f"Number of unique texts: {len(texts)}") - - # Fill cache beyond capacity - print("\n1️⃣ Filling cache with 5 texts (capacity = 3):") - for i, text in enumerate(texts): - await model.get_embedding(text) - stats = model.get_cache_stats() - print( - f" After text {i+1}: cache_size={stats['cache_size']}, " - f"hits={stats['cache_hits']}, misses={stats['cache_misses']}", - ) - - final_stats = model.get_cache_stats() - assert final_stats["cache_size"] <= 3, "Cache size should not exceed max_cache_size" - assert final_stats["cache_misses"] == 5, "Should have 5 cache misses for 5 unique texts" - - # Access the most recent entries - should be cache hits - print("\n2️⃣ Accessing recent entries (should be cached):") - recent_texts = texts[-3:] # Last 3 texts should still be in cache - - for i, text in enumerate(recent_texts): - await model.get_embedding(text) - stats = model.get_cache_stats() - print(f" Text {len(texts) - 3 + i + 1}: hits={stats['cache_hits']}") - - final_stats = model.get_cache_stats() - assert final_stats["cache_hits"] == 3, "Should have 3 cache hits for recent entries" - - # Access oldest entries - should be cache misses (evicted) - print("\n3️⃣ Accessing oldest entries (should be evicted):") - old_texts = texts[:2] # First 2 texts should have been evicted - - before_misses = final_stats["cache_misses"] - for i, text in enumerate(old_texts): - await model.get_embedding(text) - stats = model.get_cache_stats() - print(f" Text {i + 1}: misses={stats['cache_misses']}") - - final_stats = model.get_cache_stats() - assert final_stats["cache_misses"] == before_misses + 2, "Should have 2 more cache misses for evicted entries" - - await model.close() - print("\n✓ PASSED: LRU eviction works correctly") - finally: - shutil.rmtree(temp_dir, ignore_errors=True) - - -async def test_cache_stats_and_clear(): - """Test cache statistics tracking and clearing.""" - print(f"\n{'='*60}") - print("Test 4: Cache Statistics and Clearing") - print(f"{'='*60}") - - temp_dir = tempfile.mkdtemp() - try: - model = OpenAIEmbeddingModel( - model_name="text-embedding-v4", - dimensions=1024, - max_cache_size=100, - max_retries=2, - raise_exception=True, - cache_dir=temp_dir, - ) - - texts = get_test_texts() - - # Generate some cache activity - print("\n1️⃣ Generating cache activity:") - await model.get_embeddings(texts) - await model.get_embeddings(texts[:3]) # Repeat first 3 - - stats = model.get_cache_stats() - print(f" Cache size: {stats['cache_size']}") - print(f" Max cache size: {stats['max_cache_size']}") - print(f" Cache hits: {stats['cache_hits']}") - print(f" Cache misses: {stats['cache_misses']}") - print(f" Hit rate: {stats['hit_rate']:.2%}") - - assert stats["cache_size"] > 0, "Cache should not be empty" - assert stats["cache_hits"] >= 3, "Should have at least 3 cache hits" - assert "hit_rate" in stats, "Stats should include hit_rate" - - # Clear cache - print("\n2️⃣ Clearing cache:") - model.clear_cache() - stats_after_clear = model.get_cache_stats() - - print(f" Cache size after clear: {stats_after_clear['cache_size']}") - print(f" Hits after clear: {stats_after_clear['cache_hits']}") - print(f" Misses after clear: {stats_after_clear['cache_misses']}") - print(f" Hit rate after clear: {stats_after_clear['hit_rate']:.2%}") - - assert stats_after_clear["cache_size"] == 0, "Cache should be empty after clear" - assert stats_after_clear["cache_hits"] == 0, "Hits should be reset" - assert stats_after_clear["cache_misses"] == 0, "Misses should be reset" - assert stats_after_clear["hit_rate"] == 0.0, "Hit rate should be 0" - - await model.close() - print("\n✓ PASSED: Cache statistics and clearing work correctly") - finally: - shutil.rmtree(temp_dir, ignore_errors=True) - - -async def test_cache_disabled(): - """Test behavior when cache is disabled (max_cache_size=0).""" - print(f"\n{'='*60}") - print("Test 5: Cache Disabled") - print(f"{'='*60}") - - temp_dir = tempfile.mkdtemp() - try: - model = OpenAIEmbeddingModel( - model_name="text-embedding-v4", - dimensions=1024, - max_cache_size=0, # Disable cache - max_retries=2, - raise_exception=True, - cache_dir=temp_dir, - ) - - test_text = "Test text with cache disabled" - - print(f"Cache size limit: {model.max_cache_size} (disabled)") - print(f"Input text: {test_text}") - - # Call twice with same text - print("\n1️⃣ First call:") - embedding1 = await model.get_embedding(test_text) - stats1 = model.get_cache_stats() - print(f" Cache size: {stats1['cache_size']}") - print(f" Cache misses: {stats1['cache_misses']}") - - print("\n2️⃣ Second call (same text):") - embedding2 = await model.get_embedding(test_text) - stats2 = model.get_cache_stats() - print(f" Cache size: {stats2['cache_size']}") - print(f" Cache misses: {stats2['cache_misses']}") - print(f" Cache hits: {stats2['cache_hits']}") - - assert stats2["cache_size"] == 0, "Cache should remain empty when disabled" - assert stats2["cache_misses"] == 2, "Both calls should be cache misses" - assert stats2["cache_hits"] == 0, "Should have no cache hits when disabled" - assert embedding1 == embedding2, "Embeddings should still be consistent" - - await model.close() - print("\n✓ PASSED: Cache correctly disabled when max_cache_size=0") - finally: - shutil.rmtree(temp_dir, ignore_errors=True) - - -async def test_cache_performance_demo(): - """Demonstrate cache performance improvements.""" - print(f"\n{'='*60}") - print("Test 6: Cache Performance Demo") - print(f"{'='*60}") - - temp_dir = tempfile.mkdtemp() - try: - model = OpenAIEmbeddingModel( - model_name="text-embedding-v4", - dimensions=1024, - max_cache_size=1000, - max_retries=2, - raise_exception=True, - cache_dir=temp_dir, - ) - - texts = get_test_texts() - - # Create a realistic workload with many repeated queries - workload = texts * 3 # 15 queries total, 5 unique - - print(f"\nProcessing {len(workload)} queries ({len(texts)} unique texts)") - print("This simulates a realistic scenario with repeated queries\n") - - # Process all queries - for i, text in enumerate(workload, 1): - await model.get_embedding(text) - if i % 5 == 0: # Report every 5 queries - stats = model.get_cache_stats() - print( - f"After {i:2d} queries: hits={stats['cache_hits']:2d}, " - f"misses={stats['cache_misses']:2d}, " - f"hit_rate={stats['hit_rate']:5.1%}", - ) - - final_stats = model.get_cache_stats() - total_requests = final_stats["cache_hits"] + final_stats["cache_misses"] - - print(f"\n{'─'*60}") - print("📊 Final Statistics:") - print(f"{'─'*60}") - print(f" Total queries: {total_requests}") - print(f" Unique texts: {len(texts)}") - print(f" Cache hits: {final_stats['cache_hits']}") - print(f" Cache misses: {final_stats['cache_misses']}") - print(f" Hit rate: {final_stats['hit_rate']:.1%}") - print(f" Cache size: {final_stats['cache_size']}/{final_stats['max_cache_size']}") - print(f"{'─'*60}") - print( - f"💰 API calls saved: {final_stats['cache_hits']} out of {total_requests} " - f"({final_stats['cache_hits']/total_requests*100:.1f}%)", - ) - print(f"{'─'*60}") - - assert final_stats["cache_hits"] == 10, "Should have 10 cache hits (2 repeats × 5 texts)" - assert final_stats["cache_misses"] == 5, "Should have 5 cache misses (5 unique texts)" - assert final_stats["hit_rate"] > 0.6, "Hit rate should be > 60%" - - await model.close() - print("\n✓ PASSED: Cache provides significant performance improvement") - finally: - shutil.rmtree(temp_dir, ignore_errors=True) - - -async def main(): - """Run all cache tests.""" - print("\n" + "#" * 60) - print("# EMBEDDING CACHE TESTS") - print("#" * 60) - - try: - await test_cache_basic_functionality() - await test_batch_cache_efficiency() - await test_cache_lru_eviction() - await test_cache_stats_and_clear() - await test_cache_disabled() - await test_cache_performance_demo() - - print("\n" + "=" * 60) - print("✅ ALL CACHE TESTS PASSED") - print("=" * 60) - print("\nKey takeaways:") - print(" • Cache correctly tracks hits/misses") - print(" • LRU eviction works as expected") - print(" • Duplicate queries are efficiently cached") - print(" • Cache can be disabled or cleared") - print(" • Significant performance improvement with realistic workloads") - print("=" * 60 + "\n") - - except Exception as e: - print(f"\n✗ TEST FAILED: {type(e).__name__}: {e}") - raise - - -if __name__ == "__main__": - asyncio.run(main())