diff --git a/reme2/component/__init__.py b/reme2/component/__init__.py index b7dd0e95..6b3942ec 100644 --- a/reme2/component/__init__.py +++ b/reme2/component/__init__.py @@ -2,10 +2,10 @@ from . import as_llm from . import as_llm_formatter +from . import chunk_store from . import client from . import embedding from . import file_parser -from . import file_store from . import file_watcher from . import job from . import service @@ -27,10 +27,10 @@ __all__ = [ # base components "as_llm", "as_llm_formatter", + "chunk_store", "client", "embedding", "file_parser", - "file_store", "file_watcher", "job", "service", diff --git a/reme2/component/base_step.py b/reme2/component/base_step.py index 3d7bb173..52cb699e 100644 --- a/reme2/component/base_step.py +++ b/reme2/component/base_step.py @@ -8,11 +8,12 @@ from agentscope.model import ChatModelBase from agentscope.token import TokenCounterBase from .base_component import BaseComponent +from .chunk_store import BaseChunkStore from .embedding import BaseEmbeddingModel -from .file_store import BaseFileStore from .prompt_handler import PromptHandler from .runtime_context import RuntimeContext from ..enumeration import ComponentEnum +from ..schema.file_graph import FileGraph class BaseStep(BaseComponent): @@ -83,9 +84,15 @@ class BaseStep(BaseComponent): "token_counter") @property - def file_store(self) -> BaseFileStore: - name = self.kwargs.get("file_store", "default") - return name if isinstance(name, BaseFileStore) else self._get_component(ComponentEnum.FILE_STORE, name) + def chunk_store(self) -> BaseChunkStore: + name = self.kwargs.get("chunk_store", "default") + return name if isinstance(name, BaseChunkStore) else self._get_component(ComponentEnum.CHUNK_STORE, name) + + @property + def file_graph(self) -> FileGraph: + name = self.kwargs.get("file_watcher", "default") + watcher = self._get_component(ComponentEnum.FILE_WATCHER, name) + return watcher.file_graph @property def embedding(self) -> BaseEmbeddingModel: diff --git a/reme2/component/chunk_store/__init__.py b/reme2/component/chunk_store/__init__.py new file mode 100644 index 00000000..0eb1470f --- /dev/null +++ b/reme2/component/chunk_store/__init__.py @@ -0,0 +1,17 @@ +"""Chunk store module. + +Storage backends for FileChunks with vector and full-text search. +File metadata is managed by FileGraph, not by ChunkStore. +""" + +from .base_chunk_store import BaseChunkStore +from .chroma_chunk_store import ChromaChunkStore +from .local_chunk_store import LocalChunkStore +from .sqlite_chunk_store import SqliteChunkStore + +__all__ = [ + "BaseChunkStore", + "ChromaChunkStore", + "LocalChunkStore", + "SqliteChunkStore", +] diff --git a/reme2/component/file_store/base_file_store.py b/reme2/component/chunk_store/base_chunk_store.py similarity index 76% rename from reme2/component/file_store/base_file_store.py rename to reme2/component/chunk_store/base_chunk_store.py index d1981fb1..834584da 100644 --- a/reme2/component/file_store/base_file_store.py +++ b/reme2/component/chunk_store/base_chunk_store.py @@ -1,4 +1,4 @@ -"""Abstract base class for file storage backends.""" +"""Abstract base class for chunk storage backends.""" import re from abc import abstractmethod @@ -7,18 +7,18 @@ from pathlib import Path from ..base_component import BaseComponent from ..embedding import BaseEmbeddingModel from ...enumeration import ComponentEnum -from ...schema import FileChunk, FileMetadata, SearchFilter +from ...schema import ChunkFilter, FileChunk -class BaseFileStore(BaseComponent): - """Abstract base class for file storage backends. +class BaseChunkStore(BaseComponent): + """Abstract base class for chunk storage backends. - Provides embedding resolution, validation, safe embedding retrieval, - metadata caching, hybrid search, and keyword scoring utilities. - Subclasses must implement the storage-specific CRUD and search methods. + Handles chunk persistence and retrieval (vector / keyword / hybrid search). + File-level metadata and search filter resolution live in FileGraph; this + layer only consumes a compiled ChunkFilter (path set) for restricting search. """ - component_type = ComponentEnum.FILE_STORE + component_type = ComponentEnum.CHUNK_STORE def __init__( self, @@ -45,7 +45,6 @@ class BaseFileStore(BaseComponent): raise ValueError("At least one of embedding_model or fts_enabled must be set.") async def _start(self): - """Resolve embedding model from app_context.""" if not self._embedding_model_name: return assert self.app_context is not None, "app_context must be provided" @@ -139,7 +138,7 @@ class BaseFileStore(BaseComponent): limit: int, vector_weight: float = 0.7, candidate_multiplier: float = 3.0, - search_filter: SearchFilter | None = None, + chunk_filter: ChunkFilter | None = None, ) -> list[FileChunk]: """Perform hybrid search combining vector and keyword results.""" assert 0.0 <= vector_weight <= 1.0 @@ -148,8 +147,8 @@ class BaseFileStore(BaseComponent): text_weight = 1.0 - vector_weight if self.vector_enabled and self.fts_enabled: - keyword_results = await self.keyword_search(query, candidates, search_filter) - vector_results = await self.vector_search(query, candidates, search_filter) + keyword_results = await self.keyword_search(query, candidates, chunk_filter) + vector_results = await self.vector_search(query, candidates, chunk_filter) if not keyword_results: return vector_results[:limit] @@ -164,9 +163,9 @@ class BaseFileStore(BaseComponent): ) return merged[:limit] elif self.vector_enabled: - return await self.vector_search(query, limit, search_filter) + return await self.vector_search(query, limit, chunk_filter) elif self.fts_enabled: - return await self.keyword_search(query, limit, search_filter) + return await self.keyword_search(query, limit, chunk_filter) return [] @staticmethod @@ -199,24 +198,11 @@ class BaseFileStore(BaseComponent): # -- Filter utility ----------------------------------------------------- - def _apply_filter( - self, - chunks: list[FileChunk], - search_filter: SearchFilter | None, - file_metadata: dict[str, FileMetadata] | None = None, - ) -> list[FileChunk]: - """Apply search filter to a list of chunks. - - Args: - chunks: Candidate chunks to filter. - search_filter: Filter conditions. - file_metadata: File-level metadata lookup (path -> FileMetadata). - Used for tag filtering since tags are file-level, not chunk-level. - """ - if not search_filter or search_filter.is_empty(): + @staticmethod + def _apply_filter(chunks: list[FileChunk], chunk_filter: ChunkFilter | None) -> list[FileChunk]: + if chunk_filter is None or chunk_filter.resolved_paths is None: return chunks - fm = file_metadata or {} - return [c for c in chunks if search_filter.match(c.path, fm[c.path].metadata if c.path in fm else None)] + return [c for c in chunks if chunk_filter.match_path(c.path)] # -- Abstract methods --------------------------------------------------- @@ -225,34 +211,31 @@ class BaseFileStore(BaseComponent): """Clear all indexed data.""" @abstractmethod - async def upsert_file(self, file_meta: FileMetadata, chunks: list[FileChunk]): - """Insert or update a file and its chunks.""" + async def upsert_chunks(self, path: str, chunks: list[FileChunk]): + """Insert or update all chunks for a file path.""" @abstractmethod - async def delete_file(self, path: str): - """Delete a file and all its chunks.""" + async def delete_chunks(self, path: str): + """Delete all chunks for a file path.""" @abstractmethod - async def list_files(self) -> list[str]: - """List all indexed file paths.""" + async def get_chunks(self, path: str) -> list[FileChunk]: + """Get all chunks for a file path.""" @abstractmethod - async def get_file_chunks(self, path: str) -> list[FileChunk]: - """Get all chunks for a file.""" - - @abstractmethod - async def get_file_metadata(self, path: str) -> FileMetadata | None: - """Get metadata for a specific file.""" - - @abstractmethod - async def vector_search(self, query: str, limit: int, search_filter: SearchFilter | None = None) -> list[FileChunk]: - """Perform vector similarity search.""" + async def vector_search( + self, + query: str, + limit: int, + chunk_filter: ChunkFilter | None = None, + ) -> list[FileChunk]: + """Perform vector similarity search, optionally restricted by chunk_filter.""" @abstractmethod async def keyword_search( self, query: str, limit: int, - search_filter: SearchFilter | None = None, + chunk_filter: ChunkFilter | None = None, ) -> list[FileChunk]: - """Perform full-text/keyword search.""" + """Perform full-text/keyword search, optionally restricted by chunk_filter.""" diff --git a/reme2/component/file_store/chroma_file_store.py b/reme2/component/chunk_store/chroma_chunk_store.py similarity index 53% rename from reme2/component/file_store/chroma_file_store.py rename to reme2/component/chunk_store/chroma_chunk_store.py index 35df0b7d..44c17966 100644 --- a/reme2/component/file_store/chroma_file_store.py +++ b/reme2/component/chunk_store/chroma_chunk_store.py @@ -1,12 +1,10 @@ -"""ChromaDB storage backend for file store.""" +"""ChromaDB chunk storage backend.""" -import json import time -from pathlib import Path -from .base_file_store import BaseFileStore +from .base_chunk_store import BaseChunkStore from ..component_registry import R -from ...schema import FileChunk, FileMetadata, SearchFilter +from ...schema import ChunkFilter, FileChunk try: import chromadb @@ -20,12 +18,11 @@ except Exception as e: @R.register("chroma") -class ChromaFileStore(BaseFileStore): - """ChromaDB file storage with vector and full-text search. +class ChromaChunkStore(BaseChunkStore): + """ChromaDB chunk storage with vector and full-text search. Uses ChromaDB's native vector search and `where_document` $contains - for keyword matching. File metadata is persisted to a JSON file - alongside the ChromaDB database. + for keyword matching. """ def __init__(self, **kwargs): @@ -34,42 +31,11 @@ class ChromaFileStore(BaseFileStore): super().__init__(**kwargs) self.client: "chromadb.ClientAPI | None" = None self.chunks_collection: "chromadb.Collection | None" = None - self._metadata_file: Path = self.db_path / f"{self.store_name}_file_metadata.json" - self._metadata_cache: dict[str, FileMetadata] = {} @property def collection_name(self) -> str: return f"chunks_{self.store_name}" - # -- Persistence helpers ------------------------------------------------ - - async def _load_metadata(self) -> None: - if not self._metadata_file.exists(): - return - try: - data = self._metadata_file.read_text(encoding="utf-8") - raw = json.loads(data) - self._metadata_cache = {path: FileMetadata(**meta) for path, meta in raw.items()} - except Exception as e: - self.logger.warning(f"Failed to load metadata: {e}") - - async def _save_metadata(self) -> None: - try: - raw = { - path: meta.model_dump(mode="json") - for path, meta in self._metadata_cache.items() - } - data = json.dumps(raw, indent=2, ensure_ascii=False) - temp = self._metadata_file.with_suffix(".tmp") - temp.write_text(data, encoding="utf-8") - temp.replace(self._metadata_file) - except Exception as e: - self.logger.error(f"Failed to save metadata: {e}") - raise - finally: - if temp.exists(): - temp.unlink() - # -- Lifecycle ---------------------------------------------------------- async def _start(self, app_context=None) -> None: @@ -81,76 +47,77 @@ class ChromaFileStore(BaseFileStore): name=self.collection_name, metadata={"hnsw:space": "cosine"}, ) - await self._load_metadata() - self.logger.info( - f"ChromaFileStore '{self.store_name}' ready: " - f"collection={self.collection_name}, metadata at {self._metadata_file}", - ) + self.logger.info(f"ChromaChunkStore '{self.store_name}' ready: collection={self.collection_name}") await super()._start(app_context) async def _close(self) -> None: - await self._save_metadata() self.client = None self.chunks_collection = None await super()._close() + # -- Filter helper ------------------------------------------------------ + + @staticmethod + def _path_where(chunk_filter: ChunkFilter | None) -> dict | None: + if chunk_filter is None or chunk_filter.resolved_paths is None: + return None + paths = chunk_filter.resolved_paths + if not paths: + return {"path": "__nonexistent__"} + if len(paths) == 1: + return {"path": next(iter(paths))} + return {"path": {"$in": list(paths)}} + # -- Write operations --------------------------------------------------- - async def upsert_file(self, file_meta: FileMetadata, chunks: list[FileChunk]) -> None: - # Always delete existing data for this file first - if file_meta.path: - await self.delete_file(file_meta.path) + async def upsert_chunks(self, path: str, chunks: list[FileChunk]) -> None: + await self.delete_chunks(path) - if chunks: - chunks = await self.get_chunk_embeddings(chunks) + if not chunks: + return - ids, documents, embeddings, metadatas = [], [], [], [] - now = int(time.time() * 1000) - for chunk in chunks: - ids.append(chunk.id) - documents.append(chunk.text) - embeddings.append(chunk.embedding if chunk.embedding else [0.0] * self.embedding_dim) - metadatas.append({ - "path": file_meta.path, - "start_line": chunk.start_line, - "end_line": chunk.end_line, - "hash": chunk.hash, - "updated_at": now, - }) + chunks = await self.get_chunk_embeddings(chunks) - self.chunks_collection.upsert( - ids=ids, - documents=documents, - embeddings=embeddings, - metadatas=metadatas, - ) + ids, documents, embeddings, metadatas = [], [], [], [] + now = int(time.time() * 1000) + for chunk in chunks: + ids.append(chunk.id) + documents.append(chunk.text) + embeddings.append(chunk.embedding if chunk.embedding else [0.0] * self.embedding_dim) + metadatas.append({ + "path": path, + "start_line": chunk.start_line, + "end_line": chunk.end_line, + "hash": chunk.hash, + "updated_at": now, + }) - if file_meta.path: - self._metadata_cache[file_meta.path] = FileMetadata( - modified_time=file_meta.modified_time, - path=file_meta.path, - metadata=file_meta.metadata, - ) + self.chunks_collection.upsert( + ids=ids, + documents=documents, + embeddings=embeddings, + metadatas=metadatas, + ) - async def delete_file(self, path: str) -> None: + async def delete_chunks(self, path: str) -> None: results = self.chunks_collection.get(where={"path": path}, include=[]) if results["ids"]: self.chunks_collection.delete(ids=results["ids"]) - self._metadata_cache.pop(path, None) # -- Read operations ---------------------------------------------------- - async def list_files(self) -> list[str]: - """List all indexed file paths.""" - return list(self._metadata_cache.keys()) - - async def get_file_metadata(self, path: str) -> FileMetadata | None: - """Get metadata for a specific file.""" - return self._metadata_cache.get(path) + async def get_chunks(self, path: str) -> list[FileChunk]: + results = self.chunks_collection.get(where={"path": path}, include=["documents", "metadatas"]) + chunks: list[FileChunk] = [] + for cid, md, text in zip(results["ids"], results["metadatas"], results["documents"]): + chunks.append(self._chunk_from_chroma(cid, md, text)) + chunks.sort(key=lambda c: c.start_line) + return chunks # -- Search helpers ----------------------------------------------------- - def _chunk_from_chroma(self, chunk_id: str, md: dict, text: str, embedding=None) -> FileChunk: + @staticmethod + def _chunk_from_chroma(chunk_id: str, md: dict, text: str, embedding=None) -> FileChunk: return FileChunk( id=chunk_id, path=md["path"], @@ -163,7 +130,12 @@ class ChromaFileStore(BaseFileStore): # -- Search operations -------------------------------------------------- - async def vector_search(self, query: str, limit: int, search_filter: SearchFilter | None = None) -> list[FileChunk]: + async def vector_search( + self, + query: str, + limit: int, + chunk_filter: ChunkFilter | None = None, + ) -> list[FileChunk]: if not self.vector_enabled or not query: return [] @@ -175,6 +147,7 @@ class ChromaFileStore(BaseFileStore): results = self.chunks_collection.query( query_embeddings=[query_embedding], n_results=limit, + where=self._path_where(chunk_filter), include=["documents", "metadatas", "distances"], ) except Exception as e: @@ -192,7 +165,6 @@ class ChromaFileStore(BaseFileStore): chunk.scores = {"vector": score, "score": score} chunks.append(chunk) - chunks = self._apply_filter(chunks, search_filter) chunks.sort(key=lambda c: c.score, reverse=True) return chunks[:limit] @@ -200,9 +172,8 @@ class ChromaFileStore(BaseFileStore): self, query: str, limit: int, - search_filter: SearchFilter | None = None, + chunk_filter: ChunkFilter | None = None, ) -> list[FileChunk]: - """Keyword search via ChromaDB $contains with case variants.""" if not self.fts_enabled or not query: return [] @@ -210,7 +181,6 @@ class ChromaFileStore(BaseFileStore): if not words: return [] - # Generate case variants for case-insensitive matching word_variants = set() for word in words: word_variants.add(word) @@ -225,6 +195,7 @@ class ChromaFileStore(BaseFileStore): where_document = {"$or": [{"$contains": w} for w in variants_list]} results = self.chunks_collection.get( + where=self._path_where(chunk_filter), where_document=where_document, include=["documents", "metadatas"], ) @@ -241,7 +212,6 @@ class ChromaFileStore(BaseFileStore): chunk.scores = {"keyword": score, "score": score} chunks.append(chunk) - chunks = self._apply_filter(chunks, search_filter) chunks.sort(key=lambda c: c.score, reverse=True) return chunks[:limit] @@ -253,6 +223,4 @@ class ChromaFileStore(BaseFileStore): name=self.collection_name, metadata={"hnsw:space": "cosine"}, ) - self._metadata_cache = {} - await self._save_metadata() - self.logger.info(f"Cleared all data from ChromaFileStore '{self.store_name}'") + self.logger.info(f"Cleared all data from ChromaChunkStore '{self.store_name}'") diff --git a/reme2/component/file_store/local_file_store.py b/reme2/component/chunk_store/local_chunk_store.py similarity index 62% rename from reme2/component/file_store/local_file_store.py rename to reme2/component/chunk_store/local_chunk_store.py index 468bc43f..062758e8 100644 --- a/reme2/component/file_store/local_file_store.py +++ b/reme2/component/chunk_store/local_chunk_store.py @@ -1,31 +1,25 @@ -"""Pure-Python file storage with JSONL persistence.""" +"""Pure-Python chunk storage with JSONL persistence.""" import json from pathlib import Path import numpy as np -from .base_file_store import BaseFileStore +from .base_chunk_store import BaseChunkStore from ..component_registry import R -from ...schema import FileChunk, FileMetadata, SearchFilter +from ...schema import ChunkFilter, FileChunk from ...utils import batch_cosine_similarity @R.register("local") -class LocalFileStore(BaseFileStore): - """In-memory file storage with JSONL disk persistence. - - No external database required. All data lives in Python dicts; - writes are flushed to JSONL files on disk and survive restarts. - """ +class LocalChunkStore(BaseChunkStore): + """In-memory chunk storage with JSONL disk persistence.""" def __init__(self, encoding: str = "utf-8", **kwargs): super().__init__(**kwargs) self._encoding: str = encoding self._chunks: dict[str, FileChunk] = {} - self._files: dict[str, FileMetadata] = {} self._chunks_file: Path = self.db_path / f"{self.store_name}_chunks.jsonl" - self._metadata_file: Path = self.db_path / f"{self.store_name}_file_metadata.json" # -- Persistence helpers ------------------------------------------------ @@ -59,88 +53,40 @@ class LocalFileStore(BaseFileStore): if temp_path.exists(): temp_path.unlink() - async def _load_metadata(self) -> None: - """Load file metadata from JSON file into memory.""" - if not self._metadata_file.exists(): - return - try: - data = self._metadata_file.read_text(encoding=self._encoding) - raw: dict = json.loads(data) - self._files = {path: FileMetadata(**meta) for path, meta in raw.items()} - except Exception as e: - self.logger.warning(f"Failed to load metadata: {e}") - - async def _save_metadata(self) -> None: - """Persist file metadata to JSON file with atomic write.""" - raw = {path: meta.model_dump(mode="json") for path, meta in self._files.items()} - content = json.dumps(raw, indent=2, ensure_ascii=False) - temp_path = self._metadata_file.with_suffix(".tmp") - try: - temp_path.write_text(content, encoding=self._encoding) - temp_path.replace(self._metadata_file) - except Exception as e: - self.logger.error(f"Failed to save metadata: {e}") - raise - finally: - if temp_path.exists(): - temp_path.unlink() - # -- Lifecycle ---------------------------------------------------------- async def _start(self) -> None: """Load persisted data into memory.""" - await self._load_metadata() await self._load_chunks() - self.logger.info( - f"LocalFileStore '{self.store_name}' ready: " - f"{len(self._chunks)} chunks, metadata at {self._metadata_file}", - ) + self.logger.info(f"LocalChunkStore '{self.store_name}' ready: {len(self._chunks)} chunks") await super()._start() async def _close(self) -> None: """Flush state to disk and clear memory.""" - await self._save_metadata() await self._save_chunks() self._chunks.clear() - self._files.clear() await super()._close() # -- Write operations --------------------------------------------------- - async def upsert_file(self, file_meta: FileMetadata, chunks: list[FileChunk]) -> None: + async def upsert_chunks(self, path: str, chunks: list[FileChunk]) -> None: """Insert or update a file and its chunks.""" - await self.delete_file(file_meta.path) + await self.delete_chunks(path) + if not chunks: + return + chunks = await self.get_chunk_embeddings(chunks) + for chunk in chunks: + self._chunks[chunk.id] = chunk - if chunks: - chunks = await self.get_chunk_embeddings(chunks) - for chunk in chunks: - self._chunks[chunk.id] = chunk - - if file_meta.path: - self._files[file_meta.path] = FileMetadata( - modified_time=file_meta.modified_time, - path=file_meta.path, - metadata=file_meta.metadata, - ) - - async def delete_file(self, path: str) -> None: + async def delete_chunks(self, path: str) -> None: """Delete a file and all its chunks.""" to_delete = [cid for cid, chunk in self._chunks.items() if chunk.path == path] for cid in to_delete: del self._chunks[cid] - self._files.pop(path, None) # -- Read operations ---------------------------------------------------- - async def list_files(self) -> list[str]: - """List all indexed file paths.""" - return list(self._files.keys()) - - async def get_file_metadata(self, path: str) -> FileMetadata | None: - """Get metadata for a specific file.""" - return self._files.get(path) - - async def get_file_chunks(self, path: str) -> list[FileChunk]: + async def get_chunks(self, path: str) -> list[FileChunk]: """Get all chunks for a file, sorted by start_line.""" chunks = [chunk for chunk in self._chunks.values() if chunk.path == path] chunks.sort(key=lambda c: c.start_line) @@ -148,7 +94,12 @@ class LocalFileStore(BaseFileStore): # -- Search operations -------------------------------------------------- - async def vector_search(self, query: str, limit: int, search_filter: SearchFilter | None = None) -> list[FileChunk]: + async def vector_search( + self, + query: str, + limit: int, + chunk_filter: ChunkFilter | None = None, + ) -> list[FileChunk]: """Cosine-similarity vector search over in-memory embeddings.""" if not self.vector_enabled or not query: return [] @@ -159,14 +110,12 @@ class LocalFileStore(BaseFileStore): candidates = self._apply_filter( [c for c in self._chunks.values() if c.embedding], - search_filter, + chunk_filter, ) if not candidates: return [] expected_dim = self.embedding_dim - - # Validate and align embedding dimensions valid_embeddings = [] for chunk in candidates: emb = chunk.embedding @@ -201,17 +150,16 @@ class LocalFileStore(BaseFileStore): self, query: str, limit: int, - search_filter: SearchFilter | None = None, + chunk_filter: ChunkFilter | None = None, ) -> list[FileChunk]: """Keyword search via substring matching.""" if not self.fts_enabled or not query: return [] - words = query.split() - if not words: + if not query.split(): return [] - filtered_chunks = self._apply_filter(list(self._chunks.values()), search_filter) + filtered_chunks = self._apply_filter(list(self._chunks.values()), chunk_filter) results = [] for chunk in filtered_chunks: @@ -236,7 +184,5 @@ class LocalFileStore(BaseFileStore): async def clear_all(self) -> None: """Clear all indexed data from memory and disk.""" self._chunks.clear() - self._files.clear() await self._save_chunks() - await self._save_metadata() - self.logger.info(f"Cleared all data from LocalFileStore '{self.store_name}'") + self.logger.info(f"Cleared all data from LocalChunkStore '{self.store_name}'") diff --git a/reme2/component/file_store/sqlite_file_store.py b/reme2/component/chunk_store/sqlite_chunk_store.py similarity index 74% rename from reme2/component/file_store/sqlite_file_store.py rename to reme2/component/chunk_store/sqlite_chunk_store.py index fbebaac2..bb01dd86 100644 --- a/reme2/component/file_store/sqlite_file_store.py +++ b/reme2/component/chunk_store/sqlite_chunk_store.py @@ -1,18 +1,18 @@ -"""SQLite storage backend for file store.""" +"""SQLite chunk storage backend.""" import json import sqlite3 import struct import time -from .base_file_store import BaseFileStore +from .base_chunk_store import BaseChunkStore from ..component_registry import R -from ...schema import FileChunk, FileMetadata, SearchFilter +from ...schema import ChunkFilter, FileChunk @R.register("sqlite") -class SqliteFileStore(BaseFileStore): - """SQLite file storage with vector and full-text search. +class SqliteChunkStore(BaseChunkStore): + """SQLite chunk storage with vector and full-text search. Uses sqlite-vec for vector similarity search and FTS5 with trigram tokenizer for keyword search. Falls back to LIKE-based substring @@ -30,10 +30,6 @@ class SqliteFileStore(BaseFileStore): def chunks_table(self) -> str: return f"chunks_{self.store_name}" - @property - def files_table(self) -> str: - return f"files_{self.store_name}" - @property def vector_table(self) -> str: return f"chunks_vec_{self.store_name}" @@ -88,23 +84,12 @@ class SqliteFileStore(BaseFileStore): self.conn.enable_load_extension(False) await self._create_tables() - self.logger.info( - f"SqliteFileStore '{self.store_name}' ready: " - f"db={self.db_path / 'reme.db'}", - ) + self.logger.info(f"SqliteChunkStore '{self.store_name}' ready: db={self.db_path / 'reme.db'}") await super()._start(app_context) async def _create_tables(self) -> None: cursor = self.conn.cursor() try: - cursor.execute(f""" - CREATE TABLE IF NOT EXISTS {self.files_table} ( - path TEXT PRIMARY KEY, - mtime REAL, - metadata TEXT - ) - """) - cursor.execute(f""" CREATE TABLE IF NOT EXISTS {self.chunks_table} ( id TEXT PRIMARY KEY, @@ -117,6 +102,10 @@ class SqliteFileStore(BaseFileStore): updated_at INTEGER ) """) + cursor.execute( + f"CREATE INDEX IF NOT EXISTS idx_{self.chunks_table}_path " + f"ON {self.chunks_table}(path)", + ) if self.vector_enabled: cursor.execute(f""" @@ -153,52 +142,17 @@ class SqliteFileStore(BaseFileStore): self.conn = None await super()._close() - # -- Metadata helpers --------------------------------------------------- - - async def _get_all_metadata(self) -> dict[str, FileMetadata]: - """Load all file metadata from SQL into a dict for tag filtering.""" - cursor = self.conn.cursor() - try: - cursor.execute(f"SELECT path, mtime, metadata FROM {self.files_table}") - result = {} - for path, mtime, meta_str in cursor.fetchall(): - metadata = json.loads(meta_str) if meta_str else {} - result[path] = FileMetadata( - modified_time=mtime, - path=path, - metadata=metadata, - ) - return result - except Exception as e: - self.logger.error(f"Failed to load all metadata: {e}") - return {} - finally: - cursor.close() - # -- Write operations --------------------------------------------------- - async def upsert_file(self, file_meta: FileMetadata, chunks: list[FileChunk]) -> None: + async def upsert_chunks(self, path: str, chunks: list[FileChunk]) -> None: cursor = self.conn.cursor() try: cursor.execute("BEGIN") - # Upsert file metadata - cursor.execute( - f"""INSERT OR REPLACE INTO {self.files_table} - (path, mtime, metadata) - VALUES (?, ?, ?)""", - ( - file_meta.path, - file_meta.modified_time, - json.dumps(file_meta.metadata, ensure_ascii=False) if file_meta.metadata else None, - ), - ) - - # Delete old chunks and vectors/fts for this file old_ids = [ row[0] for row in cursor.execute( f"SELECT id FROM {self.chunks_table} WHERE path = ?", - (file_meta.path,), + (path,), ).fetchall() ] if old_ids: @@ -208,9 +162,8 @@ class SqliteFileStore(BaseFileStore): for oid in old_ids: cursor.execute(f"DELETE FROM {self.vector_table} WHERE id = ?", (oid,)) if self.fts_enabled: - cursor.execute(f"DELETE FROM {self.fts_table} WHERE path = ?", (file_meta.path,)) + cursor.execute(f"DELETE FROM {self.fts_table} WHERE path = ?", (path,)) - # Insert new chunks if chunks: chunks = await self.get_chunk_embeddings(chunks) now = int(time.time() * 1000) @@ -221,7 +174,7 @@ class SqliteFileStore(BaseFileStore): VALUES (?, ?, ?, ?, ?, ?, ?, ?)""", ( chunk.id, - file_meta.path, + path, chunk.start_line, chunk.end_line, chunk.hash, @@ -242,18 +195,18 @@ class SqliteFileStore(BaseFileStore): f"""INSERT INTO {self.fts_table} (text, id, path, start_line, end_line) VALUES (?, ?, ?, ?, ?)""", - (chunk.text, chunk.id, file_meta.path, chunk.start_line, chunk.end_line), + (chunk.text, chunk.id, path, chunk.start_line, chunk.end_line), ) cursor.execute("COMMIT") except Exception as e: cursor.execute("ROLLBACK") - self.logger.error(f"Failed to upsert file {file_meta.path}: {e}") + self.logger.error(f"Failed to upsert chunks for {path}: {e}") raise finally: cursor.close() - async def delete_file(self, path: str) -> None: + async def delete_chunks(self, path: str) -> None: cursor = self.conn.cursor() try: cursor.execute("BEGIN") @@ -272,53 +225,18 @@ class SqliteFileStore(BaseFileStore): cursor.execute(f"DELETE FROM {self.fts_table} WHERE path = ?", (path,)) cursor.execute(f"DELETE FROM {self.chunks_table} WHERE path = ?", (path,)) - cursor.execute(f"DELETE FROM {self.files_table} WHERE path = ?", (path,)) cursor.execute("COMMIT") except Exception as e: cursor.execute("ROLLBACK") - self.logger.error(f"Failed to delete file {path}: {e}") + self.logger.error(f"Failed to delete chunks for {path}: {e}") raise finally: cursor.close() - # -- Read operations (SQL overrides) ------------------------------------ + # -- Read operations ---------------------------------------------------- - async def list_files(self) -> list[str]: - cursor = self.conn.cursor() - try: - cursor.execute(f"SELECT path FROM {self.files_table}") - return [row[0] for row in cursor.fetchall()] - except Exception as e: - self.logger.error(f"Failed to list files: {e}") - return [] - finally: - cursor.close() - - async def get_file_metadata(self, path: str) -> FileMetadata | None: - cursor = self.conn.cursor() - try: - cursor.execute( - f"SELECT mtime, metadata FROM {self.files_table} WHERE path = ?", - (path,), - ) - row = cursor.fetchone() - if not row: - return None - mtime, meta_str = row - metadata = json.loads(meta_str) if meta_str else {} - return FileMetadata( - modified_time=mtime, - path=path, - metadata=metadata, - ) - except Exception as e: - self.logger.error(f"Failed to get file metadata for {path}: {e}") - return None - finally: - cursor.close() - - async def get_file_chunks(self, path: str) -> list[FileChunk]: + async def get_chunks(self, path: str) -> list[FileChunk]: cursor = self.conn.cursor() try: cursor.execute( @@ -347,7 +265,7 @@ class SqliteFileStore(BaseFileStore): )) return chunks except Exception as e: - self.logger.error(f"Failed to get file chunks for {path}: {e}") + self.logger.error(f"Failed to get chunks for {path}: {e}") return [] finally: cursor.close() @@ -364,9 +282,25 @@ class SqliteFileStore(BaseFileStore): cleaned = cleaned.replace(ch, " ") return " ".join(cleaned.split()) + @staticmethod + def _path_filter_clause(chunk_filter: ChunkFilter | None, column: str = "path") -> tuple[str, list]: + """Build a WHERE clause fragment for ChunkFilter; returns (sql_fragment, params).""" + if chunk_filter is None or chunk_filter.resolved_paths is None: + return "", [] + paths = chunk_filter.resolved_paths + if not paths: + return "0", [] # filter excludes everything + placeholders = ",".join("?" * len(paths)) + return f"{column} IN ({placeholders})", list(paths) + # -- Search operations -------------------------------------------------- - async def vector_search(self, query: str, limit: int, search_filter: SearchFilter | None = None) -> list[FileChunk]: + async def vector_search( + self, + query: str, + limit: int, + chunk_filter: ChunkFilter | None = None, + ) -> list[FileChunk]: if not self.vector_enabled or not query: return [] @@ -391,7 +325,7 @@ class SqliteFileStore(BaseFileStore): chunks = [] for cid, path, start, end, text, dist in cursor.fetchall(): score = max(0.0, 1.0 - dist / 2.0) - chunk = FileChunk( + chunks.append(FileChunk( id=cid, path=path, start_line=start, @@ -399,12 +333,9 @@ class SqliteFileStore(BaseFileStore): text=text, hash="", scores={"vector": score, "score": score}, - ) - chunks.append(chunk) + )) - # Apply filter with SQL-based metadata for tag support - file_meta = await self._get_all_metadata() if (search_filter and search_filter.tags) else None - chunks = self._apply_filter(chunks, search_filter, file_meta) + chunks = self._apply_filter(chunks, chunk_filter) chunks.sort(key=lambda c: c.score, reverse=True) return chunks[:limit] except Exception as e: @@ -417,7 +348,7 @@ class SqliteFileStore(BaseFileStore): self, query: str, limit: int, - search_filter: SearchFilter | None = None, + chunk_filter: ChunkFilter | None = None, ) -> list[FileChunk]: if not self.fts_enabled or not query: return [] @@ -430,18 +361,14 @@ class SqliteFileStore(BaseFileStore): if not words: return [] - file_meta = await self._get_all_metadata() if (search_filter and search_filter.tags) else None - - # FTS5 trigram requires all terms >= 3 chars if all(len(w) >= 3 for w in words): results = await self._fts_trigram_search(words, limit) if results: - return self._apply_filter(results, search_filter, file_meta)[:limit] + return self._apply_filter(results, chunk_filter)[:limit] return self._apply_filter( await self._like_search(cleaned, words, limit), - search_filter, - file_meta, + chunk_filter, )[:limit] async def _fts_trigram_search(self, words: list[str], limit: int) -> list[FileChunk]: @@ -464,7 +391,7 @@ class SqliteFileStore(BaseFileStore): chunks = [] for cid, path, start, end, text, rank in cursor.fetchall(): score = max(0.0, 1.0 / (1.0 + abs(rank))) - chunk = FileChunk( + chunks.append(FileChunk( id=cid, path=path, start_line=start, @@ -472,8 +399,7 @@ class SqliteFileStore(BaseFileStore): text=text, hash="", scores={"keyword": score, "score": score}, - ) - chunks.append(chunk) + )) chunks.sort(key=lambda c: c.score, reverse=True) return chunks except Exception as e: @@ -511,7 +437,7 @@ class SqliteFileStore(BaseFileStore): if score == 0.0: continue - chunk = FileChunk( + chunks.append(FileChunk( id=cid, path=path, start_line=start, @@ -519,8 +445,7 @@ class SqliteFileStore(BaseFileStore): text=text, hash="", scores={"keyword": score, "score": score}, - ) - chunks.append(chunk) + )) chunks.sort(key=lambda c: c.score, reverse=True) return chunks[:limit] @@ -536,7 +461,6 @@ class SqliteFileStore(BaseFileStore): cursor = self.conn.cursor() try: cursor.execute("BEGIN") - cursor.execute(f"DELETE FROM {self.files_table}") cursor.execute(f"DELETE FROM {self.chunks_table}") if self.vector_enabled: cursor.execute(f"DELETE FROM {self.vector_table}") @@ -549,4 +473,4 @@ class SqliteFileStore(BaseFileStore): raise finally: cursor.close() - self.logger.info(f"Cleared all data from SqliteFileStore '{self.store_name}'") + self.logger.info(f"Cleared all data from SqliteChunkStore '{self.store_name}'") diff --git a/reme2/component/file_store/__init__.py b/reme2/component/file_store/__init__.py deleted file mode 100644 index 28616593..00000000 --- a/reme2/component/file_store/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -"""File store module for persistent memory management. - -Provides storage backends for memory chunks and file metadata with -vector and full-text search capabilities. -""" - -from .base_file_store import BaseFileStore -from .chroma_file_store import ChromaFileStore -from .local_file_store import LocalFileStore -from .sqlite_file_store import SqliteFileStore - -__all__ = [ - "BaseFileStore", - "ChromaFileStore", - "LocalFileStore", - "SqliteFileStore", -] diff --git a/reme2/component/file_watcher/base_file_watcher.py b/reme2/component/file_watcher/base_file_watcher.py index da296a09..d0740b79 100644 --- a/reme2/component/file_watcher/base_file_watcher.py +++ b/reme2/component/file_watcher/base_file_watcher.py @@ -6,10 +6,10 @@ from pathlib import Path from watchfiles import Change, awatch from ..base_component import BaseComponent +from ..chunk_store import BaseChunkStore from ..file_parser import BaseFileParser -from ..file_store import BaseFileStore from ...enumeration import ComponentEnum -from ...file_graph import FileGraph +from ...schema.file_graph import FileGraph class BaseFileWatcher(BaseComponent): @@ -31,16 +31,16 @@ class BaseFileWatcher(BaseComponent): debounce: int = 2000, chunk_tokens: int = 400, chunk_overlap: int = 80, - file_store: str = "default", + chunk_store: str = "default", default_parser: str | None = None, rebuild_index_on_start: bool = False, poll_delay_ms: int = 2000, **kwargs, ): super().__init__(**kwargs) - self._file_store_name: str = file_store + self._chunk_store_name: str = chunk_store self._default_parser_name: str | None = default_parser - self.file_store: BaseFileStore | None = None + self.chunk_store: BaseChunkStore | None = None self._suffix_to_parser: dict[str, BaseFileParser] = {} self._default_parser: BaseFileParser | None = None self.watch_path: str = watch_path @@ -66,17 +66,17 @@ class BaseFileWatcher(BaseComponent): return self._meta_path / "file_graph.json" async def _start(self): - """Resolve file_store, load or build file_graph, and start watching.""" - if self._file_store_name: + """Resolve chunk_store, load or build file_graph, and start watching.""" + if self._chunk_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__}") - self.file_store = store + stores = self.app_context.components.get(ComponentEnum.CHUNK_STORE, {}) + if self._chunk_store_name not in stores: + raise ValueError(f"Chunk store '{self._chunk_store_name}' not found.") + store = stores[self._chunk_store_name] + if not isinstance(store, BaseChunkStore): + raise TypeError(f"Expected BaseChunkStore, got {type(store).__name__}") + self.chunk_store = store parsers = self.app_context.components.get(ComponentEnum.FILE_PARSER, {}) for parser in parsers.values(): @@ -127,13 +127,13 @@ class BaseFileWatcher(BaseComponent): self._watch_task = None self._stop_event.clear() - self.file_store = None + self.chunk_store = None self._suffix_to_parser.clear() self._default_parser = None self.logger.info("Stopped watching") async def _scan_existing_files(self) -> None: - if not self.file_store: + if not self.chunk_store: return watch_path = Path(self.watch_path) @@ -202,7 +202,7 @@ class BaseFileWatcher(BaseComponent): return self._suffix_to_parser.get(suffix, self._default_parser) async def on_changes(self, changes: set[tuple[Change, str]]) -> None: - if not self.file_store: + if not self.chunk_store: self.logger.warning("File store not initialized, skipping changes") return @@ -227,8 +227,8 @@ class BaseFileWatcher(BaseComponent): self.logger.debug(f"No parser for {path}, skipping") return file_meta, chunks = await parser.parse(path) - await self.file_store.upsert_file(file_meta, chunks) - self.file_graph.add(file_meta) + await self.chunk_store.upsert_chunks(path, chunks) + self.file_graph.create(file_meta) self.logger.info(f"Added {path} ({len(chunks)} chunks)") async def _on_modified(self, path: str) -> None: @@ -237,11 +237,11 @@ class BaseFileWatcher(BaseComponent): self.logger.debug(f"No parser for {path}, skipping") return file_meta, chunks = await parser.parse(path) - await self.file_store.upsert_file(file_meta, chunks) - self.file_graph.update(file_meta) + await self.chunk_store.upsert_chunks(path, chunks) + self.file_graph.create(file_meta) self.logger.info(f"Modified {path} ({len(chunks)} chunks)") async def _on_deleted(self, path: str) -> None: - await self.file_store.delete_file(path) - self.file_graph.remove(path) + await self.chunk_store.delete_chunks(path) + self.file_graph.delete(path) self.logger.info(f"Deleted {path}") diff --git a/reme2/config/paw.yaml b/reme2/config/paw.yaml index 7d604ddd..999ac9ed 100644 --- a/reme2/config/paw.yaml +++ b/reme2/config/paw.yaml @@ -48,7 +48,7 @@ components: default: backend: default - file_store: + chunk_store: default: backend: local embedding_model: default @@ -58,7 +58,7 @@ components: file_watcher: default: backend: full - file_store: default + chunk_store: default default_parser: default watch_paths: [ "./test_data" ] recursive: true diff --git a/reme2/enumeration/component_enum.py b/reme2/enumeration/component_enum.py index d3426ef7..7f74a167 100644 --- a/reme2/enumeration/component_enum.py +++ b/reme2/enumeration/component_enum.py @@ -26,7 +26,7 @@ class ComponentEnum(str, Enum): FILE_PARSER = "file_parser" - FILE_STORE = "file_store" + CHUNK_STORE = "chunk_store" FILE_WATCHER = "file_watcher" diff --git a/reme2/file_based/memory_search.py b/reme2/file_based/memory_search.py index 0d062bdd..4abb7507 100644 --- a/reme2/file_based/memory_search.py +++ b/reme2/file_based/memory_search.py @@ -4,7 +4,6 @@ import json from ..component import R from ..component.base_step import BaseStep -from ..schema import SearchFilter @R.register("memory_search") @@ -38,22 +37,20 @@ class MemorySearch(BaseStep): isinstance(max_results, int) and max_results > 0 ), f"max_results must be a positive integer, got {max_results}" - filter_paths: list[str] | None = self.context.get("paths") or None - filter_tags: list[str] | None = self.context.get("tags") or None - exclude_paths: list[str] | None = self.context.get("exclude_paths") or None - search_filter = None - if filter_paths or filter_tags or exclude_paths: - search_filter = SearchFilter(paths=filter_paths, tags=filter_tags, exclude_paths=exclude_paths) + chunk_filter = self.file_graph.filter( + paths=self.context.get("paths") or None, + tags=self.context.get("tags") or None, + exclude_paths=self.context.get("exclude_paths") or None, + ) - results = await self.file_store.hybrid_search( + results = await self.chunk_store.hybrid_search( query=query, limit=max_results, vector_weight=self.vector_weight, candidate_multiplier=self.candidate_multiplier, - search_filter=search_filter, + chunk_filter=chunk_filter, ) - # Filter by min_score results = [r for r in results if r.score >= min_score] return json.dumps([result.model_dump(exclude_none=True) for result in results], indent=2, ensure_ascii=False) diff --git a/reme2/schema/__init__.py b/reme2/schema/__init__.py index 7940a1be..b0c9310f 100644 --- a/reme2/schema/__init__.py +++ b/reme2/schema/__init__.py @@ -3,11 +3,11 @@ from .application_config import ApplicationConfig, ComponentConfig, JobConfig from .as_msg_stat import AsBlockStat, AsMsgStat from .base_node import BaseNode +from .chunk_filter import ChunkFilter from .file_chunk import FileChunk from .file_metadata import FileMetadata from .request import Request from .response import Response -from .search_filter import SearchFilter from .stream_chunk import StreamChunk __all__ = [ @@ -17,10 +17,10 @@ __all__ = [ "AsBlockStat", "AsMsgStat", "BaseNode", + "ChunkFilter", "FileChunk", "FileMetadata", "Request", "Response", - "SearchFilter", "StreamChunk", ] diff --git a/reme2/schema/chunk_filter.py b/reme2/schema/chunk_filter.py new file mode 100644 index 00000000..9cd4870f --- /dev/null +++ b/reme2/schema/chunk_filter.py @@ -0,0 +1,49 @@ +"""Filter for chunk-store search. + +User-facing fields (paths/tags/exclude_paths) describe metadata-level intent. +FileGraph compiles them into `resolved_paths` (the concrete path set that +ChunkStore actually consumes for filtering). +""" + +from pydantic import BaseModel, Field + + +class ChunkFilter(BaseModel): + """User-facing search filter, resolved by FileGraph into a path set. + + User input fields: + paths: include only paths starting with any of these prefixes + tags: include only files whose metadata contains ALL these tags + exclude_paths: exclude paths starting with any of these prefixes + + Compiled field (set by FileGraph.filter): + resolved_paths: concrete path set to restrict ChunkStore search. + None = no restriction (all chunks). + Empty set = no chunks match (search returns empty). + """ + + paths: list[str] | None = Field(default=None) + tags: list[str] | None = Field(default=None) + exclude_paths: list[str] | None = Field(default=None) + resolved_paths: set[str] | None = Field(default=None) + + def is_empty(self) -> bool: + return not self.paths and not self.tags and not self.exclude_paths + + def match_metadata(self, path: str, metadata: dict | None = None) -> bool: + """Match a single path+metadata against user-input conditions.""" + if self.paths and not any(path.startswith(p) for p in self.paths): + return False + if self.exclude_paths and any(path.startswith(p) for p in self.exclude_paths): + return False + if self.tags: + file_tags = set((metadata or {}).get("tags", [])) + if not all(t in file_tags for t in self.tags): + return False + return True + + def match_path(self, path: str) -> bool: + """Match a path against the resolved path set (used by ChunkStore).""" + if self.resolved_paths is None: + return True + return path in self.resolved_paths diff --git a/reme2/file_graph.py b/reme2/schema/file_graph.py similarity index 80% rename from reme2/file_graph.py rename to reme2/schema/file_graph.py index 8c7764ff..0dc7a65c 100644 --- a/reme2/file_graph.py +++ b/reme2/schema/file_graph.py @@ -2,7 +2,8 @@ import json from collections import defaultdict from pathlib import Path -from reme2.schema.file_metadata import FileMetadata +from .chunk_filter import ChunkFilter +from .file_metadata import FileMetadata class FileGraph: @@ -13,7 +14,7 @@ class FileGraph: # -- CRUD ---------------------------------------------------------------- - def add(self, *metadatas: FileMetadata) -> None: + def create(self, *metadatas: FileMetadata) -> None: for metadata in metadatas: path = metadata.path if path in self._nodes: @@ -22,6 +23,9 @@ class FileGraph: for metadata in metadatas: self._add_forward(metadata) + def read(self, path: str) -> FileMetadata | None: + return self._nodes.get(path) + def update(self, path: str, **fields) -> FileMetadata | None: metadata = self._nodes.get(path) if metadata is None: @@ -32,7 +36,7 @@ class FileGraph: self._add_forward(updated) return updated - def remove(self, path: str) -> FileMetadata | None: + def delete(self, path: str) -> FileMetadata | None: metadata = self._nodes.pop(path, None) if metadata is None: return None @@ -40,9 +44,6 @@ class FileGraph: self._backlinks.pop(path, None) return metadata - def get(self, path: str) -> FileMetadata | None: - return self._nodes.get(path) - # -- Link queries -------------------------------------------------------- def get_links(self, path: str) -> list[FileMetadata]: @@ -58,6 +59,21 @@ class FileGraph: if src in self._nodes ] + def filter( + self, + paths: list[str] | None = None, + tags: list[str] | None = None, + exclude_paths: list[str] | None = None, + ) -> ChunkFilter: + cf = ChunkFilter(paths=paths, tags=tags, exclude_paths=exclude_paths) + if cf.is_empty(): + return cf + cf.resolved_paths = { + path for path, meta in self._nodes.items() + if cf.match_metadata(path, meta.metadata) + } + return cf + # -- Index helpers ------------------------------------------------------- def _add_forward(self, metadata: FileMetadata) -> None: @@ -87,7 +103,7 @@ class FileGraph: return graph raw: dict = json.loads(path.read_text(encoding="utf-8")) nodes = [FileMetadata(**meta) for meta in raw.values()] - graph.add(*nodes) + graph.create(*nodes) return graph # -- Dunder -------------------------------------------------------------- diff --git a/reme2/schema/search_filter.py b/reme2/schema/search_filter.py deleted file mode 100644 index bc0a2ad5..00000000 --- a/reme2/schema/search_filter.py +++ /dev/null @@ -1,36 +0,0 @@ -"""Search filter schema for constraining search results.""" - -from pydantic import BaseModel, Field - - -class SearchFilter(BaseModel): - """Filter conditions for search operations. - - All specified conditions are combined with AND logic. - Within paths/exclude_paths, items are combined with OR logic. - Within tags, items are combined with AND logic (all must match). - """ - - paths: list[str] | None = Field( - default=None, - description="Include only chunks whose path starts with any of these prefixes", - ) - tags: list[str] | None = Field(default=None, description="Include only chunks containing ALL specified tags") - exclude_paths: list[str] | None = Field( - default=None, - description="Exclude chunks whose path starts with any of these prefixes", - ) - - def is_empty(self) -> bool: - return not self.paths and not self.tags and not self.exclude_paths - - def match(self, path: str, metadata: dict | None = None) -> bool: - if self.paths and not any(path.startswith(p) for p in self.paths): - return False - if self.exclude_paths and any(path.startswith(p) for p in self.exclude_paths): - return False - if self.tags: - chunk_tags = set((metadata or {}).get("tags", [])) - if not all(t in chunk_tags for t in self.tags): - return False - return True