This commit is contained in:
jinli.yl 2026-05-17 13:25:11 +08:00
parent 0b3c124c79
commit b8d42d1ae5
19 changed files with 1005 additions and 73 deletions

View file

@ -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
]

View file

@ -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):

View file

@ -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<predicate>[A-Za-z][\w-]*)\s*::\s*)?"
r"\[\[\s*(?P<target>[^\[\]|#]+?)"
r"(?:#(?P<anchor>[^\[\]|]+?))?"
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."""

View file

@ -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:

View file

@ -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

View file

@ -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]):

View file

@ -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."""

View file

@ -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"

View file

@ -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")

View file

@ -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",

View file

@ -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

View file

@ -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")

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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",

View file

@ -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

View file

@ -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所有测试通过!")

View file

@ -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所有测试通过!")