diff --git a/reme2/component/embedding/base_embedding_model.py b/reme2/component/embedding/base_embedding_model.py index 8ea7da96..8844e891 100644 --- a/reme2/component/embedding/base_embedding_model.py +++ b/reme2/component/embedding/base_embedding_model.py @@ -11,7 +11,7 @@ import numpy as np from ..base_component import BaseComponent from ...enumeration import ComponentEnum -from ...schema import BaseNode +from ...schema import EmbNode class BaseEmbeddingModel(BaseComponent): @@ -204,7 +204,7 @@ class BaseEmbeddingModel(BaseComponent): 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 = texts_to_compute[i: i + self.max_batch_size] batch_indices = [idx for idx, _ in batch] batch_texts = [text for _, text in batch] @@ -225,7 +225,7 @@ class BaseEmbeddingModel(BaseComponent): return results - async def get_node_embeddings(self, nodes: list[BaseNode], **kwargs) -> list[BaseNode]: + 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) diff --git a/reme2/component/file_parser/base_file_parser.py b/reme2/component/file_parser/base_file_parser.py index a79706fe..3a891e9e 100644 --- a/reme2/component/file_parser/base_file_parser.py +++ b/reme2/component/file_parser/base_file_parser.py @@ -1,98 +1,46 @@ -"""Abstract base class for file parsers. - -Single-pass parser: `parse(path, existing_chunks)` returns a `ParsedFile` -carrying metadata + chunks (with embeddings attached) + edges. - -The parser owns the embedding pipeline. To avoid re-embedding unchanged -blocks, the watcher passes in the file's prior chunks; the parser -hash-diffs and only calls the embedding API for blocks whose hash is -new. Vanished hashes simply drop out (the file_store's upsert is a -delete-and-insert). - -If `existing_chunks` is None / empty, every chunk is embedded fresh. -""" +"""Base file parser interface.""" from abc import abstractmethod from ..base_component import BaseComponent -from ..embedding import BaseEmbeddingModel -from ...enumeration import ComponentEnum, FileSuffixEnum -from ...schema import FileChunk, ParsedFile +from ...enumeration import ComponentEnum +from ...schema import FileChunk, FileNode class BaseFileParser(BaseComponent): - """Single-pass parser producing an embedded `ParsedFile`.""" + """Abstract base class for file parsers. + + Subclasses must implement `_parse` to extract file content into chunks. + """ component_type = ComponentEnum.FILE_PARSER - suffixes: list[FileSuffixEnum] = [] - def __init__(self, embedding_model: str = "", **kwargs): - super().__init__(**kwargs) - self._embedding_model_name: str = embedding_model - self.embedding_model: BaseEmbeddingModel | None = None + async def parse(self, path: str, cache_chunks: list[FileChunk] | None = None) -> tuple[FileNode, list[FileChunk]]: + """Parse a file and optionally reuse cached chunks by hash. - async def _start(self) -> None: - if not self._embedding_model_name: - return - assert self.app_context is not None, "app_context must be provided" - models = self.app_context.components.get(ComponentEnum.EMBEDDING_MODEL, {}) - if self._embedding_model_name not in models: - raise ValueError(f"Embedding model '{self._embedding_model_name}' not found.") - model = models[self._embedding_model_name] - if not isinstance(model, BaseEmbeddingModel): - raise TypeError(f"Expected BaseEmbeddingModel, got {type(model).__name__}") - self.embedding_model = model + Args: + path: Path to the file to parse. + cache_chunks: Previously parsed chunks to reuse when hashes match. - async def _close(self) -> None: - self.embedding_model = None - - async def _embed_chunks(self, chunks: list[FileChunk]) -> list[FileChunk]: - """Attach embeddings to chunks if an embedding model is configured. - - Failures (missing embeddings) leave `chunk.embedding=None`; the - file_store can still keyword-search them. + Returns: + Tuple of file metadata and parsed chunks. """ - if not chunks or self.embedding_model is None: - return chunks - try: - await self.embedding_model.get_node_embeddings(chunks) - except Exception as e: - self.logger.warning(f"embedding chunks failed: {e}") - return chunks - - @staticmethod - def _hash_diff_attach( - chunks: list[FileChunk], - existing_chunks: list[FileChunk] | None, - ) -> list[FileChunk]: - """Attach cached embeddings to chunks whose hash already exists. - - Returns the dirty subset (chunks still needing embeddings). Mutates - the input chunks in place: matched chunks get their `embedding` - field set from the cache. - """ - if not existing_chunks: - return list(chunks) - cached = {c.hash: c.embedding for c in existing_chunks if c.embedding} - if not cached: - return list(chunks) - dirty: list[FileChunk] = [] - for c in chunks: - cached_emb = cached.get(c.hash) - if cached_emb is not None: - c.embedding = cached_emb - else: - dirty.append(c) - return dirty + file_node, chunks = await self._parse(path) + if cache_chunks: + cache_chunk_dict = {chunk.hash: chunk for chunk in cache_chunks} + for i in range(len(chunks)): + chunk = chunks[i] + if chunk.hash in cache_chunk_dict: + chunks[i] = cache_chunk_dict[chunk.hash] + return file_node, chunks @abstractmethod - async def parse( - self, - path: str, - existing_chunks: list[FileChunk] | None = None, - ) -> ParsedFile: - """Parse + embed the file, returning a fully populated `ParsedFile`. + async def _parse(self, path: str) -> tuple[FileNode, list[FileChunk]]: + """Parse a file into metadata and chunks. Subclasses must implement this. - `existing_chunks`, when provided, lets the parser reuse cached - embeddings for blocks whose hash hasn't changed. + Args: + path: Path to the file to parse. + + Returns: + Tuple of file metadata and parsed chunks. """ diff --git a/reme2/schema/__init__.py b/reme2/schema/__init__.py index 37751258..cee3737b 100644 --- a/reme2/schema/__init__.py +++ b/reme2/schema/__init__.py @@ -2,11 +2,11 @@ from .application_config import ApplicationConfig, ComponentConfig, JobConfig from .as_msg_stat import AsBlockStat, AsMsgStat -from .base_node import BaseNode +from .emb_node import EmbNode from .chunk_filter import ChunkFilter from .file_chunk import FileChunk from .file_edge import FileEdge -from .file_metadata import FileMetadata +from .file_node import FileNode from .parsed_file import ParsedFile from .request import Request from .response import Response @@ -18,11 +18,11 @@ __all__ = [ "JobConfig", "AsBlockStat", "AsMsgStat", - "BaseNode", + "EmbNode", "ChunkFilter", "FileChunk", "FileEdge", - "FileMetadata", + "FileNode", "ParsedFile", "Request", "Response", diff --git a/reme2/schema/base_node.py b/reme2/schema/emb_node.py similarity index 74% rename from reme2/schema/base_node.py rename to reme2/schema/emb_node.py index 47510567..36011d23 100644 --- a/reme2/schema/base_node.py +++ b/reme2/schema/emb_node.py @@ -3,9 +3,7 @@ from uuid import uuid4 from pydantic import BaseModel, Field -class BaseNode(BaseModel): - """Base node model for graph and document structures.""" - +class EmbNode(BaseModel): id: str = Field(default_factory=lambda: uuid4().hex) text: str = Field(default="") embedding: list[float] | None = Field(default=None) diff --git a/reme2/schema/file_chunk.py b/reme2/schema/file_chunk.py index 5c29a1bd..1b46addc 100644 --- a/reme2/schema/file_chunk.py +++ b/reme2/schema/file_chunk.py @@ -1,16 +1,12 @@ from pydantic import Field -from .base_node import BaseNode +from .emb_node import EmbNode -class FileChunk(BaseNode): - """File content chunk with positional and scoring metadata.""" - - path: str = Field(...) - start_line: int = Field(...) - end_line: int = Field(...) - - hash: str = Field(...) +class FileChunk(EmbNode): + path: str = Field(default="") + start_line: int = Field(default=0) + end_line: int = Field(default=0) scores: dict[str, float] = Field(default_factory=dict) @property @@ -18,5 +14,5 @@ class FileChunk(BaseNode): return self.scores.get("score", 0.0) @property - def unique_key(self) -> str: - return f"{self.path}:{self.start_line}:{self.end_line}" + def hash(self) -> str: + return "_".join([self.id, self.path, str(self.start_line), str(self.end_line)])