diff --git a/reme4/components/file_graph/local_file_graph.py b/reme4/components/file_graph/local_file_graph.py index 221b022a..2d6f58bb 100644 --- a/reme4/components/file_graph/local_file_graph.py +++ b/reme4/components/file_graph/local_file_graph.py @@ -9,7 +9,7 @@ from ...schema import FileLink, FileNode @R.register("local") class LocalFileGraph(BaseFileGraph): - """Dict-backed file graph; uses FileLink.path for adjacency.""" + """Dict-backed file graph; uses FileLink.target_path for adjacency.""" def __init__(self, **kwargs): super().__init__(**kwargs) @@ -81,12 +81,12 @@ class LocalFileGraph(BaseFileGraph): old = self._nodes.get(path) if old is not None: for link in old.links: - if link.path: - self._remove_edge(path, link.path) + if link.target_path: + self._remove_edge(path, link.target_path) self._nodes[path] = node for link in node.links: - if link.path: - self._add_edge(path, link.path) + if link.target_path: + self._add_edge(path, link.target_path) # Promote pending edges that now target a real node. promoted = self._pending.pop(path, None) if promoted: @@ -98,8 +98,8 @@ class LocalFileGraph(BaseFileGraph): if node is None: continue for link in node.links: - if link.path: - self._remove_edge(path, link.path) + if link.target_path: + self._remove_edge(path, link.target_path) # Demote inbound edges to pending (sources still reference this path). demoted = self._inverse.pop(path, None) if demoted: @@ -116,8 +116,8 @@ class LocalFileGraph(BaseFileGraph): self._pending.clear() for src, node in self._nodes.items(): for link in node.links: - if link.path: - self._add_edge(src, link.path) + if link.target_path: + self._add_edge(src, link.target_path) async def clear(self): self._nodes.clear() @@ -130,9 +130,11 @@ class LocalFileGraph(BaseFileGraph): node = self._nodes.get(path) if node is None: return [] - return [lnk for lnk in node.links if lnk.path and lnk.path in self._nodes] + return [lnk for lnk in node.links if lnk.target_path and lnk.target_path in self._nodes] async def get_inlinks(self, path: str) -> list[FileLink]: if path not in self._nodes: return [] - return [link for src in self._inverse.get(path, ()) for link in self._nodes[src].links if link.path == path] + return [ + link for src in self._inverse.get(path, ()) for link in self._nodes[src].links if link.target_path == path + ] diff --git a/reme4/components/file_graph/nx_file_graph.py b/reme4/components/file_graph/nx_file_graph.py index dcb608d6..485212c5 100644 --- a/reme4/components/file_graph/nx_file_graph.py +++ b/reme4/components/file_graph/nx_file_graph.py @@ -15,7 +15,7 @@ from ...schema import FileLink, FileNode @R.register("nx") class NxFileGraph(BaseFileGraph): - """Networkx-backed file graph; uses FileLink.path for adjacency.""" + """Networkx-backed file graph; uses FileLink.target_path for adjacency.""" def __init__(self, **kwargs): super().__init__(**kwargs) @@ -70,7 +70,7 @@ class NxFileGraph(BaseFileGraph): self._graph.remove_edges_from(list(self._graph.out_edges(path, keys=True))) self._graph.add_node(path, node=node) # promotes virtual node if present # Missing targets become attr-less virtual nodes. - self._graph.add_edges_from((path, lnk.path, {"link": lnk}) for lnk in node.links if lnk.path) + self._graph.add_edges_from((path, lnk.target_path, {"link": lnk}) for lnk in node.links if lnk.target_path) async def delete_nodes(self, paths: list[str]) -> None: for path in paths: @@ -94,10 +94,10 @@ class NxFileGraph(BaseFileGraph): virtual = [n for n, d in self._graph.nodes(data=True) if "node" not in d] self._graph.remove_nodes_from(virtual) self._graph.add_edges_from( - (path, lnk.path, {"link": lnk}) + (path, lnk.target_path, {"link": lnk}) for path, data in self._graph.nodes(data=True) for lnk in data["node"].links - if lnk.path + if lnk.target_path ) async def clear(self): diff --git a/reme4/components/file_parser/default_file_parser.py b/reme4/components/file_parser/default_file_parser.py index d64b0234..69a7c694 100644 --- a/reme4/components/file_parser/default_file_parser.py +++ b/reme4/components/file_parser/default_file_parser.py @@ -1,5 +1,6 @@ """Default file parser with byte-based overlapping chunking.""" +import re from bisect import bisect_right from pathlib import Path @@ -8,7 +9,20 @@ import yaml from .base_file_parser import BaseFileParser from ..component_registry import R -from ...schema import FileChunk, FileNode, FileFrontMatter +from ...schema import FileChunk, FileFrontMatter, FileLink, FileNode + +# Single-pass wikilink + optional Dataview predicate. +# Covers: [[X]] / [[X#h]] / [[X|alias]] / pred:: [[X]] / [pred:: [[X]]] +# - predicate group: optional leading '[' (Dataview inline-bracket form), an identifier, +# then '::' β€” the whole prefix is non-capturing-optional so bare wikilinks still match. +# - target / anchor: target stops before '#', '|', '[', ']'; anchor stops before '|', '[', ']'. +# - alias '|...': consumed but not captured (we don't need display text). +_LINK_RE = re.compile( + r"(?:\[?\s*(?P[A-Za-z][\w-]*)\s*::\s*)?" + r"\[\[\s*(?P[^\[\]|#]+?)" + r"(?:#(?P[^\[\]|]+?))?" + r"\s*(?:\|[^\[\]]*?)?\s*\]\]", +) @R.register("default") @@ -21,6 +35,25 @@ class DefaultFileParser(BaseFileParser): self.chunk_byte_size = max(100, chunk_byte_size) self.overlap_byte_size = max(4, overlap_byte_size) + @staticmethod + def parse_links(content: str, source_path: str) -> list[FileLink]: + """Extract wikilinks with optional Dataview predicate as outgoing FileLinks.""" + links: list[FileLink] = [] + for m in _LINK_RE.finditer(content): + target = m["target"].strip() + if not target: + continue + anchor = m["anchor"] + links.append( + FileLink( + source_path=source_path, + target_path=target, + target_anchor=anchor.strip() if anchor else None, + predicate=m["predicate"], + ), + ) + return links + @staticmethod def _parse_front_matter(text: str) -> tuple[FileFrontMatter, str]: """Parse YAML front matter delimited by ---, return (front_matter, remaining).""" @@ -51,9 +84,19 @@ class DefaultFileParser(BaseFileParser): if not content: return FileNode(path=rel_path, st_mtime=stat.st_mtime, front_matter=front_matter), [] + links = self.parse_links(content, rel_path) chunks = self._chunk_content(content, rel_path) chunk_ids = [c.id for c in chunks] - return FileNode(path=rel_path, st_mtime=stat.st_mtime, front_matter=front_matter, chunk_ids=chunk_ids), chunks + return ( + FileNode( + path=rel_path, + st_mtime=stat.st_mtime, + front_matter=front_matter, + links=links, + chunk_ids=chunk_ids, + ), + chunks, + ) def _chunk_content(self, content: str, rel_path: str) -> list[FileChunk]: """Split content into overlapping byte-range chunks with line numbers.""" diff --git a/reme4/components/file_store/base_file_store.py b/reme4/components/file_store/base_file_store.py index 770d043e..5cc6edd9 100644 --- a/reme4/components/file_store/base_file_store.py +++ b/reme4/components/file_store/base_file_store.py @@ -81,6 +81,12 @@ class BaseFileStore(BaseComponent): raise RuntimeError("file_graph is required for delete_by_path") return await self.file_graph.rebuild_links() + async def get_nodes(self, paths: list[str]) -> list[FileNode]: + """Return file nodes for the given paths (missing paths are skipped).""" + if not self.file_graph: + raise RuntimeError("file_graph is required for get_nodes") + return await self.file_graph.get_nodes(paths) + async def get_outlinks(self, path: str) -> list[FileLink]: """Return outgoing links for *path*.""" if not self.file_graph: diff --git a/reme4/components/file_store/local_file_store.py b/reme4/components/file_store/local_file_store.py index 052694e4..c02dfa5f 100644 --- a/reme4/components/file_store/local_file_store.py +++ b/reme4/components/file_store/local_file_store.py @@ -46,7 +46,7 @@ class LocalFileStore(BaseFileStore): self.logger.exception(f"Failed to load {self.chunks_path}: {e}") async def dump(self) -> None: - """Persist chunks to JSONL via atomic rename.""" + """Persist chunks to JSONL via atomic rename, then cascade to keyword_index and file_graph.""" try: tmp = self.chunks_path.with_suffix(".tmp") async with aiofiles.open(tmp, "w", encoding=self.encoding) as f: @@ -54,6 +54,10 @@ class LocalFileStore(BaseFileStore): tmp.replace(self.chunks_path) except Exception as e: self.logger.exception(f"Failed to write {self.chunks_path}: {e}") + if self.keyword_index: + await self.keyword_index.dump() + if self.file_graph: + await self.file_graph.dump() # Base class interface diff --git a/reme4/components/file_watcher/base_file_watcher.py b/reme4/components/file_watcher/base_file_watcher.py index b246537c..3c6fbc9e 100644 --- a/reme4/components/file_watcher/base_file_watcher.py +++ b/reme4/components/file_watcher/base_file_watcher.py @@ -94,26 +94,16 @@ class BaseFileWatcher(BaseComponent): files[self._get_relative_path(p)] = p.absolute() return files - async def clear_store(self): - """Remove all entries from the file store.""" - if self.file_store is None: - raise ValueError("file_store is not initialized!") - await self.file_store.clear() - - async def reset_store(self): - """Clear the store and re-index all existing files.""" - if self.file_store is None: - raise ValueError("file_store is not initialized!") - await self.file_store.clear() - await self.on_added(list((await self.scan_existing_files()).keys())) - @abstractmethod async def watch_loop(self): """Watch for file changes and dispatch events.""" @abstractmethod - async def update_store(self): - """Sync the store with the current state of watch_paths.""" + async def update_store(self, dump: bool = True) -> dict[str, int]: + """Sync the store with watch_paths; dump store if any changes and dump=True. + + Returns counts {"added": int, "modified": int, "deleted": int}. + """ @abstractmethod async def on_added(self, path: str | list[str]): diff --git a/reme4/components/file_watcher/lite_file_watcher.py b/reme4/components/file_watcher/lite_file_watcher.py index e1fea327..da48b2b3 100644 --- a/reme4/components/file_watcher/lite_file_watcher.py +++ b/reme4/components/file_watcher/lite_file_watcher.py @@ -70,7 +70,7 @@ class LiteFileWatcher(BaseFileWatcher): self.logger.info(f"Detected {len(buckets[change])} {label} file(s)") await handler(buckets[change]) - async def update_store(self): + async def update_store(self, dump: bool = True) -> dict[str, int]: if self.file_store is None: raise ValueError("file_store is not initialized!") @@ -92,8 +92,13 @@ class LiteFileWatcher(BaseFileWatcher): if to_add: self.logger.info(f"Indexing {len(to_add)} new file(s)") await self.on_added(to_add) - if not to_modify and not to_delete and not to_add: + + changed = bool(to_add or to_modify or to_delete) + if not changed: self.logger.info("Store is up to date") + if dump and changed: + await self.file_store.dump() + return {"added": len(to_add), "modified": len(to_modify), "deleted": len(to_delete)} async def _parse_and_upsert(self, paths: list[str], action: str): """Parse files and upsert into store. Shared by on_added / on_modified.""" diff --git a/reme4/config/default.yaml b/reme4/config/default.yaml index 8bcfd5c0..5e753124 100644 --- a/reme4/config/default.yaml +++ b/reme4/config/default.yaml @@ -49,6 +49,44 @@ jobs: steps: - backend: help_step + - backend: base + name: search + description: "hybrid search over file_store: vector + keyword fused via RRF" + parameters: + type: object + properties: + query: + type: string + description: "search query" + limit: + type: integer + description: "max results to return" + default: 5 + min_score: + type: number + description: "minimum fused score threshold (RRF scores are small; default 0 disables filter)" + default: 0.0 + vector_weight: + type: number + description: "weight for vector results in [0, 1]; keyword weight = 1 - vector_weight" + default: 0.7 + candidate_multiplier: + type: number + description: "candidate pool multiplier per branch (capped at 200)" + default: 3.0 + expand_links: + type: boolean + description: "attach outlinks/inlinks (with neighbor meta) to each result" + default: true + max_links_per_direction: + type: integer + description: "max neighbors shown per direction per result" + default: 10 + required: + - query + steps: + - backend: search_step + - backend: stream name: stream_demo description: "stream demo job: repeat query 10x and stream char-by-char" diff --git a/reme4/schema/file_link.py b/reme4/schema/file_link.py index b93cf1ca..b029d5a9 100644 --- a/reme4/schema/file_link.py +++ b/reme4/schema/file_link.py @@ -4,19 +4,15 @@ from pydantic import BaseModel, ConfigDict, Field class FileLink(BaseModel): - """A parsed wikilink with optional anchor and predicate.""" + """file link + [[target_path]] + [[target_path#target_anchor]] + predicate:: [[target_*]] + [predicate:: [[target_*]]] + """ model_config = ConfigDict(extra="forbid") - - path: str = Field( - default=..., - description="Wikilink target β€” raw text pre-resolution, vault-relative path after.", - ) - anchor: str | None = Field( - default=None, - description="Heading or block anchor (text after '#'); None if absent.", - ) - predicate: str | None = Field( - default=None, - description="Dataview-style typed-link predicate; None for bare [[X]].", - ) + source_path: str = Field(default=..., description="source file path relative to working dir") + target_path: str = Field(default=..., description="target file path relative to working dir") + target_anchor: str | None = Field(default=None, description="Heading or block anchor (text after '#')") + predicate: str | None = Field(default=None, description="Dataview-style typed-link predicate") diff --git a/reme4/steps/common/__init__.py b/reme4/steps/common/__init__.py index e31ba302..e09c2073 100644 --- a/reme4/steps/common/__init__.py +++ b/reme4/steps/common/__init__.py @@ -3,6 +3,8 @@ from .demo import DemoEchoStep1, DemoEchoStep2 from .health_check import HealthCheckStep from .help import HelpStep +from .reindex import ReindexStep +from .search import SearchStep from .stream_demo import StreamDemoStep1, StreamDemoStep2 from .version import VersionStep @@ -11,6 +13,8 @@ __all__ = [ "DemoEchoStep2", "HealthCheckStep", "HelpStep", + "ReindexStep", + "SearchStep", "StreamDemoStep1", "StreamDemoStep2", "VersionStep", diff --git a/reme4/steps/common/health_check.py b/reme4/steps/common/health_check.py index 43a40573..fa7d55c0 100644 --- a/reme4/steps/common/health_check.py +++ b/reme4/steps/common/health_check.py @@ -124,7 +124,7 @@ def _is_status_healthy(ctype: ComponentEnum, status: dict) -> bool: @R.register("health_check_step") class HealthCheckStep(BaseStep): - """Collect a concise health-check snapshot of the relevant components.""" + """Collect a concise health check snapshot of the relevant components.""" async def execute(self): assert self.context is not None @@ -145,6 +145,7 @@ class HealthCheckStep(BaseStep): health = {"version": __version__, "healthy": healthy, "components": components} self.logger.info(f"[{self.name}] health collected: {health}") - self.context.response.answer = f"ReMe v{__version__} - {'healthy' if healthy else 'unhealthy'}" + status_emoji = "βœ…" if healthy else "❌" + self.context.response.answer = f"{status_emoji} ReMe v{__version__} - {'healthy' if healthy else 'unhealthy'}" self.context.response.metadata["health"] = health return self.context.response diff --git a/reme4/steps/common/help.py b/reme4/steps/common/help.py index 782ea877..6c850772 100644 --- a/reme4/steps/common/help.py +++ b/reme4/steps/common/help.py @@ -4,27 +4,27 @@ from ..base_step import BaseStep from ...components import R -def _format_params(parameters: dict) -> str: - props = (parameters or {}).get("properties") or {} - if not props: - return "no args" - required = set((parameters or {}).get("required") or []) - parts = [] - for pname, pschema in props.items(): - ptype = pschema.get("type", "any") - if pname in required: - parts.append(f"{pname}:{ptype}*") - elif "default" in pschema: - parts.append(f"{pname}:{ptype}={pschema['default']}") - else: - parts.append(f"{pname}:{ptype}") - return ", ".join(parts) - - @R.register("help_step") class HelpStep(BaseStep): """List all registered jobs (excluding self) as compact one-liners for an LLM.""" + @staticmethod + def _format_params(parameters: dict) -> str: + props = (parameters or {}).get("properties") or {} + if not props: + return "no args" + required = set((parameters or {}).get("required") or []) + parts = [] + for pname, pschema in props.items(): + ptype = pschema.get("type", "any") + if pname in required: + parts.append(f"{pname}:{ptype}*") + elif "default" in pschema: + parts.append(f"{pname}:{ptype}={pschema['default']}") + else: + parts.append(f"{pname}:{ptype}") + return ", ".join(parts) + async def execute(self): assert self.context is not None @@ -33,7 +33,7 @@ class HelpStep(BaseStep): for name, job in self.app_context.jobs.items(): if name == "help": continue - lines.append(f"πŸ› οΈ `{name}` β€” {job.description} πŸ“₯ {_format_params(job.parameters)}") + lines.append(f"πŸ› οΈ `{name}` β€” {job.description} πŸ“₯ {self._format_params(job.parameters)}") self.logger.info(f"[{self.name}] returning {len(lines)} jobs") diff --git a/reme4/steps/common/reindex.py b/reme4/steps/common/reindex.py new file mode 100644 index 00000000..f4d280e9 --- /dev/null +++ b/reme4/steps/common/reindex.py @@ -0,0 +1,24 @@ +"""Wipe the file store and rebuild it from the watcher's tracked files.""" + +from ..base_step import BaseStep +from ...components import R + + +@R.register("reindex_step") +class ReindexStep(BaseStep): + """Full re-index: stop watcher, clear store, sync from disk, then restart.""" + + async def execute(self): + assert self.context is not None + + await self.file_watcher.close() + try: + await self.file_store.clear() + counts = await self.file_watcher.update_store() + finally: + await self.file_watcher.start() + + self.logger.info(f"[{self.name}] reindexed {counts}") + self.context.response.answer = f"πŸ”„ Reindexed {counts['added']} file(s)" + self.context.response.metadata["counts"] = counts + return self.context.response diff --git a/reme4/steps/common/search.py b/reme4/steps/common/search.py new file mode 100644 index 00000000..deded425 --- /dev/null +++ b/reme4/steps/common/search.py @@ -0,0 +1,224 @@ +"""Hybrid search over file_store using RRF fusion of vector + keyword results.""" + +import asyncio + +from ..base_step import BaseStep +from ...components import R +from ...schema import FileChunk, FileLink, FileNode + +_RRF_K = 60 +_MAX_CANDIDATES = 200 + + +@R.register("search_step") +class SearchStep(BaseStep): + """Hybrid search: run vector + keyword in parallel, fuse via RRF, filter, truncate.""" + + @staticmethod + def _rrf_merge( + vector: list[FileChunk], + keyword: list[FileChunk], + vector_weight: float, + ) -> list[FileChunk]: + """Fuse two ranked lists with Reciprocal Rank Fusion, keyed by chunk.id.""" + text_weight = 1.0 - vector_weight + merged: dict[str, FileChunk] = {} + + for rank, chunk in enumerate(vector, start=1): + contrib = vector_weight / (_RRF_K + rank) + c = chunk.model_copy(deep=False) + c.scores = {**chunk.scores, "vector": chunk.scores.get("vector", chunk.score), "score": contrib} + merged[c.id] = c + + for rank, chunk in enumerate(keyword, start=1): + contrib = text_weight / (_RRF_K + rank) + existing = merged.get(chunk.id) + if existing is not None: + existing.scores = { + **existing.scores, + "keyword": chunk.scores.get("keyword", chunk.score), + "score": existing.scores["score"] + contrib, + } + else: + c = chunk.model_copy(deep=False) + c.scores = {**chunk.scores, "keyword": chunk.scores.get("keyword", chunk.score), "score": contrib} + merged[c.id] = c + + results = list(merged.values()) + results.sort(key=lambda r: r.score, reverse=True) + return results + + @staticmethod + def _format_scores(scores: dict[str, float], hybrid: bool) -> str: + """Format scores for the answer line: always show fused; show per-branch when hybrid.""" + parts = [f"score={scores.get('score', 0.0):.4f}"] + if hybrid: + for k in ("vector", "keyword"): + v = scores.get(k) + parts.append(f"{k}={v:.4f}" if v is not None else f"{k}=-") + return " ".join(parts) + + @staticmethod + def _group_by_neighbor(links: list[FileLink], key_attr: str) -> dict[str, list[dict]]: + """Group edges by neighbor path (insertion-ordered), each value a list of {predicate, anchor}.""" + out: dict[str, list[dict]] = {} + for lnk in links: + neighbor = getattr(lnk, key_attr) + if not neighbor: + continue + out.setdefault(neighbor, []).append( + {"predicate": lnk.predicate, "anchor": lnk.target_anchor}, + ) + return out + + @staticmethod + def _node_meta(node: FileNode | None) -> dict: + """Extract a compact meta dict (title/description/tags) from a FileNode.""" + if node is None: + return {} + fm = node.front_matter + meta: dict = {} + if fm.title: + meta["title"] = fm.title + if fm.description: + meta["description"] = fm.description + if fm.tags: + meta["tags"] = list(fm.tags) + return meta + + @staticmethod + def _format_meta_inline(meta: dict) -> str: + """One-line render of node meta for the answer; '(no meta)' when empty.""" + parts = [] + if "title" in meta: + parts.append(f'title="{meta["title"]}"') + if "tags" in meta: + parts.append(f"tags={meta['tags']}") + return " ".join(parts) if parts else "(no meta)" + + @staticmethod + def _format_via(edge: dict) -> str: + """Render a single (predicate, anchor) edge as a 'via ...' descriptor.""" + bits = [] + if edge.get("predicate"): + bits.append(f"predicate={edge['predicate']}") + if edge.get("anchor"): + bits.append(f"anchor=#{edge['anchor']}") + return ", ".join(bits) if bits else "plain" + + async def _expand_links( + self, + chunk_paths: list[str], + max_per_direction: int, + ) -> dict[str, dict]: + """Fetch out/in links for each chunk path; attach neighbor meta. Returns per-path expansion.""" + if not chunk_paths: + return {} + + out_lists, in_lists = await asyncio.gather( + asyncio.gather(*(self.file_store.get_outlinks(p) for p in chunk_paths)), + asyncio.gather(*(self.file_store.get_inlinks(p) for p in chunk_paths)), + ) + + # Pre-group + cap per direction so we only fetch meta for displayed neighbors. + out_grouped = [ + dict(list(self._group_by_neighbor(outs, "target_path").items())[:max_per_direction]) for outs in out_lists + ] + in_grouped = [ + dict(list(self._group_by_neighbor(ins, "source_path").items())[:max_per_direction]) for ins in in_lists + ] + + neighbor_paths = sorted({n for g in out_grouped for n in g} | {n for g in in_grouped for n in g}) + nodes = await self.file_store.get_nodes(neighbor_paths) if neighbor_paths else [] + meta_by_path = {n.path: self._node_meta(n) for n in nodes} + + def _attach(grouped: dict[str, list[dict]]) -> list[dict]: + return [ + {"path": npath, "meta": meta_by_path.get(npath, {}), "edges": edges} for npath, edges in grouped.items() + ] + + return { + cp: {"outlinks": _attach(og), "inlinks": _attach(ig)} + for cp, og, ig in zip(chunk_paths, out_grouped, in_grouped) + } + + @classmethod + def _render_expansion_lines(cls, expansion: dict) -> list[str]: + """Render outlinks/inlinks blocks for one chunk path; return zero or more indented lines.""" + lines: list[str] = [] + for direction, arrow, items in ( + ("outlinks", "β†’", expansion.get("outlinks") or []), + ("inlinks", "←", expansion.get("inlinks") or []), + ): + if not items: + continue + lines.append(f" {direction} ({len(items)}):") + for item in items: + lines.append(f" {arrow} {item['path']} {cls._format_meta_inline(item['meta'])}") + for edge in item["edges"]: + lines.append(f" via {cls._format_via(edge)}") + return lines + + async def execute(self): + assert self.context is not None + query: str = (self.context.get("query", "") or "").strip() + limit: int = int(self.context.get("limit", 5)) + min_score: float = float(self.context.get("min_score", 0.0)) + vector_weight: float = float(self.context.get("vector_weight", 0.7)) + candidate_multiplier: float = float(self.context.get("candidate_multiplier", 3.0)) + expand_links: bool = bool(self.context.get("expand_links", True)) + max_links_per_direction: int = int(self.context.get("max_links_per_direction", 10)) + + assert query, "query cannot be empty" + assert 0.0 <= vector_weight <= 1.0, f"vector_weight must be in [0, 1], got {vector_weight}" + assert limit > 0, f"limit must be positive, got {limit}" + + candidates = min(_MAX_CANDIDATES, max(1, int(limit * candidate_multiplier))) + search_filter: dict = self.context.get("search_filter", {}) or {} + + vector_results, keyword_results = await asyncio.gather( + self.file_store.vector_search(query, candidates, search_filter), + self.file_store.keyword_search(query, candidates, search_filter), + ) + + self.logger.info( + f"[{self.name}] query={query!r} candidates={candidates} " + f"vector_hits={len(vector_results)} keyword_hits={len(keyword_results)}", + ) + + hybrid = bool(vector_results) and bool(keyword_results) + if not vector_results and not keyword_results: + fused: list[FileChunk] = [] + elif not keyword_results: + fused = vector_results + elif not vector_results: + fused = keyword_results + else: + fused = self._rrf_merge(vector_results, keyword_results, vector_weight) + + if min_score > 0.0: + fused = [c for c in fused if c.score >= min_score] + fused = fused[:limit] + + unique_paths = list(dict.fromkeys(c.path for c in fused)) + link_expansion: dict[str, dict] = ( + await self._expand_links(unique_paths, max_links_per_direction) if expand_links else {} + ) + + answer_lines: list[str] = [] + for c in fused: + answer_lines.append( + f"{c.path}:{c.start_line}-{c.end_line} [{self._format_scores(c.scores, hybrid)}] {c.text}", + ) + answer_lines.extend(self._render_expansion_lines(link_expansion.get(c.path, {}))) + + self.context.response.answer = "\n".join(answer_lines) + self.context.response.metadata["results"] = [c.model_dump(exclude_none=True) for c in fused] + self.context.response.metadata["link_expansion"] = link_expansion + self.context.response.metadata["counts"] = { + "vector": len(vector_results), + "keyword": len(keyword_results), + "returned": len(fused), + "hybrid": hybrid, + } + return self.context.response diff --git a/reme4/steps/common/version.py b/reme4/steps/common/version.py index 477802e6..e43aa698 100644 --- a/reme4/steps/common/version.py +++ b/reme4/steps/common/version.py @@ -1,7 +1,7 @@ """Return the package version.""" from ..base_step import BaseStep -from ... import __version__ + from ...components import R @@ -11,8 +11,9 @@ class VersionStep(BaseStep): async def execute(self): assert self.context is not None - self.logger.info(f"[{self.name}] version={__version__}") + from ... import __version__ + self.logger.info(f"[{self.name}] version={__version__}") self.context.response.answer = __version__ self.context.response.metadata["version"] = __version__ return self.context.response diff --git a/reme4/utils/__init__.py b/reme4/utils/__init__.py index c9e2d275..45a5bfb9 100644 --- a/reme4/utils/__init__.py +++ b/reme4/utils/__init__.py @@ -1,6 +1,12 @@ """Utility modules.""" -from .common_utils import hash_text, execute_stream_task +from .common_utils import ( + hash_text, + execute_stream_task, + mock_reme_server, + call_action, + call_and_check, +) from .env_utils import load_env from .logger_utils import get_logger from .logo_utils import print_logo @@ -10,6 +16,9 @@ from .similarity_utils import cosine_similarity, batch_cosine_similarity __all__ = [ "hash_text", "execute_stream_task", + "mock_reme_server", + "call_action", + "call_and_check", "load_env", "get_logger", "print_logo", diff --git a/reme4/utils/common_utils.py b/reme4/utils/common_utils.py index 46d6f97e..53996187 100644 --- a/reme4/utils/common_utils.py +++ b/reme4/utils/common_utils.py @@ -2,10 +2,17 @@ import asyncio import hashlib -from collections.abc import AsyncGenerator +import json +import socket +import subprocess +import sys +import time +from collections.abc import AsyncGenerator, Callable +from contextlib import asynccontextmanager from typing import Any, Literal from .logger_utils import get_logger +from ..constants import REME_DEFAULT_HOST, REME_DEFAULT_PORT from ..enumeration import ChunkEnum from ..schema import StreamChunk @@ -103,3 +110,139 @@ async def execute_stream_task( await task except asyncio.CancelledError: pass + + +def _pick_free_port(host: str = REME_DEFAULT_HOST) -> int: + """Bind to port 0 and return the OS-assigned free port.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind((host, 0)) + return s.getsockname()[1] + + +async def _wait_reme_ready(host: str, port: int, timeout: float) -> None: + """Poll find_reme until it reports 'reme' or timeout elapses.""" + from .service_utils import find_reme + + deadline = time.time() + timeout + while time.time() < deadline: + status = await find_reme(host, port) + if status == "reme": + return + await asyncio.sleep(0.2) + raise TimeoutError(f"ReMe service did not become ready at {host}:{port} within {timeout}s") + + +@asynccontextmanager +async def mock_reme_server( + host: str = REME_DEFAULT_HOST, + port: int | None = None, + config: str | None = None, + extra_args: list[str] | None = None, + startup_timeout: float = 30.0, + shutdown_timeout: float = 10.0, + log_to_file: bool = False, + enable_logo: bool = False, +): + """Spawn `reme4 start` as a subprocess and yield (host, port) once ready. + + Auto-picks a free port when port is None. Subprocess is terminated on exit. + """ + logger = get_logger() + if port is None: + port = _pick_free_port(host) + + cmd: list[str] = [ + sys.executable, + "-m", + "reme4.reme", + "start", + f"service.host={host}", + f"service.port={port}", + f"log_to_file={'true' if log_to_file else 'false'}", + f"enable_logo={'true' if enable_logo else 'false'}", + ] + if config: + cmd.append(f"config={config}") + if extra_args: + cmd.extend(extra_args) + + logger.info(f"Launching mock reme server: {' '.join(cmd)}") + proc = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + try: + await _wait_reme_ready(host, port, startup_timeout) + yield host, port + except Exception: + # Capture early-exit output for diagnostics. + if proc.poll() is not None and proc.stdout is not None: + tail = proc.stdout.read() + logger.error(f"reme server exited early. output:\n{tail}") + raise + finally: + if proc.poll() is None: + proc.terminate() + try: + proc.wait(timeout=shutdown_timeout) + except subprocess.TimeoutExpired: + logger.warning("reme server did not terminate gracefully, killing") + proc.kill() + proc.wait(timeout=shutdown_timeout) + if proc.stdout is not None: + try: + proc.stdout.close() + except Exception: + pass + + +async def call_action( + action: str, + host: str = REME_DEFAULT_HOST, + port: int = REME_DEFAULT_PORT, + timeout: float = 30.0, + **kwargs, +) -> dict | str: + """POST to /{action}; return parsed JSON (dict) for JSON endpoints, raw text for SSE.""" + from ..components.client.http_client import HttpClient + + pieces: list[str] = [] + async with HttpClient(action=action, host=host, port=port, timeout=timeout, **kwargs) as client: + async for chunk in client(): + pieces.append(chunk) + raw = "".join(pieces) + try: + return json.loads(raw) + except (ValueError, json.JSONDecodeError): + return raw + + +async def call_and_check( + action: str, + host: str = REME_DEFAULT_HOST, + port: int = REME_DEFAULT_PORT, + validator: Callable[[Any], bool] | None = None, + expected: Any = None, + timeout: float = 30.0, + **kwargs, +) -> Any: + """Call action and verify response. Raises AssertionError on mismatch. + + - validator(result) -> bool: custom predicate. + - expected: deep-equality target (compared to result, or to result[key] when expected is dict). + """ + result = await call_action(action, host=host, port=port, timeout=timeout, **kwargs) + if validator is not None and not validator(result): + raise AssertionError(f"validator rejected response for action={action!r}: {result!r}") + if expected is not None: + if isinstance(expected, dict) and isinstance(result, dict): + for k, v in expected.items(): + if result.get(k) != v: + raise AssertionError( + f"action={action!r} expected {k}={v!r}, got {result.get(k)!r} (full: {result!r})", + ) + elif result != expected: + raise AssertionError(f"action={action!r} expected {expected!r}, got {result!r}") + return result diff --git a/tests4/unittest/test_common_steps.py b/tests4/unittest/test_common_steps.py new file mode 100644 index 00000000..a2f60559 --- /dev/null +++ b/tests4/unittest/test_common_steps.py @@ -0,0 +1,260 @@ +"""End-to-end tests for reme4 common steps: spawn `reme4 start`, drive via HTTP, +verify responses, then shut down. Each test uses an isolated cwd so the working_dir +(.reme by default) does not collide. +""" + +import asyncio +import os +import tempfile +import warnings + +from reme4 import __version__ as REME_VERSION +from reme4.utils import call_action, call_and_check, mock_reme_server + +warnings.filterwarnings("ignore", category=DeprecationWarning, module="jieba") +warnings.filterwarnings("ignore", category=DeprecationWarning, module="pkg_resources") + + +class _temp_chdir: + """chdir to path for the duration of the block; restore on exit.""" + + def __init__(self, path): + self.path = path + self._old = None + + def __enter__(self): + self._old = os.getcwd() + os.chdir(self.path) + return self + + def __exit__(self, *exc): + os.chdir(self._old) + + +def _run(coro): + """Run an async coroutine on a fresh isolated event loop.""" + asyncio.run(coro) + + +# --------------------------------------------------------------------------- +# Individual job tests +# --------------------------------------------------------------------------- + + +def test_version_job(): + """version job should return the package version string.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, _temp_chdir(tmp): + async with mock_reme_server() as (host, port): + await call_and_check( + "version", + host=host, + port=port, + validator=lambda r: ( + isinstance(r, dict) + and r.get("success") is True + and r.get("answer") == REME_VERSION + and r.get("metadata", {}).get("version") == REME_VERSION + ), + ) + print("βœ“ test_version_job passed") + + _run(run()) + + +def test_help_job(): + """help job should list jobs except itself.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, _temp_chdir(tmp): + async with mock_reme_server() as (host, port): + result = await call_and_check( + "help", + host=host, + port=port, + validator=lambda r: ( + isinstance(r, dict) + and r.get("success") is True + and isinstance(r.get("answer"), str) + and r.get("metadata", {}).get("job_count", 0) > 0 + and "help" not in r["answer"] + ), + ) + # Spot-check that a couple of known jobs appear in the listing. + answer = result["answer"] + for expected_job in ("version", "health_check", "search"): + if expected_job not in answer: + raise AssertionError(f"help output missing job {expected_job!r}: {answer!r}") + print("βœ“ test_help_job passed") + + _run(run()) + + +def test_health_check_job(): + """health_check job should return a structured health snapshot.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, _temp_chdir(tmp): + async with mock_reme_server() as (host, port): + result = await call_and_check( + "health_check", + host=host, + port=port, + validator=lambda r: ( + isinstance(r, dict) + and r.get("success") is True + and isinstance(r.get("metadata"), dict) + and isinstance(r["metadata"].get("health"), dict) + and r["metadata"]["health"].get("version") == REME_VERSION + and isinstance(r["metadata"]["health"].get("components"), dict) + ), + ) + # Validate that each expected component type is in the snapshot. + components = result["metadata"]["health"]["components"] + for ctype in ( + "embedding_model", + "file_graph", + "file_store", + "file_watcher", + "keyword_index", + ): + if ctype not in components: + raise AssertionError(f"health snapshot missing component {ctype!r}: {components!r}") + print("βœ“ test_health_check_job passed") + + _run(run()) + + +def test_search_job_empty_store(): + """search on an empty store should return successfully with zero results.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, _temp_chdir(tmp): + async with mock_reme_server() as (host, port): + await call_and_check( + "search", + host=host, + port=port, + query="hello world", + limit=5, + validator=lambda r: ( + isinstance(r, dict) + and r.get("success") is True + and isinstance(r.get("metadata"), dict) + and isinstance(r["metadata"].get("counts"), dict) + and r["metadata"]["counts"].get("returned", -1) == 0 + ), + ) + print("βœ“ test_search_job_empty_store passed") + + _run(run()) + + +def test_search_job_missing_query(): + """search without a query should surface the assertion error in `answer`.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, _temp_chdir(tmp): + async with mock_reme_server() as (host, port): + result = await call_action("search", host=host, port=port, query="") + if not isinstance(result, dict): + raise AssertionError(f"expected dict response, got {result!r}") + if "query" not in str(result.get("answer", "")).lower(): + raise AssertionError(f"expected query-related error in answer, got {result!r}") + print("βœ“ test_search_job_missing_query passed") + + _run(run()) + + +def test_demo_job(): + """demo job should echo back the normalized query and adjusted min_score.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, _temp_chdir(tmp): + async with mock_reme_server() as (host, port): + await call_and_check( + "demo", + host=host, + port=port, + query=" Hello World ", + min_score=0.8, + validator=lambda r: ( + isinstance(r, dict) + and r.get("success") is True + and "hello world" in str(r.get("answer", "")) + and abs(r.get("metadata", {}).get("adjusted_min_score", 0) - 0.72) < 1e-6 + ), + ) + print("βœ“ test_demo_job passed") + + _run(run()) + + +# --------------------------------------------------------------------------- +# Aggregate test: reuse one server instance for all jobs (faster). +# --------------------------------------------------------------------------- + + +def test_all_jobs_one_server(): + """Run every common job against a single shared server for efficiency.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, _temp_chdir(tmp): + async with mock_reme_server() as (host, port): + # version + await call_and_check( + "version", + host=host, + port=port, + validator=lambda r: isinstance(r, dict) and r.get("answer") == REME_VERSION, + ) + # help + await call_and_check( + "help", + host=host, + port=port, + validator=lambda r: isinstance(r, dict) and r.get("metadata", {}).get("job_count", 0) > 0, + ) + # health_check + await call_and_check( + "health_check", + host=host, + port=port, + validator=lambda r: isinstance(r, dict) + and isinstance( + r.get("metadata", {}).get("health"), + dict, + ), + ) + # search (empty store) + await call_and_check( + "search", + host=host, + port=port, + query="anything", + validator=lambda r: isinstance(r, dict) and r.get("success") is True, + ) + # demo + await call_and_check( + "demo", + host=host, + port=port, + query="Foo", + validator=lambda r: isinstance(r, dict) and "foo" in str(r.get("answer", "")), + ) + print("βœ“ test_all_jobs_one_server passed") + + _run(run()) + + +if __name__ == "__main__": + print("\n=== reme4 common steps E2E tests ===") + test_version_job() + test_help_job() + test_health_check_job() + test_search_job_empty_store() + test_search_job_missing_query() + test_demo_job() + test_all_jobs_one_server() + print("\nζ‰€ζœ‰ζ΅‹θ―•ι€šθΏ‡!") diff --git a/tests4/unittest/test_default_file_parser.py b/tests4/unittest/test_default_file_parser.py index 75f68ba5..5730078b 100644 --- a/tests4/unittest/test_default_file_parser.py +++ b/tests4/unittest/test_default_file_parser.py @@ -169,6 +169,175 @@ def test_file_chunk_properties(): asyncio.run(run()) +def test_parse_links_bare(): + """Bare wikilink: [[target]].""" + links = DefaultFileParser.parse_links("see [[note]]", "src.md") + assert len(links) == 1 + link = links[0] + assert link.source_path == "src.md" + assert link.target_path == "note" + assert link.target_anchor is None + assert link.predicate is None + print("βœ“ test_parse_links_bare passed") + + +def test_parse_links_with_anchor(): + """Wikilink with anchor: [[target#anchor]].""" + links = DefaultFileParser.parse_links("see [[note#section A]]", "src.md") + assert len(links) == 1 + assert links[0].target_path == "note" + assert links[0].target_anchor == "section A" + assert links[0].predicate is None + print("βœ“ test_parse_links_with_anchor passed") + + +def test_parse_links_alias_dropped(): + """Alias after '|' is consumed but not captured as anchor.""" + links = DefaultFileParser.parse_links("see [[note|display text]]", "src.md") + assert len(links) == 1 + assert links[0].target_path == "note" + assert links[0].target_anchor is None + print("βœ“ test_parse_links_alias_dropped passed") + + +def test_parse_links_anchor_and_alias(): + """[[target#anchor|alias]] β€” anchor captured, alias dropped.""" + links = DefaultFileParser.parse_links("see [[note#sec|disp]]", "src.md") + assert len(links) == 1 + assert links[0].target_path == "note" + assert links[0].target_anchor == "sec" + print("βœ“ test_parse_links_anchor_and_alias passed") + + +def test_parse_links_predicate_simple(): + """Dataview inline: predicate:: [[target]].""" + links = DefaultFileParser.parse_links("author:: [[Alice]]", "src.md") + assert len(links) == 1 + assert links[0].predicate == "author" + assert links[0].target_path == "Alice" + assert links[0].target_anchor is None + print("βœ“ test_parse_links_predicate_simple passed") + + +def test_parse_links_predicate_bracketed(): + """Dataview inline-bracket: [predicate:: [[target]]].""" + links = DefaultFileParser.parse_links("text [author:: [[Alice]]] more", "src.md") + assert len(links) == 1 + assert links[0].predicate == "author" + assert links[0].target_path == "Alice" + print("βœ“ test_parse_links_predicate_bracketed passed") + + +def test_parse_links_predicate_bracketed_with_anchor(): + """[predicate:: [[target_path#target_anchor]]] β€” combined form.""" + links = DefaultFileParser.parse_links( + "[predicate:: [[target_path#target_anchor]]]", + "src.md", + ) + assert len(links) == 1 + link = links[0] + assert link.source_path == "src.md" + assert link.predicate == "predicate" + assert link.target_path == "target_path" + assert link.target_anchor == "target_anchor" + print("βœ“ test_parse_links_predicate_bracketed_with_anchor passed") + + +def test_parse_links_predicate_sticks_to_first(): + """Predicate attaches only to the immediately following wikilink.""" + links = DefaultFileParser.parse_links("pred:: [[a]] and bare [[b]]", "src.md") + assert len(links) == 2 + assert links[0].predicate == "pred" and links[0].target_path == "a" + assert links[1].predicate is None and links[1].target_path == "b" + print("βœ“ test_parse_links_predicate_sticks_to_first passed") + + +def test_parse_links_multiple_on_one_line(): + """Multiple bare wikilinks on the same line are all captured.""" + links = DefaultFileParser.parse_links("see [[x]] and [[y#h]]", "src.md") + assert [(link.target_path, link.target_anchor) for link in links] == [ + ("x", None), + ("y", "h"), + ] + print("βœ“ test_parse_links_multiple_on_one_line passed") + + +def test_parse_links_no_match(): + """Strings without [[]] yield no links, even if '::' appears.""" + assert len(DefaultFileParser.parse_links("no link here :: foo", "src.md")) == 0 + assert len(DefaultFileParser.parse_links("plain text without brackets", "src.md")) == 0 + assert len(DefaultFileParser.parse_links("", "src.md")) == 0 + print("βœ“ test_parse_links_no_match passed") + + +def test_parse_links_predicate_with_dash_and_digits(): + """Predicate identifier accepts letters, digits, underscore, dash.""" + links = DefaultFileParser.parse_links("see-also-2:: [[target]]", "src.md") + assert len(links) == 1 + assert links[0].predicate == "see-also-2" + assert links[0].target_path == "target" + print("βœ“ test_parse_links_predicate_with_dash_and_digits passed") + + +def test_parse_links_in_file(): + """Integration: parse() populates FileNode.links from file content.""" + + async def run(): + content = ( + "---\n" + "title: demo\n" + "---\n" + "\n" + "Intro paragraph with [[alpha]] and [[beta#h2]].\n" + "author:: [[Alice]]\n" + "[ref:: [[paper#chapter 1]]]\n" + ) + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".md") as f: + f.write(content) + temp_path = f.name + + try: + parser = DefaultFileParser() + file_node, _ = await parser.parse(temp_path) + triples = {(link.predicate, link.target_path, link.target_anchor) for link in file_node.links} + assert (None, "alpha", None) in triples + assert (None, "beta", "h2") in triples + assert ("author", "Alice", None) in triples + assert ("ref", "paper", "chapter 1") in triples + assert all(link.source_path == file_node.path for link in file_node.links) + print("βœ“ test_parse_links_in_file passed") + finally: + os.unlink(temp_path) + + asyncio.run(run()) + + +def test_parse_links_empty_when_no_content(): + """Empty file and front-matter-only file both yield no links.""" + + async def run(): + # Empty file + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".md") as f: + empty_path = f.name + # Front-matter-only file + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".md") as f: + f.write("---\ntitle: x\n---\n") + fm_only_path = f.name + + try: + parser = DefaultFileParser() + node1, _ = await parser.parse(empty_path) + node2, _ = await parser.parse(fm_only_path) + assert node1.links == [] + assert node2.links == [] + print("βœ“ test_parse_links_empty_when_no_content passed") + finally: + os.unlink(empty_path) + os.unlink(fm_only_path) + + asyncio.run(run()) + + def test_min_chunk_and_overlap_size(): """Test that minimum chunk and overlap sizes are enforced.""" @@ -201,5 +370,18 @@ if __name__ == "__main__": test_parse_with_custom_encoding() test_file_node_properties() test_file_chunk_properties() + test_parse_links_bare() + test_parse_links_with_anchor() + test_parse_links_alias_dropped() + test_parse_links_anchor_and_alias() + test_parse_links_predicate_simple() + test_parse_links_predicate_bracketed() + test_parse_links_predicate_bracketed_with_anchor() + test_parse_links_predicate_sticks_to_first() + test_parse_links_multiple_on_one_line() + test_parse_links_no_match() + test_parse_links_predicate_with_dash_and_digits() + test_parse_links_in_file() + test_parse_links_empty_when_no_content() test_min_chunk_and_overlap_size() print("\nζ‰€ζœ‰ζ΅‹θ―•ι€šθΏ‡!")