From 5162873c3ca57be2682898321bd1b32c88380ef5 Mon Sep 17 00:00:00 2001 From: xinyao1994 Date: Wed, 5 Aug 2026 20:52:03 +0800 Subject: [PATCH 1/3] feat(memory): add wikimem implementation --- .gitignore | 1 + agent_plugin/wikimem/__init__.py | 45 + agent_plugin/wikimem/team_sync.py | 486 +++++ evaluation/wikimem/__init__.py | 259 +++ evaluation/wikimem/ama_bench.py | 680 ++++++ evaluation/wikimem/example_eval.py | 1561 ++++++++++++++ evaluation/wikimem/llm_semantics.py | 314 +++ evaluation/wikimem/locomo_refined.py | 1214 +++++++++++ evaluation/wikimem/longmemeval.py | 563 +++++ evaluation/wikimem/qmd_consensus.py | 1282 ++++++++++++ evaluation/wikimem/retained_eval.py | 1828 +++++++++++++++++ evaluation/wikimem/retrieval_profile.py | 1763 ++++++++++++++++ evaluation/wikimem/wiki_builder.py | 586 ++++++ src/construction/evolver_impl/__init__.py | 1 + .../evolver_impl/wikimem_baseline_evolver.py | 150 ++ src/construction/extractor_impl/__init__.py | 1 + .../wikimem_baseline_extractor.py | 324 +++ src/retrieval/recaller_impl/__init__.py | 1 + .../recaller_impl/wikimem_memdir_recaller.py | 87 + src/retrieval/wikimem_memdir.py | 523 +++++ src/retrieval/wikimem_options.py | 129 ++ .../agent_plugin/test_wikimem_team_sync.py | 227 ++ .../test_wikimem_baseline_evolver.py | 92 + .../test_wikimem_baseline_extractor.py | 92 + .../unit/evaluation/test_wikimem_ama_bench.py | 378 ++++ .../evaluation/test_wikimem_example_eval.py | 1345 ++++++++++++ .../evaluation/test_wikimem_longmemeval.py | 216 ++ .../evaluation/test_wikimem_qmd_consensus.py | 189 ++ .../evaluation/test_wikimem_retained_eval.py | 568 +++++ .../test_wikimem_retrieval_profile.py | 490 +++++ tests/unit/retrieval/test_wikimem_memdir.py | 147 ++ .../retrieval/test_wikimem_memdir_recaller.py | 60 + tests/unit/retrieval/test_wikimem_options.py | 69 + 33 files changed, 15671 insertions(+) create mode 100644 agent_plugin/wikimem/__init__.py create mode 100644 agent_plugin/wikimem/team_sync.py create mode 100644 evaluation/wikimem/__init__.py create mode 100644 evaluation/wikimem/ama_bench.py create mode 100644 evaluation/wikimem/example_eval.py create mode 100644 evaluation/wikimem/llm_semantics.py create mode 100644 evaluation/wikimem/locomo_refined.py create mode 100644 evaluation/wikimem/longmemeval.py create mode 100644 evaluation/wikimem/qmd_consensus.py create mode 100644 evaluation/wikimem/retained_eval.py create mode 100644 evaluation/wikimem/retrieval_profile.py create mode 100644 evaluation/wikimem/wiki_builder.py create mode 100644 src/construction/evolver_impl/wikimem_baseline_evolver.py create mode 100644 src/construction/extractor_impl/wikimem_baseline_extractor.py create mode 100644 src/retrieval/recaller_impl/wikimem_memdir_recaller.py create mode 100644 src/retrieval/wikimem_memdir.py create mode 100644 src/retrieval/wikimem_options.py create mode 100644 tests/unit/agent_plugin/test_wikimem_team_sync.py create mode 100644 tests/unit/construction/test_wikimem_baseline_evolver.py create mode 100644 tests/unit/construction/test_wikimem_baseline_extractor.py create mode 100644 tests/unit/evaluation/test_wikimem_ama_bench.py create mode 100644 tests/unit/evaluation/test_wikimem_example_eval.py create mode 100644 tests/unit/evaluation/test_wikimem_longmemeval.py create mode 100644 tests/unit/evaluation/test_wikimem_qmd_consensus.py create mode 100644 tests/unit/evaluation/test_wikimem_retained_eval.py create mode 100644 tests/unit/evaluation/test_wikimem_retrieval_profile.py create mode 100644 tests/unit/retrieval/test_wikimem_memdir.py create mode 100644 tests/unit/retrieval/test_wikimem_memdir_recaller.py create mode 100644 tests/unit/retrieval/test_wikimem_options.py diff --git a/.gitignore b/.gitignore index c100c783..25b0f21c 100644 --- a/.gitignore +++ b/.gitignore @@ -24,6 +24,7 @@ logs/ automated_testing/data/ evaluation/benchmark/data/ results.json +tmp/ # 预下载的本地模型（GB 级二进制；用 deploy/docker/download-models.sh 自行拉取，不入库） deploy/docker/models/ diff --git a/agent_plugin/wikimem/__init__.py b/agent_plugin/wikimem/__init__.py new file mode 100644 index 00000000..8abd4647 --- /dev/null +++ b/agent_plugin/wikimem/__init__.py @@ -0,0 +1,45 @@ +"""wikimem agent adapters.""" + +from __future__ import annotations + +from agent_plugin.wikimem.team_sync import ( + FetchOutcome, + HashesProbe, + LocalTeamMemorySnapshot, + PullOutcome, + PushOutcomeSummary, + PutOutcome, + RemoteTeamMemoryData, + SkippedSecretFile, + SyncState, + TeamMemoryFailureKind, + TeamMemoryRemote, + TeamMemorySyncFailure, + batch_delta_by_bytes, + hash_content, + pull_team_memory, + push_team_memory, + read_local_team_memory, + validate_relative_team_memory_key, +) + +__all__ = [ + "FetchOutcome", + "HashesProbe", + "LocalTeamMemorySnapshot", + "PullOutcome", + "PushOutcomeSummary", + "PutOutcome", + "RemoteTeamMemoryData", + "SkippedSecretFile", + "SyncState", + "TeamMemoryFailureKind", + "TeamMemoryRemote", + "TeamMemorySyncFailure", + "batch_delta_by_bytes", + "hash_content", + "pull_team_memory", + "push_team_memory", + "read_local_team_memory", + "validate_relative_team_memory_key", +] diff --git a/agent_plugin/wikimem/team_sync.py b/agent_plugin/wikimem/team_sync.py new file mode 100644 index 00000000..088aad06 --- /dev/null +++ b/agent_plugin/wikimem/team_sync.py @@ -0,0 +1,486 @@ +"""wikimem team memory synchronization adapter.""" + +from __future__ import annotations + +import hashlib +import json +import os +import urllib.parse +from dataclasses import dataclass, field +from enum import Enum +from pathlib import Path +from typing import Protocol + +MAX_FILE_SIZE_BYTES = 250_000 +MAX_PUT_BODY_BYTES = 200_000 +MAX_CONFLICT_RETRIES = 2 + +_SECRET_MARKERS = ( + "ghp_", + "gho_", + "ghu_", + "ghs_", + "ghr_", + "github_pat_", + "sk-ant-", + "sk-proj-", + "sk-svcacct-", + "sk-admin-", + "xoxb-", + "xoxp-", + "xoxe-", + "xapp-", + "akia", + "asia", + "abia", + "acca", +) + + +@dataclass +class SyncState: + last_known_checksum: str | None = None + server_checksums: dict[str, str] = field(default_factory=dict) + server_max_entries: int | None = None + + +@dataclass(frozen=True) +class SkippedSecretFile: + path: str + reason: str + + +@dataclass(frozen=True) +class LocalTeamMemorySnapshot: + entries: dict[str, str] + skipped_secrets: list[SkippedSecretFile] + + +@dataclass(frozen=True) +class RemoteTeamMemoryData: + checksum: str | None + entries: dict[str, str] + entry_checksums: dict[str, str] + + +@dataclass(frozen=True) +class FetchOutcome: + kind: str + checksum: str | None = None + data: RemoteTeamMemoryData | None = None + + @classmethod + def not_modified(cls, checksum: str | None = None) -> FetchOutcome: + return cls(kind="not_modified", checksum=checksum) + + @classmethod + def empty(cls) -> FetchOutcome: + return cls(kind="empty") + + @classmethod + def data(cls, data: RemoteTeamMemoryData) -> FetchOutcome: + return cls(kind="data", data=data) + + +@dataclass(frozen=True) +class HashesProbe: + checksum: str | None + entry_checksums: dict[str, str] + + +@dataclass(frozen=True) +class PutOutcome: + kind: str + checksum: str | None = None + max_entries: int | None = None + received_entries: int | None = None + + @classmethod + def success(cls, checksum: str | None = None) -> PutOutcome: + return cls(kind="success", checksum=checksum) + + @classmethod + def conflict(cls) -> PutOutcome: + return cls(kind="conflict") + + @classmethod + def too_many_entries(cls, max_entries: int, received_entries: int) -> PutOutcome: + return cls( + kind="too_many_entries", + max_entries=max_entries, + received_entries=received_entries, + ) + + +class TeamMemoryFailureKind(str, Enum): + AUTH = "auth" + TIMEOUT = "timeout" + NETWORK = "network" + CONFLICT = "conflict" + NO_OAUTH = "no_oauth" + NO_REPO = "no_repo" + PARSE = "parse" + UNKNOWN = "unknown" + + +@dataclass(frozen=True) +class TeamMemorySyncFailure: + kind: TeamMemoryFailureKind + message: str + http_status: int | None = None + permanent: bool = False + + +@dataclass(frozen=True) +class PullOutcome: + success: bool + checksum: str | None + not_modified: bool + is_empty: bool + files_written: int + failure: TeamMemorySyncFailure | None = None + + +@dataclass(frozen=True) +class PushOutcomeSummary: + success: bool + files_uploaded: int + checksum: str | None + conflict: bool + skipped_secrets: list[SkippedSecretFile] + server_max_entries: int | None + failure: TeamMemorySyncFailure | None = None + + +class TeamMemoryRemote(Protocol): + async def fetch(self, repo_slug: str, if_none_match: str | None) -> FetchOutcome: + """Fetch all remote team memory entries.""" + + async def fetch_hashes(self, repo_slug: str) -> HashesProbe: + """Fetch remote entry checksums.""" + + async def put_entries( + self, repo_slug: str, if_match: str | None, entries: dict[str, str] + ) -> PutOutcome: + """Upload a batch of team memory entries.""" + + +async def pull_team_memory( + remote: TeamMemoryRemote, + state: SyncState, + team_memory_root: str | os.PathLike[str], + repo_slug: str, + skip_etag_cache: bool = False, +) -> PullOutcome: + if_none_match = None if skip_etag_cache else state.last_known_checksum + try: + fetch = await remote.fetch(repo_slug, if_none_match) + except TeamMemorySyncFailure as failure: + return _pull_failure(failure) + except Exception as error: # pragma: no cover - defensive adapter boundary + return _pull_failure(_unknown_failure(str(error))) + + if fetch.kind == "not_modified": + return PullOutcome( + success=True, + checksum=fetch.checksum, + not_modified=True, + is_empty=False, + files_written=0, + ) + if fetch.kind == "empty": + state.last_known_checksum = None + state.server_checksums.clear() + return PullOutcome( + success=True, + checksum=None, + not_modified=False, + is_empty=True, + files_written=0, + ) + if fetch.kind != "data" or fetch.data is None: + return _pull_failure(_unknown_failure(f"Unexpected fetch outcome: {fetch.kind}")) + + try: + files_written = _write_remote_entries(Path(team_memory_root), fetch.data.entries) + except Exception as error: + return _pull_failure(_unknown_failure(f"Failed to write pulled team memory: {error}")) + + state.last_known_checksum = fetch.data.checksum + if fetch.data.entry_checksums: + state.server_checksums = dict(fetch.data.entry_checksums) + else: + state.server_checksums = { + key: hash_content(value) for key, value in fetch.data.entries.items() + } + return PullOutcome( + success=True, + checksum=fetch.data.checksum, + not_modified=False, + is_empty=False, + files_written=files_written, + ) + + +async def push_team_memory( + remote: TeamMemoryRemote, + state: SyncState, + team_memory_root: str | os.PathLike[str], + repo_slug: str, +) -> PushOutcomeSummary: + try: + local = read_local_team_memory(team_memory_root, state.server_max_entries) + except Exception as error: + return PushOutcomeSummary( + success=False, + files_uploaded=0, + checksum=None, + conflict=False, + skipped_secrets=[], + server_max_entries=state.server_max_entries, + failure=_unknown_failure(f"Failed to read local team memory: {error}"), + ) + + local_hashes = {key: hash_content(value) for key, value in local.entries.items()} + conflict_attempts = 0 + + while True: + delta = _compute_delta(local.entries, local_hashes, state.server_checksums) + if not delta: + return PushOutcomeSummary( + success=True, + files_uploaded=0, + checksum=state.last_known_checksum, + conflict=False, + skipped_secrets=local.skipped_secrets, + server_max_entries=state.server_max_entries, + ) + + files_uploaded = 0 + needs_retry = False + for batch in batch_delta_by_bytes(delta, MAX_PUT_BODY_BYTES): + try: + outcome = await remote.put_entries(repo_slug, state.last_known_checksum, batch) + except TeamMemorySyncFailure as failure: + return PushOutcomeSummary( + success=False, + files_uploaded=files_uploaded, + checksum=state.last_known_checksum, + conflict=conflict_attempts > 0, + skipped_secrets=local.skipped_secrets, + server_max_entries=state.server_max_entries, + failure=failure, + ) + except Exception as error: # pragma: no cover - defensive adapter boundary + return PushOutcomeSummary( + success=False, + files_uploaded=files_uploaded, + checksum=state.last_known_checksum, + conflict=conflict_attempts > 0, + skipped_secrets=local.skipped_secrets, + server_max_entries=state.server_max_entries, + failure=_unknown_failure(str(error)), + ) + + if outcome.kind == "success": + if outcome.checksum is not None: + state.last_known_checksum = outcome.checksum + for key in batch: + state.server_checksums[key] = local_hashes[key] + files_uploaded += len(batch) + continue + + if outcome.kind == "too_many_entries": + state.server_max_entries = outcome.max_entries + return PushOutcomeSummary( + success=False, + files_uploaded=files_uploaded, + checksum=state.last_known_checksum, + conflict=False, + skipped_secrets=local.skipped_secrets, + server_max_entries=state.server_max_entries, + failure=TeamMemorySyncFailure( + kind=TeamMemoryFailureKind.UNKNOWN, + message="Server rejected team memory: too many entries", + http_status=413, + permanent=True, + ), + ) + + if outcome.kind == "conflict": + if conflict_attempts >= MAX_CONFLICT_RETRIES: + return PushOutcomeSummary( + success=False, + files_uploaded=files_uploaded, + checksum=state.last_known_checksum, + conflict=True, + skipped_secrets=local.skipped_secrets, + server_max_entries=state.server_max_entries, + failure=TeamMemorySyncFailure( + kind=TeamMemoryFailureKind.CONFLICT, + message="Team memory push conflicted after retries", + http_status=412, + ), + ) + try: + hashes = await remote.fetch_hashes(repo_slug) + except TeamMemorySyncFailure as failure: + return PushOutcomeSummary( + success=False, + files_uploaded=files_uploaded, + checksum=state.last_known_checksum, + conflict=True, + skipped_secrets=local.skipped_secrets, + server_max_entries=state.server_max_entries, + failure=failure, + ) + state.last_known_checksum = hashes.checksum + state.server_checksums = dict(hashes.entry_checksums) + conflict_attempts += 1 + needs_retry = True + break + + return PushOutcomeSummary( + success=False, + files_uploaded=files_uploaded, + checksum=state.last_known_checksum, + conflict=conflict_attempts > 0, + skipped_secrets=local.skipped_secrets, + server_max_entries=state.server_max_entries, + failure=_unknown_failure(f"Unexpected put outcome: {outcome.kind}"), + ) + + if needs_retry: + continue + + return PushOutcomeSummary( + success=True, + files_uploaded=files_uploaded, + checksum=state.last_known_checksum, + conflict=False, + skipped_secrets=local.skipped_secrets, + server_max_entries=state.server_max_entries, + ) + + +def read_local_team_memory( + team_memory_root: str | os.PathLike[str], + server_max_entries: int | None = None, +) -> LocalTeamMemorySnapshot: + root = Path(team_memory_root) + entries: dict[str, str] = {} + skipped_secrets: list[SkippedSecretFile] = [] + if not root.exists(): + return LocalTeamMemorySnapshot(entries=entries, skipped_secrets=skipped_secrets) + + files = (candidate for candidate in root.rglob("*") if candidate.is_file()) + for path in sorted(files, key=str): + if path.stat().st_size > MAX_FILE_SIZE_BYTES: + continue + relative = path.relative_to(root).as_posix() + content = path.read_text(encoding="utf-8") + if _has_potential_secret(content): + skipped_secrets.append(SkippedSecretFile(path=relative, reason="potential_secret")) + continue + entries[relative] = content + + if server_max_entries is not None: + entries = dict(sorted(entries.items())[:server_max_entries]) + else: + entries = dict(sorted(entries.items())) + return LocalTeamMemorySnapshot(entries=entries, skipped_secrets=skipped_secrets) + + +def batch_delta_by_bytes(delta: dict[str, str], max_body_bytes: int) -> list[dict[str, str]]: + if not delta: + return [] + + batches: list[dict[str, str]] = [] + current: dict[str, str] = {} + for key, value in sorted(delta.items()): + current[key] = value + if _estimate_put_body_bytes(current) > max_body_bytes and len(current) > 1: + current.pop(key) + batches.append(current) + current = {key: value} + if current: + batches.append(current) + return batches + + +def hash_content(content: str) -> str: + digest = hashlib.sha256(content.encode("utf-8")).hexdigest() + return f"sha256:{digest}" + + +def validate_relative_team_memory_key(key: str) -> str: + if "\0" in key or "\\" in key or key.startswith("/") or _has_windows_prefix(key): + raise ValueError(f"Invalid team memory key: {key}") + + decoded = urllib.parse.unquote(key) + if decoded != key and (".." in decoded or "/" in decoded or "\\" in decoded): + raise ValueError(f"Invalid team memory key: {key}") + + parts = key.split("/") + if not parts or any(part in {"", ".", ".."} for part in parts): + raise ValueError(f"Invalid team memory key: {key}") + return "/".join(parts) + + +def _write_remote_entries(team_memory_root: Path, entries: dict[str, str]) -> int: + team_memory_root.mkdir(parents=True, exist_ok=True) + files_written = 0 + for key, content in sorted(entries.items()): + relative = validate_relative_team_memory_key(key) + path = team_memory_root.joinpath(*relative.split("/")) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(content, encoding="utf-8") + files_written += 1 + return files_written + + +def _compute_delta( + local_entries: dict[str, str], + local_hashes: dict[str, str], + server_checksums: dict[str, str], +) -> dict[str, str]: + return { + key: value + for key, value in sorted(local_entries.items()) + if server_checksums.get(key) != local_hashes.get(key) + } + + +def _estimate_put_body_bytes(entries: dict[str, str]) -> int: + body = json.dumps({"entries": dict(sorted(entries.items()))}, separators=(",", ":")) + return len(body.encode("utf-8")) + + +def _has_potential_secret(value: str) -> bool: + lower = value.lower() + if "-----begin" in lower and "private key-----" in lower: + return True + return any(marker in lower for marker in _SECRET_MARKERS) + + +def _has_windows_prefix(key: str) -> bool: + return len(key) >= 2 and key[1] == ":" and key[0].isalpha() + + +def _pull_failure(failure: TeamMemorySyncFailure) -> PullOutcome: + return PullOutcome( + success=False, + checksum=None, + not_modified=False, + is_empty=False, + files_written=0, + failure=failure, + ) + + +def _unknown_failure(message: str) -> TeamMemorySyncFailure: + return TeamMemorySyncFailure( + kind=TeamMemoryFailureKind.UNKNOWN, + message=message, + permanent=False, + ) diff --git a/evaluation/wikimem/__init__.py b/evaluation/wikimem/__init__.py new file mode 100644 index 00000000..a6558dc4 --- /dev/null +++ b/evaluation/wikimem/__init__.py @@ -0,0 +1,259 @@ +"""wikimem retained evaluation compatibility surface.""" + +from __future__ import annotations + +from evaluation.wikimem.retained_eval import ( + CaseScore, + CategoryScoreSummary, + ConversationRecord, + EvalCase, + EvalHarnessConfig, + EvalOutput, + EvalSummary, + LoCoMoQuestion, + ObservationNote, + PreparedSample, + ProgressUpdate, + RetrievalCoverageSummary, + SessionEvents, + StageProfileArtifact, + StageProfiler, + StageTimingRecord, + WikiMode, + build_eval_cases, + build_retained_memory_files, + format_progress_message, + parse_retrieval_plugin_list, + parse_sample_filter, + prepare_locomo_samples, + normalize_wiki_mode, + run_python_locomo_retrieval_eval, + run_retained_qmd_eval, + summarize_scores, + summarize_scores_by_locomo_category, + write_harness_artifacts, +) +from evaluation.wikimem.longmemeval import ( + LongMemEvalCaseScore, + LongMemEvalRetrievalResults, + LongMemEvalSummary, + adapt_longmemeval_to_locomo_samples, + load_longmemeval_samples, + run_python_longmemeval_retrieval_eval, +) +from evaluation.wikimem.ama_bench import ( + AmaEpisode, + AmaMethodQuestionResult, + AmaMethodSummary, + AmaQuestion, + AmaTurn, + ProxyRetrievalMetrics, + aggregate_ama_method_summaries, + build_ama_memory_files, + build_ama_precision_comparison, + derive_proxy_gold_turn_ids, + load_ama_dataset_split, + load_ama_episodes_from_jsonl, + run_python_ama_retrieval_eval, + run_python_wikimem_qmd_retrieval, + score_ama_retrieval_proxy, + select_ama_lexical_turn_ids, +) +from evaluation.wikimem.example_eval import ( + build_mem_gallery_python_workspaces, + build_meta_crag_python_workspaces, +) +from evaluation.wikimem.locomo_refined import ( + MultimodalArtifactSummary, + MultimodalBuildOptions, + MultimodalRecallHit, + OfflineExportSummary, + VisionEnrichmentConfig, + build_multimodal_memory_artifacts, + export_offline_locomo_refined_dataset, + recall_multimodal_artifacts_for_question, + rescore_case_with_multimodal_artifacts, +) +from evaluation.wikimem.qmd_consensus import ( + CachedFileLexicalFeatures, + CachedFileLineLexicalFeatures, + CandidateFile, + CandidateProposal, + QmdConsensusFileMetrics, + QueryAugmentation, + QuestionProfile, + RerankProposal, + RetrievedMemoryFile, + apply_query_augmentation, + build_cached_file_lexical_features, + build_qmd_consensus_augmentation, + build_qmd_consensus_candidate_proposals, + build_qmd_consensus_late_bridge_proposals, + build_qmd_consensus_rerank_proposals, + build_question_profile, + candidate_query_hits_with_features, + keyword_ngrams, + normalize_memory_path, + qmd_consensus_is_conservative, + score_line, + significant_phrases, + tokenize_fuzzy_query, + tokenize_query, +) +from evaluation.wikimem.retrieval_profile import ( + RetrievalProfileCoverage, + RetrievalProfileResult, + SessionSourceFile, + best_corpus_correction, + build_corpus_consensus_augmentation, + build_session_source_files, + collect_session_source_companions, + infer_session_number_from_path, + rank_global_session_sources, + retrieve_qmd_consensus_files, + scoped_budgets, + select_diverse_session_sources, + select_scoped_candidate_files, + session_source_token_document_frequency, + source_companion_budget, + source_injection_budget, +) +from evaluation.wikimem.llm_semantics import ( + MEMORY_KINDS, + QueryUnderstanding, + SemanticEntity, + SemanticMemory, + SemanticRelation, + SemanticSource, + extract_semantic_memories, + understand_query, +) +from evaluation.wikimem.wiki_builder import ( + EntityResolver, + MemoryConsolidator, + WikiBuildDiagnostics, + WikiBuildResult, + WikiBuilder, + WikiBuilderMode, + TemplateExtractor, +) + +__all__ = [ + "CaseScore", + "CategoryScoreSummary", + "ConversationRecord", + "EvalCase", + "EvalHarnessConfig", + "EvalOutput", + "EvalSummary", + "LoCoMoQuestion", + "ObservationNote", + "PreparedSample", + "ProgressUpdate", + "RetrievalCoverageSummary", + "SessionEvents", + "StageProfileArtifact", + "StageProfiler", + "StageTimingRecord", + "WikiMode", + "build_eval_cases", + "build_retained_memory_files", + "format_progress_message", + "parse_retrieval_plugin_list", + "parse_sample_filter", + "prepare_locomo_samples", + "normalize_wiki_mode", + "run_python_locomo_retrieval_eval", + "run_retained_qmd_eval", + "summarize_scores", + "summarize_scores_by_locomo_category", + "write_harness_artifacts", + "LongMemEvalCaseScore", + "LongMemEvalRetrievalResults", + "LongMemEvalSummary", + "adapt_longmemeval_to_locomo_samples", + "load_longmemeval_samples", + "run_python_longmemeval_retrieval_eval", + "AmaEpisode", + "AmaMethodQuestionResult", + "AmaMethodSummary", + "AmaQuestion", + "AmaTurn", + "ProxyRetrievalMetrics", + "aggregate_ama_method_summaries", + "build_ama_memory_files", + "build_ama_precision_comparison", + "derive_proxy_gold_turn_ids", + "load_ama_dataset_split", + "load_ama_episodes_from_jsonl", + "run_python_ama_retrieval_eval", + "run_python_wikimem_qmd_retrieval", + "score_ama_retrieval_proxy", + "select_ama_lexical_turn_ids", + "build_mem_gallery_python_workspaces", + "build_meta_crag_python_workspaces", + "MultimodalArtifactSummary", + "MultimodalBuildOptions", + "MultimodalRecallHit", + "OfflineExportSummary", + "VisionEnrichmentConfig", + "build_multimodal_memory_artifacts", + "export_offline_locomo_refined_dataset", + "recall_multimodal_artifacts_for_question", + "rescore_case_with_multimodal_artifacts", + "CachedFileLexicalFeatures", + "CachedFileLineLexicalFeatures", + "CandidateFile", + "CandidateProposal", + "QmdConsensusFileMetrics", + "QueryAugmentation", + "QuestionProfile", + "RerankProposal", + "RetrievedMemoryFile", + "apply_query_augmentation", + "build_cached_file_lexical_features", + "build_qmd_consensus_augmentation", + "build_qmd_consensus_candidate_proposals", + "build_qmd_consensus_late_bridge_proposals", + "build_qmd_consensus_rerank_proposals", + "build_question_profile", + "candidate_query_hits_with_features", + "keyword_ngrams", + "normalize_memory_path", + "qmd_consensus_is_conservative", + "score_line", + "significant_phrases", + "tokenize_fuzzy_query", + "tokenize_query", + "RetrievalProfileCoverage", + "RetrievalProfileResult", + "SessionSourceFile", + "best_corpus_correction", + "build_corpus_consensus_augmentation", + "build_session_source_files", + "collect_session_source_companions", + "infer_session_number_from_path", + "rank_global_session_sources", + "retrieve_qmd_consensus_files", + "scoped_budgets", + "select_diverse_session_sources", + "select_scoped_candidate_files", + "session_source_token_document_frequency", + "source_companion_budget", + "source_injection_budget", + "MEMORY_KINDS", + "QueryUnderstanding", + "SemanticEntity", + "SemanticMemory", + "SemanticRelation", + "SemanticSource", + "extract_semantic_memories", + "understand_query", + "EntityResolver", + "MemoryConsolidator", + "WikiBuildDiagnostics", + "WikiBuildResult", + "WikiBuilder", + "WikiBuilderMode", + "TemplateExtractor", +] diff --git a/evaluation/wikimem/ama_bench.py b/evaluation/wikimem/ama_bench.py new file mode 100644 index 00000000..6d8cc5b8 --- /dev/null +++ b/evaluation/wikimem/ama_bench.py @@ -0,0 +1,680 @@ +"""AMA-Bench retrieval-only helpers for the wikimem Python migration.""" + +from __future__ import annotations + +import json +import re +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +from evaluation.wikimem.qmd_consensus import RetrievedMemoryFile +from evaluation.wikimem.retained_eval import WikiMode, normalize_wiki_mode +from evaluation.wikimem.retrieval_profile import retrieve_qmd_consensus_files + + +@dataclass(frozen=True) +class AmaTurn: + turn_idx: int + action: str = "" + observation: str = "" + + +@dataclass(frozen=True) +class AmaQuestion: + question: str + answer: str + qa_type: str = "" + options: list[str] | None = None + + +@dataclass(frozen=True) +class AmaEpisode: + episode_id: str + task: str + task_type: str + domain: str + trajectory: list[AmaTurn] + qa_pairs: list[AmaQuestion] + + +@dataclass(frozen=True) +class ProxyRetrievalMetrics: + proxy_gold_turn_ids: list[int] + retrieved_turn_ids: list[int] + hit_turn_ids: list[int] + proxy_recall_at_k: float + proxy_precision_at_k: float + proxy_hit_at_k: float + answer_support_coverage: float + + +@dataclass(frozen=True) +class AmaMethodQuestionResult: + method_name: str + episode_id: str + domain: str + task_type: str + qa_type: str + question_index: int + question: str + golden_answer: str + proxy_metrics: ProxyRetrievalMetrics + retrieved_context: str + retrieved_turn_ids: list[int] + retrieved_file_paths: list[str] + retrieval_notes: list[str] + + +@dataclass(frozen=True) +class AmaMethodSummary: + method_name: str + questions: int + proxy_recall_at_k: float + proxy_precision_at_k: float + proxy_hit_at_k: float + answer_support_coverage: float + judged_questions: int = 0 + judge_accuracy: float | None = None + + +def load_ama_episodes_from_jsonl( + path: str | Path, + limit: int | None = None, +) -> list[AmaEpisode]: + episodes = [] + buffer = "" + start_line = 1 + lines = Path(path).read_text(encoding="utf-8").splitlines(keepends=True) + for line_number, line in enumerate(lines, 1): + if not buffer and not line.strip(): + continue + if not buffer: + start_line = line_number + buffer += line + stripped = buffer.strip() + try: + raw = json.loads(_escape_json_string_newlines(stripped)) + except json.JSONDecodeError as error: + if line_number != len(lines): + continue + raise ValueError(f"failed to parse AMA episode line {start_line}: {error}") from error + episodes.append(_parse_ama_episode(raw)) + buffer = "" + if limit is not None and len(episodes) >= limit: + break + if buffer.strip(): + raise ValueError(f"failed to parse AMA episode line {start_line}: incomplete JSON record") + return episodes + + +def load_ama_dataset_split( + dataset_root: str | Path, + subset: str = "open_end", + sample_limit: int | None = None, + question_limit: int | None = None, +) -> list[AmaEpisode]: + file_name = "open_end_qa_set.jsonl" if subset == "open_end" else "mcq_set.jsonl" + episodes = load_ama_episodes_from_jsonl( + Path(dataset_root) / "test" / file_name, + limit=sample_limit, + ) + if question_limit is not None: + episodes = [ + AmaEpisode( + episode_id=episode.episode_id, + task=episode.task, + task_type=episode.task_type, + domain=episode.domain, + trajectory=episode.trajectory, + qa_pairs=episode.qa_pairs[:question_limit], + ) + for episode in episodes + ] + return episodes + + +def score_ama_retrieval_proxy( + episode: AmaEpisode, + question: AmaQuestion, + retrieved_turn_ids: list[int], +) -> ProxyRetrievalMetrics: + gold_turn_ids = derive_proxy_gold_turn_ids(episode, question) + gold_set = set(gold_turn_ids) + unique_retrieved = _unique_preserve_order(retrieved_turn_ids) + hit_turn_ids = [turn_id for turn_id in unique_retrieved if turn_id in gold_set] + recall = _ratio(len(hit_turn_ids), len(gold_turn_ids)) + precision = _ratio(len(hit_turn_ids), len(unique_retrieved)) + return ProxyRetrievalMetrics( + proxy_gold_turn_ids=gold_turn_ids, + retrieved_turn_ids=unique_retrieved, + hit_turn_ids=hit_turn_ids, + proxy_recall_at_k=_round4(recall), + proxy_precision_at_k=_round4(precision), + proxy_hit_at_k=1.0 if hit_turn_ids else 0.0, + answer_support_coverage=_round4(recall), + ) + + +def derive_proxy_gold_turn_ids(episode: AmaEpisode, question: AmaQuestion) -> list[int]: + answer_tokens = _normalize_token_set(question.answer) + question_tokens = _normalize_token_set(question.question) + scored = [] + for turn in episode.trajectory: + turn_tokens = _normalize_token_set(f"{turn.action} {turn.observation}") + answer_overlap = len(answer_tokens & turn_tokens) + question_overlap = len(question_tokens & turn_tokens) + scored.append(((answer_overlap * 2) + question_overlap, answer_overlap, turn.turn_idx)) + scored.sort(key=lambda item: (-item[0], -item[1], -item[2])) + positives = [ + turn_idx + for score, answer_overlap, turn_idx in scored + if score >= 2 and answer_overlap >= 1 + ] + if not positives: + positives = [ + turn_idx + for _, answer_overlap, turn_idx in scored + if answer_overlap >= 1 + ][:1] + return sorted(set(positives)) + + +def run_python_wikimem_qmd_retrieval( + episode: AmaEpisode, + question: AmaQuestion, + *, + top_k: int, + method_name: str = "wikimem_qmd", + wiki_mode: WikiMode = "text", +) -> AmaMethodQuestionResult: + wiki_mode = normalize_wiki_mode(wiki_mode) + if method_name == "wikimem_qmd_ama": + lexical_turn_ids = select_ama_lexical_turn_ids( + episode.trajectory, + question.question, + top_k=top_k, + ) + files = _build_ama_lexical_files(episode, lexical_turn_ids) + metrics = score_ama_retrieval_proxy(episode, question, lexical_turn_ids) + return AmaMethodQuestionResult( + method_name=method_name, + episode_id=episode.episode_id, + domain=episode.domain, + task_type=episode.task_type, + qa_type=question.qa_type, + question_index=0, + question=question.question, + golden_answer=question.answer, + proxy_metrics=metrics, + retrieved_context=_format_retrieved_context(files), + retrieved_turn_ids=_unique_preserve_order(lexical_turn_ids), + retrieved_file_paths=[file.file_path for file in files], + retrieval_notes=[ + "retrieval_plugins=qmd_consensus,ama_anchor_consensus", + "ama_merge_mode=hybrid_fusion", + "proxy_turn_source=ama_lexical_primary", + "qmd_context_supplement=skipped_in_retrieval_only_mode", + ], + ) + + files = build_ama_memory_files(episode, wiki_mode=wiki_mode) + root_files = [file for file in files if "/wiki/sources/" in file.file_path] + result = retrieve_qmd_consensus_files( + question=question.question, + files=files, + root_files=root_files, + entity_names=_candidate_entity_names(episode, question), + top_k=top_k, + ) + retrieved_turn_ids = [ + turn_id + for file in result.files + if (turn_id := _extract_turn_id_from_path(file.file_path)) is not None + ] + metrics = score_ama_retrieval_proxy(episode, question, retrieved_turn_ids) + return AmaMethodQuestionResult( + method_name=method_name, + episode_id=episode.episode_id, + domain=episode.domain, + task_type=episode.task_type, + qa_type=question.qa_type, + question_index=0, + question=question.question, + golden_answer=question.answer, + proxy_metrics=metrics, + retrieved_context=_format_retrieved_context(result.files), + retrieved_turn_ids=_unique_preserve_order(retrieved_turn_ids), + retrieved_file_paths=[file.file_path for file in result.files], + retrieval_notes=[ + "retrieval_plugins=qmd_consensus", + f"coverage_late_bridge={len(result.coverage.late_bridge_file_paths)}", + ], + ) + + +def select_ama_lexical_turn_ids( + turns: list[AmaTurn], + question: str, + top_k: int, +) -> list[int]: + top_k = max(top_k, 1) + query_tokens = _normalize_token_set(question) + scored = [ + ( + len(query_tokens & _normalize_token_set(_format_ama_turn(turn))), + turn.turn_idx, + ) + for turn in turns + ] + scored.sort(key=lambda item: (-item[0], item[1])) + selected = [turn_idx for score, turn_idx in scored[:top_k] if score > 0] + if selected: + return selected + return [turn_idx for _, turn_idx in scored[:top_k]] + + +def build_ama_memory_files( + episode: AmaEpisode, + *, + wiki_mode: WikiMode = "text", +) -> list[RetrievedMemoryFile]: + """Build AMA's text-only wiki; multimodal mode is accepted for API parity. + + AMA-Bench trajectories currently contain action/observation text only, so + selecting ``multimodal`` does not synthesize image artifacts or alter the + searchable corpus. + """ + normalize_wiki_mode(wiki_mode) + source_content = [ + "## Summary", + f"AMA episode {episode.episode_id}: {episode.task}", + "", + "## Turn Index", + ] + files = [] + for turn in episode.trajectory: + source_content.append(f"- [turn {turn.turn_idx}](../turns/T{turn.turn_idx}.md)") + files.append( + RetrievedMemoryFile( + filename=f"T{turn.turn_idx}.md", + file_path=f"/ama/{episode.episode_id}/wiki/turns/T{turn.turn_idx}.md", + mtime_ms=turn.turn_idx, + content=( + f"- Session: D1\n- Evidence: T{turn.turn_idx}\n" + f"Action: {turn.action}\nObservation: {turn.observation}" + ), + ) + ) + files.append( + RetrievedMemoryFile( + filename=f"D1_T{turn.turn_idx}_obs.md", + file_path=( + f"/ama/{episode.episode_id}/wiki/observations/" + f"D1_T{turn.turn_idx}_obs.md" + ), + mtime_ms=turn.turn_idx, + content=( + f"- Session: D1\n- Evidence: T{turn.turn_idx}\n" + f"{turn.action}\n{turn.observation}" + ), + ) + ) + + files.insert( + 0, + RetrievedMemoryFile( + filename="session_1.md", + file_path=f"/ama/{episode.episode_id}/wiki/sources/session_1.md", + mtime_ms=0, + content="\n".join(source_content), + ), + ) + return files + + +def aggregate_ama_method_summaries( + results: list[AmaMethodQuestionResult], +) -> list[AmaMethodSummary]: + grouped: dict[str, list[AmaMethodQuestionResult]] = {} + for result in results: + grouped.setdefault(result.method_name, []).append(result) + return [ + AmaMethodSummary( + method_name=method_name, + questions=len(rows), + proxy_recall_at_k=_round4( + _mean([row.proxy_metrics.proxy_recall_at_k for row in rows]) + ), + proxy_precision_at_k=_round4( + _mean([row.proxy_metrics.proxy_precision_at_k for row in rows]) + ), + proxy_hit_at_k=_round4(_mean([row.proxy_metrics.proxy_hit_at_k for row in rows])), + answer_support_coverage=_round4( + _mean([row.proxy_metrics.answer_support_coverage for row in rows]) + ), + ) + for method_name, rows in sorted(grouped.items()) + ] + + +def build_ama_precision_comparison( + *, + python_summary: AmaMethodSummary | dict[str, Any], + rust_baseline: AmaMethodSummary | dict[str, Any], + tolerance: float = 0.0, +) -> dict[str, Any]: + python_values = _summary_mapping(python_summary) + baseline_values = _summary_mapping(rust_baseline) + metrics = [ + "proxy_recall_at_k", + "proxy_precision_at_k", + "proxy_hit_at_k", + "answer_support_coverage", + ] + deltas = {} + statuses = {} + for metric in metrics: + if metric not in python_values or metric not in baseline_values: + continue + delta = _round4(float(python_values[metric]) - float(baseline_values[metric])) + deltas[metric] = delta + statuses[metric] = "pass" if delta + tolerance >= 0.0 else "regression" + + return { + "method_name": python_values.get("method_name", baseline_values.get("method_name", "")), + "questions": python_values.get("questions"), + "baseline_questions": baseline_values.get("questions"), + "status": "pass" + if all(status == "pass" for status in statuses.values()) + else "regression", + "tolerance": tolerance, + "python_summary": python_values, + "rust_baseline": baseline_values, + "metric_deltas": deltas, + "metric_status": statuses, + } + + +def run_python_ama_retrieval_eval( + *, + dataset_root: str | Path, + output_dir: str | Path, + subset: str = "open_end", + top_k: int = 24, + sample_limit: int | None = None, + question_limit: int | None = None, + methods: list[str] | None = None, + rust_baseline: AmaMethodSummary | dict[str, Any] | None = None, + tolerance: float = 0.0, + wiki_mode: WikiMode = "text", +) -> dict[str, Any]: + wiki_mode = normalize_wiki_mode(wiki_mode) + episodes = load_ama_dataset_split( + dataset_root, + subset, + sample_limit=sample_limit, + question_limit=question_limit, + ) + method_names = methods or ["wikimem_qmd"] + results = [] + for method_name in method_names: + for episode in episodes: + for question_index, question in enumerate(episode.qa_pairs): + result = run_python_wikimem_qmd_retrieval( + episode, + question, + top_k=top_k, + method_name=method_name, + wiki_mode=wiki_mode, + ) + results.append( + AmaMethodQuestionResult( + method_name=result.method_name, + episode_id=result.episode_id, + domain=result.domain, + task_type=result.task_type, + qa_type=result.qa_type, + question_index=question_index, + question=result.question, + golden_answer=result.golden_answer, + proxy_metrics=result.proxy_metrics, + retrieved_context=result.retrieved_context, + retrieved_turn_ids=result.retrieved_turn_ids, + retrieved_file_paths=result.retrieved_file_paths, + retrieval_notes=result.retrieval_notes, + ) + ) + + method_summaries = aggregate_ama_method_summaries(results) + stem = "openend" if subset == "open_end" else subset + config = { + "dataset_root": str(dataset_root), + "subset": subset, + "mode": "retrieval_only", + "output_dir": str(output_dir), + "sample_limit": sample_limit, + "question_limit": question_limit, + "methods": method_names, + "top_k": top_k, + "wiki_mode": wiki_mode, + "effective_wiki_mode": "text", + } + report = { + "config": config, + "metric_label_source": "answer_derived_proxy", + "total_episodes": len(episodes), + "total_questions": len(results), + "method_summaries": [_jsonable(summary) for summary in method_summaries], + "question_results": [_jsonable(result) for result in results], + } + summary = { + "config": config, + "metric_label_source": "answer_derived_proxy", + "total_episodes": len(episodes), + "total_questions": len(results), + "method_summaries": [_jsonable(summary) for summary in method_summaries], + "domain_summaries": _aggregate_ama_breakdown_summaries(results, "domain"), + "task_type_summaries": _aggregate_ama_breakdown_summaries(results, "task_type"), + "qa_type_summaries": _aggregate_ama_breakdown_summaries(results, "qa_type"), + } + comparison = ( + build_ama_precision_comparison( + python_summary=method_summaries[0], + rust_baseline=rust_baseline, + tolerance=tolerance, + ) + if rust_baseline is not None and method_summaries + else None + ) + + output = Path(output_dir) + _write_json(output / f"report_{stem}.json", report) + _write_json(output / f"summary_{stem}.json", summary) + if comparison is not None: + _write_json(output / f"comparison_{stem}.json", comparison) + return { + "report": report, + "summary": summary, + "comparison": comparison, + } + + +def _parse_ama_episode(raw: dict[str, Any]) -> AmaEpisode: + return AmaEpisode( + episode_id=str(raw.get("episode_id", "")), + task=str(raw.get("task", "")), + task_type=str(raw.get("task_type", "")), + domain=str(raw.get("domain", "")), + trajectory=[ + AmaTurn( + turn_idx=int(turn.get("turn_idx", 0)), + action=str(turn.get("action", "")), + observation=str(turn.get("observation", "")), + ) + for turn in raw.get("trajectory", []) + if isinstance(turn, dict) + ], + qa_pairs=[ + AmaQuestion( + question=str(question.get("question", "")), + answer=str(question.get("answer", "")), + qa_type=str(question.get("type", "")), + options=[str(item) for item in question.get("options", [])], + ) + for question in raw.get("qa_pairs", []) + if isinstance(question, dict) + ], + ) + + +def _candidate_entity_names(episode: AmaEpisode, question: AmaQuestion) -> list[str]: + text = f"{episode.task} {question.question}" + return re.findall(r"\b[A-Z][a-zA-Z]{2,}\b", text) + + +def _extract_turn_id_from_path(path: str) -> int | None: + normalized = path.replace("\\", "/") + match = re.search(r"/turns/t?(\d+)\.md$", normalized, flags=re.IGNORECASE) + if match: + return int(match.group(1)) + match = re.search(r"_t(\d+)_obs\.md$", normalized, flags=re.IGNORECASE) + return int(match.group(1)) if match else None + + +def _format_retrieved_context(files: list[RetrievedMemoryFile]) -> str: + return "\n\n".join(f"## {file.file_path}\n{file.content}" for file in files) + + +def _build_ama_lexical_files( + episode: AmaEpisode, + turn_ids: list[int], +) -> list[RetrievedMemoryFile]: + turns_by_id = {turn.turn_idx: turn for turn in episode.trajectory} + files = [] + for turn_id in turn_ids: + turn = turns_by_id.get(turn_id) + if turn is None: + continue + files.append( + RetrievedMemoryFile( + filename=f"T{turn.turn_idx}.md", + file_path=f"ama_turns/T{turn.turn_idx}.md", + mtime_ms=turn.turn_idx, + content=_format_ama_turn(turn), + ) + ) + return files + + +def _format_ama_turn(turn: AmaTurn) -> str: + return ( + f"Turn {turn.turn_idx}:\n" + f"Action: {turn.action}\n" + f"Observation: {turn.observation}" + ) + + +def _normalize_token_set(text: str) -> set[str]: + return { + token.lower() + for token in re.split(r"[^0-9A-Za-z]+", text) + if token.strip() + } + + +def _unique_preserve_order(values: list[int]) -> list[int]: + seen = set() + unique = [] + for value in values: + if value not in seen: + unique.append(value) + seen.add(value) + return unique + + +def _summary_mapping(summary: AmaMethodSummary | dict[str, Any]) -> dict[str, Any]: + if isinstance(summary, AmaMethodSummary): + return { + "method_name": summary.method_name, + "questions": summary.questions, + "proxy_recall_at_k": summary.proxy_recall_at_k, + "proxy_precision_at_k": summary.proxy_precision_at_k, + "proxy_hit_at_k": summary.proxy_hit_at_k, + "answer_support_coverage": summary.answer_support_coverage, + "judged_questions": summary.judged_questions, + "judge_accuracy": summary.judge_accuracy, + } + return dict(summary) + + +def _aggregate_ama_breakdown_summaries( + results: list[AmaMethodQuestionResult], + field_name: str, +) -> list[dict[str, Any]]: + grouped: dict[tuple[str, str], list[AmaMethodQuestionResult]] = {} + for result in results: + group_name = str(getattr(result, field_name)) + grouped.setdefault((result.method_name, group_name), []).append(result) + summaries = [] + for (method_name, group_name), rows in sorted(grouped.items()): + summary = aggregate_ama_method_summaries(rows)[0] + values = _summary_mapping(summary) + values["group_name"] = group_name + values["method_name"] = method_name + summaries.append(values) + return summaries + + +def _write_json(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(_jsonable(value), ensure_ascii=False, indent=2), + encoding="utf-8", + ) + + +def _jsonable(value: Any) -> Any: + if hasattr(value, "__dataclass_fields__"): + return asdict(value) + if isinstance(value, list): + return [_jsonable(item) for item in value] + if isinstance(value, dict): + return {key: _jsonable(item) for key, item in value.items()} + return value + + +def _escape_json_string_newlines(text: str) -> str: + escaped = [] + in_string = False + escape_next = False + for char in text: + if escape_next: + escaped.append(char) + escape_next = False + continue + if char == "\\": + escaped.append(char) + escape_next = True + continue + if char == '"': + escaped.append(char) + in_string = not in_string + continue + if char == "\n" and in_string: + escaped.append("\\n") + continue + escaped.append(char) + return "".join(escaped) + + +def _ratio(numerator: int, denominator: int) -> float: + return 0.0 if denominator == 0 else numerator / denominator + + +def _mean(values: list[float]) -> float: + return sum(values) / len(values) if values else 0.0 + + +def _round4(value: float) -> float: + return round(value + 0.0, 4) diff --git a/evaluation/wikimem/example_eval.py b/evaluation/wikimem/example_eval.py new file mode 100644 index 00000000..e0582641 --- /dev/null +++ b/evaluation/wikimem/example_eval.py @@ -0,0 +1,1561 @@ +"""Lightweight example dataset runners for wikimem migration.""" + +from __future__ import annotations + +import json +import math +import posixpath +import re +from functools import lru_cache +from pathlib import Path, PurePosixPath +from typing import Any + +from common.llm.base import LLM + +from evaluation.wikimem.qmd_consensus import RetrievedMemoryFile +from evaluation.wikimem.locomo_refined import ( + VisionEnrichmentConfig, + run_locomo_refined_offline_eval as _run_locomo_refined_multimodal_eval, +) +from evaluation.wikimem.retained_eval import WikiMode, normalize_wiki_mode +from evaluation.wikimem.wiki_builder import WikiBuilderMode +from evaluation.wikimem.retrieval_profile import retrieve_qmd_consensus_files + +_EVERMEM_STOPWORDS = { + "about", + "after", + "and", + "before", + "did", + "during", + "final", + "for", + "from", + "had", + "has", + "have", + "her", + "his", + "how", + "she", + "that", + "the", + "their", + "they", + "this", + "was", + "were", + "what", + "when", + "where", + "which", + "who", + "with", +} + +_MEM_GALLERY_STOPWORDS = { + "answer", + "find", + "for", + "from", + "help", + "image", + "into", + "matches", + "mentions", + "option", + "picture", + "please", + "question", + "same", + "shown", + "that", + "the", + "this", + "using", + "what", + "which", + "with", + "your", +} + +_STALE_RETRIEVAL_CASE_FIELDS = { + "baseline_retrieved_file_paths", + "baseline_retrieved_ids", + "final_retrieved_file_paths", + "final_retrieved_ids", + "final_ranked_clues", + "ranked_clues", + "retrieved_clue_ids", + "retrieved_file_paths", +} + +_METRIC_LABEL_SOURCES = { + "evermembench": "dataset_reference", + "mem_gallery": "human_annotated_clue", + "meta_crag": "answer_derived_proxy", +} + +_WORKSPACE_PROVENANCE_FILE = ".wikimem-workspace.json" +_WORKSPACE_PRODUCER = "mem2.0.wikimem.python" + + +def build_mem_gallery_python_workspaces( + *, + dialog_root: str | Path, + workspace_root: str | Path, + image_root: str | Path | None = None, + dataset_names: list[str] | None = None, + case_limit: int | None = None, + wiki_mode: WikiMode = "multimodal", +) -> list[dict[str, Any]]: + """Build Mem-Gallery wiki workspaces directly from raw dialog JSON.""" + wiki_mode = normalize_wiki_mode(wiki_mode) + dialog_root = Path(dialog_root) + workspace_root = Path(workspace_root) + image_root = Path(image_root) if image_root is not None else None + selected = set(dataset_names or []) + cases: list[dict[str, Any]] = [] + for path in sorted(dialog_root.glob("*.json")): + dataset_name = path.stem + if selected and dataset_name not in selected: + continue + entry = json.loads(path.read_text(encoding="utf-8")) + root = workspace_root / _sanitize_mem_gallery_component(dataset_name, keep_dash=True) + turns = _mem_gallery_turn_documents(entry, dataset_name, image_root, wiki_mode=wiki_mode) + _write_mem_gallery_python_workspace(root, turns, wiki_mode=wiki_mode) + for index, qa in enumerate(_mem_gallery_qas(entry), start=1): + cases.append( + { + "case_id": f"{dataset_name}::q{index}", + "question": str(qa.get("question") or ""), + "silver_evidence_ids": _normalize_mem_gallery_clue_ids(qa.get("clue") or []), + "knowledge_base_root": str(root), + "question_image_caption": str(qa.get("image_caption") or "").strip(), + "source_session_ids": _string_list(qa.get("session_id") or []), + } + ) + if case_limit is not None and len(cases) >= case_limit: + return cases + return cases + + +def build_meta_crag_python_workspaces( + *, + data_root: str | Path, + workspace_root: str | Path, + dataset_variant: str = "all", + case_limit: int | None = None, + wiki_mode: WikiMode = "multimodal", +) -> list[dict[str, Any]]: + """Build Meta-CRAG wiki workspaces from raw JSON/JSONL/parquet rows.""" + wiki_mode = normalize_wiki_mode(wiki_mode) + data_root = Path(data_root) + workspace_root = Path(workspace_root) + cases: list[dict[str, Any]] = [] + for variant, root in _meta_crag_variant_roots(data_root, dataset_variant): + for row in _meta_crag_rows(root): + for sample in _meta_crag_samples_from_row(row, variant, wiki_mode=wiki_mode): + root = workspace_root / _sanitize_mem_gallery_component(sample["sample_id"], keep_dash=True) + _write_meta_crag_python_workspace(root, sample, wiki_mode=wiki_mode) + cases.append( + { + "case_id": sample["sample_id"], + "question": sample["question"], + "answer": sample["answer"], + # Silver labels are derived once from the authoritative + # multimodal artifact set; changing wiki_mode must not + # silently change the evaluation labels. + "silver_evidence_ids": list(sample["silver_evidence_ids"]), + "knowledge_base_root": str(root), + } + ) + if case_limit is not None and len(cases) >= case_limit: + return cases + return cases + + +def run_locomo_refined_offline_eval( + *, + dataset_path: str | Path, + output_dir: str | Path, + workspace_root: str | Path, + top_k: int = 24, + question_limit: int | None = None, + wiki_mode: WikiMode = "multimodal", + multimodal_top_k: int = 6, + download_images: bool = True, + request_timeout_secs: int = 8, + redownload_missing_only: bool = True, + proxy_url: str | None = None, + vision_config: VisionEnrichmentConfig | None = None, + sample_filter: set[str] | None = None, + offline_export_root: str | Path | None = None, + llm: LLM | None = None, + wiki_builder_mode: WikiBuilderMode = "llm", + query_llm: LLM | None = None, +) -> dict[str, Any]: + return _run_locomo_refined_multimodal_eval( + dataset_path=dataset_path, + output_dir=output_dir, + workspace_root=workspace_root, + top_k=top_k, + question_limit=question_limit, + wiki_mode=wiki_mode, + multimodal_top_k=multimodal_top_k, + download_images=download_images, + request_timeout_secs=request_timeout_secs, + redownload_missing_only=redownload_missing_only, + proxy_url=proxy_url, + vision_config=vision_config, + sample_filter=sample_filter, + offline_export_root=offline_export_root, + llm=llm, + wiki_builder_mode=wiki_builder_mode, + query_llm=query_llm, + ) + + +def run_filesystem_proxy_eval( + *, + dataset_name: str, + cases_path: str | Path, + output_dir: str | Path, + top_k: int, + base_root: str | Path | None = None, + case_limit: int | None = None, + allow_stale_retrieval_fields: bool = False, + mem_gallery_dialog_root: str | Path | None = None, + wiki_mode: WikiMode | None = None, +) -> dict[str, Any]: + cases = json.loads(Path(cases_path).read_text(encoding="utf-8")) + cases = cases.get("cases", []) if isinstance(cases, dict) else cases + if case_limit is not None: + cases = cases[:case_limit] + if not allow_stale_retrieval_fields: + _reject_stale_retrieval_case_fields(cases, cases_path) + if dataset_name == "mem_gallery" and mem_gallery_dialog_root is not None: + _fill_mem_gallery_question_metadata(cases, Path(mem_gallery_dialog_root)) + workspace_cache: dict[str, tuple[list[RetrievedMemoryFile], dict[str, str]]] = {} + scored = [ + _score_filesystem_case(dataset_name, case, top_k, Path(base_root) if base_root else None, workspace_cache) + for case in cases + ] + result = { + "summary": _proxy_summary(dataset_name, scored, top_k, wiki_mode=wiki_mode), + "cases": scored, + } + _write_json(Path(output_dir) / f"{dataset_name}_proxy_eval.json", result) + return result + + +def _reject_stale_retrieval_case_fields( + cases: list[dict[str, Any]], + cases_path: str | Path, +) -> None: + for index, case in enumerate(cases): + stale = sorted(_STALE_RETRIEVAL_CASE_FIELDS & case.keys()) + if stale: + raise ValueError( + f"stale retrieval fields in {cases_path} " + f"case {case.get('case_id', index)}: {', '.join(stale)}" + ) + + +def _mem_gallery_qas(entry: dict[str, Any]) -> list[dict[str, Any]]: + return entry.get("human-annotated QAs") or entry.get("human_annotated_qas") or [] + + +def _meta_crag_variant_roots(data_root: Path, dataset_variant: str) -> list[tuple[str, Path]]: + if dataset_variant == "all": + return [ + ("single_turn", data_root / "single_turn" / "data"), + ("multi_turn", data_root / "multi_turn" / "data"), + ] + if (data_root / "data").is_dir(): + return [(dataset_variant, data_root / "data")] + return [(dataset_variant, data_root)] + + +def _meta_crag_rows(root: Path) -> list[dict[str, Any]]: + files = sorted(path for path in root.iterdir() if path.is_file() and path.name.startswith("validation")) + if not files: + files = sorted(path for path in root.iterdir() if path.suffix in {".json", ".jsonl", ".parquet"}) + rows: list[dict[str, Any]] = [] + for path in files: + if path.suffix == ".json": + payload = json.loads(path.read_text(encoding="utf-8")) + rows.extend(payload if isinstance(payload, list) else [payload]) + elif path.suffix == ".jsonl": + rows.extend(json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()) + elif path.suffix == ".parquet": + rows.extend(_meta_crag_parquet_rows(path)) + return rows + + +def _meta_crag_parquet_rows(path: Path) -> list[dict[str, Any]]: + try: + import pyarrow.parquet as pq # type: ignore[import-not-found] + except ModuleNotFoundError as exc: + raise RuntimeError("Meta-CRAG parquet input needs pyarrow; run with `uv run --with pyarrow ...`") from exc + return [dict(row) for batch in pq.ParquetFile(path).iter_batches() for row in batch.to_pylist()] + + +def _meta_crag_samples_from_row( + row: dict[str, Any], + variant: str, + *, + wiki_mode: WikiMode, +) -> list[dict[str, Any]]: + turns = row.get("turns") or {} + queries = _string_list(turns.get("query") or []) + searches = _string_list(turns.get("search_query") or []) + answers = _string_list(turns.get("answer") or []) + full_answers = _string_list((row.get("answers") or {}).get("ans_full") or []) + session_id = str(row.get("session_id") or "") + out = [] + all_turns = [ + { + "turn_index": index, + "query": query, + "search_query": searches[index] if index < len(searches) else query, + "answer": (answers[index] if index < len(answers) else "") or (full_answers[index] if index < len(full_answers) else ""), + } + for index, query in enumerate(queries) + ] + for turn in all_turns: + answer = full_answers[turn["turn_index"]] if turn["turn_index"] < len(full_answers) else turn["answer"] + raw_image = str(row.get("image_url") or "").strip() + if not raw_image and row.get("image"): + raw_image = "embedded-image" + sample = { + "sample_id": f"{session_id}::q{turn['turn_index'] + 1}", + "session_id": session_id, + "dataset_variant": variant, + "question_turn_index": turn["turn_index"], + "question": turn["query"], + "answer": answer, + "image_path": "" if raw_image.startswith(("http://", "https://")) else raw_image, + "image_url": raw_image if raw_image.startswith(("http://", "https://")) else "", + "history_turns": [item for item in all_turns if item["turn_index"] < turn["turn_index"]], + "current_turn": turn, + } + authoritative_artifacts = _meta_crag_artifacts(sample, wiki_mode="multimodal") + sample["silver_evidence_ids"] = _meta_crag_supported_ids( + sample, + artifacts=authoritative_artifacts, + ) + sample["artifact_documents"] = _meta_crag_artifacts(sample, wiki_mode=wiki_mode) + out.append(sample) + return out + + +def _meta_crag_artifacts( + sample: dict[str, Any], + *, + wiki_mode: WikiMode, +) -> list[dict[str, Any]]: + artifacts = [ + { + "artifact_id": f"turn-{turn['turn_index']}", + "evidence_id": f"turn-{turn['turn_index']}", + "modality": "text", + "title": f"Turn {turn['turn_index']}", + "text": f"Question: {turn['query']}\nSearch query: {turn['search_query']}\nAnswer: {turn['answer']}", + "turn_index": turn["turn_index"], + "search_query": turn["search_query"], + } + for turn in sample["history_turns"] + ] + current = sample["current_turn"] + artifacts.append( + { + "artifact_id": f"turn-{current['turn_index']}", + "evidence_id": f"turn-{current['turn_index']}", + "modality": "text", + "title": f"Question turn {current['turn_index']}", + "text": f"Question: {current['query']}\nSearch query: {current['search_query']}", + "turn_index": current["turn_index"], + "search_query": current["search_query"], + } + ) + if wiki_mode == "multimodal" and (sample.get("image_path") or sample.get("image_url")): + artifacts.append( + { + "artifact_id": "image-main", + "evidence_id": "image-main", + "modality": "image", + "title": f"Session image for {sample['session_id']}", + "text": f"Session image for {sample['session_id']}", + "turn_index": None, + "search_query": current["search_query"], + } + ) + return artifacts + + +def _write_meta_crag_python_workspace( + root: Path, + sample: dict[str, Any], + *, + wiki_mode: WikiMode, +) -> None: + for directory in [ + root / ".kb-research" / "manifests", + root / ".kb-research" / "retrieval", + root / "raw" / "multimodal", + root / "wiki" / "memories", + root / "wiki" / "sources", + ]: + directory.mkdir(parents=True, exist_ok=True) + _write_workspace_provenance(root, "meta_crag", wiki_mode=wiki_mode) + index_lines = ["# MEMORY", ""] + for artifact in sample["artifact_documents"]: + artifact_id = artifact["artifact_id"] + evidence_id = artifact["evidence_id"] + searchable = _meta_crag_searchable_text(artifact) + memory_path = root / "wiki" / "memories" / f"{artifact_id}.md" + _write_json(root / "raw" / "multimodal" / f"{artifact_id}.json", artifact) + _write_json( + root / ".kb-research" / "retrieval" / f"{artifact_id}.json", + { + "artifact_id": artifact_id, + "evidence_id": evidence_id, + "title": artifact["title"], + "searchable_text": searchable, + "memory_path": memory_path.as_posix(), + "modality": artifact["modality"], + }, + ) + memory_path.write_text( + f"# {artifact['title']}\nEvidence: {evidence_id}\nevidence_id: {evidence_id}\n\n{searchable}\n", + encoding="utf-8", + ) + (root / "wiki" / "sources" / f"{artifact_id}.md").write_text( + f"# Source Snapshot\n\n- title: {artifact['title']}\n- modality: {artifact['modality']}\n", + encoding="utf-8", + ) + index_lines.append(f"- [{artifact['title']}](wiki/memories/{artifact_id}.md)") + (root / "MEMORY.md").write_text("\n".join(index_lines), encoding="utf-8") + + +def _meta_crag_supported_ids( + sample: dict[str, Any], + *, + artifacts: list[dict[str, Any]] | None = None, +) -> list[str]: + artifacts = artifacts if artifacts is not None else sample["artifact_documents"] + answer = str(sample.get("answer") or "") + supported = [ + artifact["evidence_id"] + for artifact in artifacts + if _meta_crag_answer_supported(answer, artifact) + ] + if supported: + return supported + best = max( + artifacts, + key=lambda artifact: _meta_crag_support_proxy_score(sample, artifact), + default=None, + ) + if best and _meta_crag_support_proxy_score(sample, best) > 0: + return [best["evidence_id"]] + return [] + + +def _meta_crag_answer_supported(answer: str, artifact: dict[str, Any]) -> bool: + normalized_answer = _normalize_support_text(answer) + text = _normalize_support_text(f"{artifact['title']}\n{artifact['text']}") + answer_tokens = _support_tokens(answer) + return bool(normalized_answer and normalized_answer in text) or _support_tokens_match(answer_tokens, _support_tokens(text)) + + +def _meta_crag_support_proxy_score(sample: dict[str, Any], artifact: dict[str, Any]) -> int: + if artifact["evidence_id"] == f"turn-{sample['question_turn_index']}": + return 0 + candidate = set(_support_tokens(f"{artifact['title']}\n{artifact['text']}")) + answer_tokens = [token for token in _support_tokens(sample.get("answer") or "") if len(token) > 2] + question_tokens = [token for token in _support_tokens(sample.get("question") or "") if len(token) > 2] + score = 4 * sum(1 for token in answer_tokens if token in candidate) + score += sum(1 for token in question_tokens if token in candidate) + if artifact["evidence_id"] == "image-main" and (sample.get("image_path") or sample.get("image_url")): + score += 16 + if artifact["evidence_id"].startswith("turn-"): + score += 1 + return score + + +def _meta_crag_searchable_text(artifact: dict[str, Any]) -> str: + return f"{artifact['title']}\n{artifact['text']}\nSearch query: {artifact.get('search_query') or ''}" + + +def _normalize_support_text(text: str) -> str: + return " ".join(text.split()).strip().lower() + + +def _support_tokens(text: str) -> list[str]: + return [token.lower() for token in re.split(r"[^0-9A-Za-z]+", text) if token] + + +def _support_tokens_match(answer_tokens: list[str], candidate_tokens: list[str]) -> bool: + significant = [token for token in answer_tokens if len(token) > 2] + candidate = set(candidate_tokens) + return bool(significant) and all(token in candidate for token in significant) + + +def _mem_gallery_turn_documents( + entry: dict[str, Any], + dataset_name: str, + image_root: Path | None, + *, + wiki_mode: WikiMode, +) -> list[dict[str, Any]]: + speaker = str((entry.get("character_profile") or {}).get("name") or dataset_name) + turns = [] + for session in entry.get("multi_session_dialogues") or []: + session_id = str(session.get("session_id") or "") + for index, dialogue in enumerate(session.get("dialogues") or [], start=1): + clue_id = str(dialogue.get("round") or index) + image_path = _first(dialogue.get("input_image") or []) if wiki_mode == "multimodal" else "" + image_caption = _first(dialogue.get("image_caption") or []) if wiki_mode == "multimodal" else "" + image_id = _first(dialogue.get("image_id") or []) if wiki_mode == "multimodal" else "" + lines = [] + user = str(dialogue.get("user") or "").strip() + assistant = str(dialogue.get("assistant") or "").strip() + if user: + lines.append(f"user ({speaker}): {user}") + if assistant: + lines.append(f"assistant: {assistant}") + turns.append( + { + "clue_id": clue_id, + "session_id": session_id, + "dialogue_round": index, + "timestamp": session.get("date"), + "text": "\n".join(lines), + "image_path": _resolve_mem_gallery_image_ref(image_root, dataset_name, image_path), + "image_caption": image_caption, + "image_id": image_id, + } + ) + return turns + + +def _write_mem_gallery_python_workspace( + root: Path, + turns: list[dict[str, Any]], + *, + wiki_mode: WikiMode, +) -> None: + for directory in [ + root / "raw" / "sessions", + root / "raw" / "turns", + root / "raw" / "clue_summaries", + root / "wiki" / "turns", + root / "wiki" / "observations", + root / "wiki" / "memories", + root / ".mem-gallery", + root / ".kb-research" / "retrieval", + ]: + directory.mkdir(parents=True, exist_ok=True) + _write_workspace_provenance(root, "mem_gallery", wiki_mode=wiki_mode) + for session_id in sorted({turn["session_id"] for turn in turns}): + session_turns = [turn for turn in turns if turn["session_id"] == session_id] + (root / "raw" / "sessions" / f"{session_id}.md").write_text( + _render_mem_gallery_session(session_id, session_turns), + encoding="utf-8", + ) + support = [] + for turn in turns: + raw_slug = _mem_gallery_turn_file_slug(turn) + evidence_slug = _evidence_slug(turn["clue_id"]) + raw_turn = _render_mem_gallery_turn_raw_artifact(turn) + clue_summary = _render_mem_gallery_clue_summary(turn) + searchable = _mem_gallery_turn_searchable_text(turn) + (root / "raw" / "turns" / f"{raw_slug}.md").write_text(raw_turn, encoding="utf-8") + (root / "raw" / "clue_summaries" / f"{raw_slug}.md").write_text(clue_summary, encoding="utf-8") + (root / "wiki" / "turns" / f"{evidence_slug}.md").write_text( + f"# Turn {turn['clue_id']}\nEvidence: {turn['clue_id']}\nSession: {turn['session_id']}\n\n{searchable}\n", + encoding="utf-8", + ) + observation_path = ( + root + / "wiki" + / "observations" + / f"{evidence_slug}_obs_{_observation_suffix(turn['clue_id'])}.md" + ) + observation_path.write_text( + f"# Observation {turn['clue_id']}\nEvidence: {turn['clue_id']}\nSession: {turn['session_id']}\n\n{searchable}\n", + encoding="utf-8", + ) + memory_path = root / "wiki" / "memories" / f"clue-summary-{evidence_slug}.md" + memory_path.write_text( + f"# Clue summary {turn['clue_id']}\nEvidence: {turn['clue_id']}\nSession: {turn['session_id']}\n\n{clue_summary}\n", + encoding="utf-8", + ) + retrieval_path = root / ".kb-research" / "retrieval" / f"{evidence_slug}.json" + _write_json( + retrieval_path, + { + "artifact_id": f"turn-{evidence_slug}", + "evidence_id": turn["clue_id"], + "memory_path": memory_path.as_posix(), + "title": f"Turn {turn['clue_id']}", + "searchable_text": searchable, + }, + ) + support.append({"memory_path": memory_path.as_posix(), "linked_clue_ids": [turn["clue_id"]]}) + _write_json(root / ".mem-gallery" / "artifact_support_map.json", support) + + +def _render_mem_gallery_session(session_id: str, turns: list[dict[str, Any]]) -> str: + lines = [f"# Session {session_id}", ""] + for turn in turns: + lines.extend([f"## {turn['clue_id']}", str(turn["text"])]) + if turn.get("image_caption"): + lines.append(f"Image caption: {turn['image_caption']}") + if turn.get("image_id"): + lines.append(f"Image id: {turn['image_id']}") + lines.append("") + return "\n".join(lines) + + +def _render_mem_gallery_turn_raw_artifact(turn: dict[str, Any]) -> str: + lines = [f"# Turn {turn['clue_id']}", "", str(turn["text"])] + if turn.get("image_path"): + lines.extend(["", f"Image path: {turn['image_path']}"]) + if turn.get("image_caption"): + lines.append(f"Image caption: {turn['image_caption']}") + if turn.get("image_id"): + lines.append(f"Image id: {turn['image_id']}") + return "\n".join(lines) + + +def _render_mem_gallery_clue_summary(turn: dict[str, Any]) -> str: + lines = [f"# Clue {turn['clue_id']}", "", str(turn["text"])] + if turn.get("image_id"): + lines.append(f"Linked image clue: {turn['image_id']}") + if turn.get("image_caption"): + lines.append(f"Image caption: {turn['image_caption']}") + return "\n".join(lines) + + +def _mem_gallery_turn_searchable_text(turn: dict[str, Any]) -> str: + parts = [str(turn["text"])] + if turn.get("image_caption"): + parts.append(f"image caption: {turn['image_caption']}") + if turn.get("image_id"): + parts.append(f"image id: {turn['image_id']}") + return "\n".join(parts) + + +def _resolve_mem_gallery_image_ref(image_root: Path | None, dataset_name: str, raw_path: str) -> str: + raw_path = raw_path.strip() + if not raw_path: + return "" + relative = raw_path.replace("../image/", "") + if Path(raw_path).is_absolute() or image_root is None: + return raw_path + candidate = image_root / relative + if candidate.is_file(): + return candidate.as_posix() + return (image_root / dataset_name / Path(relative).name).as_posix() + + +def _normalize_mem_gallery_clue_ids(values: list[Any]) -> list[str]: + return _string_list(values) + + +def _mem_gallery_turn_file_slug(turn: dict[str, Any]) -> str: + return f"D{_sanitize_mem_gallery_component(turn['session_id'])}_{turn['dialogue_round']}" + + +def _sanitize_mem_gallery_component(value: Any, *, keep_dash: bool = False) -> str: + allowed = {"_"} | ({"-"} if keep_dash else set()) + return "".join(ch if ch.isascii() and (ch.isalnum() or ch in allowed) else "_" for ch in str(value)) + + +def _evidence_slug(evidence_id: str) -> str: + return evidence_id.replace(":", "_") + + +def _observation_suffix(evidence_id: str) -> str: + return evidence_id.rsplit(":", 1)[-1] or evidence_id + + +def _first(values: list[Any]) -> str: + if isinstance(values, str): + return values + return str(values[0]) if values else "" + + +def _string_list(values: Any) -> list[str]: + if isinstance(values, str): + values = [values] + return [str(value).strip() for value in values if str(value).strip()] + + +def _fill_mem_gallery_question_metadata(cases: list[dict[str, Any]], dialog_root: Path) -> None: + cache: dict[str, list[dict[str, Any]]] = {} + for case in cases: + case_id = str(case.get("case_id") or "") + if "::q" not in case_id: + continue + dataset_name, question_number = case_id.rsplit("::q", 1) + if not question_number.isdigit(): + continue + if dataset_name not in cache: + path = dialog_root / f"{dataset_name}.json" + if not path.exists(): + cache[dataset_name] = [] + else: + payload = json.loads(path.read_text(encoding="utf-8")) + cache[dataset_name] = ( + payload.get("human-annotated QAs") + or payload.get("human_annotated_qas") + or [] + ) + index = int(question_number) - 1 + qas = cache[dataset_name] + if not 0 <= index < len(qas): + continue + qa = qas[index] + if qa.get("image_caption") and not case.get("question_image_caption"): + case["question_image_caption"] = qa["image_caption"] + if qa.get("session_id") and not case.get("source_session_ids"): + case["source_session_ids"] = qa["session_id"] + + +def run_evermembench_python_eval( + *, + data_root: str | Path, + output_dir: str | Path, + top_k: int = 100, + topic_names: list[str] | None = None, + topic_limit: int | None = None, + question_offset: int = 0, + question_limit: int | None = None, + wiki_mode: WikiMode = "text", +) -> dict[str, Any]: + wiki_mode = normalize_wiki_mode(wiki_mode) + topics = sorted(path for path in Path(data_root).iterdir() if path.is_dir()) + if topic_names is not None: + selected = set(topic_names) + topics = [path for path in topics if path.name in selected] + if topic_limit is not None: + topics = topics[:topic_limit] + cases = [] + for topic in topics: + cases.extend(_score_evermem_topic(topic, top_k, question_offset, question_limit)) + result = { + "summary": _proxy_summary("evermembench", cases, top_k, wiki_mode=wiki_mode), + "cases": cases, + } + result["summary"]["effective_wiki_mode"] = "text" + _write_json(Path(output_dir) / "evermembench_python_eval.json", result) + return result + + +def _score_filesystem_case( + dataset_name: str, + case: dict[str, Any], + top_k: int, + base_root: Path | None, + workspace_cache: dict[str, tuple[list[RetrievedMemoryFile], dict[str, str]]], +) -> dict[str, Any]: + root = _resolve_root(case["knowledge_base_root"], base_root) + include_retrieval_json = ( + dataset_name in {"meta_crag", "mem_gallery"} + and _has_python_workspace_provenance(root, dataset_name) + ) + cache_key = f"{root.as_posix()}::retrieval_json={include_retrieval_json}" + if cache_key not in workspace_cache: + workspace_cache[cache_key] = _read_workspace_files( + root, + include_retrieval_json=include_retrieval_json, + ) + files, evidence_by_path = workspace_cache[cache_key] + retrieved, file_paths = _retrieve_filesystem_evidence( + dataset_name, + case, + files, + evidence_by_path, + top_k, + ) + expected = _case_expected_ids(case) + hits = [item for item in retrieved if item in set(expected)] + return _case_score( + dataset_name=dataset_name, + case_id=str(case.get("case_id") or case.get("sample_id") or len(retrieved)), + question=str(case["question"]), + expected=expected, + retrieved=retrieved, + hits=hits, + file_paths=file_paths, + top_k=top_k, + ) + + +def _score_evermem_topic( + topic: Path, + top_k: int, + question_offset: int, + question_limit: int | None, +) -> list[dict[str, Any]]: + topic_id = topic.name + profiles = _evermem_profiles(topic.parent) + rows = _evermem_rows( + topic_id, + json.loads((topic / "dialogue.json").read_text(encoding="utf-8")), + profiles, + ) + files = [ + RetrievedMemoryFile( + filename=f"{_slug(row['evidence_id'])}.md", + file_path=f"/evermembench/{topic_id}/wiki/turns/{row['session_id']}_{row['message_index']}.md", + mtime_ms=index + 1, + content=( + f"# Turn {row['evidence_id']}\n" + f"Evidence: {row['evidence_id']}\n" + f"Session: {row['session_id']}\n" + f"Speaker: {row['speaker']}\n" + f"Date: {row['date']}\n" + f"Group: {row['group']}\n\n" + f"{row['profile']}\n\n" + f"{row['text']}\n" + ), + ) + for index, row in enumerate(rows) + ] + evidence_by_path = {file.file_path: row["evidence_id"] for file, row in zip(files, rows)} + token_by_path = {file.file_path: set(_content_keywords(file.content)) for file in files} + idf_by_token = _idf_by_token(token_by_path.values()) + session_by_ref = { + (row["date"], row["group"]): row["session_id"] + for row in rows + } + questions = json.loads((topic / f"qa_{topic_id}.json").read_text(encoding="utf-8")) + if question_offset: + questions = questions[question_offset:] + if question_limit is not None: + questions = questions[:question_limit] + cases = [] + for question in questions: + retrieved, file_paths = _retrieve_evermem_evidence( + str(question["Q"]), + files=files, + evidence_by_path=evidence_by_path, + rows=rows, + token_by_path=token_by_path, + idf_by_token=idf_by_token, + top_k=top_k, + ) + expected = _evermem_expected(question.get("R") or [], session_by_ref) + hits = [item for item in retrieved if item in set(expected)] + cases.append( + _case_score( + dataset_name="evermembench", + case_id=f"{topic_id}::{question.get('id')}", + question=str(question["Q"]), + expected=expected, + retrieved=retrieved, + hits=hits, + file_paths=file_paths, + top_k=top_k, + ) + ) + return cases + + +def _read_workspace_files( + root: Path, + *, + include_retrieval_json: bool = False, +) -> tuple[list[RetrievedMemoryFile], dict[str, str]]: + files = [] + evidence_by_path = {} + index = 0 + for path in sorted(root.rglob("*.md")): + content = path.read_text(encoding="utf-8", errors="ignore") + file_path = path.as_posix() + index += 1 + files.append( + RetrievedMemoryFile( + filename=path.name, + file_path=file_path, + mtime_ms=index, + content=content, + ) + ) + content_ids = _content_evidence_ids(content) + evidence_by_path[file_path] = ( + content_ids[0] if content_ids else _evidence_id_from_path(path) + ) + _merge_mem_gallery_support_map(root, evidence_by_path) + retrieval_paths = sorted((root / ".kb-research" / "retrieval").glob("*.json")) + for path in (retrieval_paths if include_retrieval_json else []): + record = json.loads(path.read_text(encoding="utf-8")) + evidence_id = str(record.get("evidence_id") or "").strip() + if not evidence_id: + memory_path = str(record.get("memory_path") or "").replace("\\", "/") + evidence_id = evidence_by_path.get(memory_path, "") + if not evidence_id: + continue + title = str(record.get("title") or "") + searchable_text = str(record.get("searchable_text") or "") + artifact_id = str(record.get("artifact_id") or path.stem) + file_path = path.as_posix() + index += 1 + files.append( + RetrievedMemoryFile( + filename=path.name, + file_path=file_path, + mtime_ms=index, + content=( + f"# {title or artifact_id}\n" + f"Artifact: {artifact_id}\n" + f"Evidence: {evidence_id}\n\n" + f"{searchable_text}\n" + ), + ) + ) + evidence_by_path[file_path] = evidence_id + return files, evidence_by_path + + +def _write_workspace_provenance( + root: Path, + dataset_name: str, + *, + wiki_mode: WikiMode | None = None, +) -> None: + provenance = { + "producer": _WORKSPACE_PRODUCER, + "dataset_name": dataset_name, + "schema_version": 1, + } + if wiki_mode is not None: + provenance["wiki_mode"] = wiki_mode + _write_json( + root / _WORKSPACE_PROVENANCE_FILE, + provenance, + ) + + +def _has_python_workspace_provenance(root: Path, dataset_name: str) -> bool: + path = root / _WORKSPACE_PROVENANCE_FILE + if not path.is_file(): + return False + try: + provenance = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError: + return False + return ( + provenance.get("producer") == _WORKSPACE_PRODUCER + and provenance.get("dataset_name") == dataset_name + and provenance.get("schema_version") == 1 + ) + + +def _merge_mem_gallery_support_map(root: Path, evidence_by_path: dict[str, str]) -> None: + path = root / ".mem-gallery" / "artifact_support_map.json" + if not path.exists(): + return + for item in json.loads(path.read_text(encoding="utf-8")): + if not isinstance(item, dict): + continue + clue_ids = [str(clue_id) for clue_id in item.get("linked_clue_ids") or [] if clue_id] + memory_path = str(item.get("memory_path") or "") + if not clue_ids or not memory_path: + continue + evidence_by_path[memory_path.replace("\\", "/")] = clue_ids[0] + marker = "/wiki/" + if marker in memory_path: + evidence_by_path[(root / ("wiki/" + memory_path.split(marker, 1)[1])).as_posix()] = clue_ids[0] + + +def _ids_for_file( + file: RetrievedMemoryFile, + evidence_by_path: dict[str, str], +) -> list[str]: + if "/wiki/memories/clue-summary-" in file.file_path.replace("\\", "/") and file.file_path in evidence_by_path: + return [evidence_by_path[file.file_path]] + ids = _content_evidence_ids(file.content) + if ids: + return ids + + parent = PurePosixPath(file.file_path).parent + ids = _ordered_unique( + evidence_by_path.get( + posixpath.normpath((parent / match.group(1).split("#", 1)[0]).as_posix()), + "", + ) + for match in re.finditer(r"\[[^\]]+\]\(([^)#]+\.md)(?:#[^)]+)?\)", file.content) + ) + return ids or [evidence_by_path[file.file_path]] + + +def _retrieve_filesystem_evidence( + dataset_name: str, + case: dict[str, Any], + files: list[RetrievedMemoryFile], + evidence_by_path: dict[str, str], + top_k: int, +) -> tuple[list[str], list[str]]: + query = str(case["question"]) + result = retrieve_qmd_consensus_files( + question=query, + files=files, + root_files=[file for file in files if file.file_path.endswith("/MEMORY.md")], + top_k=top_k, + ) + retrieved = _ordered_unique( + evidence_id + for file in result.files + for evidence_id in _ids_for_file(file, evidence_by_path) + ) + file_paths = [file.file_path for file in result.files] + if dataset_name == "meta_crag": + adjunct_ids, adjunct_paths = _retrieve_overlap_evidence( + query, + files, + evidence_by_path, + top_k, + ) + selected_paths = _ordered_unique(file_paths + adjunct_paths)[: max(top_k, 1)] + return ( + _meta_crag_ids_for_paths(selected_paths, evidence_by_path), + selected_paths, + ) + if dataset_name == "mem_gallery": + artifact_query = _case_query_text(dataset_name, case) + retrieved = [ + evidence_id + for evidence_id in retrieved + if _is_mem_gallery_clue_id(evidence_id) + ] + retrieved = _rrf_fuse_ids( + retrieved, + _retrieve_mem_gallery_artifact_ids( + artifact_query, + files, + evidence_by_path, + top_k, + ), + top_k, + ) + return retrieved, file_paths + + +def _case_query_text(dataset_name: str, case: dict[str, Any]) -> str: + query = str(case["question"]) + if dataset_name != "mem_gallery": + return query + lines = [query] + caption = str(case.get("question_image_caption") or case.get("image_caption") or "").strip() + if caption: + lines.append(f"question image caption: {caption}") + return "\n".join(lines) + + +def _meta_crag_ids_for_paths( + paths: list[str], + evidence_by_path: dict[str, str], +) -> list[str]: + selected = [] + for path in paths: + normalized = path.replace("\\", "/") + evidence_id = evidence_by_path.get(path) or evidence_by_path.get(normalized) or "" + if "example/meta-crag/" in normalized and "/wiki/memories/" not in normalized and "/.kb-research/" not in normalized: + evidence_id = _rust_stable_path_id(_rust_relative_meta_crag_path(normalized)) + selected.append(evidence_id) + return _ordered_unique(selected) + + +def _is_mem_gallery_clue_id(evidence_id: str) -> bool: + return re.match(r"^D\d+:(?:\d+|IMG_\d+)$", evidence_id) is not None + + +def _retrieve_mem_gallery_artifact_ids( + question: str, + files: list[RetrievedMemoryFile], + evidence_by_path: dict[str, str], + top_k: int, +) -> list[str]: + query_tokens = { + token + for token in _content_keywords(question) + if token not in _MEM_GALLERY_STOPWORDS + } + if not query_tokens: + return [] + candidates = [ + file + for file in files + if "/.kb-research/retrieval/" in file.file_path.replace("\\", "/") + and _is_mem_gallery_clue_id(evidence_by_path.get(file.file_path, "")) + ] + token_by_path = { + file.file_path: set(_cached_content_keywords(file.content)) + for file in candidates + } + idf_by_token = _idf_by_token(token_by_path.values()) + ranked = [] + for file in candidates: + evidence_id = evidence_by_path.get(file.file_path, "") + score = sum( + idf_by_token.get(token, 1.0) ** 1.35 + for token in query_tokens & token_by_path[file.file_path] + ) + if score: + ranked.append((-score, file.file_path, evidence_id)) + return _ordered_unique(evidence_id for _, _, evidence_id in sorted(ranked))[: max(top_k * 9, 1)] + + +@lru_cache(maxsize=16384) +def _cached_content_keywords(text: str) -> tuple[str, ...]: + return tuple(_content_keywords(text)) + + +def _rrf_fuse_ids(primary: list[str], secondary: list[str], top_k: int) -> list[str]: + scores: dict[str, float] = {} + for rank, evidence_id in enumerate(primary, start=1): + scores[evidence_id] = scores.get(evidence_id, 0.0) + 1.0 / (60 + rank) + for rank, evidence_id in enumerate(secondary, start=1): + scores[evidence_id] = scores.get(evidence_id, 0.0) + 1.0 / (60 + rank) + return [ + evidence_id + for evidence_id, _ in sorted(scores.items(), key=lambda item: (-item[1], item[0])) + ][: max(top_k, 1)] + + +def _rust_relative_meta_crag_path(path: str) -> str: + marker = "example/meta-crag/" + return marker + path.split(marker, 1)[1] if marker in path else path + + +def _rust_stable_path_id(path: str) -> str: + data = path.encode("utf-8") + b"\xff" + mask = (1 << 64) - 1 + v0 = 0x736F6D6570736575 + v1 = 0x646F72616E646F6D + v2 = 0x6C7967656E657261 + v3 = 0x7465646279746573 + + def rotl(value: int, bits: int) -> int: + return ((value << bits) & mask) | (value >> (64 - bits)) + + def sip_round() -> None: + nonlocal v0, v1, v2, v3 + v0 = (v0 + v1) & mask + v1 = rotl(v1, 13) + v1 ^= v0 + v0 = rotl(v0, 32) + v2 = (v2 + v3) & mask + v3 = rotl(v3, 16) + v3 ^= v2 + v0 = (v0 + v3) & mask + v3 = rotl(v3, 21) + v3 ^= v0 + v2 = (v2 + v1) & mask + v1 = rotl(v1, 17) + v1 ^= v2 + v2 = rotl(v2, 32) + + end = len(data) - (len(data) % 8) + for offset in range(0, end, 8): + chunk = int.from_bytes(data[offset : offset + 8], "little") + v3 ^= chunk + sip_round() + v0 ^= chunk + tail = data[end:] + last = len(data) << 56 + for index, byte in enumerate(tail): + last |= byte << (8 * index) + v3 ^= last + sip_round() + v0 ^= last + v2 ^= 0xFF + for _ in range(3): + sip_round() + return f"baseline::{(v0 ^ v1 ^ v2 ^ v3) & mask:016x}" + + +def _retrieve_overlap_evidence( + query: str, + files: list[RetrievedMemoryFile], + evidence_by_path: dict[str, str], + top_k: int, +) -> tuple[list[str], list[str]]: + query_tokens = set(_query_keywords(query)) + if not query_tokens: + return [], [] + score_by_id: dict[str, int] = {} + path_by_id: dict[str, str] = {} + best_score_by_id: dict[str, int] = {} + for file in files: + content_tokens = set(_query_keywords(file.content)) + score = len(query_tokens & content_tokens) + if score <= 0: + continue + for evidence_id in _ids_for_file(file, evidence_by_path): + score_by_id[evidence_id] = score_by_id.get(evidence_id, 0) + score + if score > best_score_by_id.get(evidence_id, -1): + best_score_by_id[evidence_id] = score + path_by_id[evidence_id] = file.file_path + ranked = sorted(score_by_id, key=lambda item: (-score_by_id[item], item)) + retrieved = ranked[: max(top_k, 1)] + return retrieved, [path_by_id[item] for item in retrieved if item in path_by_id] + + +def _retrieve_evermem_evidence( + question: str, + *, + files: list[RetrievedMemoryFile], + evidence_by_path: dict[str, str], + rows: list[dict[str, str]], + token_by_path: dict[str, set[str]] | None = None, + idf_by_token: dict[str, float] | None = None, + top_k: int, +) -> tuple[list[str], list[str]]: + scored_files = _rank_files_by_overlap( + question, + files, + token_by_path=token_by_path, + idf_by_token=idf_by_token, + ) + index_by_id = {row["evidence_id"]: index for index, row in enumerate(rows)} + file_by_id = { + evidence_by_path[file.file_path]: file + for file in files + if file.file_path in evidence_by_path + } + candidate_scores: dict[str, float] = {} + session_scores: dict[str, float] = {} + pool_limit = max(top_k * 3, 256) + for rank, (score, file) in enumerate(scored_files[:pool_limit], start=1): + evidence_id = evidence_by_path.get(file.file_path) + if not evidence_id: + continue + index = index_by_id.get(evidence_id) + if index is None: + continue + session_id = rows[index]["session_id"] + session_scores[session_id] = max(session_scores.get(session_id, 0.0), score) + candidate_scores[evidence_id] = max( + candidate_scores.get(evidence_id, 0.0), + score + 1 / (60 + rank), + ) + for distance, neighbor in _evermem_dialogue_neighbors(rows, index, radius=5): + neighbor_score = score * 0.92 - distance * 0.08 + candidate_scores[neighbor["evidence_id"]] = max( + candidate_scores.get(neighbor["evidence_id"], 0.0), + neighbor_score, + ) + for row in rows: + score = session_scores.get(row["session_id"]) + if score: + candidate_scores[row["evidence_id"]] = max( + candidate_scores.get(row["evidence_id"], 0.0), + score * 0.45, + ) + selected = [ + evidence_id + for evidence_id, _ in sorted( + candidate_scores.items(), + key=lambda item: (-item[1], item[0]), + )[: max(top_k, 1)] + ] + paths = [ + file_by_id[evidence_id].file_path + for evidence_id in selected + if evidence_id in file_by_id + ] + return selected, paths + + +def _rank_files_by_overlap( + query: str, + files: list[RetrievedMemoryFile], + token_by_path: dict[str, set[str]] | None = None, + idf_by_token: dict[str, float] | None = None, +) -> list[tuple[float, RetrievedMemoryFile]]: + query_tokens = { + token + for token in _content_keywords(query) + if token not in _EVERMEM_STOPWORDS + } + if not query_tokens: + return [] + scored = [] + for file in files: + content_tokens = ( + token_by_path[file.file_path] + if token_by_path is not None + else set(_content_keywords(file.content)) + ) + score = sum( + (idf_by_token.get(token, 1.0) if idf_by_token else 1.0) ** 1.35 + for token in query_tokens & content_tokens + ) + if score > 0: + scored.append((score, file)) + scored.sort(key=lambda item: (-item[0], item[1].file_path)) + return scored + + +def _idf_by_token(token_sets) -> dict[str, float]: + token_sets = list(token_sets) + total = len(token_sets) + document_frequency: dict[str, int] = {} + for tokens in token_sets: + for token in tokens: + document_frequency[token] = document_frequency.get(token, 0) + 1 + return { + token: math.log((total + 1) / (frequency + 1)) + 1 + for token, frequency in document_frequency.items() + } + + +def _evermem_dialogue_neighbors( + rows: list[dict[str, str]], + index: int, + *, + radius: int, +) -> list[tuple[int, dict[str, str]]]: + row = rows[index] + neighbors = [] + for distance in range(1, radius + 1): + for neighbor_index in (index - distance, index + distance): + if not 0 <= neighbor_index < len(rows): + continue + neighbor = rows[neighbor_index] + if neighbor["session_id"] == row["session_id"]: + neighbors.append((distance, neighbor)) + return neighbors + + +def _query_keywords(text: str) -> list[str]: + return _text_keywords(text, limit=18) + + +def _content_keywords(text: str) -> list[str]: + return _text_keywords(text, limit=None) + + +def _text_keywords(text: str, *, limit: int | None) -> list[str]: + keywords = _ordered_unique( + token.lower() + for token in re.split(r"[^0-9A-Za-z]+", text) + if len(token) >= 3 + ) + return keywords[:limit] if limit is not None else keywords + + +def _content_evidence_ids(content: str) -> list[str]: + explicit_ids = [ + _clean_evidence_id(match.group(1)) + for match in re.finditer(r"(?im)^(?:evidence|evidence_id):\s*(.+)$", content) + ] + clue_ids = [ + f"{match.group(2)}:{match.group(3)}" + for match in re.finditer(r"\bclue:([^:\s\"']+):([^:\s\"']+):([^:\s\"']+)", content) + ] + return _ordered_unique(explicit_ids + clue_ids) + + +def _clean_evidence_id(value: str) -> str: + return value.strip().strip("'\"") + + +def _case_expected_ids(case: dict[str, Any]) -> list[str]: + for key in ("silver_evidence_ids", "gold_clue_ids", "expected_evidence_ids"): + if case.get(key): + return [str(item) for item in case[key]] + return [] + + +def _proxy_summary( + dataset_name: str, + cases: list[dict[str, Any]], + top_k: int, + *, + wiki_mode: WikiMode | None = None, +) -> dict[str, Any]: + metric_cases = ( + [case for case in cases if case["expected_evidence_ids"]] + if dataset_name == "meta_crag" + else cases + ) + summary = { + "dataset_name": dataset_name, + "metric_label_source": _METRIC_LABEL_SOURCES.get(dataset_name, "unspecified_proxy"), + "total_cases": len(cases), + "evidence_labeled_cases": len(metric_cases), + "top_k": top_k, + "precision_at_k": _round(_mean(case["precision_at_k"] for case in metric_cases)), + "recall_at_k": _round(_mean(case["recall_at_k"] for case in metric_cases)), + "hitrate_at_k": _round( + _mean(1.0 if case["hit_evidence_ids"] else 0.0 for case in metric_cases) + ), + } + if wiki_mode is not None: + summary["wiki_mode"] = normalize_wiki_mode(wiki_mode) + return summary + + +def _case_score( + *, + dataset_name: str, + case_id: str, + question: str, + expected: list[str], + retrieved: list[str], + hits: list[str], + file_paths: list[str], + top_k: int, +) -> dict[str, Any]: + return { + "dataset_name": dataset_name, + "case_id": case_id, + "question": question, + "expected_evidence_ids": expected, + "retrieved_evidence_ids": retrieved, + "hit_evidence_ids": hits, + "retrieved_file_paths": file_paths, + "top_k": top_k, + "precision_at_k": _round(len(hits) / len(retrieved)) if retrieved else 0.0, + "recall_at_k": _round(len(hits) / len(expected)) if expected else 0.0, + } + + +def _evermem_profiles(data_root: Path) -> dict[str, str]: + path = data_root / "profiles.json" + if not path.exists(): + return {} + profiles = json.loads(path.read_text(encoding="utf-8")) + return { + str(profile.get("Name", "")): _render_evermem_profile(profile) + for profile in profiles + if profile.get("Name") + } + + +def _render_evermem_profile(profile: dict[str, Any]) -> str: + skills = ", ".join( + str(item.get("skill", "")) + for item in profile.get("Skills_List", []) + if isinstance(item, dict) and item.get("skill") + ) + interests = ", ".join(str(item) for item in profile.get("Interests", []) if item) + fields = [ + ("Profile", profile.get("Name")), + ("Dept", profile.get("Dept")), + ("Title", profile.get("Title")), + ("Major", profile.get("Major")), + ("Skills", skills), + ("Interests", interests), + ] + return "\n".join(f"{key}: {value}" for key, value in fields if value) + + +def _evermem_rows( + topic_id: str, + days: list[dict[str, Any]], + profiles: dict[str, str] | None = None, +) -> list[dict[str, str]]: + rows = [] + session_number = 0 + for day in days: + date = str(day["date"]) + for group in sorted(day.get("dialogues", {})): + messages = day.get("dialogues", {}).get(group) or [] + if not messages: + continue + session_number += 1 + session_id = f"D{session_number}" + for message in messages or []: + index = str(message["message_index"]) + speaker = str(message.get("speaker", "")) + rows.append( + { + "evidence_id": f"{session_id}:{index}", + "session_id": session_id, + "message_index": index, + "date": date, + "group": str(group), + "speaker": speaker, + "profile": (profiles or {}).get(speaker, ""), + "text": str(message.get("dialogue", "")), + } + ) + return rows + + +def _evermem_expected( + refs: list[dict[str, Any]], + session_by_ref: dict[tuple[str, str], str], +) -> list[str]: + ids = [] + for ref in refs: + session_id = session_by_ref.get((str(ref["date"]), str(ref["group"]))) + if not session_id: + continue + for index in _expand_indices(str(ref["message_index"])): + ids.append(f"{session_id}:{index}") + return _ordered_unique(ids) + + +def _expand_indices(raw: str) -> list[str]: + values = [] + for part in raw.split(","): + part = part.strip() + if "-" in part: + start, end = [int(item.strip()) for item in part.split("-", 1)] + values.extend(str(item) for item in range(start, end + 1)) + elif part: + values.append(str(int(part))) + return values + + +def _evidence_id_from_path(path: Path) -> str: + stem = path.stem + match = re.match(r"^(D\d+)_(\d+)", stem, flags=re.IGNORECASE) + if match: + return f"{match.group(1).upper()}:{match.group(2)}" + return stem + + +def _resolve_root(raw: str, base_root: Path | None) -> Path: + path = Path(raw) + if path.exists() or base_root is None or path.is_absolute(): + return path + return base_root / raw + + +def _ordered_unique(items) -> list[str]: + seen = set() + result = [] + for item in items: + if item and item not in seen: + result.append(item) + seen.add(item) + return result + + +def _mean(values) -> float: + values = list(values) + return sum(values) / len(values) if values else 0.0 + + +def _round(value: float) -> float: + return round(value, 4) + + +def _slug(value: str) -> str: + return re.sub(r"[^0-9A-Za-z]+", "_", value).strip("_").lower() + + +def _write_json(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8") diff --git a/evaluation/wikimem/llm_semantics.py b/evaluation/wikimem/llm_semantics.py new file mode 100644 index 00000000..b5184b39 --- /dev/null +++ b/evaluation/wikimem/llm_semantics.py @@ -0,0 +1,314 @@ +"""LLM-backed semantic compilation primitives for the wikimem Wiki. + +The module is deliberately independent from any dataset adapter. It accepts +source records and returns a small canonical JSON-compatible ontology. A +caller may pass ``llm=None``; in that case no network call is made and the +caller can use its deterministic fallback. +""" + +from __future__ import annotations + +import hashlib +import json +import re +from dataclasses import dataclass, field +from typing import Any, Iterable + +from common.llm.base import LLM +from common.type_def.chat import ChatMessage + + +MEMORY_KINDS = ( + "entity", + "fact", + "event", + "preference", + "skill", + "relationship", + "decision", + "constraint", + "context", + "artifact", +) +MemoryKind = str + + +@dataclass(frozen=True) +class SemanticSource: + """A dataset-independent source block supplied to the memory compiler.""" + + source_id: str + text: str + conversation_id: str = "" + session_id: str = "" + speaker: str = "" + timestamp: str = "" + metadata: dict[str, str] = field(default_factory=dict) + + +@dataclass(frozen=True) +class SemanticEntity: + name: str + entity_type: str = "thing" + description: str = "" + aliases: tuple[str, ...] = () + + +@dataclass(frozen=True) +class SemanticRelation: + subject: str + predicate: str + object: str + + +@dataclass(frozen=True) +class SemanticMemory: + memory_id: str + kind: MemoryKind + content: str + source_id: str + evidence: str = "" + entities: tuple[SemanticEntity, ...] = () + relations: tuple[SemanticRelation, ...] = () + timestamp: str = "" + confidence: float = 0.0 + tags: tuple[str, ...] = () + metadata: dict[str, str] = field(default_factory=dict) + + +@dataclass(frozen=True) +class QueryUnderstanding: + intent: str = "recall" + entities: tuple[str, ...] = () + relation: str = "" + time_expression: str = "" + expanded_terms: tuple[str, ...] = () + memory_kinds: tuple[str, ...] = () + + +_EXTRACTION_SYSTEM_PROMPT = """You are a long-term memory compiler. +Read source blocks and return ONLY a JSON array. Extract only information +explicitly stated in each source; do not use answers, evaluation labels, or +outside knowledge. Keep every item traceable to exactly one source_id. + +Each item must have: +{"source_id":"...","kind":"entity|fact|event|preference|skill|relationship| +decision|constraint|context|artifact", +"content":"one self-contained memory statement in the source language", +"evidence":"the smallest verbatim source span that supports it", +"entities":[{"name":"...","entity_type":"...","description":"...","aliases":["..."]}], +"relations":[{"subject":"...","predicate":"...","object":"..."}], +"timestamp":"ISO-8601 or empty","confidence":0.0,"tags":["..."],"metadata":{}} + +Rules: preserve negation, quantities, names, dates, and modality; do not merge +unrelated facts; emit no item for greetings or unsupported speculation; use a +confidence between 0 and 1; keep content and evidence in the source language. +""" + +_QUERY_SYSTEM_PROMPT = """You are a query understanding module for a long-term memory system. +Return ONLY one JSON object with keys: +{"intent":"recall|compare|decision|preference|procedure|event|profile", +"entities":["..."],"relation":"...","time_expression":"...", +"expanded_terms":["..."],"memory_kinds":["entity|fact|event|preference|skill|relationship| +decision|constraint|context|artifact"]} +Expand paraphrases conservatively. Do not invent entities or facts. +""" + + +def extract_semantic_memories( + llm: LLM, + sources: Iterable[SemanticSource], + *, + batch_size: int = 8, + max_tokens: int = 4096, +) -> list[SemanticMemory]: + """Extract ontology records with a bounded number of LLM calls.""" + + materialized = [source for source in sources if source.text.strip()] + result: list[SemanticMemory] = [] + for start in range(0, len(materialized), max(1, batch_size)): + batch = materialized[start : start + max(1, batch_size)] + source_text = "\n".join( + "---\n" + f"source_id: {source.source_id}\n" + f"conversation_id: {source.conversation_id}\n" + f"session_id: {source.session_id}\n" + f"speaker: {source.speaker}\n" + f"timestamp: {source.timestamp}\n" + f"metadata: {json.dumps(source.metadata, ensure_ascii=False)}\n" + f"text: {source.text}\n---" + for source in batch + ) + response = llm.chat( + [ + ChatMessage(role="system", content=_EXTRACTION_SYSTEM_PROMPT), + ChatMessage(role="user", content=source_text), + ], + temperature=0.0, + max_tokens=max_tokens, + ) + payload = _parse_json_payload(response) + if not isinstance(payload, list): + raise ValueError("LLM memory extraction must return a JSON array") + valid_ids = {source.source_id for source in batch} + for item in payload: + memory = _parse_memory_item(item, valid_ids) + if memory is not None: + result.append(memory) + return result + + +def understand_query( + llm: LLM, + question: str, + *, + known_entities: Iterable[str] = (), + max_tokens: int = 768, +) -> QueryUnderstanding: + """Add semantic intent and conservative query expansions.""" + + known = ", ".join(str(value).strip() for value in known_entities if str(value).strip()) + response = llm.chat( + [ + ChatMessage(role="system", content=_QUERY_SYSTEM_PROMPT), + ChatMessage( + role="user", + content=f"Known entities: {known or '(none)'}\nQuestion: {question}", + ), + ], + temperature=0.0, + max_tokens=max_tokens, + ) + payload = _parse_json_payload(response) + if not isinstance(payload, dict): + raise ValueError("LLM query understanding must return a JSON object") + return QueryUnderstanding( + intent=_clean_scalar(payload.get("intent"), "recall"), + entities=_clean_string_tuple(payload.get("entities")), + relation=_clean_scalar(payload.get("relation")), + time_expression=_clean_scalar(payload.get("time_expression")), + expanded_terms=_clean_string_tuple(payload.get("expanded_terms")), + memory_kinds=tuple( + value + for value in _clean_string_tuple(payload.get("memory_kinds")) + if value in MEMORY_KINDS + ), + ) + + +def _parse_memory_item(item: Any, valid_ids: set[str]) -> SemanticMemory | None: + if not isinstance(item, dict): + return None + source_id = _clean_scalar(item.get("source_id")) + content = _clean_scalar(item.get("content")) + if not source_id or source_id not in valid_ids or not content: + return None + kind = _clean_scalar(item.get("kind"), "context").lower() + if kind not in MEMORY_KINDS: + kind = "context" + confidence = _safe_confidence(item.get("confidence")) + entities = tuple(_parse_entity(value) for value in _as_list(item.get("entities"))) + entities = tuple(value for value in entities if value is not None) + relations = tuple(_parse_relation(value) for value in _as_list(item.get("relations"))) + relations = tuple(value for value in relations if value is not None) + stable_key = f"{source_id}:{kind}:{content}".encode("utf-8") + memory_id = _clean_scalar(item.get("memory_id")) or ( + f"{source_id}:{kind}:{hashlib.sha1(stable_key).hexdigest()[:16]}" + ) + return SemanticMemory( + memory_id=memory_id, + kind=kind, + content=content, + source_id=source_id, + evidence=_clean_scalar(item.get("evidence")), + entities=entities, + relations=relations, + timestamp=_clean_scalar(item.get("timestamp")), + confidence=confidence, + tags=_clean_string_tuple(item.get("tags"))[:8], + metadata={ + str(key): _clean_scalar(value) + for key, value in item.get("metadata", {}).items() + } + if isinstance(item.get("metadata"), dict) + else {}, + ) + + +def _parse_entity(value: Any) -> SemanticEntity | None: + if isinstance(value, str): + name = value.strip() + return SemanticEntity(name=name) if name else None + if not isinstance(value, dict): + return None + name = _clean_scalar(value.get("name")) + if not name: + return None + return SemanticEntity( + name=name, + entity_type=_clean_scalar(value.get("entity_type"), "thing"), + description=_clean_scalar(value.get("description")), + aliases=_clean_string_tuple(value.get("aliases"))[:8], + ) + + +def _parse_relation(value: Any) -> SemanticRelation | None: + if not isinstance(value, dict): + return None + subject = _clean_scalar(value.get("subject")) + predicate = _clean_scalar(value.get("predicate")) + obj = _clean_scalar(value.get("object")) + if not subject or not predicate or not obj: + return None + return SemanticRelation(subject=subject, predicate=predicate, object=obj) + + +def _parse_json_payload(response: str) -> Any: + text = str(response or "").strip() + if text.startswith("```"): + text = re.sub( + r"^```(?:json)?\s*|\s*```$", + "", + text, + flags=re.IGNORECASE | re.DOTALL, + ).strip() + try: + return json.loads(text) + except json.JSONDecodeError: + starts = [index for index in (text.find("["), text.find("{")) if index >= 0] + if not starts: + raise + start = min(starts) + for end in range(len(text), start, -1): + try: + return json.loads(text[start:end]) + except json.JSONDecodeError: + continue + raise + + +def _as_list(value: Any) -> list[Any]: + if isinstance(value, list): + return value + return [] + + +def _clean_scalar(value: Any, default: str = "") -> str: + if value is None: + return default + return str(value).strip() or default + + +def _clean_string_tuple(value: Any) -> tuple[str, ...]: + return tuple( + item + for item in (_clean_scalar(raw) for raw in _as_list(value)) + if item + ) + + +def _safe_confidence(value: Any) -> float: + try: + return min(1.0, max(0.0, float(value))) + except (TypeError, ValueError): + return 0.5 diff --git a/evaluation/wikimem/locomo_refined.py b/evaluation/wikimem/locomo_refined.py new file mode 100644 index 00000000..b2e23fdf --- /dev/null +++ b/evaluation/wikimem/locomo_refined.py @@ -0,0 +1,1214 @@ +"""Rust-compatible LoCoMo_refined multimodal adjunct evaluation. + +The retained evaluator builds the ordinary text wiki first. This module then +implements the example-local Rust adjunct: one multimodal artifact per turn, +optional image materialization, optional OpenAI-compatible vision enrichment, +lexical adjunct retrieval, and post-retrieval evidence rescoring. +""" + +from __future__ import annotations + +import base64 +import json +import mimetypes +import re +import shutil +import urllib.parse +import urllib.request +from collections import defaultdict +from dataclasses import asdict, dataclass, replace +from pathlib import Path +from typing import Any + +from common.llm.base import LLM + +from evaluation.wikimem.retained_eval import ( + CaseScore, + EvalOutput, + PreparedSample, + WikiMode, + _extract_evidence_ids, + _ordered_unique, + prepare_locomo_samples, + normalize_wiki_mode, + run_retained_qmd_eval, + summarize_scores, + summarize_scores_by_locomo_category, + write_harness_artifacts, +) +from evaluation.wikimem.wiki_builder import WikiBuilderMode + + +@dataclass(frozen=True) +class VisionEnrichmentConfig: + api_key: str + model: str + base_url: str + + +@dataclass(frozen=True) +class MultimodalBuildOptions: + download_images: bool = True + request_timeout_secs: int = 8 + redownload_missing_only: bool = True + proxy_url: str | None = None + vision: VisionEnrichmentConfig | None = None + + +@dataclass(frozen=True) +class MultimodalArtifactSummary: + artifact_count: int + downloaded_images: int + failed_images: int + + +@dataclass(frozen=True) +class MultimodalRecallHit: + artifact_id: str + evidence_id: str + file_path: str + score: float + + +@dataclass(frozen=True) +class OfflineExportSummary: + export_root: str + dataset_path: str + manifest_path: str + missing_images_path: str + sample_count: int + total_image_references: int + downloaded_images: int + unresolved_images: int + + +@dataclass(frozen=True) +class _DownloadedImageAsset: + source_url: str + resolved_url: str | None + local_path: str | None + mime_type: str | None + download_status: str + strategy: str | None + error: str | None + + +@dataclass(frozen=True) +class _VisionSummary: + status: str = "disabled" + summary: str = "" + entities: tuple[str, ...] = () + actions: tuple[str, ...] = () + attributes: tuple[str, ...] = () + keywords: tuple[str, ...] = () + error: str | None = None + + +def build_multimodal_memory_artifacts( + sample: PreparedSample, + kb_root: str | Path, + options: MultimodalBuildOptions | None = None, +) -> MultimodalArtifactSummary: + """Build the same raw/manifest/retrieval/wiki adjunct layout as Rust.""" + + options = options or MultimodalBuildOptions() + kb_root = Path(kb_root) + manifests_dir = kb_root / ".kb-research" / "manifests" + retrieval_dir = kb_root / ".kb-research" / "retrieval" + raw_dir = kb_root / "raw" / "multimodal" + memories_dir = kb_root / "wiki" / "memories" + assets_dir = raw_dir / "assets" + for directory in (manifests_dir, retrieval_dir, raw_dir, memories_dir, assets_dir): + directory.mkdir(parents=True, exist_ok=True) + _write_json( + kb_root / ".wikimem-workspace.json", + { + "producer": "mem2.0.wikimem.python", + "dataset_name": "locomo_refined", + "schema_version": 1, + }, + ) + + turns = _collect_multimodal_turns(sample.raw_sample) + downloaded_count = 0 + failed_count = 0 + for turn in turns: + raw_path = raw_dir / f"{turn['artifact_id']}.json" + memory_path = memories_dir / f"{turn['artifact_id']}.md" + manifest_path = manifests_dir / f"{turn['artifact_id']}.json" + retrieval_path = retrieval_dir / f"{turn['artifact_id']}.json" + image_assets = _materialize_image_assets(turn, assets_dir, options) + downloaded_count += sum(item.download_status == "downloaded" for item in image_assets) + failed_count += sum(item.download_status == "failed" for item in image_assets) + vision = _enrich_vision_summary(turn, image_assets, options) + linked_topics = _collect_topics(turn, vision) + manifest = { + "type": "memory", + "modality": "image", + "date": turn["session_date"], + "updated": "2026-04-21", + "tags": ["locomo_refined", "multimodal"], + "aliases": [turn["evidence_id"]], + "sources": ( + [f"locomo_refined:{turn['evidence_id']}"] + if not turn["images"] + else turn["images"] + ), + "maturity": "compiled", + "artifact_id": turn["artifact_id"], + "linked_entities": [turn["speaker"]], + "linked_topics": linked_topics, + "title": f"{turn['speaker']} multimodal memory {turn['evidence_id']}", + "evidence_id": turn["evidence_id"], + "session_number": turn["session_number"], + "speaker": turn["speaker"], + "raw_path": str(raw_path), + "memory_path": str(memory_path), + } + searchable_text = _build_searchable_text(turn, image_assets, vision) + retrieval = { + "artifact_id": turn["artifact_id"], + "evidence_id": turn["evidence_id"], + "title": manifest["title"], + "searchable_text": searchable_text, + "memory_path": str(memory_path), + "vision_status": vision.status, + "vision_summary": vision.summary or None, + "vision_keywords": list(vision.keywords), + } + _write_json( + raw_path, + { + "artifact_id": turn["artifact_id"], + "evidence_id": turn["evidence_id"], + "speaker": turn["speaker"], + "session_number": turn["session_number"], + "session_date": turn["session_date"], + "text": turn["text"], + "caption": turn["caption"], + "query": turn["query"], + "images": turn["images"], + "downloaded_images": [asdict(item) for item in image_assets], + "vision_status": vision.status, + "vision_summary": vision.summary, + "vision_entities": list(vision.entities), + "vision_actions": list(vision.actions), + "vision_attributes": list(vision.attributes), + "vision_keywords": list(vision.keywords), + "keywords": list(vision.keywords), + "vision_error": vision.error, + }, + ) + _write_json(manifest_path, manifest) + _write_json(retrieval_path, retrieval) + memory_path.write_text( + _render_memory_page(manifest, turn, image_assets, vision), + encoding="utf-8", + ) + return MultimodalArtifactSummary( + artifact_count=len(turns), + downloaded_images=downloaded_count, + failed_images=failed_count, + ) + + +def recall_multimodal_artifacts_for_question( + kb_root: str | Path, + question: str, + top_k: int, +) -> list[MultimodalRecallHit]: + provenance_path = Path(kb_root) / ".wikimem-workspace.json" + try: + provenance = json.loads(provenance_path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return [] + if provenance != { + "producer": "mem2.0.wikimem.python", + "dataset_name": "locomo_refined", + "schema_version": 1, + }: + return [] + retrieval_dir = Path(kb_root) / ".kb-research" / "retrieval" + if not retrieval_dir.exists(): + return [] + hits: list[MultimodalRecallHit] = [] + for path in sorted(retrieval_dir.glob("*.json")): + try: + record = json.loads(path.read_text(encoding="utf-8")) + score = _token_score(question, str(record.get("searchable_text") or "")) + if score > 0.0: + hits.append( + MultimodalRecallHit( + artifact_id=str(record["artifact_id"]), + evidence_id=str(record["evidence_id"]), + file_path=str(record["memory_path"]), + score=score, + ) + ) + except (OSError, KeyError, TypeError, ValueError): + continue + hits.sort(key=lambda item: (-item.score, item.artifact_id)) + return hits[: max(top_k, 1)] + + +def rescore_case_with_multimodal_artifacts( + baseline: CaseScore, + adjunct_hits: list[MultimodalRecallHit], +) -> CaseScore: + retrieved_file_paths = list(baseline.retrieved_file_paths) + for hit in adjunct_hits: + if hit.file_path not in retrieved_file_paths: + retrieved_file_paths.append(hit.file_path) + retrieved_evidence = list(baseline.retrieved_evidence) + for path in retrieved_file_paths: + try: + retrieved_evidence.extend( + _extract_evidence_ids(Path(path).read_text(encoding="utf-8")) + ) + except OSError: + continue + retrieved_evidence = _ordered_unique(retrieved_evidence) + expected = _ordered_unique(baseline.expected_evidence) + expected_set = set(expected) + hit_evidence = [item for item in retrieved_evidence if item in expected_set] + retrieved_count = len(retrieved_evidence) + hit_count = len(hit_evidence) + expected_count = len(expected) + return replace( + baseline, + expected_evidence=expected, + retrieved_evidence=retrieved_evidence, + hit_evidence=hit_evidence, + retrieved_file_paths=retrieved_file_paths, + retrieved_file_count=len(retrieved_file_paths), + evidence_precision=hit_count / retrieved_count if retrieved_count else 0.0, + evidence_recall=hit_count / expected_count if expected_count else 0.0, + full_evidence_hit=bool(expected_count) and hit_count == expected_count, + ) + + +def run_locomo_refined_offline_eval( + *, + dataset_path: str | Path, + output_dir: str | Path, + workspace_root: str | Path, + top_k: int = 24, + question_limit: int | None = None, + wiki_mode: WikiMode = "multimodal", + multimodal_top_k: int = 6, + download_images: bool = True, + request_timeout_secs: int = 8, + redownload_missing_only: bool = True, + proxy_url: str | None = None, + vision_config: VisionEnrichmentConfig | None = None, + sample_filter: set[str] | None = None, + offline_export_root: str | Path | None = None, + llm: LLM | None = None, + wiki_builder_mode: WikiBuilderMode = "llm", + query_llm: LLM | None = None, +) -> dict[str, Any]: + """Run Rust's text baseline plus multimodal adjunct rescoring.""" + + wiki_mode = normalize_wiki_mode(wiki_mode) + payloads = json.loads(Path(dataset_path).read_text(encoding="utf-8")) + samples = prepare_locomo_samples( + payloads, + sample_filter=sample_filter, + include_multimodal_context=True, + ) + output_dir = Path(output_dir) + workspace_root = Path(workspace_root) + baseline = run_retained_qmd_eval( + dataset_name="locomo_refined", + samples=samples, + workspace_root=workspace_root, + top_k=top_k, + question_limit=question_limit, + harness_root=output_dir / "harness_baseline", + wiki_mode="text", + llm=llm, + wiki_builder_mode=wiki_builder_mode, + query_llm=query_llm, + ) + output_dir.mkdir(parents=True, exist_ok=True) + _write_json(output_dir / "locomo_refined_retrieval_eval_baseline.json", _eval_output_dict(baseline)) + + summaries: list[MultimodalArtifactSummary] = [] + if wiki_mode == "multimodal": + options = MultimodalBuildOptions( + download_images=download_images, + request_timeout_secs=request_timeout_secs, + redownload_missing_only=redownload_missing_only, + proxy_url=proxy_url, + vision=vision_config, + ) + for sample in samples: + summaries.append( + build_multimodal_memory_artifacts( + sample, + workspace_root / sample.sample_id, + options, + ) + ) + + offline_export = None + if offline_export_root is not None: + offline_export = export_offline_locomo_refined_dataset( + samples=samples, + workspace_root=workspace_root, + export_root=offline_export_root, + ) + + final_cases: list[CaseScore] = [] + multimodal_cases_with_hits = 0 + for case in baseline.cases: + hits = ( + recall_multimodal_artifacts_for_question( + case.knowledge_base_root, + case.question, + multimodal_top_k, + ) + if wiki_mode == "multimodal" + else [] + ) + if hits: + multimodal_cases_with_hits += 1 + final_cases.append(rescore_case_with_multimodal_artifacts(case, hits)) + final = EvalOutput( + summary=summarize_scores("locomo_refined", final_cases), + cases=final_cases, + stage_profile=baseline.stage_profile, + ) + _write_json(output_dir / "locomo_refined_retrieval_eval.json", _eval_output_dict(final)) + write_harness_artifacts( + output_dir / "harness", + final, + config=_harness_config( + workspace_root=workspace_root, + top_k=top_k, + question_limit=question_limit, + wiki_mode=wiki_mode, + vision_config=vision_config, + sample_filter=sample_filter, + wiki_builder_mode=wiki_builder_mode, + llm=llm, + query_llm=query_llm, + ), + ) + download_report_path: Path | None = None + if wiki_mode == "multimodal": + download_report_path = output_dir / "harness" / "download_source_report.json" + _write_json( + download_report_path, + _build_download_source_report(workspace_root), + ) + report = { + "baseline_summary": asdict(baseline.summary), + "final_summary": asdict(final.summary), + "category_breakdown": [ + asdict(item) for item in summarize_scores_by_locomo_category(final.cases) + ], + "multimodal_artifact_count": sum(item.artifact_count for item in summaries), + "multimodal_cases_with_hits": multimodal_cases_with_hits, + "downloaded_images": sum(item.downloaded_images for item in summaries), + "failed_images": sum(item.failed_images for item in summaries), + "download_report_path": str(download_report_path) if download_report_path else None, + "offline_export": asdict(offline_export) if offline_export else None, + "baseline_output_path": str(output_dir / "locomo_refined_retrieval_eval_baseline.json"), + "final_output_path": str(output_dir / "locomo_refined_retrieval_eval.json"), + "harness_dir": str(output_dir / "harness"), + } + _write_json(output_dir / "harness" / "locomo_refined_run_report.json", report) + # Keep the legacy runner contract (summary/cases/stage_profile) while + # exposing the new multimodal run report fields to callers. + return _jsonable( + report + | { + "summary": asdict(final.summary), + "cases": [asdict(case) for case in final.cases], + "stage_profile": asdict(final.stage_profile), + } + ) + + +def _build_download_source_report(workspace_root: Path) -> dict[str, Any]: + """Aggregate image outcomes by source domain like Rust's report.""" + + successful: dict[str, int] = defaultdict(int) + failed: dict[str, dict[str, Any]] = defaultdict( + lambda: {"failed_count": 0, "example_errors": []} + ) + for raw_path in sorted(workspace_root.glob("*/raw/multimodal/*_multimodal.json")): + try: + payload = json.loads(raw_path.read_text(encoding="utf-8")) + except (OSError, ValueError): + continue + for item in payload.get("downloaded_images") or []: + if not isinstance(item, dict): + continue + source_url = str(item.get("source_url") or "") + domain = urllib.parse.urlparse(source_url).hostname or "(unknown)" + status = item.get("download_status") + if status == "downloaded": + successful[domain] += 1 + elif status == "failed": + entry = failed[domain] + entry["failed_count"] += 1 + error = str(item.get("error") or "unknown_error") + if error not in entry["example_errors"] and len(entry["example_errors"]) < 3: + entry["example_errors"].append(error) + successful_domains = [ + {"domain": domain, "downloaded_count": count} + for domain, count in sorted(successful.items(), key=lambda pair: (-pair[1], pair[0])) + ] + failed_domains = [ + {"domain": domain, **value} + for domain, value in sorted( + failed.items(), key=lambda pair: (-pair[1]["failed_count"], pair[0]) + ) + ] + return { + "total_downloaded_images": sum(successful.values()), + "total_failed_images": sum(item["failed_count"] for item in failed.values()), + "successful_domains": successful_domains, + "failed_domains": failed_domains, + } + + +def export_offline_locomo_refined_dataset( + *, + samples: list[PreparedSample], + workspace_root: str | Path, + export_root: str | Path, +) -> OfflineExportSummary: + """Rewrite image references to a self-contained offline dataset. + + Only assets recorded by the Python-generated multimodal artifacts are + copied. Missing or failed downloads remain as their original URLs and are + listed in ``missing_images.json``. + """ + + workspace_root = Path(workspace_root) + export_root = Path(export_root) + export_root.mkdir(parents=True, exist_ok=True) + exported_samples: list[dict[str, Any]] = [] + manifest_entries: list[dict[str, Any]] = [] + missing_entries: list[dict[str, Any]] = [] + total_image_references = 0 + downloaded_images = 0 + for sample in samples: + sample_value = json.loads(json.dumps(sample.raw_sample, ensure_ascii=False)) + asset_map = _load_sample_asset_map( + workspace_root / sample.sample_id / "raw" / "multimodal" + ) + sample_total, sample_downloaded = _rewrite_sample_images_for_offline_export( + sample_id=sample.sample_id, + sample_value=sample_value, + export_root=export_root, + asset_map=asset_map, + manifest_entries=manifest_entries, + missing_entries=missing_entries, + ) + total_image_references += sample_total + downloaded_images += sample_downloaded + exported_samples.append(sample_value) + dataset_path = export_root / "locomo_refined_offline.json" + manifest_path = export_root / "manifest.json" + missing_images_path = export_root / "missing_images.json" + _write_json(dataset_path, exported_samples) + _write_json(manifest_path, {"entries": manifest_entries}) + _write_json(missing_images_path, {"entries": missing_entries}) + return OfflineExportSummary( + export_root=str(export_root), + dataset_path=str(dataset_path), + manifest_path=str(manifest_path), + missing_images_path=str(missing_images_path), + sample_count=len(samples), + total_image_references=total_image_references, + downloaded_images=downloaded_images, + unresolved_images=len(missing_entries), + ) + + +def _load_sample_asset_map(raw_dir: Path) -> dict[str, list[dict[str, Any]]]: + asset_map: dict[str, list[dict[str, Any]]] = {} + if not raw_dir.is_dir(): + return asset_map + for path in sorted(raw_dir.glob("*_multimodal.json")): + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + continue + evidence_id = payload.get("evidence_id") + if isinstance(evidence_id, str): + records = payload.get("downloaded_images") + asset_map[evidence_id] = records if isinstance(records, list) else [] + return asset_map + + +def _rewrite_sample_images_for_offline_export( + *, + sample_id: str, + sample_value: dict[str, Any], + export_root: Path, + asset_map: dict[str, list[dict[str, Any]]], + manifest_entries: list[dict[str, Any]], + missing_entries: list[dict[str, Any]], +) -> tuple[int, int]: + conversation = sample_value.get("conversation") + if not isinstance(conversation, dict): + raise ValueError(f"sample {sample_id} conversation must be an object") + session_keys = sorted( + ( + int(key.removeprefix("session_")), + key, + ) + for key in conversation + if key.startswith("session_") and key.removeprefix("session_").isdigit() + ) + total = 0 + downloaded = 0 + for _, session_key in session_keys: + turns = conversation.get(session_key) + if not isinstance(turns, list): + continue + for turn in turns: + if not isinstance(turn, dict): + continue + dia_id = str(turn.get("dia_id") or "") + if not dia_id: + continue + source_urls = _normalize_multimodal_image_list(turn.get("img_url")) + if not source_urls: + continue + speaker = str(turn.get("speaker") or "") + turn_text = str(turn.get("text") or "") + records = asset_map.get(dia_id, []) + rewritten_urls: list[str] = [] + original_urls: list[str] = [] + local_paths: list[str | None] = [] + statuses: list[str] = [] + errors: list[str | None] = [] + for image_index, source_url in enumerate(source_urls): + total += 1 + asset = ( + records[image_index] + if image_index < len(records) + and isinstance(records[image_index], dict) + else None + ) + resolved = _resolve_offline_image( + sample_id=sample_id, + dia_id=dia_id, + image_index=image_index, + source_url=source_url, + asset=asset, + export_root=export_root, + ) + if resolved["local_relative_path"] is not None: + downloaded += 1 + entry = { + "sample_id": sample_id, + "dia_id": dia_id, + "speaker": speaker, + "turn_text": turn_text, + "image_index": image_index, + "source_url": source_url, + "rewritten_img_url": resolved["rewritten_img_url"], + "local_relative_path": resolved["local_relative_path"], + "download_status": resolved["download_status"], + "error": resolved["error"], + } + manifest_entries.append(entry) + if entry["local_relative_path"] is None: + missing_entries.append(entry) + rewritten_urls.append(resolved["rewritten_img_url"]) + original_urls.append(source_url) + local_paths.append(resolved["local_relative_path"]) + statuses.append(resolved["download_status"]) + errors.append(resolved["error"]) + turn["img_url"] = rewritten_urls + turn["img_url_original"] = original_urls + turn["img_local_path"] = local_paths + turn["img_download_status"] = statuses + turn["img_download_error"] = errors + return total, downloaded + + +def _resolve_offline_image( + *, + sample_id: str, + dia_id: str, + image_index: int, + source_url: str, + asset: dict[str, Any] | None, + export_root: Path, +) -> dict[str, Any]: + if asset is None: + return { + "rewritten_img_url": source_url, + "local_relative_path": None, + "download_status": "missing_artifact", + "error": "missing multimodal artifact", + } + local_path = asset.get("local_path") + if isinstance(local_path, str) and local_path: + source_path = Path(local_path) + if source_path.exists(): + extension = source_path.suffix.lstrip(".") or "img" + relative_path = ( + Path("images") + / sample_id + / f"{dia_id.replace(':', '_')}_{image_index:02d}.{extension}" + ).as_posix() + destination = export_root / Path(relative_path) + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source_path, destination) + return { + "rewritten_img_url": relative_path, + "local_relative_path": relative_path, + "download_status": "downloaded", + "error": None, + } + status = str(asset.get("download_status") or "unresolved") + return { + "rewritten_img_url": source_url, + "local_relative_path": None, + "download_status": status, + "error": str(asset.get("error") or "local downloaded asset missing on disk"), + } + + +def _collect_multimodal_turns(raw_sample: dict[str, Any]) -> list[dict[str, Any]]: + conversation = raw_sample.get("conversation") + if not isinstance(conversation, dict): + raise ValueError("raw_sample.conversation must be an object") + sessions = sorted( + ( + int(key.removeprefix("session_")), + value, + ) + for key, value in conversation.items() + if key.startswith("session_") + and key.removeprefix("session_").isdigit() + ) + turns: list[dict[str, Any]] = [] + for session_number, messages in sessions: + if not isinstance(messages, list): + continue + date = conversation.get(f"session_{session_number}_date_time") + session_date = date if isinstance(date, str) else "" + for message in messages: + if not isinstance(message, dict): + continue + images = _normalize_multimodal_image_list(message.get("img_url")) + caption = message.get("blip_caption") + query = message.get("query") + if caption is not None and not isinstance(caption, str): + caption = "" + if query is not None and not isinstance(query, str): + query = "" + caption = (caption or "").strip() + query = (query or "").strip() + if not images and not caption and not query: + continue + evidence_id = message.get("dia_id") + if not isinstance(evidence_id, str): + raise ValueError("multimodal message missing dia_id") + turns.append( + { + "artifact_id": f"{evidence_id.replace(':', '_')}_multimodal", + "evidence_id": evidence_id, + "speaker": str(message.get("speaker") or ""), + "session_number": session_number, + "session_date": session_date, + "text": str(message.get("text") or "").strip(), + "caption": caption, + "query": query, + "images": images, + } + ) + return turns + + +def _normalize_multimodal_image_list(raw: Any) -> list[str]: + """Match Rust artifact collection: ignore non-string array members.""" + + if isinstance(raw, str): + value = raw.strip() + return [value] if value else [] + if isinstance(raw, list): + return [item.strip() for item in raw if isinstance(item, str) and item.strip()] + return [] + + +def _materialize_image_assets( + turn: dict[str, Any], + assets_dir: Path, + options: MultimodalBuildOptions, +) -> list[_DownloadedImageAsset]: + if not options.download_images: + return [ + _DownloadedImageAsset(url, None, None, None, "skipped", "disabled", None) + for url in turn["images"] + ] + opener = _build_url_opener(options.proxy_url) + assets: list[_DownloadedImageAsset] = [] + for index, source_url in enumerate(turn["images"]): + local_path = _resolve_asset_path(assets_dir, turn["artifact_id"], index, source_url) + if options.redownload_missing_only and local_path.exists(): + assets.append( + _DownloadedImageAsset( + source_url, + source_url, + str(local_path), + _infer_mime_from_path(local_path), + "downloaded", + "cache_reuse", + None, + ) + ) + continue + if source_url.startswith("data:"): + try: + mime = _materialize_data_url_asset(source_url, local_path) + assets.append( + _DownloadedImageAsset( + source_url, + source_url, + str(local_path), + mime, + "downloaded", + "data_url_inline", + None, + ) + ) + except Exception as exc: + assets.append( + _DownloadedImageAsset( + source_url, + source_url, + None, + None, + "failed", + "data_url_inline", + str(exc), + ) + ) + continue + last_error = "download did not start" + last_strategy = "initial" + last_url = source_url + for url, strategy, browser_headers, identity in _build_download_attempt_plans(source_url): + last_strategy, last_url = strategy, url + try: + mime, resolved = _execute_download_attempt( + opener, + url, + strategy, + browser_headers, + identity, + local_path, + source_url, + options.request_timeout_secs, + ) + assets.append( + _DownloadedImageAsset( + source_url, + resolved, + str(local_path), + mime, + "downloaded", + strategy, + None, + ) + ) + last_error = "" + break + except Exception as exc: + last_error = str(exc) + if last_error: + assets.append( + _DownloadedImageAsset( + source_url, + last_url, + None, + None, + "failed", + last_strategy, + last_error, + ) + ) + return assets + + +def _build_url_opener(proxy_url: str | None, *, disable_proxy: bool = False): + if disable_proxy: + return urllib.request.build_opener(urllib.request.ProxyHandler({})) + if proxy_url and proxy_url.strip(): + return urllib.request.build_opener(urllib.request.ProxyHandler({"http": proxy_url, "https": proxy_url})) + return urllib.request.build_opener() + + +def _build_download_attempt_plans(source_url: str) -> list[tuple[str, str, bool, bool]]: + plans: list[tuple[str, str, bool, bool]] = [] + candidates = [ + (source_url, "direct", False, False), + (source_url, "browser_headers", True, False), + (source_url, "browser_headers_identity", True, True), + ] + parsed = urllib.parse.urlparse(source_url) + if parsed.scheme == "http": + candidates.append((urllib.parse.urlunparse(parsed._replace(scheme="https")), "scheme_swap_browser_identity", True, True)) + if parsed.hostname and parsed.hostname.lower() == "imgur.com": + candidates.append((urllib.parse.urlunparse(parsed._replace(netloc="i.imgur.com")), "normalized_direct_browser_identity", True, True)) + if parsed.hostname and parsed.hostname.lower() == "i.redd.it": + candidates.append(("https://www.reddit.com/media?url=" + urllib.parse.quote(source_url, safe=""), "reddit_media_browser_identity", True, True)) + seen: set[tuple[str, bool, bool]] = set() + for item in candidates: + key = (item[0], item[2], item[3]) + if key not in seen: + seen.add(key) + plans.append(item) + return plans + + +def _execute_download_attempt( + opener: Any, + url: str, + strategy: str, + browser_headers: bool, + force_identity: bool, + local_path: Path, + source_url: str, + timeout_secs: int, +) -> tuple[str | None, str]: + headers = {} + if browser_headers: + headers.update( + { + "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 Chrome/135.0 Safari/537.36", + "Accept": "image/avif,image/webp,image/apng,image/*,*/*;q=0.8", + "Referer": _derive_referer(source_url), + } + ) + if force_identity: + headers["Accept-Encoding"] = "identity" + request = urllib.request.Request(url, headers=headers) + with opener.open(request, timeout=max(timeout_secs, 1)) as response: + status = getattr(response, "status", response.getcode()) + if status < 200 or status >= 300: + raise RuntimeError(f"http {status}") + data = response.read() + mime = response.headers.get("Content-Type") + local_path.parent.mkdir(parents=True, exist_ok=True) + local_path.write_bytes(data) + return mime, url + + +def _resolve_asset_path(assets_dir: Path, artifact_id: str, index: int, source_url: str) -> Path: + extension = Path(urllib.parse.urlparse(source_url).path).suffix.lstrip(".") + if not extension or len(extension) > 8 or not extension.isalnum(): + extension = "img" + return assets_dir / f"{artifact_id}_{index}.{extension}" + + +def _derive_referer(source_url: str) -> str: + parsed = urllib.parse.urlparse(source_url) + return f"{parsed.scheme}://{parsed.netloc}/" if parsed.scheme and parsed.netloc else "https://www.google.com/" + + +def _materialize_data_url_asset(source_url: str, local_path: Path) -> str | None: + metadata, encoded = source_url.split(",", 1) + if not metadata.endswith(";base64"): + raise ValueError("invalid data url: only base64 payloads are supported") + mime = metadata.removeprefix("data:").removesuffix(";base64") or None + local_path.parent.mkdir(parents=True, exist_ok=True) + local_path.write_bytes(base64.b64decode(encoded)) + return mime + + +def _infer_mime_from_path(path: Path) -> str | None: + return mimetypes.guess_type(path.name)[0] + + +def _enrich_vision_summary( + turn: dict[str, Any], + image_assets: list[_DownloadedImageAsset], + options: MultimodalBuildOptions, +) -> _VisionSummary: + if options.vision is None: + return _VisionSummary() + local_assets = [item for item in image_assets if item.local_path] + if not local_assets: + return _VisionSummary(status="skipped", error="no_local_images") + merged_summary: list[str] = [] + entities: list[str] = [] + actions: list[str] = [] + attributes: list[str] = [] + keywords: list[str] = [] + for asset in local_assets: + try: + data = Path(asset.local_path).read_bytes() + mime = asset.mime_type or "image/jpeg" + data_url = f"data:{mime};base64,{base64.b64encode(data).decode('ascii')}" + prompt = ( + "You are building a retrieval memory for a conversation benchmark. " + "Return strict JSON with keys summary, entities, actions, attributes, keywords. " + "Be concise and concrete.\n" + f"Speaker: {turn['speaker']}\nTurn text: {turn['text']}\n" + f"Caption hint: {turn['caption']}\nQuery hint: {turn['query']}" + ) + body = { + "model": options.vision.model, + "stream": False, + "temperature": 0.0, + "messages": [{ + "role": "user", + "content": [ + {"type": "text", "text": prompt}, + {"type": "image_url", "image_url": {"url": data_url}}, + ], + }], + } + content = _vision_request(options.vision, body, options.request_timeout_secs) + parsed = _parse_vision_output(content) + if parsed["summary"]: + merged_summary.append(parsed["summary"].strip()) + entities.extend(parsed["entities"]) + actions.extend(parsed["actions"]) + attributes.extend(parsed["attributes"]) + keywords.extend(parsed["keywords"]) + except Exception as exc: + return _VisionSummary( + status="failed", + summary=" | ".join(item for item in merged_summary if item), + entities=tuple(_ordered_unique(entities)), + actions=tuple(_ordered_unique(actions)), + attributes=tuple(_ordered_unique(attributes)), + keywords=tuple(_ordered_unique(keywords)), + error=str(exc), + ) + return _VisionSummary( + status="enriched", + summary=" | ".join(item for item in merged_summary if item), + entities=tuple(_ordered_unique(entities)), + actions=tuple(_ordered_unique(actions)), + attributes=tuple(_ordered_unique(attributes)), + keywords=tuple(_ordered_unique(keywords)), + ) + + +def _vision_request(config: VisionEnrichmentConfig, body: dict[str, Any], timeout_secs: int) -> str: + endpoint = _normalize_openai_base_url(config.base_url) + request = urllib.request.Request( + endpoint, + data=json.dumps(body).encode("utf-8"), + headers={ + "Authorization": f"Bearer {config.api_key}", + "Content-Type": "application/json", + }, + method="POST", + ) + opener = _build_url_opener(None, disable_proxy=True) + with opener.open(request, timeout=max(timeout_secs, 1)) as response: + status = getattr(response, "status", response.getcode()) + if status < 200 or status >= 300: + raise RuntimeError(f"http {status}") + payload = json.loads(response.read().decode("utf-8")) + choices = payload.get("choices") if isinstance(payload, dict) else None + message = choices[0].get("message") if isinstance(choices, list) and choices else None + content = message.get("content") if isinstance(message, dict) else "" + if isinstance(content, str): + return content + if isinstance(content, list): + return "\n".join(str(item.get("text")) for item in content if isinstance(item, dict) and item.get("text")) + return "" + + +def _normalize_openai_base_url(base_url: str) -> str: + trimmed = base_url.strip().rstrip("/") + if not trimmed: + return "https://api.openai.com/v1/chat/completions" + parsed = urllib.parse.urlparse(trimmed) + path = parsed.path.rstrip("/") + if path.endswith("/chat/completions"): + return trimmed + if path in {"", "/", "/v1", "/api/v1", "/v1beta", "/api/v1beta"}: + path = "/v1/chat/completions" if path in {"", "/"} else f"{path}/chat/completions" + return urllib.parse.urlunparse(parsed._replace(path=path)) + return trimmed + + +def _parse_vision_output(content: str) -> dict[str, Any]: + cleaned = content.strip() + if not cleaned: + raise ValueError("empty vision content") + for candidate in (cleaned, cleaned.removeprefix("```json").removeprefix("```").removesuffix("```").strip()): + try: + parsed = json.loads(candidate) + if isinstance(parsed, dict) and isinstance(parsed.get("summary"), str): + return { + "summary": parsed["summary"], + "entities": _string_list(parsed.get("entities")), + "actions": _string_list(parsed.get("actions")), + "attributes": _string_list(parsed.get("attributes")), + "keywords": _string_list(parsed.get("keywords")), + } + except json.JSONDecodeError: + pass + return {"summary": cleaned, "entities": [], "actions": [], "attributes": [], "keywords": _tokenize(cleaned)} + + +def _string_list(value: Any) -> list[str]: + if not isinstance(value, list): + return [] + return [item.strip() for item in value if isinstance(item, str) and item.strip()] + + +def _collect_topics(turn: dict[str, Any], vision: _VisionSummary) -> list[str]: + values = [turn["query"], turn["caption"], vision.summary, *vision.keywords] + topics = _ordered_unique(item for item in values if item) + return topics or ["multimodal"] + + +def _build_searchable_text( + turn: dict[str, Any], + image_assets: list[_DownloadedImageAsset], + vision: _VisionSummary, +) -> str: + local_assets = " ".join(item.local_path for item in image_assets if item.local_path) + return "\n".join( + [ + turn["speaker"], + turn["text"], + turn["caption"], + turn["query"], + turn["session_date"], + " ".join(turn["images"]), + local_assets, + vision.summary, + " ".join(vision.entities), + " ".join(vision.actions), + " ".join(vision.attributes), + " ".join(vision.keywords), + vision.status, + ] + ) + + +def _render_memory_page( + manifest: dict[str, Any], + turn: dict[str, Any], + image_assets: list[_DownloadedImageAsset], + vision: _VisionSummary, +) -> str: + image_lines = "(none)" if not image_assets else "\n".join( + f"- source: {item.source_url}\n local: {item.local_path or '(none)'}\n status: {item.download_status}" + for item in image_assets + ) + if not any((vision.summary, vision.entities, vision.actions, vision.attributes, vision.keywords)): + vision_section = f"## Vision Summary\nstatus: {vision.status}\n" + else: + vision_section = ( + f"## Vision Summary\nstatus: {vision.status}\nsummary: {vision.summary}\n" + f"entities: {', '.join(vision.entities) or '(none)'}\n" + f"actions: {', '.join(vision.actions) or '(none)'}\n" + f"attributes: {', '.join(vision.attributes) or '(none)'}\n" + f"keywords: {', '.join(vision.keywords) or '(none)'}\n" + ) + quote = lambda value: str(value).replace('"', '\\"') + sources = ", ".join(f'"{quote(item)}"' for item in manifest["sources"]) + topics = ", ".join(f'"{quote(item)}"' for item in manifest["linked_topics"]) + return ( + f'---\ntype: "{quote(manifest["type"])}"\nmodality: "{quote(manifest["modality"])}"\n' + f'date: "{quote(manifest["date"])}"\nupdated: "{quote(manifest["updated"])}"\n' + f'tags: ["locomo_refined", "multimodal"]\naliases: ["{quote(turn["evidence_id"])}"]\n' + f"sources: [{sources}]\nmaturity: \"compiled\"\nartifact_id: \"{quote(turn['artifact_id'])}\"\n" + f'linked_entities: ["{quote(turn["speaker"])}"]\nlinked_topics: [{topics}]\n---\n\n' + f"# {manifest['title']}\n\n- Evidence: {turn['evidence_id']}\n- Session: D{turn['session_number']}\n" + f"- Speaker: {turn['speaker']}\n- Query: {turn['query']}\n\n## Turn Text\n{turn['text']}\n\n" + f"## Caption\n{turn['caption'] or '(none)'}\n\n## Images\n{image_lines}\n\n{vision_section}" + ) + + +def _token_score(query: str, text: str) -> float: + query_tokens = _tokenize(query) + text_tokens = set(_tokenize(text)) + if not query_tokens or not text_tokens: + return 0.0 + hits = sum(token in text_tokens for token in query_tokens) + phrase_bonus = 2.0 if query.lower() in text.lower() else 0.0 + return float(hits) + phrase_bonus if hits else 0.0 + + +def _tokenize(text: str) -> list[str]: + values = [item for item in re.split(r"[^0-9A-Za-z]+", text) if item] + result = [] + for value in values: + lower = value.lower() + if lower.endswith("ing") and len(lower) > 5: + lower = lower[:-3] + elif lower.endswith("ed") and len(lower) > 4: + lower = lower[:-2] + elif lower.endswith("s") and len(lower) > 3: + lower = lower[:-1] + if len(lower) > 1: + result.append(lower) + return result + + +def _write_json(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(value, ensure_ascii=False, indent=2), encoding="utf-8") + + +def _eval_output_dict(output: EvalOutput) -> dict[str, Any]: + return { + "summary": asdict(output.summary), + "cases": [asdict(case) for case in output.cases], + "stage_profile": asdict(output.stage_profile), + } + + +def _harness_config( + *, + workspace_root: Path, + top_k: int, + question_limit: int | None, + wiki_mode: str, + vision_config: VisionEnrichmentConfig | None, + sample_filter: set[str] | None, + wiki_builder_mode: WikiBuilderMode = "deterministic", + llm: LLM | None = None, + query_llm: LLM | None = None, +): + from evaluation.wikimem.retained_eval import EvalHarnessConfig + + return EvalHarnessConfig( + dataset_name="locomo_refined", + samples=",".join(sorted(sample_filter)) if sample_filter else None, + question_limit=question_limit, + top_k=top_k, + workspace_root=str(workspace_root), + llm_provider=( + type(llm).__name__ + if llm is not None + else type(query_llm).__name__ + if query_llm is not None + else ("nvidia" if vision_config else None) + ), + retrieval_plugins=["qmd_consensus"], + wiki_mode=wiki_mode, + wiki_builder_mode=wiki_builder_mode, + ) + + +def _jsonable(value: Any) -> Any: + if hasattr(value, "__dataclass_fields__"): + return {key: _jsonable(item) for key, item in asdict(value).items()} + if isinstance(value, dict): + return {str(key): _jsonable(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_jsonable(item) for item in value] + return value diff --git a/evaluation/wikimem/longmemeval.py b/evaluation/wikimem/longmemeval.py new file mode 100644 index 00000000..177ecf2d --- /dev/null +++ b/evaluation/wikimem/longmemeval.py @@ -0,0 +1,563 @@ +"""LongMemEval retrieval-only helpers for the wikimem Python migration.""" + +from __future__ import annotations + +import json +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +from common.llm.base import LLM + +from evaluation.wikimem.retained_eval import ( + ConversationRecord, + LoCoMoQuestion, + ObservationNote, + PreparedSample, + SessionEvents, + WikiMode, + run_retained_qmd_eval, +) +from evaluation.wikimem.wiki_builder import WikiBuilderMode + + +@dataclass(frozen=True) +class LongMemEvalSummary: + dataset_name: str + granularity: str + total_cases: int + evaluated_cases: int + skipped_abstention_cases: int + skipped_no_target_cases: int + averaged_metrics: dict[str, dict[str, float]] + + +@dataclass(frozen=True) +class LongMemEvalRetrievalResults: + query: str + granularity: str + ranked_items: list[dict[str, Any]] + ranked_item_count: int + metrics: dict[str, dict[str, float]] + + +@dataclass(frozen=True) +class LongMemEvalCaseScore: + question_id: str + question_type: str + question: str + question_date: str | None + answer: Any + abstention_case: bool + evaluated: bool + expected_session_ids: list[str] + expected_turn_ids: list[str] + retrieval_results: LongMemEvalRetrievalResults + knowledge_base_root: str + kb_fallback_used: bool + kb_warnings: list[str] = field(default_factory=list) + + +def load_longmemeval_samples( + path: str | Path, + sample_filter: set[str] | None = None, +) -> list[dict[str, Any]]: + rows = json.loads(Path(path).read_text(encoding="utf-8")) + samples = [] + for raw in rows: + question_id = str(raw.get("question_id", "")) + if sample_filter is not None and question_id not in sample_filter: + continue + samples.append(_prepare_longmemeval_sample(raw)) + return samples + + +def run_python_longmemeval_retrieval_eval( + *, + dataset_path: str | Path, + output_dir: str | Path, + workspace_root: str | Path, + top_k: int = 24, + granularity: str = "turn", + sample_limit: int | None = None, + wiki_mode: WikiMode = "text", + llm: LLM | None = None, + wiki_builder_mode: WikiBuilderMode = "llm", + query_llm: LLM | None = None, +) -> dict[str, Any]: + samples = load_longmemeval_samples(dataset_path) + if sample_limit is not None: + samples = samples[:sample_limit] + locomo_samples = adapt_longmemeval_to_locomo_samples(samples) + output_dir = Path(output_dir) + retained_output = run_retained_qmd_eval( + dataset_name="longmemeval", + samples=locomo_samples, + workspace_root=workspace_root, + top_k=top_k, + question_limit=1, + harness_root=output_dir / "harness", + wiki_mode=wiki_mode, + llm=llm, + wiki_builder_mode=wiki_builder_mode, + query_llm=query_llm, + ) + result = convert_retained_output_to_longmemeval( + samples=samples, + retained_cases=retained_output.cases, + granularity=granularity, + workspace_root=workspace_root, + ) + output_dir.mkdir(parents=True, exist_ok=True) + (output_dir / "longmemeval_retrieval_eval.json").write_text( + json.dumps(_jsonable(result), ensure_ascii=False, indent=2), + encoding="utf-8", + ) + return _jsonable(result) + + +def adapt_longmemeval_to_locomo_samples( + samples: list[dict[str, Any]], +) -> list[PreparedSample]: + prepared = [] + for sample in samples: + session_number_by_id = { + session_id: index + 1 + for index, session_id in enumerate(sample.get("haystack_session_ids", [])) + } + records: list[ConversationRecord] = [] + session_summaries: dict[int, str] = {} + session_events: dict[int, SessionEvents] = {} + observations: dict[int, list[ObservationNote]] = {} + support_by_path: dict[str, dict[str, list[str]]] = {} + turn_ids_by_speaker: dict[str, list[str]] = {} + session_ids_by_speaker: dict[str, list[str]] = {} + for session_index, session_id in enumerate(sample.get("haystack_session_ids", [])): + session_number = session_number_by_id[session_id] + session_token = f"D{session_number}" + date = _get_index(sample.get("haystack_dates", []), session_index) + user_ordinal = 0 + user_texts = [] + for turn_index, turn in enumerate( + _get_index(sample.get("haystack_sessions", []), session_index) or [] + ): + role = str(turn.get("role", "")) + content = str(turn.get("content", "")).strip() + if role.lower() == "user": + user_ordinal += 1 + evidence_id = f"D{session_number}:{user_ordinal}" + turn_id = f"{session_id}_{turn_index + 1}" + user_texts.append(content) + observations.setdefault(session_number, []).append( + ObservationNote(speaker="User", evidence_id=evidence_id, text=content) + ) + turn_ids_by_speaker.setdefault("User", []).append(turn_id) + session_ids_by_speaker.setdefault("User", []).append(session_id) + # Keep the conversion support map aligned with the Rust + # LongMemEval adapter. The retained evaluator extracts + # evidence IDs from selected turn/event files, but those + # files are not registered as support pages by the Rust + # adapter's fallback path. Only source/topic/entity and + # fallback-observation pages participate in ranked-page + # metrics; otherwise Python would count an extra page + # type and shift every later rank. + else: + evidence_id = f"A{session_number}_{turn_index + 1}" + records.append( + ConversationRecord( + dia_id=evidence_id, + session_id=session_token, + speaker=_normalize_role(role), + text=content, + ) + ) + session_summaries[session_number] = " ".join(user_texts) + session_events[session_number] = SessionEvents( + date=str(date) if date is not None else None, + items_by_speaker={ + "User": [ + f"Evidence {note.evidence_id}: {note.text}" + for note in observations.get(session_number, []) + ] + }, + ) + session_turn_ids = [ + f"{session_id}_{index + 1}" + for index, turn in enumerate( + _get_index(sample.get("haystack_sessions", []), session_index) or [] + ) + if str(turn.get("role", "")).lower() == "user" + ] + session_support = {"turn_ids": session_turn_ids, "session_ids": [session_id]} + support_by_path[f"wiki/sources/session_{session_number}.md"] = session_support + support_by_path[ + f"wiki/topics/session_{session_number}_observations.md" + ] = session_support + support_by_path[f"wiki/topics/session_{session_number}_events.md"] = session_support + for ordinal, note in enumerate(observations.get(session_number, []), start=1): + turn_support = { + "turn_ids": session_turn_ids[ordinal - 1 : ordinal], + "session_ids": [session_id], + } + support_by_path[ + f"wiki/observations/{_evidence_slug(note.evidence_id)}_obs_{ordinal}.md" + ] = turn_support + # Event pages intentionally have no direct support mapping in + # the Rust fallback adapter. Their evidence IDs are still + # extracted by the retained evaluator, but are appended as + # evidence support after ranked file pages during conversion. + + for speaker, turn_ids in turn_ids_by_speaker.items(): + support_by_path[f"wiki/entities/{_entity_slug(speaker)}.md"] = { + "turn_ids": _unique(turn_ids), + "session_ids": _unique(session_ids_by_speaker.get(speaker, [])), + } + sample["_locomo_support_by_path"] = support_by_path + + evidence = [ + sample.get("_locomo_evidence_by_turn_id", {}).get(turn_id) + or _locomo_evidence_id_for_turn_id(turn_id, session_number_by_id) + for turn_id in sample.get("answer_turn_ids", []) + ] + evidence = [item for item in evidence if item] + prepared.append( + PreparedSample( + sample_id=str(sample.get("question_id", "")), + raw_sample=sample, + records=records, + questions=[ + LoCoMoQuestion( + question=str(sample.get("question", "")), + answer=sample.get("answer"), + evidence=evidence, + category=None, + ) + ], + session_datetimes={ + session_number_by_id[session_id]: str(date) + for session_id, date in zip( + sample.get("haystack_session_ids", []), + sample.get("haystack_dates", []), + ) + if session_id in session_number_by_id + }, + session_summaries=session_summaries, + event_summaries=session_events, + observations=observations, + ) + ) + return prepared + + +def convert_retained_output_to_longmemeval( + *, + samples: list[dict[str, Any]], + retained_cases: list[Any], + granularity: str, + workspace_root: str | Path, +) -> dict[str, Any]: + sample_by_id = {str(sample.get("question_id", "")): sample for sample in samples} + cases = [] + for score in retained_cases: + sample = sample_by_id[score.sample_id] + expected_turn_ids = list(sample.get("answer_turn_ids", [])) + expected_session_ids = list(sample.get("answer_session_ids", [])) + ranked_support_pages = _ranked_support_pages(score, sample) + abstention_case = _is_longmemeval_abstention(sample) + evaluated = _is_longmemeval_evaluated( + granularity=granularity, + abstention_case=abstention_case, + expected_turn_ids=expected_turn_ids, + expected_session_ids=expected_session_ids, + ) + metrics = ( + _longmemeval_metrics( + granularity=granularity, + expected_turn_ids=expected_turn_ids, + expected_session_ids=expected_session_ids, + ranked_support_pages=ranked_support_pages, + ) + if evaluated + else {"turn": {}, "session": {}} + ) + ranked_items = [ + { + "corpus_id": _relative_retrieved_path(path, score.knowledge_base_root), + "text": "", + "timestamp": None, + } + for path in score.retrieved_file_paths + ] + cases.append( + LongMemEvalCaseScore( + question_id=score.sample_id, + question_type=str(sample.get("question_type", "")), + question=score.question, + question_date=sample.get("question_date"), + answer=sample.get("answer"), + abstention_case=abstention_case, + evaluated=evaluated, + expected_session_ids=expected_session_ids, + expected_turn_ids=expected_turn_ids, + retrieval_results=LongMemEvalRetrievalResults( + query=score.question, + granularity=granularity, + ranked_items=ranked_items, + ranked_item_count=len(ranked_items), + metrics=metrics, + ), + knowledge_base_root=str(Path(workspace_root) / score.sample_id), + kb_fallback_used=False, + kb_warnings=["locomo_qmdconsensus python migration"], + ) + ) + + evaluated = [case for case in cases if case.evaluated] + summary = LongMemEvalSummary( + dataset_name="longmemeval", + granularity=granularity, + total_cases=len(cases), + evaluated_cases=len(evaluated), + skipped_abstention_cases=sum(1 for case in cases if case.abstention_case), + skipped_no_target_cases=sum( + 1 for case in cases if not case.evaluated and not case.abstention_case + ), + averaged_metrics={ + "turn": _average_metric_bucket( + [case.retrieval_results.metrics["turn"] for case in evaluated] + ), + "session": _average_metric_bucket( + [case.retrieval_results.metrics["session"] for case in evaluated] + ), + }, + ) + return {"summary": summary, "cases": cases} + + +def _is_longmemeval_abstention(sample: dict[str, Any]) -> bool: + return "_abs" in str(sample.get("question_id", "")) + + +def _is_longmemeval_evaluated( + *, + granularity: str, + abstention_case: bool, + expected_turn_ids: list[str], + expected_session_ids: list[str], +) -> bool: + if abstention_case: + return False + if granularity == "session": + return bool(expected_session_ids) + return bool(expected_turn_ids) + + +def _prepare_longmemeval_sample(raw: dict[str, Any]) -> dict[str, Any]: + sample = dict(raw) + session_documents = [] + turn_documents = [] + answer_turn_ids = [] + locomo_evidence_by_turn_id = {} + for session_index, session_id in enumerate(sample.get("haystack_session_ids", [])): + turns = _get_index(sample.get("haystack_sessions", []), session_index) or [] + date = _get_index(sample.get("haystack_dates", []), session_index) + user_ordinal = 0 + session_texts = [] + for raw_turn_index, turn in enumerate(turns, start=1): + if str(turn.get("role", "")).lower() != "user": + continue + user_ordinal += 1 + content = str(turn.get("content", "")).strip() + turn_id = f"{session_id}_{raw_turn_index}" + locomo_evidence_by_turn_id[turn_id] = f"D{session_index + 1}:{user_ordinal}" + session_texts.append(content) + turn_documents.append( + { + "corpus_id": turn_id, + "session_id": session_id, + "turn_index": raw_turn_index, + "text": content, + "timestamp": date, + } + ) + if turn.get("has_answer") is True: + answer_turn_ids.append(turn_id) + session_documents.append( + { + "corpus_id": session_id, + "session_id": session_id, + "turn_index": None, + "text": " ".join(session_texts), + "timestamp": date, + } + ) + sample["session_documents"] = session_documents + sample["turn_documents"] = turn_documents + sample["answer_turn_ids"] = answer_turn_ids + sample["_locomo_evidence_by_turn_id"] = locomo_evidence_by_turn_id + sample["has_user_answer_turn"] = bool(answer_turn_ids) + return sample + + +def _locomo_evidence_id_for_turn_id( + turn_id: str, + session_number_by_id: dict[str, int], +) -> str | None: + session_id, _, ordinal = turn_id.rpartition("_") + if session_id not in session_number_by_id or not ordinal.isdigit(): + return None + return f"D{session_number_by_id[session_id]}:{int(ordinal)}" + + +def _longmemeval_metrics( + *, + granularity: str, + expected_turn_ids: list[str], + expected_session_ids: list[str], + ranked_support_pages: list[dict[str, list[str]]], +) -> dict[str, dict[str, float]]: + if granularity == "session": + return { + "turn": {}, + "session": _page_support_metrics( + expected_session_ids, + ranked_support_pages, + "session_ids", + ), + } + return { + "turn": _page_support_metrics(expected_turn_ids, ranked_support_pages, "turn_ids"), + "session": _page_support_metrics( + expected_session_ids, + ranked_support_pages, + "session_ids", + ), + } + + +def _ranked_support_pages(score: Any, sample: dict[str, Any]) -> list[dict[str, list[str]]]: + support_by_path = sample.get("_locomo_support_by_path", {}) + pages = [] + for path in score.retrieved_file_paths: + relative = _relative_retrieved_path(path, score.knowledge_base_root) + support = support_by_path.get(relative) + if support: + pages.append(support) + + turn_id_by_evidence = { + evidence_id: turn_id + for turn_id, evidence_id in sample.get("_locomo_evidence_by_turn_id", {}).items() + } + session_ids = sample.get("haystack_session_ids", []) + for evidence_id in score.retrieved_evidence: + turn_id = turn_id_by_evidence.get(evidence_id) + session_number = evidence_id.removeprefix("D").split(":", maxsplit=1)[0] + session_id = ( + session_ids[int(session_number) - 1] + if session_number.isdigit() and 0 < int(session_number) <= len(session_ids) + else None + ) + if turn_id or session_id: + pages.append( + { + "turn_ids": [turn_id] if turn_id else [], + "session_ids": [session_id] if session_id else [], + } + ) + return pages + + +def _relative_retrieved_path(path: str, knowledge_root: str) -> str: + normalized = path.replace("\\", "/") + root = knowledge_root.replace("\\", "/").rstrip("/") + prefix = f"{root}/" + return normalized[len(prefix) :] if normalized.startswith(prefix) else normalized.lstrip("/") + + +def _page_support_metrics( + expected: list[str], + pages: list[dict[str, list[str]]], + support_key: str, +) -> dict[str, float]: + expected_unique = _unique(expected) + metrics = {} + for k in (1, 3, 5, 10, 30, 50): + first_rank: dict[str, int] = {} + for rank, page in enumerate(pages[:k], start=1): + for item in page[support_key]: + first_rank.setdefault(item, rank) + hits = [item for item in expected_unique if item in first_rank] + best_rank = min((first_rank[item] for item in hits), default=None) + metrics[f"recall_any@{k}"] = 1.0 if hits else 0.0 + metrics[f"recall_all@{k}"] = ( + 1.0 if expected_unique and len(hits) == len(expected_unique) else 0.0 + ) + # The Rust comparison table calls this ``recall_avg``: average the + # fraction of expected evidence items recovered by each case. It is + # distinct from the binary any/all indicators above. + metrics[f"recall_avg@{k}"] = ( + len(hits) / len(expected_unique) if expected_unique else 0.0 + ) + metrics[f"ndcg_any@{k}"] = ( + _round4(_discounted_gain_for_rank(best_rank)) if best_rank is not None else 0.0 + ) + return metrics + + +def _discounted_gain_for_rank(rank: int) -> float: + if rank <= 1: + return 1.0 + import math + + return 1.0 / math.log2(rank + 1) + + +def _average_metric_bucket(rows: list[dict[str, float]]) -> dict[str, float]: + if not rows: + return {} + keys = sorted({key for row in rows for key in row}) + return { + key: _round4(sum(row.get(key, 0.0) for row in rows) / len(rows)) + for key in keys + } + + +def _normalize_role(role: str) -> str: + return "User" if role.lower() == "user" else "Assistant" + + +def _get_index(values: list[Any], index: int) -> Any: + return values[index] if 0 <= index < len(values) else None + + +def _unique(values: list[str]) -> list[str]: + seen = set() + result = [] + for value in values: + if value not in seen: + result.append(value) + seen.add(value) + return result + + +def _evidence_slug(evidence_id: str) -> str: + return evidence_id.replace(":", "_") + + +def _entity_slug(value: str) -> str: + slug = "".join(ch if ch.isalnum() or ch in "_.-" else "_" for ch in value) + return slug.strip("_") or "entity" + + +def _round4(value: float) -> float: + return round(value + 0.0, 4) + + +def _jsonable(value: Any) -> Any: + if hasattr(value, "__dataclass_fields__"): + return asdict(value) + if isinstance(value, list): + return [_jsonable(item) for item in value] + if isinstance(value, dict): + return {key: _jsonable(item) for key, item in value.items()} + return value diff --git a/evaluation/wikimem/qmd_consensus.py b/evaluation/wikimem/qmd_consensus.py new file mode 100644 index 00000000..39fcc342 --- /dev/null +++ b/evaluation/wikimem/qmd_consensus.py @@ -0,0 +1,1282 @@ +"""Deterministic qmd_consensus retrieval helpers ported from wikimem.""" + +from __future__ import annotations + +import posixpath +import re +from dataclasses import dataclass, replace +from functools import lru_cache +from pathlib import PurePosixPath + + +_QUERY_STOPWORDS = { + "a", + "an", + "and", + "are", + "at", + "be", + "by", + "did", + "do", + "for", + "how", + "in", + "is", + "it", + "of", + "on", + "or", + "the", + "to", + "was", + "what", + "when", + "where", + "which", + "who", + "why", + "would", +} + +_QMD_FOCUS_STOPWORDS = { + "what", + "when", + "where", + "which", + "would", + "could", + "pursue", + "should", + "might", + "likely", + "still", + "there", + "their", + "about", + "after", + "before", + "around", +} + +_EXPANSION_STOPWORDS = { + "about", + "after", + "before", + "could", + "might", + "still", + "their", + "there", + "these", + "those", + "would", +} + + +@dataclass(frozen=True) +class RetrievedMemoryFile: + filename: str + file_path: str + mtime_ms: int + content: str + description: str | None = None + memory_type: str | None = None + scope: str = "auto" + + +@dataclass(frozen=True) +class QuestionProfile: + question: str + query_tokens: list[str] + query_fuzzy_tokens: list[str] + expansion_tokens: list[str] + expansion_fuzzy_tokens: list[str] + expansion_phrases: list[str] + named_entities: list[str] + temporal: bool + location: bool + identity: bool + hypothetical: bool + relational: bool + aggregate: bool + + +@dataclass(frozen=True) +class QueryAugmentation: + tokens: list[str] + fuzzy_tokens: list[str] + phrases: list[str] + + +@dataclass(frozen=True) +class CachedFileLineLexicalFeatures: + lower_text: str + tokens: list[str] + fuzzy_tokens: list[str] + normalized_tokens: list[str] + normalized_token_set: set[str] + normalized_prefix3_set: set[str] + token_set: set[str] + fuzzy_token_set: set[str] + + +@dataclass(frozen=True) +class CachedFileLexicalFeatures: + lower_content: str + content_tokens: list[str] + content_fuzzy_tokens: list[str] + path_tokens: list[str] + content_token_set: set[str] + content_fuzzy_token_set: set[str] + path_token_set: set[str] + bridge_target_paths: list[str] + has_session_marker: bool + has_evidence_marker: bool + line_features: list[CachedFileLineLexicalFeatures] + + +@dataclass(frozen=True) +class QmdConsensusFileMetrics: + query_hits: int + best_line_score: float + support_density: float + + +@dataclass(frozen=True) +class CandidateProposal: + file_path: str + query_hits: int + seed_boost: float + + +@dataclass(frozen=True) +class RerankProposal: + file_path: str + query_hits: int + seed_boost: float + + +@dataclass(frozen=True) +class CandidateFile: + file: RetrievedMemoryFile + query_hits: int + seed_boost: float + + +def build_question_profile(question: str, entity_names: list[str]) -> QuestionProfile: + lower = question.lower() + named_entities = [ + name for name in entity_names if _word_boundary_contains(lower, name.lower()) + ] + return QuestionProfile( + question=question, + query_tokens=tokenize_query(question), + query_fuzzy_tokens=tokenize_fuzzy_query(question), + expansion_tokens=[], + expansion_fuzzy_tokens=[], + expansion_phrases=[], + named_entities=named_entities, + temporal=_contains_any_phrase( + lower, + ["when", "how long", "what year", "what month", "what day"], + ), + location=_contains_any_phrase(lower, ["where", "which park", "which place"]), + identity=_contains_any_phrase( + lower, + ["identity", "who is", "relationship status", "member of", "ally"], + ), + hypothetical=_contains_any_phrase( + lower, + ["would", "likely", "plan", "planning", "pursue"], + ), + relational=_contains_any_phrase( + lower, + ["relationship", "friends", "family", "mentor", "support"], + ), + aggregate=_contains_any_phrase( + lower, + [ + "what activities", + "what events", + "what books", + "what artists", + "what subjects", + "what items", + "what are some", + "in what ways", + "what has", + "what kind of art", + "what musical", + ], + ), + ) + + +def qmd_consensus_is_conservative(profile: QuestionProfile) -> bool: + return ( + profile.temporal + and not profile.location + and not profile.identity + and not profile.hypothetical + and not profile.relational + and not profile.aggregate + ) + + +@lru_cache(maxsize=8192) +def build_cached_file_lexical_features(file: RetrievedMemoryFile) -> CachedFileLexicalFeatures: + lower_content = file.content.lower() + normalized_path = normalize_memory_path(file.file_path) + line_features = [ + _build_line_features(line.strip()) + for line in file.content.splitlines() + if line.strip() + ] + content_tokens = tokenize_query(file.content) + content_fuzzy_tokens = tokenize_fuzzy_query(file.content) + path_tokens = tokenize_query(normalized_path) + return CachedFileLexicalFeatures( + lower_content=lower_content, + content_tokens=content_tokens, + content_fuzzy_tokens=content_fuzzy_tokens, + path_tokens=path_tokens, + content_token_set=set(content_tokens), + content_fuzzy_token_set=set(content_fuzzy_tokens), + path_token_set=set(path_tokens), + bridge_target_paths=_collect_bridge_target_paths(file.file_path, file.content), + has_session_marker="- Session: D" in file.content, + has_evidence_marker="- Evidence: D" in file.content, + line_features=line_features, + ) + + +def build_qmd_consensus_augmentation( + question: str, + profile: QuestionProfile, + root_files: list[RetrievedMemoryFile], +) -> QueryAugmentation: + if qmd_consensus_is_conservative(profile): + return QueryAugmentation(tokens=[], fuzzy_tokens=[], phrases=[]) + + significant = significant_phrases(question) + ngrams = keyword_ngrams(question) + focused = _qmd_focus_profile(profile) + token_scores: dict[str, float] = {} + token_support: dict[str, set[str]] = {} + phrase_scores: dict[str, float] = {} + phrase_support: dict[str, set[str]] = {} + + for file in _preferred_seed_files(root_files): + normalized_path = normalize_memory_path(file.file_path) + path_weight = _qmd_seed_path_weight(normalized_path) + for line, line_score in _best_seed_lines(file.content, focused, question): + if not _qmd_seed_line_has_anchor_overlap(line, focused): + continue + weighted_score = line_score * path_weight + for token in tokenize_query(line): + if _should_keep_expansion_token( + token, + focused.query_tokens, + focused.named_entities, + significant, + ngrams, + ): + token_scores[token] = token_scores.get(token, 0.0) + weighted_score + token_support.setdefault(token, set()).add(normalized_path) + for phrase in _expansion_phrases_from_line(line): + if phrase in significant: + continue + phrase_scores[phrase] = phrase_scores.get(phrase, 0.0) + weighted_score + phrase_support.setdefault(phrase, set()).add(normalized_path) + + tokens = [ + token + for token, _ in sorted( + ( + (token, score) + for token, score in token_scores.items() + if len(token_support.get(token, set())) >= 2 or score >= 7.0 + ), + key=lambda item: ( + -len(token_support.get(item[0], set())), + -item[1], + item[0], + ), + )[:6] + ] + phrases = [ + phrase + for phrase, _ in sorted( + ( + (phrase, score) + for phrase, score in phrase_scores.items() + if (len(phrase_support.get(phrase, set())) >= 1 and score >= 6.0) + or len(phrase_support.get(phrase, set())) >= 2 + ), + key=lambda item: ( + -len(phrase_support.get(item[0], set())), + -item[1], + item[0], + ), + )[:4] + ] + return QueryAugmentation( + tokens=tokens, + fuzzy_tokens=tokenize_fuzzy_query(" ".join(tokens)), + phrases=phrases, + ) + + +def apply_query_augmentation( + base_profile: QuestionProfile, + augmentation: QueryAugmentation, +) -> QuestionProfile: + return replace( + base_profile, + expansion_tokens=_merge_unique(base_profile.expansion_tokens, augmentation.tokens), + expansion_fuzzy_tokens=_merge_unique( + base_profile.expansion_fuzzy_tokens, + augmentation.fuzzy_tokens, + ), + expansion_phrases=_merge_unique(base_profile.expansion_phrases, augmentation.phrases), + ) + + +def build_qmd_consensus_candidate_proposals( + question: str, + profile: QuestionProfile, + files: list[RetrievedMemoryFile], +) -> list[CandidateProposal]: + if qmd_consensus_is_conservative(profile): + return [] + + cached_files, cached_features = _cache_files(files) + focused = _qmd_focus_profile(profile) + significant = significant_phrases(question) + ngrams = keyword_ngrams(question) + named_lower = [name.lower() for name in focused.named_entities] + metrics_cache: dict[str, QmdConsensusFileMetrics] = {} + source_view: list[tuple[str, int, float]] = [] + anchor_view: list[tuple[str, int, float]] = [] + + for normalized_path, file in cached_files.items(): + kind = _qmd_candidate_view_kind(normalized_path) + if kind is None: + continue + features = cached_features[normalized_path] + metrics = _metrics_for( + metrics_cache, + normalized_path, + file, + features, + focused, + significant, + ) + score = _score_cached_file_with_metrics( + features, + normalized_path, + focused, + significant, + ngrams, + named_lower, + metrics, + ) + if kind == "source": + score += metrics.best_line_score * 0.65 + score += metrics.support_density * 0.3 + score += metrics.query_hits * 0.25 + if score >= 5.8: + source_view.append((normalized_path, max(metrics.query_hits, 1), score)) + else: + score += metrics.best_line_score * 0.2 + score += _candidate_path_weight(normalized_path) * 0.4 + score += metrics.query_hits * 0.15 + if score >= 6.2: + anchor_view.append((normalized_path, max(metrics.query_hits, 1), score)) + + _sort_qmd_ranked_items(source_view) + _sort_qmd_ranked_items(anchor_view) + linked_view = _build_qmd_linked_candidate_view( + question, + focused, + cached_files, + cached_features, + source_view, + anchor_view, + ) + return [ + CandidateProposal(file_path=path, query_hits=query_hits, seed_boost=boost) + for path, query_hits, boost in _fuse_ranked_views( + [(2.0, source_view), (1.5, anchor_view), (1.0, linked_view)], + 12, + ) + ] + + +def build_qmd_consensus_rerank_proposals( + question: str, + profile: QuestionProfile, + ranked_candidates: list[CandidateFile], + files: list[RetrievedMemoryFile], +) -> list[RerankProposal]: + if qmd_consensus_is_conservative(profile): + return [] + + cached_files, cached_features = _cache_files(files) + focused = _qmd_focus_profile(profile) + significant = significant_phrases(question) + metrics_cache: dict[str, QmdConsensusFileMetrics] = {} + + source_view = [] + anchor_view = [] + for candidate in ranked_candidates: + normalized_path = normalize_memory_path(candidate.file.file_path) + features = cached_features.get(normalized_path) + if features is None: + continue + metrics = _metrics_for( + metrics_cache, + normalized_path, + candidate.file, + features, + focused, + significant, + ) + item = ( + normalized_path, + max(candidate.query_hits, 1), + candidate.seed_boost + metrics.best_line_score, + ) + if "/wiki/sources/" in normalized_path: + source_view.append(item) + elif _qmd_candidate_view_kind(normalized_path) == "anchor": + anchor_view.append(item) + + _sort_qmd_ranked_items(source_view) + _sort_qmd_ranked_items(anchor_view) + source_view = source_view[:6] + anchor_view = anchor_view[:6] + seed_files = [candidate.file for candidate in ranked_candidates[:6]] + linked_view = [ + (proposal.file_path, max(proposal.query_hits, 1), proposal.seed_boost) + for proposal in _build_bridge_proposals( + question, + focused, + seed_files, + cached_files, + cached_features, + 8, + 4.2, + ) + ] + return [ + RerankProposal(file_path=path, query_hits=query_hits, seed_boost=boost) + for path, query_hits, boost in _fuse_ranked_views( + [(2.0, source_view), (1.5, anchor_view), (1.0, linked_view)], + 12, + ) + ] + + +def build_qmd_consensus_late_bridge_proposals( + question: str, + profile: QuestionProfile, + seed_files: list[RetrievedMemoryFile], + files: list[RetrievedMemoryFile], +) -> list[RerankProposal]: + if qmd_consensus_is_conservative(profile): + return [] + + cached_files, cached_features = _cache_files(files) + focused = _qmd_focus_profile(profile) + return _build_bridge_proposals( + question, + focused, + seed_files, + cached_files, + cached_features, + 8, + 4.2, + ) + + +def score_line(text: str, profile: QuestionProfile, question: str) -> float: + return _score_line_with_phrases(text, profile, significant_phrases(question)) + + +def candidate_query_hits_with_features( + file_features: CachedFileLexicalFeatures, + profile: QuestionProfile, +) -> int: + return _token_overlap_with_set(profile.query_tokens, file_features.content_token_set) + ( + _token_overlap_with_set(profile.query_fuzzy_tokens, file_features.content_fuzzy_token_set) + ) + + +def normalize_memory_path(path: str) -> str: + return path.replace("\\", "/").lower() + + +def tokenize_query(text: str) -> list[str]: + tokens = [] + for token in re.split(r"[^0-9A-Za-z]+", text): + lowered = token.strip().lower() + if len(lowered) <= 1 or lowered in _QUERY_STOPWORDS: + continue + tokens.append(lowered) + return tokens + + +def tokenize_fuzzy_query(text: str) -> list[str]: + features = [] + seen = set() + for token in tokenize_query(text): + for feature in _fuzzy_token_features(token): + if len(feature) > 2 and feature not in seen: + features.append(feature) + seen.add(feature) + return features + + +def significant_phrases(question: str) -> list[str]: + return [phrase for phrase in keyword_ngrams(question) if len(phrase.split()) >= 2] + + +def keyword_ngrams(question: str) -> list[str]: + tokens = tokenize_query(question.lower()) + seen = set() + ngrams = [] + for size in (2, 3): + for index in range(0, max(0, len(tokens) - size + 1)): + phrase = " ".join(tokens[index : index + size]) + if phrase not in seen: + ngrams.append(phrase) + seen.add(phrase) + return ngrams + + +def _build_line_features(line: str) -> CachedFileLineLexicalFeatures: + tokens = tokenize_query(line) + fuzzy_tokens = tokenize_fuzzy_query(line) + normalized_tokens = [_strip_common_token_suffixes(token) for token in tokens] + return CachedFileLineLexicalFeatures( + lower_text=line.lower(), + tokens=tokens, + fuzzy_tokens=fuzzy_tokens, + normalized_tokens=normalized_tokens, + normalized_token_set=set(normalized_tokens), + normalized_prefix3_set={token[:3] for token in normalized_tokens if len(token) >= 3}, + token_set=set(tokens), + fuzzy_token_set=set(fuzzy_tokens), + ) + + +def _cache_files( + files: list[RetrievedMemoryFile], +) -> tuple[dict[str, RetrievedMemoryFile], dict[str, CachedFileLexicalFeatures]]: + cached_files = {normalize_memory_path(file.file_path): file for file in files} + cached_features = { + normalize_memory_path(file.file_path): build_cached_file_lexical_features(file) + for file in files + } + return cached_files, cached_features + + +def _qmd_focus_profile(profile: QuestionProfile) -> QuestionProfile: + focused_tokens = [ + token + for token in profile.query_tokens + if _qmd_keeps_focus_token(token, profile.named_entities) + ] + if not focused_tokens: + return profile + focused_expansion = [ + token + for token in profile.expansion_tokens + if _qmd_keeps_focus_token(token, profile.named_entities) + ] + return replace( + profile, + query_tokens=focused_tokens, + query_fuzzy_tokens=tokenize_fuzzy_query(" ".join(focused_tokens)), + expansion_tokens=focused_expansion, + expansion_fuzzy_tokens=tokenize_fuzzy_query(" ".join(focused_expansion)), + ) + + +def _qmd_keeps_focus_token(token: str, named_entities: list[str]) -> bool: + if len(token) < 4: + return False + if any(entity.lower() == token.lower() for entity in named_entities): + return False + return token not in _QMD_FOCUS_STOPWORDS + + +def _metrics_for( + metrics_cache: dict[str, QmdConsensusFileMetrics], + normalized_path: str, + _file: RetrievedMemoryFile, + features: CachedFileLexicalFeatures, + profile: QuestionProfile, + significant: list[str], +) -> QmdConsensusFileMetrics: + if normalized_path not in metrics_cache: + line_scores = [ + _score_line_features_with_phrases(line, profile, significant) + for line in features.line_features + ] + best_line_score, support_density = _summarize_qmd_line_scores(line_scores) + metrics_cache[normalized_path] = QmdConsensusFileMetrics( + query_hits=candidate_query_hits_with_features(features, profile), + best_line_score=best_line_score, + support_density=support_density, + ) + return metrics_cache[normalized_path] + + +def _score_cached_file_with_metrics( + features: CachedFileLexicalFeatures, + normalized_path: str, + profile: QuestionProfile, + significant: list[str], + ngrams: list[str], + named_entities_lower: list[str], + metrics: QmdConsensusFileMetrics, +) -> float: + score = metrics.best_line_score + score += _candidate_path_weight(normalized_path) * 1.8 + score += _token_overlap_with_set(profile.query_tokens, features.path_token_set) * 1.8 + score += sum(2.3 for phrase in significant if phrase in features.lower_content) + score += sum(1.2 for phrase in ngrams if phrase in features.lower_content) + score += sum(1.5 for entity in named_entities_lower if entity in features.lower_content) + if features.has_session_marker: + score += 1.0 + if features.has_evidence_marker: + score += 4.0 + return score + + +def _build_qmd_linked_candidate_view( + question: str, + profile: QuestionProfile, + cached_files: dict[str, RetrievedMemoryFile], + cached_features: dict[str, CachedFileLexicalFeatures], + source_view: list[tuple[str, int, float]], + anchor_view: list[tuple[str, int, float]], +) -> list[tuple[str, int, float]]: + significant = significant_phrases(question) + metrics_cache: dict[str, QmdConsensusFileMetrics] = {} + direct_atomic: list[tuple[str, int, float]] = [] + for normalized_path, file in cached_files.items(): + if ( + "/wiki/turns/" not in normalized_path + and "/wiki/memories/" not in normalized_path + and "/wiki/memory/" not in normalized_path + ): + continue + features = cached_features[normalized_path] + metrics = _metrics_for( + metrics_cache, + normalized_path, + file, + features, + profile, + significant, + ) + direct_score = ( + metrics.best_line_score * 1.15 + + metrics.support_density * 0.2 + + metrics.query_hits * 0.3 + ) + if direct_score >= 5.0: + _push_qmd_ranked_top_k( + direct_atomic, + (normalized_path, max(metrics.query_hits, 1), direct_score), + 6, + ) + + seed_specs = [ + (path, hits, score, True) for path, hits, score in source_view[:8] + ] + [(path, hits, score, False) for path, hits, score in anchor_view[:8]] + seed_specs.sort(key=lambda item: (-item[2], -item[1], item[0])) + proposals: dict[str, tuple[int, float]] = {} + + for seed_path, seed_query_hits, seed_score, seed_is_source in seed_specs: + if seed_score < 6.0: + continue + seed = cached_files.get(seed_path) + seed_features = cached_features.get(seed_path) + if seed is None or seed_features is None: + continue + seed_session = _infer_session_number(seed.file_path) + same_session: list[tuple[str, int, float]] = [] + cross_session: list[tuple[str, int, float]] = [] + for target in _extract_candidate_targets_with_features(seed_features, cached_files): + normalized_target = normalize_memory_path(target.file_path) + target_features = cached_features.get(normalized_target) + if target_features is None: + continue + metrics = _metrics_for( + metrics_cache, + normalized_target, + target, + target_features, + profile, + significant, + ) + target_score = metrics.best_line_score + if target_score < 4.4: + continue + query_hits = max(metrics.query_hits, 1) + proposal_score = ( + target_score * 0.9 + + min(seed_score, 10.0) * 0.22 + + _candidate_path_weight(normalized_target) * 0.3 + + seed_query_hits * 0.08 + ) + same_session_target = ( + seed_is_source + and seed_session is not None + and seed_session == _infer_session_number(target.file_path) + ) + if same_session_target: + if not _qmd_allows_same_session_source_target( + normalized_target, + target_score, + query_hits, + ): + continue + same_session.append((normalized_target, query_hits, proposal_score + 0.15)) + else: + cross_session.append((normalized_target, query_hits, proposal_score)) + + _sort_qmd_ranked_items(same_session) + _sort_qmd_ranked_items(cross_session) + target_limit = 2 if seed_is_source else 3 + for file_path, query_hits, seed_boost in same_session[:2] + cross_session[:target_limit]: + existing = proposals.get(file_path) + if existing is None: + proposals[file_path] = (query_hits, seed_boost) + else: + proposals[file_path] = (max(existing[0], query_hits), max(existing[1], seed_boost)) + + linked_view = [ + (file_path, query_hits, seed_boost) + for file_path, (query_hits, seed_boost) in proposals.items() + ] + _sort_qmd_ranked_items(linked_view) + linked_view = linked_view[:8] + return _fuse_ranked_views([(1.2, direct_atomic), (1.0, linked_view)], 8) + + +def _build_bridge_proposals( + question: str, + profile: QuestionProfile, + seed_files: list[RetrievedMemoryFile], + cached_files: dict[str, RetrievedMemoryFile], + cached_features: dict[str, CachedFileLexicalFeatures], + max_candidates: int, + min_target_score: float, +) -> list[RerankProposal]: + significant = significant_phrases(question) + proposals: dict[str, RerankProposal] = {} + for seed in _preferred_seed_files(seed_files): + normalized_seed = normalize_memory_path(seed.file_path) + seed_weight = _candidate_path_weight(normalized_seed) + seed_session = _infer_session_number(seed.file_path) + seed_is_source = "/wiki/sources/" in normalized_seed + seed_features = cached_features.get(normalized_seed) + targets = ( + _extract_candidate_targets_with_features(seed_features, cached_files) + if seed_features is not None + else [] + ) + for target in targets: + key = normalize_memory_path(target.file_path) + if ( + seed_is_source + and seed_session is not None + and seed_session == _infer_session_number(target.file_path) + ): + continue + target_features = cached_features.get(key) + if target_features is None: + continue + score = _score_file_lines_with_features(target_features, profile, significant) + if score < min_target_score: + continue + query_hits = max(candidate_query_hits_with_features(target_features, profile), 1) + proposal = RerankProposal( + file_path=key, + query_hits=query_hits, + seed_boost=score * 0.8 + seed_weight * 1.4, + ) + existing = proposals.get(key) + if existing is None or (proposal.seed_boost, proposal.query_hits) > ( + existing.seed_boost, + existing.query_hits, + ): + proposals[key] = proposal + + return sorted( + proposals.values(), + key=lambda proposal: (-proposal.seed_boost, -proposal.query_hits, proposal.file_path), + )[:max_candidates] + + +def _score_file_lines_with_features( + features: CachedFileLexicalFeatures, + profile: QuestionProfile, + significant: list[str], +) -> float: + if not features.line_features: + return 0.0 + return max( + _score_line_features_with_phrases(line, profile, significant) + for line in features.line_features + ) + + +def _score_line_with_phrases( + text: str, + profile: QuestionProfile, + significant: list[str], +) -> float: + return _score_line_features_with_phrases(_build_line_features(text), profile, significant) + + +def _score_line_features_with_phrases( + line: CachedFileLineLexicalFeatures, + profile: QuestionProfile, + significant: list[str], +) -> float: + query_normalized = [_strip_common_token_suffixes(token) for token in profile.query_tokens] + query_overlap = _token_overlap_with_set(profile.query_tokens, line.token_set) + fuzzy_overlap = _token_overlap_with_set(profile.query_fuzzy_tokens, line.fuzzy_token_set) + expansion_overlap = _token_overlap_with_set(profile.expansion_tokens, line.token_set) + expansion_fuzzy = _token_overlap_with_set(profile.expansion_fuzzy_tokens, line.fuzzy_token_set) + soft_possible = _query_soft_overlap_possible(query_normalized, line) + if not (query_overlap or fuzzy_overlap or expansion_overlap or expansion_fuzzy or soft_possible): + return 0.0 + score = ( + query_overlap * 2.0 + + fuzzy_overlap * 1.3 + + expansion_overlap * 1.2 + + expansion_fuzzy * 0.7 + ) + score += sum(2.4 for phrase in significant if len(phrase) >= 6 and phrase in line.lower_text) + score += sum( + 1.6 + for phrase in profile.expansion_phrases + if len(phrase) >= 6 and phrase in line.lower_text + ) + score += sum( + 1.4 for entity in profile.named_entities if entity.lower() in line.lower_text + ) + if soft_possible: + score += _soft_token_overlap_with_line_fastpath( + profile.query_tokens, + query_normalized, + line, + ) * 0.8 + return score + + +def _summarize_qmd_line_scores(line_scores: list[float]) -> tuple[float, float]: + best = 0.0 + top_scores = [0.0, 0.0, 0.0] + for score in line_scores: + best = max(best, score) + if score < 2.4: + continue + candidate = score + for index, slot in enumerate(top_scores): + if candidate > slot: + top_scores[index], candidate = candidate, slot + density = top_scores[0] + top_scores[1] * 0.55 + top_scores[2] * 0.35 + return best, density + + +def _fuse_ranked_views( + views: list[tuple[float, list[tuple[str, int, float]]]], + limit: int, +) -> list[tuple[str, int, float]]: + fused: dict[str, tuple[int, float, int, float]] = {} + for weight, view in views: + for rank, (file_path, query_hits, local_score) in enumerate(view, start=1): + key = normalize_memory_path(file_path) + existing = fused.get(key, (0, 0.0, 10**9, 0.0)) + fused[key] = ( + max(existing[0], query_hits), + existing[1] + weight / (60.0 + rank), + min(existing[2], rank), + max(existing[3], local_score), + ) + + ranked = [] + for file_path, (query_hits, rrf_score, top_rank, local_score) in fused.items(): + if top_rank == 1: + top_rank_bonus = 0.05 + elif top_rank <= 3: + top_rank_bonus = 0.02 + else: + top_rank_bonus = 0.0 + boost = rrf_score * 12.0 + top_rank_bonus * 8.0 + min(local_score, 12.0) * 0.25 + ranked.append((file_path, query_hits, boost)) + _sort_qmd_ranked_items(ranked) + return ranked[:limit] + + +def _push_qmd_ranked_top_k( + ranked: list[tuple[str, int, float]], + candidate: tuple[str, int, float], + limit: int, +) -> None: + ranked.append(candidate) + _sort_qmd_ranked_items(ranked) + del ranked[limit:] + + +def _sort_qmd_ranked_items(ranked: list[tuple[str, int, float]]) -> None: + ranked.sort(key=lambda item: (-item[2], -item[1], item[0])) + + +def _qmd_candidate_view_kind(path: str) -> str | None: + if "/wiki/sources/" in path: + return "source" + if ( + "/wiki/entities/" in path + or "/wiki/events/" in path + or "/wiki/observations/" in path + or "/wiki/memories/" in path + or "/wiki/memory/" in path + ): + return "anchor" + return None + + +def _qmd_allows_same_session_source_target( + normalized_target: str, + target_score: float, + query_hits: int, +) -> bool: + return ( + "/wiki/turns/" in normalized_target + or "/wiki/observations/" in normalized_target + or "/wiki/memory/" in normalized_target + ) and ((query_hits >= 2 and target_score >= 5.2) or target_score >= 6.4) + + +def _preferred_seed_files(files: list[RetrievedMemoryFile]) -> list[RetrievedMemoryFile]: + preferred = [ + file + for file in files + if any( + part in normalize_memory_path(file.file_path) + for part in ( + "/wiki/entities/", + "/wiki/topics/", + "/wiki/sources/", + "/wiki/memories/", + "/wiki/memory/", + ) + ) + ] + return (preferred or files)[:6] + + +def _extract_candidate_targets_with_features( + features: CachedFileLexicalFeatures, + cached_files: dict[str, RetrievedMemoryFile], +) -> list[RetrievedMemoryFile]: + return [ + cached_files[path] + for path in features.bridge_target_paths + if path in cached_files + ] + + +def _collect_bridge_target_paths(file_path: str, content: str) -> list[str]: + if "](" not in content: + return [] + targets = [] + seen = set() + for target in re.findall(r"\]\(([^)]+)\)", content): + resolved = _resolve_relative_target(file_path, target.strip()) + if resolved is None: + continue + key = normalize_memory_path(resolved) + if _is_bridge_target_path(key) and key not in seen: + targets.append(key) + seen.add(key) + return targets + + +def _resolve_relative_target(base_path: str, target: str) -> str | None: + if target.startswith(("http://", "https://")): + return None + parent = str(PurePosixPath(base_path.replace("\\", "/")).parent) + return posixpath.normpath(posixpath.join(parent, target)) + + +def _is_bridge_target_path(path: str) -> bool: + normalized = normalize_memory_path(path) + return any( + part in normalized + for part in ( + "/wiki/observations/", + "/wiki/turns/", + "/wiki/events/", + "/wiki/memory/", + ) + ) + + +def _infer_session_number(path: str) -> int | None: + normalized = normalize_memory_path(path) + for marker in ("/wiki/sources/session_", "/wiki/events/session_"): + if marker in normalized: + tail = normalized.split(marker, maxsplit=1)[1] + digits = "".join(ch for ch in tail if ch.isdigit()) + return int(digits) if digits else None + stem = PurePosixPath(normalized).stem + if not stem.startswith("d"): + return None + digits = "" + for ch in stem[1:]: + if not ch.isdigit(): + break + digits += ch + return int(digits) if digits else None + + +def _candidate_path_weight(path: str) -> float: + if "/wiki/observations/" in path: + return 3.4 + if "/wiki/turns/" in path: + return 3.0 + if "/wiki/events/" in path: + return 2.7 + if "/wiki/topics/" in path: + return 2.4 + if "/wiki/entities/" in path: + return 2.2 + if "/wiki/memories/" in path: + return 2.1 + if "/wiki/memory/" in path: + return 2.5 + if "/wiki/sources/" in path: + return 2.0 + return 0.5 + + +def _qmd_seed_path_weight(path: str) -> float: + if "/wiki/sources/" in path: + return 1.35 + if "/wiki/entities/" in path: + return 1.2 + if "/wiki/memories/" in path: + return 1.15 + if "/wiki/memory/" in path: + return 1.25 + if "/wiki/topics/" in path: + return 1.1 + return 1.0 + + +def _best_seed_lines( + text: str, + profile: QuestionProfile, + question: str, +) -> list[tuple[str, float]]: + lines = [] + for raw in text.splitlines(): + trimmed = raw.strip() + if not trimmed or trimmed.startswith("#"): + continue + score = score_line(trimmed, profile, question) + if score > 0.0: + lines.append((trimmed, score)) + lines.sort(key=lambda item: (-item[1], item[0])) + return lines[:4] + + +def _should_keep_expansion_token( + token: str, + query_tokens: list[str], + named_entities: list[str], + significant: list[str], + ngrams: list[str], +) -> bool: + if len(token) < 4 or token in query_tokens: + return False + if any(entity.lower() == token.lower() for entity in named_entities): + return False + stripped = _strip_common_token_suffixes(token) + blocked_phrases = significant + ngrams + return ( + token not in _EXPANSION_STOPWORDS + and not any(phrase == token or phrase == stripped for phrase in blocked_phrases) + ) + + +def _expansion_phrases_from_line(line: str) -> list[str]: + tokens = tokenize_query(line) + return [ + f"{left} {right}" + for left, right in zip(tokens, tokens[1:]) + if len(f"{left} {right}") >= 9 + ] + + +def _qmd_seed_line_has_anchor_overlap(line: str, profile: QuestionProfile) -> bool: + line_tokens = tokenize_query(line) + line_fuzzy = tokenize_fuzzy_query(line) + return ( + _token_overlap(profile.query_tokens, line_tokens) > 0 + or _token_overlap(profile.query_fuzzy_tokens, line_fuzzy) > 0 + or _token_overlap(profile.expansion_tokens, line_tokens) > 0 + or _token_overlap(profile.expansion_fuzzy_tokens, line_fuzzy) > 0 + or any( + len(phrase) >= 6 and phrase in line.lower() + for phrase in profile.expansion_phrases + ) + ) + + +def _merge_unique(existing: list[str], additions: list[str]) -> list[str]: + seen = set(existing) + merged = list(existing) + for item in additions: + if item not in seen: + merged.append(item) + seen.add(item) + return merged + + +def _token_overlap(left: list[str], right: list[str]) -> int: + right_set = set(right) + return sum(1 for token in left if token in right_set) + + +def _token_overlap_with_set(left: list[str], right: set[str]) -> int: + return sum(1 for token in left if token in right) + + +def _query_soft_overlap_possible( + query_normalized: list[str], + line: CachedFileLineLexicalFeatures, +) -> bool: + for token in query_normalized: + if token in line.normalized_token_set: + return True + if len(token) >= 3 and token[:3] in line.normalized_prefix3_set: + return True + return False + + +def _soft_token_overlap_with_line_fastpath( + left: list[str], + left_normalized: list[str], + line: CachedFileLineLexicalFeatures, +) -> float: + total = 0.0 + for left_token, left_norm in zip(left, left_normalized): + if left_token in line.token_set: + best = 1.0 + elif len(left_norm) <= 3 and left_norm in line.normalized_token_set: + best = 0.92 + elif not _qmd_line_may_have_soft_match_candidate(left_norm, line): + best = 0.0 + else: + best = max( + ( + _soft_token_similarity(left_token, left_norm, right_token, right_norm) + for right_token, right_norm in zip(line.tokens, line.normalized_tokens) + ), + default=0.0, + ) + if best >= 0.72: + total += best + return total + + +def _qmd_line_may_have_soft_match_candidate( + left_norm: str, + line: CachedFileLineLexicalFeatures, +) -> bool: + if not left_norm: + return False + if len(left_norm) <= 3: + return left_norm in line.normalized_token_set + return left_norm[:3] in line.normalized_prefix3_set + + +def _soft_token_similarity( + left: str, + left_norm: str, + right: str, + right_norm: str, +) -> float: + if left == right: + return 1.0 + if left_norm == right_norm: + return 0.92 + if len(left_norm) >= 4 and len(right_norm) >= 4 and left_norm[:4] == right_norm[:4]: + return 0.82 + prefix = 0 + for left_ch, right_ch in zip(left_norm, right_norm): + if left_ch != right_ch: + break + prefix += 1 + return prefix / max(len(left_norm), len(right_norm)) + + +def _fuzzy_token_features(token: str) -> list[str]: + stripped = _strip_common_token_suffixes(token) + features = set() + for candidate in (token, stripped): + if len(candidate) >= 4: + features.add(candidate[:4]) + if len(candidate) >= 5: + features.add(candidate[:5]) + skeleton = _consonant_skeleton(candidate) + if len(skeleton) >= 4: + features.add(skeleton) + if stripped != token: + features.add(stripped) + return sorted(features) + + +def _strip_common_token_suffixes(token: str) -> str: + suffixes = [ + ("ies", "y"), + ("ions", ""), + ("tion", ""), + ("ing", ""), + ("ment", ""), + ("ness", ""), + ("ity", ""), + ("ed", ""), + ("es", ""), + ("s", ""), + ] + for suffix, replacement in suffixes: + if len(token) > len(suffix) + 2 and token.endswith(suffix): + return token[: -len(suffix)] + replacement + return token + + +def _consonant_skeleton(token: str) -> str: + if not token: + return "" + return token[0] + "".join(ch for ch in token[1:] if ch not in "aeiou") + + +def _contains_any_phrase(text: str, phrases: list[str]) -> bool: + return any(phrase in text for phrase in phrases) + + +def _word_boundary_contains(text: str, needle: str) -> bool: + return re.search(rf"(? WikiMode: + """Validate the workspace materialization mode shared by all evaluators.""" + + if value not in {"text", "multimodal"}: + raise ValueError(f"unsupported wiki_mode: {value!r}; expected 'text' or 'multimodal'") + return value # type: ignore[return-value] + + +@dataclass(frozen=True) +class LoCoMoQuestion: + question: str + answer: Any = None + adversarial_answer: str | None = None + evidence: list[str] = field(default_factory=list) + category: int | None = None + + +@dataclass(frozen=True) +class ConversationRecord: + dia_id: str + session_id: str + speaker: str + text: str + metadata: dict[str, str] = field(default_factory=dict) + + +@dataclass(frozen=True) +class ObservationNote: + speaker: str + evidence_id: str + text: str + + +@dataclass(frozen=True) +class SessionEvents: + date: str | None = None + items_by_speaker: dict[str, list[str]] = field(default_factory=dict) + + +@dataclass(frozen=True) +class PreparedSample: + sample_id: str + raw_sample: dict[str, Any] + records: list[ConversationRecord] + questions: list[LoCoMoQuestion] + session_datetimes: dict[int, str] + session_summaries: dict[int, str] + event_summaries: dict[int, SessionEvents] + observations: dict[int, list[ObservationNote]] + + +@dataclass(frozen=True) +class EvalCase: + case_id: str + sample_id: str + question_index: int + question: str + answer: str + category: int | None + expected_evidence: list[str] + + +@dataclass(frozen=True) +class RetrievalCoverageSummary: + miss_stage: str + root_hit_evidence: list[str] = field(default_factory=list) + candidate_pool_hit_evidence: list[str] = field(default_factory=list) + late_bridge_hit_evidence: list[str] = field(default_factory=list) + final_hit_evidence: list[str] = field(default_factory=list) + root_file_paths: list[str] = field(default_factory=list) + candidate_pool_file_paths: list[str] = field(default_factory=list) + late_bridge_file_paths: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class CaseScore: + case_id: str + sample_id: str + question_index: int + question: str + category: int | None + expected_evidence: list[str] + retrieved_evidence: list[str] + hit_evidence: list[str] + retrieved_record_ids: list[str] + retrieved_file_paths: list[str] + retrieved_entrypoint_paths: list[str] + knowledge_base_root: str + retrieved_record_count: int + retrieved_file_count: int + retrieved_entrypoint_count: int + evidence_precision: float + evidence_recall: float + full_evidence_hit: bool + retrieval_coverage: RetrievalCoverageSummary | None = None + + +@dataclass(frozen=True) +class EvalSummary: + dataset_name: str + total_cases: int + cases_with_evidence: int + evidence_precision_macro: float + evidence_precision_micro: float + evidence_recall_macro: float + evidence_recall_micro: float + full_evidence_hit_rate: float + + +@dataclass(frozen=True) +class CategoryScoreSummary: + category: int + label: str + count: int + evidence_precision_macro: float + evidence_recall_macro: float + full_evidence_hit_rate: float + + +@dataclass(frozen=True) +class StageTimingRecord: + stage: str + calls: int + total_ms: float + avg_ms: float + max_ms: float + + +@dataclass(frozen=True) +class StageProfileArtifact: + created_at_ms: int + total_samples: int + total_cases: int + stages: list[StageTimingRecord] + + def with_additional_stage(self, stage: str, duration_ms: float) -> StageProfileArtifact: + elapsed_ms = _round4(duration_ms) + stages = list(self.stages) + for index, record in enumerate(stages): + if record.stage != stage: + continue + calls = record.calls + 1 + total_ms = _round4(record.total_ms + elapsed_ms) + stages[index] = StageTimingRecord( + stage=record.stage, + calls=calls, + total_ms=total_ms, + avg_ms=_round4(total_ms / max(calls, 1)), + max_ms=_round4(max(record.max_ms, elapsed_ms)), + ) + break + else: + stages.append( + StageTimingRecord( + stage=stage, + calls=1, + total_ms=elapsed_ms, + avg_ms=elapsed_ms, + max_ms=elapsed_ms, + ) + ) + + return StageProfileArtifact( + created_at_ms=self.created_at_ms, + total_samples=self.total_samples, + total_cases=self.total_cases, + stages=_sort_stage_timings(stages), + ) + + +@dataclass +class _StageTimingTotals: + calls: int = 0 + total_ms: float = 0.0 + max_ms: float = 0.0 + + +class StageProfiler: + """Aggregate retained_eval timing records using the Rust artifact shape.""" + + def __init__(self) -> None: + self._stages: dict[str, _StageTimingTotals] = {} + + def record_elapsed(self, stage: str, started_at: float) -> None: + self.record_duration_ms(stage, (time.perf_counter() - started_at) * 1000.0) + + def record_duration(self, stage: str, duration_seconds: float) -> None: + self.record_duration_ms(stage, duration_seconds * 1000.0) + + def record_duration_ms(self, stage: str, duration_ms: float) -> None: + totals = self._stages.setdefault(stage, _StageTimingTotals()) + totals.calls += 1 + totals.total_ms += duration_ms + totals.max_ms = max(totals.max_ms, duration_ms) + + def snapshot( + self, + total_samples: int, + total_cases: int, + created_at_ms: int | None = None, + ) -> StageProfileArtifact: + stages = [ + StageTimingRecord( + stage=stage, + calls=totals.calls, + total_ms=_round4(totals.total_ms), + avg_ms=_round4(totals.total_ms / max(totals.calls, 1)), + max_ms=_round4(totals.max_ms), + ) + for stage, totals in self._stages.items() + ] + return StageProfileArtifact( + created_at_ms=_unix_now_ms() if created_at_ms is None else created_at_ms, + total_samples=total_samples, + total_cases=total_cases, + stages=_sort_stage_timings(stages), + ) + + +@dataclass(frozen=True) +class EvalOutput: + summary: EvalSummary + cases: list[CaseScore] + stage_profile: StageProfileArtifact + + +@dataclass(frozen=True) +class EvalHarnessConfig: + dataset_name: str + samples: str | None + question_limit: int | None + top_k: int + workspace_root: str + llm_provider: str | None + retrieval_plugins: list[str] = field(default_factory=list) + wiki_mode: WikiMode = "text" + wiki_builder_mode: WikiBuilderMode = "deterministic" + + +@dataclass(frozen=True) +class ProgressUpdate: + completed_cases: int + total_cases: int + sample_index: int + total_samples: int + sample_id: str + question_index: int + sample_question_total: int + + +def prepare_locomo_samples( + payloads: list[dict[str, Any]], + sample_filter: set[str] | None = None, + include_multimodal_context: bool = False, +) -> list[PreparedSample]: + prepared: list[PreparedSample] = [] + for sample in payloads: + if not isinstance(sample, dict) or not isinstance(sample.get("sample_id"), str): + raise ValueError("LoCoMo samples require a string sample_id") + sample_id = sample["sample_id"] + if sample_filter is not None and sample_id not in sample_filter: + continue + conversation = sample.get("conversation", {}) + qa = sample.get("qa", []) + session_summary = sample.get("session_summary", {}) + event_summary = sample.get("event_summary", {}) + observation = sample.get("observation", {}) + if not isinstance(qa, list): + raise ValueError(f"LoCoMo sample {sample_id} qa must be an array") + if not isinstance(session_summary, dict): + raise ValueError(f"LoCoMo sample {sample_id} session_summary must be an object") + if not isinstance(event_summary, dict): + raise ValueError(f"LoCoMo sample {sample_id} event_summary must be an object") + if not isinstance(observation, dict): + raise ValueError(f"LoCoMo sample {sample_id} observation must be an object") + prepared.append( + PreparedSample( + sample_id=sample_id, + raw_sample=sample, + records=extract_conversation_records( + conversation, + include_multimodal_context=include_multimodal_context, + ), + questions=_parse_questions(qa), + session_datetimes=_parse_session_datetimes(conversation), + session_summaries=_parse_session_summaries(session_summary), + event_summaries=_parse_event_summaries(event_summary), + observations=_parse_observations(observation), + ) + ) + return prepared + + +def extract_conversation_records( + conversation: dict[str, Any], + *, + include_multimodal_context: bool = False, +) -> list[ConversationRecord]: + if not isinstance(conversation, dict): + raise ValueError("conversation payload must be an object") + sessions: list[tuple[int, list[dict[str, Any]]]] = [] + for key, value in conversation.items(): + number = _session_number(key) + if number is None: + continue + if not isinstance(value, list): + raise ValueError(f"session_{number} conversation must be an array") + sessions.append((number, value)) + sessions.sort(key=lambda item: item[0]) + + records: list[ConversationRecord] = [] + for number, turns in sessions: + for turn in turns: + if not isinstance(turn, dict): + raise ValueError(f"session_{number} contains a non-object dialogue turn") + speaker = turn.get("speaker") + dia_id = turn.get("dia_id") + text_value = turn.get("text") + if not isinstance(speaker, str) or not isinstance(dia_id, str) or not isinstance( + text_value, str + ): + raise ValueError( + f"session_{number} dialogue turns require string speaker, dia_id, and text" + ) + text = text_value.strip() + query = "" + images: list[str] = [] + caption = turn.get("blip_caption") + if caption is not None and not isinstance(caption, str): + raise ValueError(f"session_{number} blip_caption must be a string or null") + if include_multimodal_context: + parts = [text] + if caption and caption.strip(): + parts.append(f"[caption: {caption.strip()}]") + raw_query = turn.get("query") + if raw_query is not None and not isinstance(raw_query, str): + raise ValueError(f"session_{number} query must be a string or null") + query = raw_query or "" + if query and query.strip(): + parts.append(f"[query: {query.strip()}]") + images = _normalize_refined_image_list(turn.get("img_url")) + if images: + parts.append(f"[images: {', '.join(images)}]") + text = " ".join(parts) + elif caption is not None: + # Rust's retained LoCoMo parser appends a present caption even + # when it trims to an empty string. Keep that text-only path + # unchanged; refined samples use the branch above. + text = f"{text} [caption: {caption.strip()}]" + records.append( + ConversationRecord( + dia_id=dia_id, + session_id=f"D{number}", + speaker=speaker, + text=text, + metadata={ + key: value + for key, value in { + "blip_caption": caption or "", + "query": query, + "img_url": ",".join(images), + }.items() + if value + }, + ) + ) + return records + + +def parse_sample_filter(raw: str | None) -> set[str] | None: + if raw is None: + return None + raw = raw.strip() + if not raw or raw.lower() == "all": + return None + + result: set[str] = set() + for token in raw.replace("/", ",").replace(";", ",").replace(" ", ",").split(","): + token = token.strip() + if not token: + continue + if token.isdigit(): + result.add(f"conv-{int(token)}") + elif token.startswith("conv-") and token[5:].isdigit(): + result.add(f"conv-{int(token[5:])}") + else: + result.add(token) + return result + + +def build_eval_cases( + samples: list[PreparedSample], + question_limit: int | None = None, +) -> list[EvalCase]: + cases: list[EvalCase] = [] + for sample in samples: + for index, question in enumerate(sample.questions): + if question_limit is not None and index >= question_limit: + break + cases.append( + EvalCase( + case_id=f"{sample.sample_id}::q{index + 1}", + sample_id=sample.sample_id, + question_index=index, + question=question.question, + answer=_stringify_answer(question), + category=question.category, + expected_evidence=list(question.evidence), + ) + ) + return cases + + +def parse_retrieval_plugin_list(value: str | None) -> list[str]: + if not value: + return [] + return [ + _normalize_plugin_name(token) + for token in value.split(",") + if _normalize_plugin_name(token) + ] + + +def summarize_scores(dataset_name: str, scores: list[CaseScore]) -> EvalSummary: + cases_with_evidence = sum(1 for score in scores if score.expected_evidence) + total_hits = sum(len(score.hit_evidence) for score in scores) + total_retrieved = sum(len(score.retrieved_evidence) for score in scores) + total_expected = sum(len(score.expected_evidence) for score in scores) + full_hit_cases = sum(1 for score in scores if score.full_evidence_hit) + return EvalSummary( + dataset_name=dataset_name, + total_cases=len(scores), + cases_with_evidence=cases_with_evidence, + evidence_precision_macro=_round4(_macro([score.evidence_precision for score in scores])), + evidence_precision_micro=_round4(_ratio(total_hits, total_retrieved)), + evidence_recall_macro=_round4(_macro([score.evidence_recall for score in scores])), + evidence_recall_micro=_round4(_ratio(total_hits, total_expected)), + full_evidence_hit_rate=_round4(_ratio(full_hit_cases, cases_with_evidence)), + ) + + +def summarize_scores_by_locomo_category(scores: list[CaseScore]) -> list[CategoryScoreSummary]: + grouped: dict[int, list[CaseScore]] = {} + for score in scores: + if score.category is not None: + grouped.setdefault(score.category, []).append(score) + + summaries: list[CategoryScoreSummary] = [] + for category in sorted(grouped): + label = _locomo_category_label(category) + if label is None: + continue + bucket = grouped[category] + cases_with_evidence = sum(1 for score in bucket if score.expected_evidence) + full_hit_cases = sum(1 for score in bucket if score.full_evidence_hit) + summaries.append( + CategoryScoreSummary( + category=category, + label=label, + count=len(bucket), + evidence_precision_macro=_round4( + _macro([score.evidence_precision for score in bucket]) + ), + evidence_recall_macro=_round4( + _macro([score.evidence_recall for score in bucket]) + ), + full_evidence_hit_rate=_round4(_ratio(full_hit_cases, cases_with_evidence)), + ) + ) + return summaries + + +def run_retained_qmd_eval( + *, + dataset_name: str, + samples: list[PreparedSample], + workspace_root: str | Path, + top_k: int = 24, + question_limit: int | None = None, + retrieval_plugins: list[str] | None = None, + harness_root: str | Path | None = None, + wiki_mode: WikiMode = "text", + llm: LLM | None = None, + wiki_builder_mode: WikiBuilderMode = "llm", + query_llm: LLM | None = None, +) -> EvalOutput: + """Build and retrieve a retained wiki with an explicit modality mode. + + ``wiki_mode="text"`` keeps the deterministic text-only workspace used for + Rust-compatible LoCoMo/LongMemEval runs. ``wiki_mode="multimodal"`` adds + one ``wiki/memories/*_multimodal.md`` page for each turn carrying image, + caption, or image-query metadata; it never downloads images or calls a + vision model. + """ + wiki_mode = normalize_wiki_mode(wiki_mode) + profiler = StageProfiler() + scores: list[CaseScore] = [] + cases = build_eval_cases(samples, question_limit) + completed_cases = 0 + for sample in samples: + sample_started = time.perf_counter() + sample_root = Path(workspace_root) / sample.sample_id + if llm is not None and wiki_builder_mode == "llm": + semantic_sources = [ + SemanticSource( + source_id=record.dia_id, + text=record.text, + conversation_id=sample.sample_id, + session_id=record.session_id, + speaker=record.speaker, + timestamp=sample.session_datetimes.get(_record_session_number(record), ""), + metadata=record.metadata, + ) + for record in sample.records + ] + build_result = WikiBuilder(llm=llm, mode="llm").build( + semantic_sources, sample_root, wiki_mode=wiki_mode + ) + files = ( + build_retained_memory_files(sample, sample_root, wiki_mode=wiki_mode) + if ( + not build_result.diagnostics.llm_used + and build_result.diagnostics.fallback_reason + ) + else build_result.files + ) + else: + files = build_retained_memory_files(sample, sample_root, wiki_mode=wiki_mode) + files = _materialize_rust_workspace_files(files) + profiler.record_elapsed("workspace_build", sample_started) + for question_index, question in enumerate(sample.questions): + if question_limit is not None and question_index >= question_limit: + break + retrieve_started = time.perf_counter() + result = retrieve_qmd_consensus_files( + question=question.question, + files=files, + entity_names=_sample_entity_names(sample, files), + top_k=top_k, + knowledge_root=Path(workspace_root) / sample.sample_id, + retrieval_plugins=retrieval_plugins, + llm=( + query_llm + if query_llm is not None + else (llm if wiki_builder_mode == "llm" else None) + ), + ) + profiler.record_elapsed("qmd_consensus_retrieve", retrieve_started) + expected = _normalize_expected_evidence(question.evidence) + retrieved = _ordered_unique( + evidence_id + for file in result.files + for evidence_id in _extract_retrieved_evidence_ids( + file, + question.question, + result.profile, + ) + ) + hits = [evidence_id for evidence_id in retrieved if evidence_id in set(expected)] + scores.append( + CaseScore( + case_id=f"{sample.sample_id}::q{question_index + 1}", + sample_id=sample.sample_id, + question_index=question_index, + question=question.question, + category=question.category, + expected_evidence=expected, + retrieved_evidence=retrieved, + hit_evidence=hits, + retrieved_record_ids=[], + retrieved_file_paths=[file.file_path for file in result.files], + retrieved_entrypoint_paths=[], + knowledge_base_root=str(Path(workspace_root) / sample.sample_id), + retrieved_record_count=0, + retrieved_file_count=len(result.files), + retrieved_entrypoint_count=0, + evidence_precision=_ratio(len(hits), len(retrieved)), + evidence_recall=_ratio(len(hits), len(expected)), + full_evidence_hit=bool(expected) and set(expected) <= set(hits), + retrieval_coverage=_coverage_from_profile_result(expected, result), + ) + ) + completed_cases += 1 + + output = EvalOutput( + summary=summarize_scores(dataset_name, scores), + cases=scores, + stage_profile=profiler.snapshot( + total_samples=len(samples), + total_cases=len(cases) if cases else completed_cases, + ), + ) + if harness_root is not None: + write_harness_artifacts( + harness_root, + output, + EvalHarnessConfig( + dataset_name=dataset_name, + samples=None, + question_limit=question_limit, + top_k=top_k, + workspace_root=str(workspace_root), + llm_provider=( + type(llm).__name__ + if llm is not None + else type(query_llm).__name__ + if query_llm is not None + else None + ), + retrieval_plugins=retrieval_plugins or ["qmd_consensus"], + wiki_mode=wiki_mode, + wiki_builder_mode=wiki_builder_mode, + ), + ) + return output + + +def run_python_locomo_retrieval_eval( + *, + dataset_path: str | Path, + output_dir: str | Path, + workspace_root: str | Path, + top_k: int = 24, + sample_limit: int | None = None, + question_limit: int | None = None, + sample_filter: set[str] | None = None, + wiki_mode: WikiMode = "text", + llm: LLM | None = None, + wiki_builder_mode: WikiBuilderMode = "llm", + query_llm: LLM | None = None, +) -> dict[str, Any]: + """Run LoCoMo retrieval with ``text`` or ``multimodal`` wiki generation.""" + payloads = json.loads(Path(dataset_path).read_text(encoding="utf-8")) + if sample_limit is not None: + payloads = payloads[:sample_limit] + samples = prepare_locomo_samples(payloads, sample_filter) + output_dir = Path(output_dir) + output = run_retained_qmd_eval( + dataset_name="locomo", + samples=samples, + workspace_root=workspace_root, + top_k=top_k, + question_limit=question_limit, + harness_root=output_dir / "harness", + wiki_mode=wiki_mode, + llm=llm, + wiki_builder_mode=wiki_builder_mode, + query_llm=query_llm, + ) + result = { + "summary": asdict(output.summary), + "cases": [asdict(case) for case in output.cases], + "stage_profile": asdict(output.stage_profile), + } + output_dir.mkdir(parents=True, exist_ok=True) + (output_dir / "locomo_retrieval_eval.json").write_text( + json.dumps(result, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + return result + + +def build_retained_memory_files( + sample: PreparedSample, + sample_root: str | Path, + include_multiview: bool = True, + wiki_mode: WikiMode = "text", +) -> list[RetrievedMemoryFile]: + """Materialize the retained wiki in text or multimodal mode. + + ``include_multiview`` remains a legacy switch for the minimal MEMORY.md + projection; when it is false, the minimal path intentionally skips all + adjunct pages. ``wiki_mode`` controls whether multimodal adjunct pages are + present in the full retained workspace. + """ + wiki_mode = normalize_wiki_mode(wiki_mode) + root = Path(sample_root).as_posix() + files: list[RetrievedMemoryFile] = [] + memory_lines = [ + f"# Memory Index for {sample.sample_id}", + "", + ( + "This memory index points to the most relevant cross-session profile, source " + "transcripts, observations, and event timelines." + ), + "", + "- [profile](wiki/synthesis/profile.md) - cross-session profile and recurring themes", + ] + for session_number in sorted(sample.session_summaries): + session_records = [ + item for item in sample.records if item.session_id == f"D{session_number}" + ] + session_summary = sample.session_summaries[session_number] + source_lines = [ + "---", + f"description: Session {session_number} summary with links to atomic dialogue turns", + "type: project", + "---", + f"# Session {session_number}", + "", + "## Session Date", + sample.session_datetimes.get(session_number, "Unknown"), + "", + "## Summary", + session_summary, + "", + "## Turn Index", + ] + for record in session_records: + slug = _evidence_slug(record.dia_id) + source_lines.append( + f"- [turn {record.dia_id}](../turns/{slug}.md) " + f"{record.speaker}: {_truncate_preview(record.text, 120)}" + ) + source_lines.append("") + files.append( + RetrievedMemoryFile( + filename=f"session_{session_number}.md", + file_path=f"{root}/wiki/sources/session_{session_number}.md", + mtime_ms=len(files) + 1, + content="\n".join(source_lines), + description=f"session {session_number} summary and turn index", + ) + ) + memory_lines.append( + f"- [session_{session_number} source](wiki/sources/session_{session_number}.md) " + "- session summary and turn index" + ) + + for record in sample.records: + slug = _evidence_slug(record.dia_id) + session_number = _record_session_number(record) + date_line = ( + f"- Session date: {sample.session_datetimes[session_number]}\n" + if session_number in sample.session_datetimes + else "" + ) + date_description = ( + f" on {sample.session_datetimes[session_number]}" + if session_number in sample.session_datetimes + else "" + ) + files.append( + RetrievedMemoryFile( + filename=f"{slug}.md", + file_path=f"{root}/wiki/turns/{slug}.md", + mtime_ms=len(files) + 1, + content=( + f"---\ndescription: Dialogue turn {record.dia_id} in session " + f"{record.session_id}{date_description} where {record.speaker} says: " + f"{_truncate_preview(record.text, 140)}\ntype: project\n---\n" + f"# Turn {record.dia_id}\n\n" + f"- Session: {record.session_id}\n" + f"{date_line}" + f"- Speaker: {record.speaker}\n" + f"- Evidence: {record.dia_id}\n\n" + f"## Content\n{record.text}\n" + ), + ) + ) + + if not include_multiview: + files.insert( + 0, + RetrievedMemoryFile( + filename="MEMORY.md", + file_path=f"{root}/MEMORY.md", + mtime_ms=0, + content="\n".join(memory_lines), + description=f"Retained memory index for {sample.sample_id}", + ), + ) + return files + + for session_number, notes in sorted(sample.observations.items()): + topic_lines = [ + "---", + ( + f"description: Observation index for session {session_number} with links to " + "atomic evidence notes" + ), + "type: project", + "---", + f"# Session {session_number} Observations", + "", + ] + for ordinal, note in enumerate(notes, start=1): + slug = _evidence_slug(note.evidence_id) + topic_lines.append( + f"- [observation {note.evidence_id}]" + f"(../observations/{slug}_obs_{ordinal}.md) " + f"{note.speaker}: {_truncate_preview(note.text, 120)}" + ) + files.append( + RetrievedMemoryFile( + filename=f"{slug}_obs_{ordinal}.md", + file_path=f"{root}/wiki/observations/{slug}_obs_{ordinal}.md", + mtime_ms=len(files) + 1, + content=( + f"---\ndescription: Observation {note.evidence_id} for session " + f"{session_number} about {note.speaker}: " + f"{_truncate_preview(note.text, 140)}\n" + f"type: project\n---\n# Observation {ordinal}\n\n" + f"- Session: D{session_number}\n- Speaker: {note.speaker}\n" + f"- Evidence: {note.evidence_id}\n\n## Note\n{note.text}\n" + ), + ) + ) + topic_lines.append("") + files.append( + RetrievedMemoryFile( + filename=f"session_{session_number}_observations.md", + file_path=f"{root}/wiki/topics/session_{session_number}_observations.md", + mtime_ms=len(files) + 1, + content="\n".join(topic_lines), + ) + ) + memory_lines.append( + f"- [session_{session_number} observations]" + f"(wiki/topics/session_{session_number}_observations.md) " + "- evidence-grounded observation index" + ) + + for session_number, events in sorted(sample.event_summaries.items()): + topic_lines = [ + "---", + ( + f"description: Event index for session {session_number} with dated atomic " + "event notes" + ), + "type: project", + "---", + f"# Session {session_number} Events", + "", + ] + if events.date is not None: + topic_lines.extend([f"Date: {events.date}", ""]) + ordinal = 0 + for speaker, entries in sorted(events.items_by_speaker.items()): + for entry in entries: + ordinal += 1 + topic_lines.append( + f"- [event {ordinal}]" + f"(../events/session_{session_number}_event_{ordinal}.md) " + f"{speaker}: {_truncate_preview(entry, 120)}" + ) + date_line = ( + f"- Date: {events.date}\n" if events.date is not None else "" + ) + files.append( + RetrievedMemoryFile( + filename=f"session_{session_number}_event_{ordinal}.md", + file_path=( + f"{root}/wiki/events/session_{session_number}_event_{ordinal}.md" + ), + mtime_ms=len(files) + 1, + content=( + f"---\ndescription: Session {session_number} event {ordinal} " + f"about {speaker}: {_truncate_preview(entry, 140)}\n" + "type: project\n---\n" + f"# Event {ordinal}\n\n- Session: D{session_number}\n" + f"{date_line}- Speaker: {speaker}\n\n## Event\n{entry}\n" + ), + ) + ) + topic_lines.append("") + files.append( + RetrievedMemoryFile( + filename=f"session_{session_number}_events.md", + file_path=f"{root}/wiki/topics/session_{session_number}_events.md", + mtime_ms=len(files) + 1, + content="\n".join(topic_lines), + ) + ) + memory_lines.append( + f"- [session_{session_number} events]" + f"(wiki/topics/session_{session_number}_events.md) - dated event index" + ) + + for speaker in _sample_entity_names(sample): + entity_slug = _entity_slug(speaker) + files.append( + RetrievedMemoryFile( + filename=f"{entity_slug}.md", + file_path=f"{root}/wiki/entities/{entity_slug}.md", + mtime_ms=len(files) + 1, + content=_render_retained_entity_page(sample, speaker), + description=f"Cross-session profile for {speaker}", + ) + ) + + evidence_lines = [ + f"- [{note.evidence_id}] {note.speaker}: {note.text}" + for notes in sample.observations.values() + for note in notes + ][:8] + files.append( + RetrievedMemoryFile( + filename="profile.md", + file_path=f"{root}/wiki/synthesis/profile.md", + mtime_ms=len(files) + 1, + content=( + f"---\ndescription: Cross-session profile for {sample.sample_id}\n" + "type: project\n---\n# Profile\n\n" + "This page summarizes recurring facts and themes across the conversation.\n\n" + "## Evidence Highlights\n" + + "\n".join(evidence_lines) + + "\n" + ), + ) + ) + + if wiki_mode == "multimodal": + for multimodal in _iter_multimodal_turns(sample): + artifact_slug = f"{_evidence_slug(multimodal['evidence_id'])}_multimodal" + memory_lines.append( + f"- [{multimodal['evidence_id']} multimodal]" + f"(wiki/memories/{artifact_slug}.md) - image/caption/query adjunct" + ) + files.append( + RetrievedMemoryFile( + filename=f"{artifact_slug}.md", + file_path=f"{root}/wiki/memories/{artifact_slug}.md", + mtime_ms=len(files) + 1, + content=_render_multimodal_memory_page(multimodal), + ) + ) + + files.extend( + [ + RetrievedMemoryFile( + filename="index.md", + file_path=f"{root}/index.md", + mtime_ms=len(files) + 1, + content=( + f"---\ndescription: Index of LoCoMo knowledge pages for " + f"{sample.sample_id}\ntype: reference\n---\n# Index\n\n" + "- [[wiki/synthesis/profile.md]]\n" + "- [[wiki/sources]] session transcript pages\n" + "- [[wiki/topics]] observation and event pages\n" + + ( + "- [[wiki/memories]] multimodal adjunct pages\n" + if wiki_mode == "multimodal" + else "" + ) + ), + ), + RetrievedMemoryFile( + filename="log.md", + file_path=f"{root}/log.md", + mtime_ms=len(files) + 2, + content=( + f"---\ndescription: Build log for {sample.sample_id}\ntype: reference\n" + "---\n- Materialized LoCoMo sample into local memory knowledge base.\n" + ), + ), + RetrievedMemoryFile( + filename=".wiki-schema.md", + file_path=f"{root}/.wiki-schema.md", + mtime_ms=0, + content=( + f"# Wiki Schema\n\n- topic: LoCoMo {sample.sample_id}\n" + "- language: en\n- structure: raw + wiki\n" + "- purpose: persistent memory retrieval evaluation\n" + ), + ), + RetrievedMemoryFile( + filename="sample.json", + file_path=f"{root}/raw/sample.json", + mtime_ms=0, + content=json.dumps(sample.raw_sample, indent=2, ensure_ascii=False), + ), + ] + ) + + files.insert( + 0, + RetrievedMemoryFile( + filename="MEMORY.md", + file_path=f"{root}/MEMORY.md", + mtime_ms=0, + content="\n".join(memory_lines) + "\n", + description=f"Retained memory index for {sample.sample_id}", + ), + ) + return _assign_rust_scaffold_mtimes(files, root) + + +def _assign_rust_scaffold_mtimes( + files: list[RetrievedMemoryFile], + root: str, +) -> list[RetrievedMemoryFile]: + """Use the same deterministic creation order as Rust's scaffold_memory.""" + + normalized = {normalize_memory_path(file.file_path): file for file in files} + ordered: list[RetrievedMemoryFile] = [] + seen: set[str] = set() + + def take(path: str) -> None: + key = normalize_memory_path(path) + file = normalized.get(key) + if file is not None and key not in seen: + ordered.append(file) + seen.add(key) + + for path in ( + f"{root}/.wiki-schema.md", + f"{root}/index.md", + f"{root}/log.md", + f"{root}/MEMORY.md", + f"{root}/raw/sample.json", + ): + take(path) + + def matching(prefix: str) -> list[RetrievedMemoryFile]: + prefix_key = normalize_memory_path(prefix) + return [ + file + for file in files + if normalize_memory_path(file.file_path).startswith(prefix_key) + and normalize_memory_path(file.file_path) not in seen + ] + + source_files = matching(f"{root}/wiki/sources/") + source_files.sort( + key=lambda file: int(re.search(r"session_(\d+)\.md$", file.file_path).group(1)) + if re.search(r"session_(\d+)\.md$", file.file_path) + else 0 + ) + for file in source_files: + take(file.file_path) + + for file in files: + if "/wiki/turns/" in normalize_memory_path(file.file_path): + take(file.file_path) + + observation_topics = matching(f"{root}/wiki/topics/") + observation_topics = [ + file for file in observation_topics if "_observations.md" in file.file_path + ] + observation_topics.sort( + key=lambda file: int(re.search(r"session_(\d+)_observations\.md$", file.file_path).group(1)) + if re.search(r"session_(\d+)_observations\.md$", file.file_path) + else 0 + ) + for topic in observation_topics: + session_match = re.search(r"session_(\d+)_observations\.md$", topic.file_path) + session_number = int(session_match.group(1)) if session_match else 0 + take(topic.file_path) + observation_files = [ + file + for file in matching(f"{root}/wiki/observations/") + if (match := re.search(r"_obs_(\d+)\.md$", file.file_path)) + and (prefix_match := re.search(r"(D\d+)_", Path(file.file_path).name)) + and int(prefix_match.group(1)[1:]) == session_number + ] + observation_files.sort( + key=lambda file: int(re.search(r"_obs_(\d+)\.md$", file.file_path).group(1)) + ) + for file in observation_files: + take(file.file_path) + + event_topics = matching(f"{root}/wiki/topics/") + event_topics = [file for file in event_topics if "_events.md" in file.file_path] + event_topics.sort( + key=lambda file: int(re.search(r"session_(\d+)_events\.md$", file.file_path).group(1)) + if re.search(r"session_(\d+)_events\.md$", file.file_path) + else 0 + ) + for topic in event_topics: + session_match = re.search(r"session_(\d+)_events\.md$", topic.file_path) + session_number = int(session_match.group(1)) if session_match else 0 + take(topic.file_path) + event_files = [ + file + for file in matching(f"{root}/wiki/events/") + if (match := re.search(r"session_(\d+)_event_(\d+)\.md$", file.file_path)) + and int(match.group(1)) == session_number + ] + event_files.sort( + key=lambda file: int(re.search(r"_event_(\d+)\.md$", file.file_path).group(1)) + ) + for file in event_files: + take(file.file_path) + + for file in sorted(matching(f"{root}/wiki/entities/"), key=lambda item: item.filename): + take(file.file_path) + take(f"{root}/wiki/synthesis/profile.md") + for file in files: + take(file.file_path) + # Rust's scan uses filesystem mtimes only as a recency tie-breaker. Keep + # a deterministic logical write clock in the same scaffold order instead + # of importing timestamps or group boundaries from a reference run. + write_clock = { + normalize_memory_path(file.file_path): index + for index, file in enumerate(ordered, start=1) + } + return [ + replace( + file, + mtime_ms=write_clock.get(normalize_memory_path(file.file_path), 0), + ) + for file in files + ] + + +def _materialize_rust_workspace_files( + files: list[RetrievedMemoryFile], +) -> list[RetrievedMemoryFile]: + """Write the virtual wiki in Rust scaffold order and retain filesystem mtimes. + + Rust's retained evaluator scans real files and uses ``metadata.modified()`` + as its recency tie-break. The logical clock remains the fallback for + callers that only construct an in-memory file list; the evaluation harness + materializes the same files before retrieval so its ordering observes the + same filesystem semantics without consulting Rust output. + """ + + ordered = sorted( + files, + key=lambda file: ( + file.mtime_ms, + normalize_memory_path(file.file_path), + ), + ) + for file in ordered: + path = Path(file.file_path) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(file.content, encoding="utf-8") + materialized: list[RetrievedMemoryFile] = [] + for file in files: + path = Path(file.file_path) + try: + mtime_ms = path.stat().st_mtime_ns // 1_000_000 + except OSError: + mtime_ms = file.mtime_ms + materialized.append(replace(file, mtime_ms=mtime_ms)) + return materialized + + +def format_progress_message(update: ProgressUpdate) -> str: + return ( + f"sample {update.sample_index}/{update.total_samples} " + f"{update.sample_id} q{update.question_index}/{update.sample_question_total} " + f"overall {update.completed_cases}/{update.total_cases}" + ) + + +def write_harness_artifacts( + harness_root: str | Path, + output: EvalOutput, + config: EvalHarnessConfig, +) -> None: + root = Path(harness_root) + cases_dir = root / "cases" + cases_dir.mkdir(parents=True, exist_ok=True) + + traces = [_case_trace_from_score(score) for score in output.cases] + for trace in traces: + _write_json(cases_dir / f"{trace['case_id'].replace('::', '__')}.json", trace) + + failures = [trace for trace in traces if trace["failure_kind"] != "full_hit"] + _write_json( + root / "run_manifest.json", + { + "created_at_ms": _unix_now_ms(), + "config": asdict(config), + "summary": asdict(output.summary), + "total_failures": len(failures), + }, + ) + _write_json(root / "category_breakdown.json", summarize_scores_by_locomo_category(output.cases)) + _write_json(root / "failure_buckets.json", _summarize_failure_buckets(failures)) + _write_json(root / "failure_report.json", failures) + (root / "failure_report.md").write_text( + _render_failure_report_markdown(output.summary, failures), + encoding="utf-8", + ) + _write_json(root / "stage_profile.json", output.stage_profile) + + +def _parse_questions(raw_questions: list[dict[str, Any]]) -> list[LoCoMoQuestion]: + questions: list[LoCoMoQuestion] = [] + for raw in raw_questions: + if not isinstance(raw, dict) or not isinstance(raw.get("question"), str): + raise ValueError("LoCoMo qa entries require a string question") + answer = raw.get("answer") + adversarial_answer = raw.get("adversarial_answer") + if adversarial_answer is not None and not isinstance(adversarial_answer, str): + raise ValueError("LoCoMo adversarial_answer must be a string or null") + evidence = raw.get("evidence", []) + if not isinstance(evidence, list) or not all(isinstance(item, str) for item in evidence): + raise ValueError("LoCoMo evidence must be an array of strings") + category = raw.get("category") + if category is not None and ( + not isinstance(category, int) or isinstance(category, bool) or category < 0 + ): + raise ValueError("LoCoMo category must be a non-negative integer or null") + questions.append( + LoCoMoQuestion( + question=raw["question"], + answer=answer, + adversarial_answer=adversarial_answer, + evidence=list(evidence), + category=category, + ) + ) + return questions + + +def _parse_session_datetimes(conversation: dict[str, Any]) -> dict[int, str]: + result: dict[int, str] = {} + for key, value in conversation.items(): + if not key.startswith("session_") or not key.endswith("_date_time"): + continue + number = key.removeprefix("session_").removesuffix("_date_time") + if number.isdigit() and isinstance(value, str) and value.strip(): + result[int(number)] = value.strip() + return dict(sorted(result.items())) + + +def _parse_session_summaries(raw: dict[str, str]) -> dict[int, str]: + if not isinstance(raw, dict): + raise ValueError("session_summary must be an object") + result: dict[int, str] = {} + for key, value in raw.items(): + if not isinstance(key, str) or not isinstance(value, str): + raise ValueError("session_summary keys and values must be strings") + suffix = key.removeprefix("session_").removesuffix("_summary") + number = int(suffix) if key.startswith("session_") and suffix.isdigit() else None + if number is not None: + result[number] = value + return dict(sorted(result.items())) + + +def _parse_event_summaries(raw: dict[str, Any]) -> dict[int, SessionEvents]: + result: dict[int, SessionEvents] = {} + for key, value in raw.items(): + number = ( + int(key.removeprefix("events_session_")) + if key.startswith("events_session_") + and key.removeprefix("events_session_").isdigit() + else None + ) + if number is None or not isinstance(value, dict): + continue + date_value = value.get("date") + date = date_value if isinstance(date_value, str) else None + items = { + str(speaker): [item for item in entries if isinstance(item, str)] + for speaker, entries in sorted(value.items()) + if speaker != "date" and isinstance(entries, list) + } + items = {speaker: entries for speaker, entries in items.items() if entries} + result[number] = SessionEvents( + date=date, + items_by_speaker=items, + ) + return dict(sorted(result.items())) + + +def _parse_observations(raw: dict[str, Any]) -> dict[int, list[ObservationNote]]: + result: dict[int, list[ObservationNote]] = {} + for key, value in raw.items(): + suffix = key.removeprefix("session_").removesuffix("_observation") + number = int(suffix) if key.startswith("session_") and suffix.isdigit() else None + if number is None or not isinstance(value, dict): + continue + notes: list[ObservationNote] = [] + # Rust deserializes JSON objects into ordered maps; iterate speakers in + # lexical order so event/observation ordinals and file names match. + for speaker, entries in sorted(value.items()): + if not isinstance(entries, list): + continue + for entry in entries: + if ( + isinstance(entry, list) + and len(entry) >= 2 + and isinstance(entry[0], str) + and isinstance(entry[1], str) + ): + notes.append( + ObservationNote( + speaker=str(speaker), + evidence_id=str(entry[1]), + text=str(entry[0]), + ) + ) + if notes: + result[number] = notes + return dict(sorted(result.items())) + + +def _sample_entity_names( + sample: PreparedSample, + files: list[RetrievedMemoryFile] | None = None, +) -> list[str]: + names = { + record.speaker.strip() + for record in sample.records + if record.speaker.strip() + } + for notes in sample.observations.values(): + names.update(note.speaker.strip() for note in notes if note.speaker.strip()) + for events in sample.event_summaries.values(): + names.update( + speaker.strip() for speaker in events.items_by_speaker if speaker.strip() + ) + if files is None: + return sorted(names) + ordered = [] + for file in sorted( + ( + file + for file in files + if "/wiki/entities/" in normalize_memory_path(file.file_path) + ), + key=lambda file: (-file.mtime_ms, file.filename), + ): + speaker = Path(file.file_path).stem + canonical_match = re.search(r"(?im)^canonical_id:\s*(.+)$", file.content) + if canonical_match: + canonical = canonical_match.group(1).strip() + if canonical and canonical not in ordered: + ordered.append(canonical) + if speaker in names and speaker not in ordered: + ordered.append(speaker) + return ordered + sorted(names.difference(ordered)) + + +def _iter_multimodal_turns(sample: PreparedSample) -> list[dict[str, Any]]: + conversation = sample.raw_sample.get("conversation", {}) + if not isinstance(conversation, dict): + return [] + turns: list[dict[str, Any]] = [] + for session_number in sorted( + number + for key in conversation + if (number := _session_number(str(key))) is not None + ): + session = conversation.get(f"session_{session_number}") + if not isinstance(session, list): + continue + date = str(conversation.get(f"session_{session_number}_date_time", "")) + for index, raw_turn in enumerate(session, start=1): + if not isinstance(raw_turn, dict): + continue + images = raw_turn.get("img_url", []) + if isinstance(images, str): + images = [images] if images.strip() else [] + elif isinstance(images, list): + images = [str(item).strip() for item in images if str(item).strip()] + else: + images = [] + caption = str(raw_turn.get("blip_caption", "") or "").strip() + query = str(raw_turn.get("query", "") or "").strip() + if not images and not caption and not query: + continue + evidence_id = str(raw_turn.get("dia_id") or f"D{session_number}:{index}") + turns.append( + { + "evidence_id": evidence_id, + "session_number": session_number, + "date": date, + "speaker": str(raw_turn.get("speaker", "")), + "text": str(raw_turn.get("text", "") or "").strip(), + "caption": caption, + "query": query, + "images": images, + } + ) + return turns + + +def _render_multimodal_memory_page(turn: dict[str, Any]) -> str: + images = turn["images"] + image_lines = ( + "\n".join(f"- source: {image}\n local: (none)\n status: skipped" for image in images) + if images + else "(none)" + ) + topics = [item for item in (turn["query"], turn["caption"]) if item] + topics_text = ", ".join(topics) if topics else "multimodal" + sources = ", ".join(f'"{item}"' for item in images) + if not sources: + sources = f'"locomo_refined:{turn["evidence_id"]}"' + return ( + "---\n" + 'type: "memory"\nmodality: "image"\n' + f'date: "{turn["date"]}"\nupdated: "2026-04-21"\n' + 'tags: ["locomo_refined", "multimodal"]\n' + f'aliases: ["{turn["evidence_id"]}"]\n' + f"sources: [{sources}]\n" + 'maturity: "compiled"\n' + f'artifact_id: "{_evidence_slug(turn["evidence_id"])}_multimodal"\n' + f'linked_entities: ["{turn["speaker"]}"]\n' + f'linked_topics: ["{topics_text}"]\n' + "---\n\n" + f"# {turn['speaker']} multimodal memory {turn['evidence_id']}\n\n" + f"- Evidence: {turn['evidence_id']}\n" + f"- Session: D{turn['session_number']}\n" + f"- Speaker: {turn['speaker']}\n" + f"- Query: {turn['query']}\n\n" + f"## Turn Text\n{turn['text']}\n\n" + f"## Caption\n{turn['caption'] or '(none)'}\n\n" + f"## Images\n{image_lines}\n\n" + "## Vision Summary\nstatus: disabled\n" + ) + + +def _render_retained_entity_page(sample: PreparedSample, speaker: str) -> str: + normalized_speaker = speaker.casefold() + records_by_session: dict[int, list[ConversationRecord]] = {} + for record in sample.records: + if record.speaker.casefold() != normalized_speaker: + continue + session_number = _record_session_number(record) + records_by_session.setdefault(session_number, []).append(record) + records = _select_evenly_spaced( + [ + record + for session_number in sorted(records_by_session) + for record in _first_and_last(records_by_session[session_number]) + ], + 10, + ) + + observations_by_session: dict[int, list[tuple[int, ObservationNote]]] = {} + for session_number, notes in sorted(sample.observations.items()): + matching = [ + (ordinal, note) + for ordinal, note in enumerate(notes, start=1) + if note.speaker.casefold() == normalized_speaker + ] + if matching: + observations_by_session[session_number] = matching + observations = _select_evenly_spaced( + [ + note + for session_number in sorted(observations_by_session) + for note in _first_and_last(observations_by_session[session_number]) + ], + 14, + ) + + events_by_session: dict[int, list[tuple[int, str]]] = {} + for session_number, events in sorted(sample.event_summaries.items()): + flattened = [] + ordinal = 0 + for event_speaker, entries in sorted(events.items_by_speaker.items()): + for entry in entries: + ordinal += 1 + if event_speaker.casefold() == normalized_speaker: + flattened.append((ordinal, entry)) + if flattened: + events_by_session[session_number] = flattened + selected_events = _select_evenly_spaced( + [ + (session_number, ordinal, entry) + for session_number in sorted(events_by_session) + for ordinal, entry in _first_and_last(events_by_session[session_number]) + ], + 10, + ) + + lines = [ + "---", + ( + f"description: Cross-session profile page for {speaker} with representative " + "evidence-linked observations, turns, and events across sessions" + ), + "type: project", + "---", + f"# {speaker}", + "", + ] + if observations: + lines.append("## Key Observations") + for ordinal, note in observations: + lines.append( + f"- [observation {note.evidence_id}]" + f"(../observations/{_evidence_slug(note.evidence_id)}_obs_{ordinal}.md) " + f"{_truncate_preview(note.text, 100)}" + ) + lines.append("") + if selected_events: + lines.append("## Related Events") + for session_number, ordinal, entry in selected_events: + lines.append( + f"- [event {ordinal}]" + f"(../events/session_{session_number}_event_{ordinal}.md) " + f"{_truncate_preview(entry, 100)}" + ) + lines.append("") + if records: + lines.append("## Related Turns") + for record in records: + lines.append( + f"- [turn {record.dia_id}]" + f"(../turns/{_evidence_slug(record.dia_id)}.md) " + f"{_truncate_preview(record.text, 100)}" + ) + lines.append("") + return "\n".join(lines) + + +def _record_session_number(record: ConversationRecord) -> int: + suffix = record.session_id.removeprefix("D") + return int(suffix) if suffix.isdigit() else 0 + + +def _first_and_last(values: list[Any]) -> list[Any]: + if len(values) <= 1: + return values.copy() + return [values[0], values[-1]] + + +def _select_evenly_spaced(values: list[Any], max_items: int) -> list[Any]: + if max_items <= 0 or not values: + return [] + if len(values) <= max_items: + return values.copy() + if max_items == 1: + return values[:1] + last_index = len(values) - 1 + return [values[(slot * last_index) // (max_items - 1)] for slot in range(max_items)] + + +def _coverage_from_profile_result( + expected: list[str], + result: Any, +) -> RetrievalCoverageSummary: + expected_set = set(expected) + final_hit_evidence = _intersect_evidence_hits( + expected_set, + result.coverage.final_file_paths, + result.files, + ) + candidate_hit_evidence = _intersect_evidence_hits( + expected_set, + result.coverage.candidate_pool_file_paths, + result.files, + ) + late_bridge_hit_evidence = _intersect_evidence_hits( + expected_set, + result.coverage.late_bridge_file_paths, + result.files, + ) + root_hit_evidence = _intersect_evidence_hits( + expected_set, + result.coverage.root_file_paths, + result.files, + ) + if not expected: + miss_stage = "no_expected_evidence" + elif set(expected) <= set(final_hit_evidence): + miss_stage = "full_hit" + elif candidate_hit_evidence or late_bridge_hit_evidence: + miss_stage = "final_assembly" + elif root_hit_evidence: + miss_stage = "candidate_ranking" + else: + miss_stage = "root" + return RetrievalCoverageSummary( + miss_stage=miss_stage, + root_hit_evidence=root_hit_evidence, + candidate_pool_hit_evidence=candidate_hit_evidence, + late_bridge_hit_evidence=late_bridge_hit_evidence, + final_hit_evidence=final_hit_evidence, + root_file_paths=result.coverage.root_file_paths, + candidate_pool_file_paths=result.coverage.candidate_pool_file_paths, + late_bridge_file_paths=result.coverage.late_bridge_file_paths, + ) + + +def _intersect_evidence_hits( + expected_set: set[str], + file_paths: list[str], + files: list[RetrievedMemoryFile], +) -> list[str]: + selected = set(file_paths) + values = [] + for file in files: + if file.file_path not in selected: + continue + values.extend(_extract_evidence_ids(file.content)) + if file.description: + values.extend(_extract_evidence_ids(file.description)) + return _ordered_unique(evidence_id for evidence_id in values if evidence_id in expected_set) + + +def _extract_evidence_ids(text: str) -> list[str]: + return _ordered_unique( + match.group(0) + for match in re.finditer(r"D\d+:\d+", text) + ) + + +def _extract_retrieved_evidence_ids( + file: RetrievedMemoryFile, + question: str, + profile: Any, +) -> list[str]: + # Rust's retained evaluator extracts every evidence id from every selected + # file (and its description), without applying a question-dependent line + # filter. Keep the Python scorer identical so recall is determined solely + # by retrieval, not by a Python-only post-processing rule. + values = _extract_evidence_ids(file.content) + if file.description: + values.extend(_extract_evidence_ids(file.description)) + return _ordered_unique(values) + + +def _is_family_entity_question(question: str) -> bool: + return any( + term in question.lower() + for term in ("kid", "kids", "child", "children", "family") + ) + + +def _normalize_expected_evidence(raw: list[str]) -> list[str]: + normalized: list[str] = [] + for item in raw: + ids = _extract_evidence_ids(item) + if ids: + normalized.extend(ids) + elif item.strip(): + normalized.append(item.strip()) + return _ordered_unique(normalized) + + +def _evidence_slug(evidence_id: str) -> str: + return evidence_id.replace(":", "_") + + +def _truncate_preview(text: str, max_chars: int) -> str: + compact = " ".join(text.split()) + return compact if len(compact) <= max_chars else compact[:max_chars] + "..." + + +def _entity_slug(name: str) -> str: + # Match Rust's speaker_file_name: only ASCII alphanumerics survive. + slug = re.sub(r"[^A-Za-z0-9]+", "_", name.strip()) + return slug.strip("_") or "entity" + + +def _ordered_unique(values: Any) -> list[Any]: + seen = set() + unique = [] + for value in values: + if value not in seen: + unique.append(value) + seen.add(value) + return unique + + +def _session_number(key: str) -> int | None: + if key.startswith("events_session_"): + suffix = key.removeprefix("events_session_") + elif key.startswith("session_"): + suffix = key.removeprefix("session_") + suffix = suffix.removesuffix("_observation") + suffix = suffix.removesuffix("_summary") + else: + return None + return int(suffix) if suffix.isdigit() else None + + +def _normalize_refined_image_list(raw: Any) -> list[str]: + """Normalize Rust refined-dataset ``img_url`` values without coercion.""" + + if raw is None: + return [] + if isinstance(raw, str): + value = raw.strip() + return [value] if value else [] + if isinstance(raw, list): + values: list[str] = [] + for item in raw: + if not isinstance(item, str): + raise ValueError("LoCoMo_refined img_url arrays require string items") + value = item.strip() + if value: + values.append(value) + return values + raise ValueError("LoCoMo_refined img_url must be a string, array, or null") + + +def _stringify_answer(question: LoCoMoQuestion) -> str: + # Rust prefers the typed `answer` field and only falls back to + # `adversarial_answer` when it is absent. + if isinstance(question.answer, str): + return question.answer + if question.answer is not None: + return json.dumps(question.answer, ensure_ascii=False, separators=(",", ":")) + if question.adversarial_answer is not None: + return question.adversarial_answer + return "" + + +def _normalize_plugin_name(value: str) -> str: + return value.strip().lower().replace("-", "_") + + +def _round4(value: float) -> float: + # Rust's f64::round rounds halfway cases away from zero; Python's round + # uses bankers rounding. All retained metrics are non-negative. + return math.floor(value * 10_000.0 + 0.5) / 10_000.0 + + +def _sort_stage_timings(stages: list[StageTimingRecord]) -> list[StageTimingRecord]: + return sorted(stages, key=lambda record: (-record.total_ms, record.stage)) + + +def _macro(values: list[float]) -> float: + if not values: + return 0.0 + return sum(values) / len(values) + + +def _ratio(numerator: int, denominator: int) -> float: + if denominator == 0: + return 0.0 + return numerator / denominator + + +def _locomo_category_label(category: int) -> str | None: + return { + 1: "1 Multi Hop", + 2: "2 Temporal", + 3: "3 Open Domain", + 4: "4 Single Hop", + 5: "5 Adversarial", + }.get(category) + + +def _case_trace_from_score(score: CaseScore) -> dict[str, Any]: + missing = sorted(set(score.expected_evidence) - set(score.hit_evidence)) + unexpected = sorted(set(score.retrieved_evidence) - set(score.expected_evidence)) + failure_kind = _failure_kind(score) + return { + **asdict(score), + "missing_expected_evidence": missing, + "unexpected_retrieved_evidence": unexpected, + "failure_kind": failure_kind, + "failure_bucket": failure_kind, + "evidence_precision": f"{score.evidence_precision:.4f}", + "evidence_recall": f"{score.evidence_recall:.4f}", + } + + +def _failure_kind(score: CaseScore) -> str: + if score.full_evidence_hit: + return "full_hit" + if not score.retrieved_evidence: + return "no_retrieval" + if not score.hit_evidence: + return "miss" + return "partial_hit" + + +def _summarize_failure_buckets(failures: list[dict[str, Any]]) -> dict[str, int]: + buckets: dict[str, int] = {} + for trace in failures: + bucket = str(trace["failure_bucket"]) + buckets[bucket] = buckets.get(bucket, 0) + 1 + return dict(sorted(buckets.items())) + + +def _render_failure_report_markdown(summary: EvalSummary, failures: list[dict[str, Any]]) -> str: + lines = [ + "# wikimem retained_eval failure report", + "", + f"- dataset: {summary.dataset_name}", + f"- total_cases: {summary.total_cases}", + f"- evidence_recall_macro: {summary.evidence_recall_macro:.4f}", + f"- full_evidence_hit_rate: {summary.full_evidence_hit_rate:.4f}", + f"- failures: {len(failures)}", + ] + for trace in failures: + lines.extend(["", f"## {trace['case_id']}", f"- failure_kind: {trace['failure_kind']}"]) + return "\n".join(lines) + "\n" + + +def _write_json(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(_jsonable(value), ensure_ascii=False, indent=2), encoding="utf-8") + + +def _jsonable(value: Any) -> Any: + if hasattr(value, "__dataclass_fields__"): + return asdict(value) + if isinstance(value, list): + return [_jsonable(item) for item in value] + if isinstance(value, dict): + return {key: _jsonable(item) for key, item in value.items()} + return value + + +def _unix_now_ms() -> int: + return int(time.time() * 1000) diff --git a/evaluation/wikimem/retrieval_profile.py b/evaluation/wikimem/retrieval_profile.py new file mode 100644 index 00000000..31b8206b --- /dev/null +++ b/evaluation/wikimem/retrieval_profile.py @@ -0,0 +1,1763 @@ +"""Retained wikimem retrieval profile assembly helpers.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field, replace +from pathlib import Path +from pathlib import PurePosixPath + +from common.llm.base import LLM + +from evaluation.wikimem.qmd_consensus import ( + CandidateFile, + QueryAugmentation, + QuestionProfile, + RetrievedMemoryFile, + apply_query_augmentation, + build_cached_file_lexical_features, + build_qmd_consensus_augmentation, + build_qmd_consensus_candidate_proposals, + build_qmd_consensus_late_bridge_proposals, + build_qmd_consensus_rerank_proposals, + build_question_profile, + candidate_query_hits_with_features, + keyword_ngrams, + normalize_memory_path, + score_line, + significant_phrases, + tokenize_fuzzy_query, + tokenize_query, +) +from evaluation.wikimem.llm_semantics import QueryUnderstanding, understand_query + + +@dataclass(frozen=True) +class RetrievalProfileCoverage: + root_file_paths: list[str] = field(default_factory=list) + source_file_paths: list[str] = field(default_factory=list) + source_companion_file_paths: list[str] = field(default_factory=list) + scoped_file_paths: list[str] = field(default_factory=list) + candidate_pool_file_paths: list[str] = field(default_factory=list) + late_bridge_file_paths: list[str] = field(default_factory=list) + final_file_paths: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class RetrievalProfileResult: + profile: QuestionProfile + files: list[RetrievedMemoryFile] + coverage: RetrievalProfileCoverage + + +@dataclass(frozen=True) +class SessionSourceFile: + session_number: int + file: RetrievedMemoryFile + search_text: str + search_tokens: list[str] + search_fuzzy_tokens: list[str] + + +def _apply_llm_query_understanding( + profile: QuestionProfile, + llm: LLM, + question: str, + known_entities: list[str], +) -> QuestionProfile: + """Merge conservative LLM intent/entity expansions into the lexical profile.""" + + try: + understanding: QueryUnderstanding = understand_query( + llm, question, known_entities=known_entities + ) + except Exception: + # Query understanding is an enhancement, never a retrieval hard dependency. + return profile + intent = understanding.intent.casefold() + extra_terms = tuple( + value + for value in ( + *understanding.expanded_terms, + understanding.relation, + understanding.time_expression, + *understanding.memory_kinds, + ) + if value + ) + extra_phrases = tuple( + value + for value in (understanding.relation, understanding.time_expression) + if value + ) + return replace( + profile, + named_entities=list(dict.fromkeys((*profile.named_entities, *understanding.entities))), + expansion_tokens=list( + dict.fromkeys( + ( + *profile.expansion_tokens, + *tokenize_query(" ".join(extra_terms)), + ) + ) + ), + expansion_fuzzy_tokens=list( + dict.fromkeys( + ( + *profile.expansion_fuzzy_tokens, + *tokenize_fuzzy_query(" ".join(extra_terms)), + ) + ) + ), + expansion_phrases=list( + dict.fromkeys( + ( + *profile.expansion_phrases, + *extra_phrases, + *[term for term in extra_terms if len(tokenize_query(term)) > 1], + ) + ) + ), + temporal=profile.temporal or bool(understanding.time_expression), + relational=( + profile.relational + or intent in {"compare", "decision", "preference"} + or bool(understanding.relation) + ), + identity=profile.identity or intent == "profile", + ) + + +def retrieve_qmd_consensus_files( + *, + question: str, + files: list[RetrievedMemoryFile], + root_files: list[RetrievedMemoryFile] | None = None, + entity_names: list[str] | None = None, + top_k: int = 24, + knowledge_root: str | Path | None = None, + retrieval_plugins: list[str] | None = None, + llm: LLM | None = None, +) -> RetrievalProfileResult: + """Assemble qmd_consensus retrieval files using the Rust retained profile order.""" + + limit = max(top_k, 1) + internal_limit = max(top_k, 8) + plugin_names = ( + ["qmd_consensus"] + if retrieval_plugins is None + else [name.strip().lower().replace("-", "_") for name in retrieval_plugins] + ) + unsupported_plugins = set(plugin_names) - {"qmd_consensus"} + if unsupported_plugins: + raise ValueError( + "Python retained profile only implements qmd_consensus; unsupported " + f"plugins: {sorted(unsupported_plugins)}" + ) + use_qmd_plugin = "qmd_consensus" in plugin_names + cached_files = [file for file in files if _is_rust_cached_retrieval_file(file)] + cached_files_by_path = { + normalize_memory_path(file.file_path): file for file in cached_files + } + profile = build_question_profile(question, entity_names or []) + if llm is not None: + profile = _apply_llm_query_understanding(profile, llm, question, entity_names or []) + roots = ( + list(root_files) + if root_files + else _rank_initial_root_files(question, files, internal_limit, knowledge_root) + ) + session_sources = build_session_source_files(files) + corpus_augmentation = build_corpus_consensus_augmentation(question, profile, session_sources) + scoped_query = _compose_augmented_retrieval_query(question, corpus_augmentation) + profile = apply_query_augmentation(profile, corpus_augmentation) + augmentation = ( + build_qmd_consensus_augmentation(question, profile, roots) + if use_qmd_plugin + else QueryAugmentation(tokens=[], fuzzy_tokens=[], phrases=[]) + ) + profile = apply_query_augmentation(profile, augmentation) + + candidate_proposals = ( + build_qmd_consensus_candidate_proposals(question, profile, cached_files) + if use_qmd_plugin + else [] + ) + plugin_candidates = [ + CandidateFile( + file=cached_files_by_path[proposal.file_path], + query_hits=proposal.query_hits, + seed_boost=proposal.seed_boost, + ) + for proposal in candidate_proposals + if proposal.file_path in cached_files_by_path + ] + ranked_candidates = _merge_candidate_files( + [CandidateFile(file=file, query_hits=2, seed_boost=0.0) for file in roots], + plugin_candidates, + ) + scoped_candidates = select_scoped_candidate_files( + scoped_query, + profile, + files, + knowledge_root=knowledge_root, + ) + ranked_candidates = _merge_candidate_files(ranked_candidates, scoped_candidates) + sources_by_session = {source.session_number: source for source in session_sources} + source_companions = collect_session_source_companions( + question, + profile, + ranked_candidates, + sources_by_session, + ) + # Rust inserts companion proposals into the main candidate map before the + # candidate ranking pass, while retaining the same proposals for the + # dedicated companion budget in final assembly. + ranked_candidates = _merge_candidate_files(ranked_candidates, source_companions) + ranked_candidates = _rank_candidate_files( + question, + profile, + ranked_candidates, + knowledge_root=knowledge_root, + ) + ranked_sources = rank_global_session_sources(question, profile, session_sources) + ranked_source_companions = _rank_candidate_files( + question, + profile, + source_companions, + knowledge_root=knowledge_root, + ) + rerank_proposals = ( + build_qmd_consensus_rerank_proposals( + question, + profile, + ranked_candidates, + cached_files, + ) + if use_qmd_plugin + else [] + ) + reranked_candidates = _rerank_candidate_files( + question, + profile, + ranked_candidates, + rerank_proposals, + cached_files_by_path, + knowledge_root=knowledge_root, + ) + late_bridge_proposals = ( + build_qmd_consensus_late_bridge_proposals( + question, + profile, + _collect_late_bridge_seed_files( + roots, + reranked_candidates, + ranked_source_companions, + ), + cached_files, + ) + if use_qmd_plugin + else [] + ) + ranked_late_bridges = _rank_candidate_files( + question, + profile, + [ + CandidateFile( + file=cached_files_by_path[proposal.file_path], + query_hits=proposal.query_hits, + seed_boost=proposal.seed_boost, + ) + for proposal in late_bridge_proposals + if proposal.file_path in cached_files_by_path + ], + knowledge_root=knowledge_root, + ) + + final_files: list[RetrievedMemoryFile] = [] + included: set[str] = set() + preserved_primary = min(limit, max(3, (limit + 1) // 2)) + + for file in roots[:preserved_primary]: + _push_unique_file(final_files, included, file, limit) + selected_sources = select_diverse_session_sources( + profile, + roots[:preserved_primary], + ranked_sources, + source_injection_budget(profile), + ) + for source in selected_sources: + _push_unique_file(final_files, included, source.file, limit) + companion_limit = source_companion_budget( + profile, + has_plugin_retrieval=use_qmd_plugin, + top_k=limit, + ) + for companion in ranked_source_companions[:companion_limit]: + _push_unique_file(final_files, included, companion.file, limit) + for candidate in reranked_candidates: + _push_unique_file(final_files, included, candidate.file, limit) + if len(final_files) >= limit: + break + for candidate in ranked_late_bridges: + _push_unique_file(final_files, included, candidate.file, limit) + if len(final_files) >= limit: + break + + return RetrievalProfileResult( + profile=profile, + files=final_files, + coverage=RetrievalProfileCoverage( + root_file_paths=[file.file_path for file in roots], + source_file_paths=[source.file.file_path for source in selected_sources], + source_companion_file_paths=[ + candidate.file.file_path for candidate in ranked_source_companions + ], + scoped_file_paths=[candidate.file.file_path for candidate in scoped_candidates], + candidate_pool_file_paths=[ + candidate.file.file_path + for candidate in ranked_candidates + ranked_source_companions + ], + late_bridge_file_paths=[proposal.file_path for proposal in late_bridge_proposals], + final_file_paths=[file.file_path for file in final_files], + ), + ) + + +def scoped_budgets(profile: QuestionProfile) -> list[tuple[str, int, float]]: + if profile.temporal: + return [ + ("wiki/sources", 3, 5.0), + ("wiki/memory", 3, 4.5), + ("wiki/observations", 3, 4.0), + ("wiki/events", 2, 3.0), + ("wiki/turns", 1, 1.0), + ("wiki/entities", 1, 2.0), + ] + if profile.identity or profile.hypothetical or profile.relational: + return [ + ("wiki/entities", 2, 4.0), + ("wiki/memory", 3, 4.5), + ("wiki/sources", 3, 5.0), + ("wiki/observations", 3, 4.0), + ("wiki/turns", 1, 2.0), + ("wiki/events", 1, 1.0), + ] + if profile.location: + return [ + ("wiki/sources", 2, 4.0), + ("wiki/memory", 3, 4.0), + ("wiki/observations", 3, 4.0), + ("wiki/events", 2, 3.0), + ("wiki/turns", 1, 1.0), + ("wiki/entities", 1, 1.0), + ] + return [ + ("wiki/sources", 2, 3.0), + ("wiki/memory", 3, 3.0), + ("wiki/observations", 3, 3.0), + ("wiki/events", 2, 2.0), + ("wiki/entities", 1, 2.0), + ("wiki/turns", 1, 1.0), + ] + + +def build_corpus_consensus_augmentation( + question: str, + base_profile: QuestionProfile, + session_sources: list[SessionSourceFile], +) -> QueryAugmentation: + if not session_sources: + return QueryAugmentation(tokens=[], fuzzy_tokens=[], phrases=[]) + + token_df = session_source_token_document_frequency(session_sources) + corrected_tokens = [ + correction + for token in base_profile.query_tokens + if ( + correction := best_corpus_correction( + token, + token_df, + base_profile.named_entities, + ) + ) + is not None + ] + correction_augmentation = QueryAugmentation( + tokens=corrected_tokens, + fuzzy_tokens=tokenize_fuzzy_query(" ".join(corrected_tokens)), + phrases=[], + ) + corrected_profile = apply_query_augmentation(base_profile, correction_augmentation) + if base_profile.temporal and not base_profile.hypothetical and not base_profile.aggregate: + return correction_augmentation + if ( + not base_profile.hypothetical + and not base_profile.identity + and not base_profile.relational + and not base_profile.aggregate + ): + return correction_augmentation + + return _extend_corpus_augmentation_with_anchors( + question, + corrected_profile, + session_sources, + token_df, + correction_augmentation, + ) + + +def session_source_token_document_frequency( + session_sources: list[SessionSourceFile], +) -> dict[str, int]: + document_frequency: dict[str, int] = {} + for source in session_sources: + for token in set(source.search_tokens): + document_frequency[token] = document_frequency.get(token, 0) + 1 + return document_frequency + + +def best_corpus_correction( + token: str, + token_df: dict[str, int], + named_entities: list[str], +) -> str | None: + if ( + len(token) < 5 + or token in token_df + or any(entity.lower() == token.lower() for entity in named_entities) + ): + return None + + query_features = tokenize_fuzzy_query(token) + candidates = [] + for candidate, df in token_df.items(): + if ( + len(candidate) < 5 + or candidate == token + or candidate[0] != token[0] + or abs(len(candidate) - len(token)) > 2 + ): + continue + distance = _bounded_edit_distance(token, candidate, 2) + if distance is None: + continue + feature_overlap = _token_overlap(query_features, tokenize_fuzzy_query(candidate)) + if distance > 1 and feature_overlap == 0: + continue + candidates.append((candidate, distance, feature_overlap, df)) + + if not candidates: + return None + candidates.sort(key=lambda item: (item[1], -item[2], -item[3], item[0])) + return candidates[0][0] + + +def _extend_corpus_augmentation_with_anchors( + question: str, + corrected_profile: QuestionProfile, + session_sources: list[SessionSourceFile], + token_df: dict[str, int], + correction_augmentation: QueryAugmentation, +) -> QueryAugmentation: + query_ngrams = keyword_ngrams(question) + significant_question_phrases = significant_phrases(question) + named_lower = [name.lower() for name in corrected_profile.named_entities] + ranked_sources = rank_global_session_sources(question, corrected_profile, session_sources) + token_scores: dict[str, float] = {} + token_support: dict[str, set[int]] = {} + phrase_scores: dict[str, float] = {} + phrase_support: dict[str, set[int]] = {} + + for rank, source in enumerate(ranked_sources[:4]): + weight = 1.35 - float(rank) * 0.2 + for fragment in _best_source_anchor_fragments( + source, + corrected_profile, + question, + correction_augmentation.tokens, + ): + for token in tokenize_query(fragment): + if not _should_keep_corpus_anchor_token( + token, + corrected_profile, + named_lower, + token_df, + ): + continue + token_scores[token] = token_scores.get(token, 0.0) + weight + token_support.setdefault(token, set()).add(source.session_number) + for phrase in keyword_ngrams(fragment): + if phrase in significant_question_phrases or phrase in query_ngrams: + continue + if not _should_keep_corpus_anchor_phrase( + phrase, + corrected_profile, + token_df, + ): + continue + phrase_scores[phrase] = phrase_scores.get(phrase, 0.0) + weight + phrase_support.setdefault(phrase, set()).add(source.session_number) + + anchor_tokens = sorted( + token_scores.items(), + key=lambda item: ( + -len(token_support.get(item[0], set())), + -item[1], + token_df.get(item[0], 1), + item[0], + ), + ) + merged_tokens = list(correction_augmentation.tokens) + seen_tokens = set(merged_tokens) + for token, score in anchor_tokens: + support = len(token_support.get(token, set())) + if support < 2 and score < 1.0: + continue + if token not in seen_tokens: + merged_tokens.append(token) + seen_tokens.add(token) + if len(merged_tokens) >= 6: + break + + token_set = set(merged_tokens) + anchor_phrases = sorted( + phrase_scores.items(), + key=lambda item: ( + -len(phrase_support.get(item[0], set())), + -item[1], + item[0], + ), + ) + phrases = [] + for phrase, score in anchor_phrases: + support = len(phrase_support.get(phrase, set())) + if support < 2 and score < 1.1: + continue + if not any(token in token_set for token in phrase.split()): + continue + phrases.append(phrase) + if len(phrases) >= 3: + break + + return QueryAugmentation( + tokens=merged_tokens, + fuzzy_tokens=tokenize_fuzzy_query(" ".join(merged_tokens)), + phrases=phrases, + ) + + +def _best_source_anchor_fragments( + source: SessionSourceFile, + profile: QuestionProfile, + question: str, + corrected_tokens: list[str], +) -> list[str]: + corrected_fuzzy = tokenize_fuzzy_query(" ".join(corrected_tokens)) + fragments = [] + for fragment in _split_anchor_fragments( + _extract_source_summary_snippet(source.file.content) + ) + _split_anchor_fragments(_extract_source_turn_index_snippet(source.file.content)): + tokens = tokenize_query(fragment) + fuzzy_tokens = tokenize_fuzzy_query(fragment) + score = ( + _token_overlap(profile.query_tokens, tokens) * 1.8 + + _token_overlap(profile.query_fuzzy_tokens, fuzzy_tokens) * 1.0 + + _token_overlap(profile.expansion_tokens, tokens) * 2.2 + + _token_overlap(profile.expansion_fuzzy_tokens, fuzzy_tokens) * 1.4 + ) + if score <= 0.0: + continue + if corrected_tokens and not ( + _token_overlap(corrected_tokens, tokens) > 0 + or _token_overlap(corrected_fuzzy, fuzzy_tokens) > 0 + ): + continue + if not profile.expansion_tokens and question.lower() not in fragment: + continue + fragments.append((fragment, score)) + fragments.sort(key=lambda item: (-item[1], -len(item[0]), item[0])) + return [fragment for fragment, _ in fragments[:3]] + + +def _split_anchor_fragments(text: str) -> list[str]: + return [ + fragment.strip().lower() + for fragment in re.split(r"[\n.!?;,]", text) + if len(fragment.strip()) >= 12 + ] + + +def _should_keep_corpus_anchor_token( + token: str, + profile: QuestionProfile, + named_lower: list[str], + token_df: dict[str, int], +) -> bool: + blocked = { + "about", + "again", + "along", + "been", + "being", + "great", + "into", + "just", + "look", + "looking", + "really", + "said", + "shared", + "talked", + "their", + "them", + "they", + } + return ( + len(token) >= 4 + and token not in blocked + and token not in profile.query_tokens + and token not in profile.expansion_tokens + and token not in named_lower + and token_df.get(token, 0) <= 8 + ) + + +def _should_keep_corpus_anchor_phrase( + phrase: str, + profile: QuestionProfile, + token_df: dict[str, int], +) -> bool: + named_lower = [name.lower() for name in profile.named_entities] + tokens = tokenize_query(phrase) + return len(tokens) >= 2 and all( + _should_keep_corpus_anchor_token(token, profile, named_lower, token_df) + or token in profile.expansion_tokens + for token in tokens + ) + + +def select_scoped_candidate_files( + question: str, + profile: QuestionProfile, + files: list[RetrievedMemoryFile], + *, + knowledge_root: str | Path | None = None, +) -> list[CandidateFile]: + candidates: list[CandidateFile] = [] + normalized_query = question.strip().lower() + query_tokens = _memory_header_tokens(normalized_query) + header_files = [_rust_header_view(file, knowledge_root) for file in files] + for relative_dir, budget, boost in scoped_budgets(profile): + scoped_files = [ + file + for file in header_files + if _relative_memory_path(file.file_path, knowledge_root).startswith( + f"{relative_dir}/" + ) + ] + projection_limit = None if relative_dir in {"wiki/observations", "wiki/turns"} else 200 + projected = ( + sorted(scoped_files, key=lambda file: (-file.mtime_ms, file.filename)) + if projection_limit is None + else sorted(scoped_files, key=lambda file: (-file.mtime_ms, file.filename))[ + :projection_limit + ] + ) + scored = [ + (score, file) + for file in projected + if ( + score := _score_memory_header(normalized_query, query_tokens, file) + ) + > 0.0 + ] + selected = _select_confident_root_files( + scored, + max(budget, 2), + 4.0, + ) + if not selected: + # memdir falls back to body scoring when no header reaches the + # confidence threshold. Keep the same projected header set and + # 4.5 minimum used by Rust's fallback selector. + body_scored = [ + (score, file) + for file in projected + if ( + score := _score_memory_body(normalized_query, query_tokens, file) + ) + > 0.0 + ] + selected = _select_confident_root_files( + body_scored, + max(budget, 2), + 4.5, + ) + candidates.extend( + CandidateFile(file=file, query_hits=1, seed_boost=boost) + for file in selected + ) + return _merge_candidate_files([], candidates) + + +def build_session_source_files(files: list[RetrievedMemoryFile]) -> list[SessionSourceFile]: + sources = [] + for file in files: + normalized_path = normalize_memory_path(file.file_path) + if "/wiki/sources/" not in normalized_path: + continue + session_number = _parse_session_source_file_number(normalized_path) + if session_number is None: + continue + search_text = _build_session_source_search_text(file.content) + source_file = replace( + file, + description=file.description or f"session {session_number} summary and turn index", + ) + sources.append( + SessionSourceFile( + session_number=session_number, + file=source_file, + search_text=search_text, + search_tokens=tokenize_query(search_text), + search_fuzzy_tokens=tokenize_fuzzy_query(search_text), + ) + ) + return sources + + +def rank_global_session_sources( + question: str, + profile: QuestionProfile, + sources: list[SessionSourceFile], +) -> list[SessionSourceFile]: + phrases = significant_phrases(question) + quoted = _exact_quoted_phrases(question) + ngrams = keyword_ngrams(question) + named_lower = [name.lower() for name in profile.named_entities] + return sorted( + sources, + key=lambda source: ( + -_score_session_source(phrases, quoted, ngrams, named_lower, profile, source), + -source.file.mtime_ms, + source.file.file_path, + ), + ) + + +def _compose_augmented_retrieval_query( + question: str, + augmentation: QueryAugmentation, +) -> str: + """Match Rust's scoped-header query composition after corpus correction.""" + + lower_question = question.lower() + parts = [question.strip()] + parts.extend( + phrase for phrase in augmentation.phrases if phrase not in lower_question + ) + parts.extend( + token + for token in augmentation.tokens + if not re.search(rf"\b{re.escape(token)}\b", lower_question) + ) + return " ".join(part for part in parts if part) + + +def source_injection_budget(profile: QuestionProfile) -> int: + if profile.aggregate: + return 8 + if profile.temporal or profile.identity or profile.hypothetical or profile.relational: + return 4 + return 3 + + +def source_companion_budget( + profile: QuestionProfile, + *, + has_plugin_retrieval: bool, + top_k: int, +) -> int: + if not has_plugin_retrieval: + budget = 2 + elif ( + profile.aggregate + or profile.identity + or profile.hypothetical + or profile.relational + or profile.location + ): + budget = 4 + else: + budget = 3 + return min(budget, max(top_k, 1)) + + +def select_diverse_session_sources( + profile: QuestionProfile, + root_files: list[RetrievedMemoryFile], + ranked_sources: list[SessionSourceFile], + budget: int, +) -> list[SessionSourceFile]: + if budget == 0 or not ranked_sources: + return [] + if ( + profile.temporal + or profile.location + or ( + not profile.hypothetical + and not profile.identity + and not profile.relational + and not profile.aggregate + ) + ): + return ranked_sources[:budget] + + root_contexts = [ + ( + infer_session_number_from_path(file.file_path), + tokenize_query(f"{file.description or ''} {file.content}"), + ) + for file in root_files + ] + max_rank = float(max(len(ranked_sources), 1)) + selected = [ranked_sources[0]] + selected_indices = {0} + selected_sessions = {ranked_sources[0].session_number} + + while len(selected) < budget and len(selected) < len(ranked_sources): + best: tuple[float, int, SessionSourceFile] | None = None + for index, source in enumerate(ranked_sources): + if index in selected_indices: + continue + rank_score = 3.0 - float(index) * (1.4 / max_rank) + query_alignment = _source_query_alignment_score(profile, source) + root_redundancy = max( + ( + _source_redundancy_penalty( + source, + session_number, + tokens, + source.session_number, + ) + for session_number, tokens in root_contexts + ), + default=0.0, + ) + selected_redundancy = max( + ( + _source_redundancy_penalty( + source, + chosen.session_number, + chosen.search_tokens, + chosen.session_number, + ) + for chosen in selected + ), + default=0.0, + ) + repeated_session_penalty = ( + 1.2 if source.session_number in selected_sessions else 0.0 + ) + anchor_miss_penalty = _source_semantic_anchor_miss_penalty( + profile, + source.search_text, + ) + score = ( + rank_score + + query_alignment + - root_redundancy + - selected_redundancy + - repeated_session_penalty + - anchor_miss_penalty + ) + # Rust's comparator prefers the later ranked index on an exact + # diversity-score tie (`right.index.cmp(left.index)`). + key = (score, index, source) + if best is None or key[:2] > best[:2]: + best = key + if best is None: + break + _, index, source = best + selected_indices.add(index) + selected_sessions.add(source.session_number) + selected.append(source) + return selected + + +def collect_session_source_companions( + question: str, + profile: QuestionProfile, + candidates: list[CandidateFile], + sources_by_session: dict[int, SessionSourceFile], +) -> list[CandidateFile]: + phrases = significant_phrases(question) + quoted = _exact_quoted_phrases(question) + ngrams = keyword_ngrams(question) + named_lower = [name.lower() for name in profile.named_entities] + companions: dict[str, CandidateFile] = {} + source_bonus_by_session: dict[int, float] = {} + + for candidate in candidates: + for session_number, query_hits, seed_boost in _infer_session_source_signals( + candidate, + profile, + phrases, + named_lower, + quoted, + ngrams, + sources_by_session, + ): + source = sources_by_session.get(session_number) + if source is None: + continue + source_bonus = source_bonus_by_session.setdefault( + session_number, + _source_companion_relevance_bonus( + source, + profile, + phrases, + quoted, + ngrams, + named_lower, + ), + ) + key = normalize_memory_path(source.file.file_path) + proposal = CandidateFile( + file=source.file, + query_hits=max(query_hits, int(source_bonus > 0.0)), + seed_boost=seed_boost + source_bonus, + ) + existing = companions.get(key) + if existing is None or (proposal.seed_boost, proposal.query_hits) > ( + existing.seed_boost, + existing.query_hits, + ): + companions[key] = proposal + + return sorted( + companions.values(), + key=lambda candidate: ( + -candidate.seed_boost, + -candidate.query_hits, + normalize_memory_path(candidate.file.file_path), + ), + ) + + +def infer_session_number_from_path(path: str) -> int | None: + normalized = path.replace("\\", "/") + for prefix in ("session_", "/D", "\\D"): + index = normalized.find(prefix) + if index < 0: + continue + suffix = normalized[index + len(prefix) :] + digits = re.match(r"\d+", suffix) + if digits: + return int(digits.group(0)) + return None + + +def _is_rust_cached_retrieval_file(file: RetrievedMemoryFile) -> bool: + """Match the Rust retained evaluator's cached retrieval file projection.""" + + path = normalize_memory_path(file.file_path) + return any( + path.startswith(marker) or f"/{marker}" in path + for marker in ( + "wiki/sources/", + "wiki/observations/", + "wiki/events/", + "wiki/entities/", + "wiki/turns/", + "wiki/memories/", + "wiki/memory/", + ) + ) + + +def _rank_initial_root_files( + question: str, + files: list[RetrievedMemoryFile], + limit: int, + knowledge_root: str | Path | None = None, +) -> list[RetrievedMemoryFile]: + # Rust builds a bounded root-header projection before selecting the initial + # files. Keep late entity/observation pages in the scoped candidate pool; + # otherwise they can displace the source/turn roots before source rescue. + # `scan_memory_directory_with_limit` only indexes markdown topic files; + # MEMORY.md is an entrypoint and raw JSON is never part of Rust's header + # projection. Excluding them here prevents Python-only root candidates. + header_files = [ + _rust_header_view(file, knowledge_root) + for file in files + if _is_rust_scanned_markdown(file, knowledge_root) + ] + # Rust orders the bounded projection by descending mtime and then relative + # filename. The Python builder stores deterministic insertion timestamps; + # retain the same comparator rather than letting absolute paths or the + # basename-only `filename` field change the projection. + projected_files = sorted( + header_files, + key=lambda file: (-file.mtime_ms, file.filename), + )[:200] + normalized_query = question.strip().lower() + query_tokens = _memory_header_tokens(normalized_query) + scored = [ + (score, file) + for file in projected_files + if (score := _score_memory_header(normalized_query, query_tokens, file)) > 0.0 + ] + selected = _select_confident_root_files(scored, limit, 4.0) + if selected: + return selected + # memdir retries with body scoring when header selection is empty. + body_scored = [ + (score, file) + for file in projected_files + if (score := _score_memory_body(normalized_query, query_tokens, file)) > 0.0 + ] + return _select_confident_root_files(body_scored, limit, 4.5) + + +def _relative_memory_path( + file_path: str, + knowledge_root: str | Path | None, +) -> str: + normalized = normalize_memory_path(file_path).rstrip("/") + if knowledge_root is None: + return normalized + root = normalize_memory_path(str(knowledge_root)).rstrip("/") + prefix = f"{root}/" + return normalized[len(prefix) :] if normalized.startswith(prefix) else normalized + + +def _rust_header_view( + file: RetrievedMemoryFile, + knowledge_root: str | Path | None, +) -> RetrievedMemoryFile: + relative_path = _relative_memory_path(file.file_path, knowledge_root) + return replace( + file, + filename=relative_path, + description=_frontmatter_value(file.content, "description") or file.description or None, + ) + + +def _score_memory_header( + normalized_query: str, + query_tokens: set[str], + file: RetrievedMemoryFile, +) -> float: + filename = file.filename.lower() + description = (file.description or _frontmatter_value(file.content, "description")).lower() + type_label = _frontmatter_value(file.content, "type") or "general" + score = 0.0 + if normalized_query in filename: + score += 4.0 + if normalized_query in description: + score += 3.5 + if filename in normalized_query: + score += 1.5 + if description in normalized_query: + score += 1.0 + score += len(query_tokens & _memory_header_tokens(filename)) * 2.5 + score += len(query_tokens & _memory_header_tokens(description)) * 2.0 + score += len(query_tokens & _memory_header_tokens(type_label)) + if _contains_reference_warning_signal(f"{filename} {description}"): + score += 2.5 + return score + + +def _score_memory_body( + normalized_query: str, + query_tokens: set[str], + file: RetrievedMemoryFile, +) -> float: + normalized_content = file.content.lower() + score = _score_memory_header(normalized_query, query_tokens, file) + if normalized_query in normalized_content: + score += 4.0 + if normalized_content in normalized_query: + score += 1.0 + score += len(query_tokens & _memory_header_tokens(normalized_content)) * 1.5 + return score + + +def _select_confident_root_files( + scored: list[tuple[float, RetrievedMemoryFile]], + limit: int, + minimum: float, + *, + prefer_recent: bool = True, +) -> list[RetrievedMemoryFile]: + scored.sort( + key=lambda item: ( + -item[0], + -item[1].mtime_ms if prefer_recent else 0, + item[1].filename, + ) + ) + if not scored or scored[0][0] < minimum: + return [] + cutoff = max(minimum, scored[0][0] * 0.45) + return [file for score, file in scored if score >= cutoff][:limit] + + +def _frontmatter_value(content: str, key: str) -> str: + match = re.search(rf"(?im)^{re.escape(key)}:\s*(.+)$", content) + return match.group(1).strip() if match else "" + + +def _is_rust_scanned_markdown( + file: RetrievedMemoryFile, + knowledge_root: str | Path | None, +) -> bool: + """Match memdir's topic-file inclusion predicate for the virtual wiki.""" + + relative = _relative_memory_path(file.file_path, knowledge_root).strip("/") + if not relative.lower().endswith(".md"): + return False + if PurePosixPath(relative).name.lower() == "memory.md": + return False + parts = relative.split("/") + if len(parts) == 4 and parts[0].lower() == "logs" and parts[3].lower().endswith(".md"): + # Rust excludes daily logs from the topic scan. The retained fixture + # does not currently create these, but keeping the predicate exact + # avoids a hidden Python-only root candidate. + year, month, filename = parts[1], parts[2], parts[3] + stem = filename[:-3] + if ( + len(year) == 4 + and year.isdigit() + and len(month) == 2 + and month.isdigit() + and re.fullmatch(r"\d{4}-\d{2}-\d{2}", stem) + ): + return False + return True + + +def _file_description(file: RetrievedMemoryFile) -> str: + # Rust gives source companions their synthesized "summary and turn index" + # description when loading them, while other cached files use parsed + # frontmatter metadata. Preserve that distinction for candidate scoring. + normalized_path = normalize_memory_path(file.file_path) + if "/wiki/sources/" in normalized_path: + return file.description or _frontmatter_value(file.content, "description") + return _frontmatter_value(file.content, "description") or file.description + + +def _memory_header_tokens(text: str) -> set[str]: + blocked = { + "a", + "an", + "and", + "at", + "do", + "for", + "help", + "how", + "i", + "if", + "in", + "is", + "it", + "me", + "my", + "of", + "on", + "or", + "our", + "please", + "should", + "that", + "the", + "this", + "to", + "use", + "what", + "when", + "with", + } + tokens = set() + for raw in re.split(r"[^0-9A-Za-z_-]+", text): + token = raw.lower() + if len(token) <= 1 or token in blocked: + continue + for suffix, replacement, minimum in ( + ("ies", "y", 4), + ("ing", "", 5), + ("ly", "", 4), + ("ed", "", 4), + ("es", "", 4), + ("s", "", 3), + ): + if token.endswith(suffix) and len(token) > minimum: + token = token[: -len(suffix)] + replacement + break + tokens.add(token) + return tokens + + +def _contains_reference_warning_signal(text: str) -> bool: + lowered = text.lower() + return any( + signal in lowered + for signal in ( + "warning", + "warn", + "gotcha", + "known issue", + "pitfall", + "danger", + "avoid", + "careful", + "caution", + ) + ) + + +def _push_unique_file( + output: list[RetrievedMemoryFile], + included: set[str], + file: RetrievedMemoryFile, + limit: int, +) -> None: + if len(output) >= limit: + return + key = normalize_memory_path(file.file_path) + if key in included: + return + included.add(key) + output.append(file) + + +def _merge_candidate_files( + base: list[CandidateFile], + additions: list[CandidateFile], +) -> list[CandidateFile]: + merged: dict[str, CandidateFile] = {} + order: list[str] = [] + for candidate in base + additions: + key = normalize_memory_path(candidate.file.file_path) + existing = merged.get(key) + if existing is None: + order.append(key) + merged[key] = candidate + continue + selected_file = ( + candidate.file + if candidate.file.mtime_ms > existing.file.mtime_ms + else existing.file + ) + merged[key] = CandidateFile( + file=selected_file, + query_hits=existing.query_hits + candidate.query_hits, + seed_boost=max(existing.seed_boost, candidate.seed_boost), + ) + return [merged[key] for key in order] + + +def _rerank_candidate_files( + question: str, + profile: QuestionProfile, + candidates: list[CandidateFile], + proposals, + files_by_path: dict[str, RetrievedMemoryFile], + *, + knowledge_root: str | Path | None = None, +) -> list[CandidateFile]: + by_path = { + normalize_memory_path(candidate.file.file_path): candidate + for candidate in candidates + } + for proposal in proposals: + existing = by_path.get(proposal.file_path) + if existing is not None: + by_path[proposal.file_path] = CandidateFile( + file=existing.file, + query_hits=max(existing.query_hits, proposal.query_hits), + seed_boost=existing.seed_boost + proposal.seed_boost, + ) + continue + file = files_by_path.get(proposal.file_path) + if file is not None: + by_path[proposal.file_path] = CandidateFile( + file=file, + query_hits=proposal.query_hits, + seed_boost=proposal.seed_boost, + ) + return _rank_candidate_files( + question, + profile, + list(by_path.values()), + knowledge_root=knowledge_root, + ) + + +def _rank_candidate_files( + question: str, + profile: QuestionProfile, + candidates: list[CandidateFile], + *, + knowledge_root: str | Path | None = None, +) -> list[CandidateFile]: + phrases = significant_phrases(question) + named_lower = [name.lower() for name in profile.named_entities] + + def score(candidate: CandidateFile) -> float: + relative_path = _relative_memory_path(candidate.file.file_path, knowledge_root) + path_markers = f"/{relative_path}" + meta_text = f"{relative_path} {_file_description(candidate.file)}".lower() + meta_tokens = set(tokenize_query(meta_text)) + features = build_cached_file_lexical_features(candidate.file) + value = ( + candidate.query_hits * 5.0 + + candidate.seed_boost + + _token_overlap(profile.query_tokens, meta_tokens) * 3.0 + + _token_overlap(profile.expansion_tokens, meta_tokens) * 1.5 + + _token_overlap(profile.query_tokens, features.content_token_set) + + _token_overlap(profile.expansion_tokens, features.content_token_set) * 0.7 + + max( + ( + _rust_candidate_line_score(line, profile, phrases) + for line in features.line_features + ), + default=0.0, + ) + ) + for name in named_lower: + if _word_boundary_contains(features.lower_content, name): + value += 3.0 + if _word_boundary_contains(meta_text, name): + value += 3.0 + if profile.temporal and any( + part in path_markers + for part in ( + "/wiki/observations/", + "/wiki/events/", + "/wiki/sources/", + "/wiki/turns/", + "/wiki/memory/", + ) + ): + value += 3.0 + if profile.identity and any( + part in path_markers + for part in ("/wiki/entities/", "/wiki/observations/", "/wiki/memory/") + ): + value += 4.0 + if (profile.hypothetical or profile.relational) and any( + part in path_markers + for part in ("/wiki/entities/", "/wiki/observations/", "/wiki/memory/") + ): + value += 2.0 + if profile.location and any( + part in path_markers + for part in ("/wiki/observations/", "/wiki/events/", "/wiki/memory/") + ): + value += 2.0 + if "/wiki/sources/" in path_markers: + value += 4.0 + if profile.temporal or profile.location: + value += 3.0 + if profile.identity or profile.hypothetical or profile.relational: + value += 2.0 + if "/wiki/turns/" in path_markers: + value += 1.0 + if profile.temporal or profile.location: + value += 1.0 + if profile.identity or profile.hypothetical or profile.relational: + value += 1.0 + for phrase in phrases: + if len(phrase) < 6: + continue + if phrase in features.lower_content: + value += 4.0 + if phrase in meta_text: + value += 3.0 + if features.has_session_marker: + value += 1.0 + if features.has_evidence_marker: + value += 4.0 + return value + + return sorted( + candidates, + key=lambda candidate: ( + -score(candidate), + -candidate.file.mtime_ms, + candidate.file.file_path, + ), + ) + + +def _rust_candidate_line_score( + line, + profile: QuestionProfile, + phrases: list[str], +) -> float: + """Mirror retained_eval's candidate-ranking line scorer (not qmd plugin scoring).""" + + score = ( + _token_overlap(profile.query_tokens, line.token_set) * 2.0 + + _token_overlap(profile.query_fuzzy_tokens, line.fuzzy_token_set) * 1.5 + + _token_overlap(profile.expansion_tokens, line.token_set) * 1.4 + + _token_overlap(profile.expansion_fuzzy_tokens, line.fuzzy_token_set) * 0.8 + ) + score += sum( + 2.0 + for entity in profile.named_entities + if _word_boundary_contains(line.lower_text, entity.lower()) + ) + score += sum( + 3.0 for phrase in phrases if len(phrase) >= 6 and phrase in line.lower_text + ) + score += sum( + 1.8 + for phrase in profile.expansion_phrases + if len(phrase) >= 6 and phrase in line.lower_text + ) + return score + + +def _word_boundary_contains(text: str, value: str) -> bool: + return re.search(rf"\b{re.escape(value)}\b", text) is not None + + +def _collect_late_bridge_seed_files( + root_files: list[RetrievedMemoryFile], + ranked_candidates: list[CandidateFile], + ranked_source_companions: list[CandidateFile], +) -> list[RetrievedMemoryFile]: + files: list[RetrievedMemoryFile] = [] + included: set[str] = set() + for file in root_files[:4]: + _push_unique_file(files, included, file, 12) + for candidate in ranked_candidates[:4]: + _push_unique_file(files, included, candidate.file, 12) + for candidate in ranked_source_companions[:4]: + _push_unique_file(files, included, candidate.file, 12) + return files + + +def _score_scoped_candidate( + file: RetrievedMemoryFile, + profile: QuestionProfile, + phrases: list[str], + boost: float, +) -> tuple[CandidateFile, float]: + features = build_cached_file_lexical_features(file) + best_line = max( + ( + _score_query_line_context( + line.lower_text, + profile, + phrases, + [name.lower() for name in profile.named_entities], + ) + for line in features.line_features + ), + default=0.0, + ) + query_hits = max(candidate_query_hits_with_features(features, profile), int(best_line > 0.0)) + if query_hits == 0 and best_line <= 0.0: + return CandidateFile(file=file, query_hits=0, seed_boost=0.0), 0.0 + score = best_line + boost + query_hits * 0.35 + return CandidateFile(file=file, query_hits=query_hits, seed_boost=boost), score + + +def _parse_session_source_file_number(normalized_path: str) -> int | None: + file_name = PurePosixPath(normalized_path).name + match = re.fullmatch(r"session_(\d+)\.md", file_name) + return int(match.group(1)) if match else None + + +def _build_session_source_search_text(content: str) -> str: + summary = _extract_source_summary_snippet(content) + turn_index = _extract_source_turn_index_snippet(content) + return f"{summary} {turn_index}".lower() + + +def _extract_source_summary_snippet(content: str) -> str: + section = _extract_section_after_heading(content, "## Summary") + text = section if section is not None else content + return text.replace("\r\n", "\n").replace("\n", " ").strip()[:700] + + +def _extract_source_turn_index_snippet(content: str) -> str: + section = _extract_section_after_heading(content, "## Turn Index") or "" + lines = [line.strip() for line in section.splitlines()] + return " ".join(line for line in lines if line.startswith("- [turn "))[:2500] + + +def _extract_section_after_heading(content: str, heading: str) -> str | None: + normalized = content.replace("\r\n", "\n") + start = normalized.find(heading) + if start < 0: + return None + tail = normalized[start + len(heading) :] + next_heading = tail.find("\n## ") + if next_heading >= 0: + tail = tail[:next_heading] + return tail.strip() + + +def _score_session_source( + phrases: list[str], + quoted: list[str], + ngrams: list[str], + named_lower: list[str], + profile: QuestionProfile, + source: SessionSourceFile, +) -> float: + score = ( + _token_overlap(profile.query_tokens, source.search_tokens) * 4.0 + + _token_overlap(profile.query_fuzzy_tokens, source.search_fuzzy_tokens) * 1.7 + + _token_overlap(profile.expansion_tokens, source.search_tokens) * 1.2 + + _token_overlap(profile.expansion_fuzzy_tokens, source.search_fuzzy_tokens) * 0.8 + ) + score += sum(8.0 for phrase in phrases if len(phrase) >= 6 and phrase in source.search_text) + score += sum(12.0 for phrase in quoted if len(phrase) >= 3 and phrase in source.search_text) + score += sum(5.0 for ngram in ngrams if len(ngram) >= 5 and ngram in source.search_text) + score += sum( + 6.0 + for phrase in profile.expansion_phrases + if len(phrase) >= 6 and phrase in source.search_text + ) + score += sum(5.0 for entity in named_lower if entity in source.search_text) + if profile.temporal: + score += 2.0 + if profile.aggregate: + score += 3.0 + return score + + +def _source_query_alignment_score(profile: QuestionProfile, source: SessionSourceFile) -> float: + score = ( + _token_overlap(profile.query_tokens, source.search_tokens) * 0.35 + + _token_overlap(profile.query_fuzzy_tokens, source.search_fuzzy_tokens) * 0.12 + + _token_overlap(profile.expansion_tokens, source.search_tokens) * 1.8 + + _token_overlap(profile.expansion_fuzzy_tokens, source.search_fuzzy_tokens) * 0.65 + ) + score += sum( + 2.0 + for phrase in profile.expansion_phrases + if len(phrase) >= 6 and phrase in source.search_text + ) + return score + + +def _source_semantic_anchor_miss_penalty( + profile: QuestionProfile, + search_text: str, +) -> float: + if not profile.expansion_tokens and not profile.expansion_phrases: + return 0.0 + search_tokens = tokenize_query(search_text) + expansion_hits = _token_overlap(profile.expansion_tokens, search_tokens) + phrase_hits = sum( + 1 + for phrase in profile.expansion_phrases + if len(phrase) >= 6 and phrase in search_text + ) + return 1.2 if expansion_hits == 0 and phrase_hits == 0 else 0.0 + + +def _source_redundancy_penalty( + source: SessionSourceFile, + other_session_number: int | None, + other_tokens: list[str], + selected_session_number: int | None, +) -> float: + shared = float(_token_overlap(source.search_tokens, other_tokens)) + normalized_overlap = shared / float(max(len(source.search_tokens), 6)) + session_penalty = 1.6 if other_session_number == selected_session_number else 0.0 + return min(shared, 4.0) * 0.35 + normalized_overlap * 1.3 + session_penalty + + +def _source_companion_relevance_bonus( + source: SessionSourceFile, + profile: QuestionProfile, + phrases: list[str], + quoted: list[str], + ngrams: list[str], + named_lower: list[str], +) -> float: + return min( + _score_session_source(phrases, quoted, ngrams, named_lower, profile, source) * 0.4, + 8.0, + ) + + +def _infer_session_source_signals( + candidate: CandidateFile, + profile: QuestionProfile, + phrases: list[str], + named_lower: list[str], + quoted: list[str], + ngrams: list[str], + sources_by_session: dict[int, SessionSourceFile], +) -> list[tuple[int, int, float]]: + session_number = infer_session_number_from_path(candidate.file.file_path) + if session_number is not None: + return [ + ( + session_number, + max(candidate.query_hits, 1), + 6.0 + min(candidate.query_hits, 2) + min(candidate.seed_boost, 2.0), + ) + ] + if not _should_infer_sessions_from_content(candidate.file.file_path): + return [] + + scored = _score_content_derived_session_mentions( + candidate.file.content, + profile, + phrases, + named_lower, + ) + if scored: + base_seed = 5.5 + min(candidate.seed_boost, 2.0) + ranked = [] + for mentioned_session, score in scored: + source = sources_by_session.get(mentioned_session) + source_bonus = ( + _source_companion_relevance_bonus( + source, + profile, + phrases, + quoted, + ngrams, + named_lower, + ) + if source is not None + else 0.0 + ) + ranked.append( + (mentioned_session, score, source_bonus, base_seed + score + source_bonus) + ) + ranked.sort(key=lambda item: (-item[3], item[0])) + selected = [ + (session_number, int(score > 0.0 or source_bonus > 0.0), base_seed + score) + for session_number, score, source_bonus, _ in ranked[:5] + ] + selected_sessions = {session_number for session_number, _, _ in selected} + bonus_backfill = [ + item + for item in ranked + if item[2] >= 4.5 and item[0] not in selected_sessions + ] + if bonus_backfill: + session_number, score, source_bonus, _ = max( + bonus_backfill, + key=lambda item: (item[2], -item[1], -item[0]), + ) + selected.append( + (session_number, int(score > 0.0 or source_bonus > 0.0), base_seed + score) + ) + return selected + + return [ + (session_number, 0, 5.5 + min(candidate.seed_boost, 2.0)) + for session_number in sorted(_collect_session_numbers_from_text(candidate.file.content))[:6] + ] + + +def _score_content_derived_session_mentions( + text: str, + profile: QuestionProfile, + phrases: list[str], + named_lower: list[str], +) -> list[tuple[int, float]]: + best_scores: dict[int, float] = {} + for line in text.splitlines(): + mentioned_sessions = _collect_session_numbers_from_text(line) + if not mentioned_sessions: + continue + score = _score_query_line_context(line, profile, phrases, named_lower) + for session_number in mentioned_sessions: + best_scores[session_number] = max(best_scores.get(session_number, score), score) + return sorted(best_scores.items(), key=lambda item: (-item[1], item[0])) + + +def _score_query_line_context( + line: str, + profile: QuestionProfile, + phrases: list[str], + named_lower: list[str], +) -> float: + lower = line.lower() + tokens = tokenize_query(line) + fuzzy_tokens = tokenize_fuzzy_query(line) + score = ( + _token_overlap(profile.query_tokens, tokens) * 2.0 + + _token_overlap(profile.query_fuzzy_tokens, fuzzy_tokens) * 1.5 + + _token_overlap(profile.expansion_tokens, tokens) * 1.4 + + _token_overlap(profile.expansion_fuzzy_tokens, fuzzy_tokens) * 0.8 + ) + score += sum( + 2.0 + for entity in named_lower + if re.search(rf"(?= 6 and phrase in lower) + score += sum( + 1.8 + for phrase in profile.expansion_phrases + if len(phrase) >= 6 and phrase in lower + ) + return score + + +def _collect_session_numbers_from_text(text: str) -> set[int]: + session_numbers = {int(match) for match in re.findall(r"session_(\d+)", text)} + session_numbers.update( + int(match) + for match in re.findall(r"\bD(\d+)(?=[:_])", text) + ) + return session_numbers + + +def _should_infer_sessions_from_content(path: str) -> bool: + normalized = normalize_memory_path(path) + return ( + normalized.endswith("/memory.md") + or "/wiki/entities/" in normalized + or "/wiki/memory/" in normalized + ) + + +def _exact_quoted_phrases(question: str) -> list[str]: + return [match.strip().lower() for match in re.findall(r'"([^"]+)"', question) if match.strip()] + + +def _token_overlap(left: list[str], right: list[str]) -> int: + right_set = set(right) + return sum(1 for token in left if token in right_set) + + +def _bounded_edit_distance(left: str, right: str, max_distance: int) -> int | None: + if abs(len(left) - len(right)) > max_distance: + return None + previous = list(range(len(right) + 1)) + for left_index, left_char in enumerate(left, start=1): + current = [left_index] + row_min = current[0] + for right_index, right_char in enumerate(right, start=1): + substitution_cost = 0 if left_char == right_char else 1 + value = min( + previous[right_index] + 1, + current[right_index - 1] + 1, + previous[right_index - 1] + substitution_cost, + ) + current.append(value) + row_min = min(row_min, value) + if row_min > max_distance: + return None + previous = current + distance = previous[-1] + return distance if distance <= max_distance else None diff --git a/evaluation/wikimem/wiki_builder.py b/evaluation/wikimem/wiki_builder.py new file mode 100644 index 00000000..dc410560 --- /dev/null +++ b/evaluation/wikimem/wiki_builder.py @@ -0,0 +1,586 @@ +"""Dataset-independent Wiki construction for wikimem. + +The builder exposes two explicit paths. ``mode="llm"`` compiles raw source +blocks into an ontology, resolves entities, and consolidates duplicate +memories. ``mode="deterministic"`` is a small, reproducible fallback for +regression and for environments where no model is configured. The benchmark +adapter keeps its historical renderer for exact Rust-compatible baselines; +this module is the open-domain construction path. +""" + +from __future__ import annotations + +import hashlib +import json +import re +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Iterable, Literal + +from common.llm.base import LLM +from common.type_def.chat import ChatMessage + +from evaluation.wikimem.llm_semantics import ( + MEMORY_KINDS, + SemanticEntity, + SemanticMemory, + SemanticSource, + _parse_json_payload, + extract_semantic_memories, +) +from evaluation.wikimem.qmd_consensus import RetrievedMemoryFile + + +WikiBuilderMode = Literal["deterministic", "llm"] + + +@dataclass(frozen=True) +class WikiBuildDiagnostics: + mode: WikiBuilderMode + llm_used: bool + source_count: int + extracted_count: int + canonical_entity_count: int + consolidated_count: int + fallback_reason: str = "" + + +@dataclass(frozen=True) +class WikiBuildResult: + files: list[RetrievedMemoryFile] + memories: list[SemanticMemory] + entities: list[SemanticEntity] + synthesis: dict[str, str] + diagnostics: WikiBuildDiagnostics + + +class EntityResolver: + """Resolve obvious aliases while retaining the original display names.""" + + _STOPWORDS = {"the", "a", "an", "inc", "corp", "corporation", "company", "co", "model"} + + def __init__(self) -> None: + self._canonical: dict[str, SemanticEntity] = {} + self._aliases: dict[str, str] = {} + + def resolve( + self, + entities: Iterable[SemanticEntity], + ) -> tuple[list[SemanticEntity], dict[str, str]]: + for entity in entities: + key = self._key(entity.name) + if not key: + continue + canonical = self._aliases.get(key) + if canonical is None: + canonical = f"entity.{self._slug(key)}" + self._aliases[key] = canonical + self._canonical[canonical] = replace( + entity, + name=canonical, + aliases=tuple(dict.fromkeys((entity.name, *entity.aliases))), + ) + else: + current = self._canonical[canonical] + description = current.description or entity.description + entity_type = ( + current.entity_type + if current.entity_type != "thing" + else entity.entity_type + ) + self._canonical[canonical] = replace( + current, + entity_type=entity_type, + description=description, + aliases=tuple(dict.fromkeys((*current.aliases, entity.name, *entity.aliases))), + ) + return list(self._canonical.values()), dict(self._aliases) + + def rewrite_memory(self, memory: SemanticMemory, aliases: dict[str, str]) -> SemanticMemory: + entities = tuple( + replace(entity, name=aliases.get(self._key(entity.name), entity.name)) + for entity in memory.entities + ) + relations = tuple( + replace( + relation, + subject=aliases.get(self._key(relation.subject), relation.subject), + object=aliases.get(self._key(relation.object), relation.object), + ) + for relation in memory.relations + ) + return replace(memory, entities=entities, relations=relations) + + @classmethod + def _key(cls, value: str) -> str: + words = re.findall(r"[\w]+", value.casefold(), flags=re.UNICODE) + return " ".join(word for word in words if word not in cls._STOPWORDS) + + @staticmethod + def _slug(value: str) -> str: + slug = re.sub(r"[^\w-]+", "-", value.casefold(), flags=re.UNICODE).strip("-")[:80] + if slug: + return slug + return f"unknown-{hashlib.sha1(value.encode('utf-8')).hexdigest()[:12]}" + + +class MemoryConsolidator: + """Merge duplicate semantic records and optionally synthesize summaries.""" + + def consolidate(self, memories: Iterable[SemanticMemory]) -> list[SemanticMemory]: + merged: dict[tuple[str, str], SemanticMemory] = {} + for memory in memories: + entity_key = ",".join(sorted(entity.name for entity in memory.entities)) + content_key = re.sub(r"\s+", " ", memory.content.casefold()).strip() + key = (memory.kind, entity_key + "|" + content_key) + previous = merged.get(key) + if previous is None: + merged[key] = memory + continue + evidence = previous.evidence or memory.evidence + source_id = previous.source_id if previous.source_id == memory.source_id else ( + f"{previous.source_id},{memory.source_id}" + ) + merged[key] = replace( + previous, + source_id=source_id, + evidence=evidence, + confidence=max(previous.confidence, memory.confidence), + tags=tuple(dict.fromkeys((*previous.tags, *memory.tags)))[:8], + ) + return list(merged.values()) + + +class TemplateExtractor: + """Deterministic source-to-memory extractor kept for regression fallback.""" + + def extract(self, sources: Iterable[SemanticSource]) -> list[SemanticMemory]: + result: list[SemanticMemory] = [] + for source in sources: + if not source.text.strip(): + continue + digest = hashlib.sha1( + f"{source.source_id}:{source.text}".encode("utf-8") + ).hexdigest()[:16] + result.append( + SemanticMemory( + memory_id=f"{source.source_id}:context:{digest}", + kind="context", + content=source.text.strip(), + source_id=source.source_id, + evidence=source.text.strip(), + timestamp=source.timestamp, + confidence=1.0, + metadata=dict(source.metadata), + ) + ) + return result + + +class WikiBuilder: + """Compile source blocks into a canonical, provenance-preserving Wiki.""" + + def __init__( + self, + llm: LLM | None = None, + *, + mode: WikiBuilderMode = "llm", + batch_size: int = 8, + max_tokens: int = 4096, + allow_fallback: bool = True, + ) -> None: + if mode not in {"llm", "deterministic"}: + raise ValueError(f"unsupported wiki builder mode: {mode!r}") + self.llm = llm + self.mode = mode + self.batch_size = batch_size + self.max_tokens = max_tokens + self.allow_fallback = allow_fallback + + def build( + self, + sources: Iterable[SemanticSource], + sample_root: str | Path, + *, + wiki_mode: str = "text", + ) -> WikiBuildResult: + source_list = [source for source in sources if source.text.strip()] + fallback_reason = "" + llm_used = self.mode == "llm" and self.llm is not None + if llm_used: + try: + memories = extract_semantic_memories( + self.llm, source_list, batch_size=self.batch_size, max_tokens=self.max_tokens + ) + except Exception as exc: + if not self.allow_fallback: + raise + memories = TemplateExtractor().extract(source_list) + llm_used = False + fallback_reason = f"llm extraction failed: {type(exc).__name__}: {exc}" + else: + memories = TemplateExtractor().extract(source_list) + if self.mode == "llm" and self.llm is None: + fallback_reason = "no LLM configured" + + source_metadata = {source.source_id: source.metadata for source in source_list} + memories = [ + replace( + memory, + metadata={ + **source_metadata.get(memory.source_id.split(",", 1)[0], {}), + **memory.metadata, + }, + ) + for memory in memories + ] + + resolver = EntityResolver() + extracted_count = len(memories) + all_entities = [entity for memory in memories for entity in memory.entities] + entities, aliases = resolver.resolve(all_entities) + memories = [resolver.rewrite_memory(memory, aliases) for memory in memories] + memories = MemoryConsolidator().consolidate(memories) + synthesis: dict[str, str] = {} + if llm_used and self.llm is not None: + try: + synthesis = self._synthesize(memories) + except Exception as exc: + if not self.allow_fallback: + raise + fallback_reason = fallback_reason or ( + f"llm synthesis failed: {type(exc).__name__}: {exc}" + ) + files = self._render( + source_list, + memories, + entities, + sample_root, + wiki_mode=wiki_mode, + synthesis=synthesis, + ) + diagnostics = WikiBuildDiagnostics( + mode=self.mode, + llm_used=llm_used, + source_count=len(source_list), + extracted_count=extracted_count, + canonical_entity_count=len(entities), + consolidated_count=len(memories), + fallback_reason=fallback_reason, + ) + return WikiBuildResult( + files=files, + memories=memories, + entities=entities, + synthesis=synthesis, + diagnostics=diagnostics, + ) + + def _synthesize(self, memories: list[SemanticMemory]) -> dict[str, str]: + if self.llm is None or not memories: + return {} + payload = json.dumps( + [ + { + "kind": memory.kind, + "content": memory.content, + "entities": [entity.name for entity in memory.entities], + "timestamp": memory.timestamp, + "provenance": memory.source_id, + } + for memory in memories + ], + ensure_ascii=False, + ) + response = self.llm.chat( + [ + ChatMessage( + role="system", + content=( + "You synthesize long-term memory. Return only JSON with string keys " + "profile, timeline, decisions. Do not add facts not present in input." + ), + ), + ChatMessage(role="user", content=payload), + ], + temperature=0.0, + max_tokens=min(self.max_tokens, 4096), + ) + parsed = _parse_json_payload(response) + if not isinstance(parsed, dict): + raise ValueError("LLM memory synthesis must return a JSON object") + return { + key: str(parsed.get(key, "")).strip() + for key in ("profile", "timeline", "decisions") + if str(parsed.get(key, "")).strip() + } + + def _render( + self, + sources: list[SemanticSource], + memories: list[SemanticMemory], + entities: list[SemanticEntity], + sample_root: str | Path, + *, + wiki_mode: str, + synthesis: dict[str, str], + ) -> list[RetrievedMemoryFile]: + root = Path(sample_root).as_posix().rstrip("/") + files: list[RetrievedMemoryFile] = [] + source_by_id = {source.source_id: source for source in sources} + index_lines = [ + "# Memory Index", + "", + "This Wiki is compiled from source blocks with explicit provenance.", + "", + ] + for memory in memories: + slug = self._slug(memory.memory_id) + path = f"{root}/wiki/memory/{memory.kind}/{slug}.md" + source_ids = memory.source_id.split(",") + source = source_by_id.get(source_ids[0], SemanticSource(source_ids[0], "")) + entity_names = ", ".join(entity.name for entity in memory.entities) or "(none)" + metadata_lines = [f"{key}: {value}" for key, value in sorted(memory.metadata.items())] + relation_lines = [ + f"- {item.subject} --{item.predicate}--> {item.object}" + for item in memory.relations + ] + related_lines = [ + f"- [source {source_id}](../../sources/{self._slug(source_id)}.md)" + for source_id in source_ids + ] + [ + f"- [entity {entity.name}](../../entities/{self._slug(entity.name)}.md)" + for entity in memory.entities + ] + body = [ + "---", + f"memory_type: {memory.kind}", + f"confidence: {memory.confidence:.3f}", + f"entities: {entity_names}", + f"provenance: {', '.join(source_ids)}", + *metadata_lines, + "---", + f"# {memory.kind.title()} memory", + "", + memory.content, + "", + f"Evidence: {memory.evidence or '(not supplied)'}", + f"Source: {source.source_id} ({source.conversation_id or 'unknown conversation'})", + ] + if relation_lines: + body.extend(["", "## Relations", *relation_lines]) + if related_lines: + body.extend(["", "## Provenance links", *related_lines]) + files.append( + RetrievedMemoryFile( + filename=f"{slug}.md", + file_path=path, + mtime_ms=len(files) + 1, + content="\n".join(body) + "\n", + description=memory.content[:160], + memory_type=memory.kind, + ) + ) + index_lines.append( + f"- [{memory.kind}]({path[len(root) + 1:]}) - {memory.content[:120]}" + ) + + for entity in entities: + slug = self._slug(entity.name) + related = [ + memory + for memory in memories + if any(item.name == entity.name for item in memory.entities) + ] + content = [ + "---", + f"entity_type: {entity.entity_type}", + f"canonical_id: {entity.name}", + f"aliases: {', '.join(entity.aliases) or entity.name}", + "---", + f"# Entity {entity.name}", + "", + entity.description or "Canonical entity resolved from source mentions.", + "", + "## Memories", + *[ + f"- [{memory.kind} memory](../memory/{memory.kind}/" + f"{self._slug(memory.memory_id)}.md): " + f"{memory.content}" + for memory in related + ], + ] + files.append( + RetrievedMemoryFile( + filename=f"{slug}.md", + file_path=f"{root}/wiki/entities/{slug}.md", + mtime_ms=len(files) + 1, + content="\n".join(content) + "\n", + description=f"canonical entity {entity.name}", + memory_type="entity", + ) + ) + + source_lines = ["# Sources", ""] + session_groups: dict[str, list[SemanticSource]] = {} + for source in sources: + slug = self._slug(source.source_id) + session_groups.setdefault(source.session_id or "unknown", []).append(source) + source_lines.append(f"- [{source.source_id}](wiki/sources/{slug}.md)") + files.append( + RetrievedMemoryFile( + filename=f"{slug}.md", + file_path=f"{root}/wiki/sources/{slug}.md", + mtime_ms=len(files) + 1, + content=( + f"---\nsource_id: {source.source_id}\n" + f"conversation_id: {source.conversation_id}\n" + f"session_id: {source.session_id}\nspeaker: {source.speaker}\n" + + "".join( + f"{key}: {value}\n" + for key, value in sorted(source.metadata.items()) + ) + + f"---\n# Source {source.source_id}\n\n{source.text}\n" + ), + description=f"raw source {source.source_id}", + ) + ) + for session_id, session_sources in sorted(session_groups.items()): + session_number = self._session_number(session_id) + if session_number is None: + continue + session_links = [ + f"- [{source.source_id}](../sources/{self._slug(source.source_id)}.md)" + for source in session_sources + ] + files.append( + RetrievedMemoryFile( + filename=f"session_{session_number}.md", + file_path=f"{root}/wiki/sources/session_{session_number}.md", + mtime_ms=len(files) + 1, + content=( + f"---\ndescription: Session {session_number} source index\n" + f"type: project\nsession_id: {session_id}\n---\n" + f"# Session {session_number}\n\n" + "\n".join(session_links) + "\n" + ), + description=f"session {session_number} source index", + ) + ) + + files.extend( + [ + RetrievedMemoryFile( + filename="MEMORY.md", + file_path=f"{root}/MEMORY.md", + mtime_ms=0, + content="\n".join(index_lines + source_lines) + "\n", + description="canonical semantic memory index", + ), + RetrievedMemoryFile( + filename="index.md", + file_path=f"{root}/index.md", + mtime_ms=len(files) + 1, + content="\n".join(index_lines + source_lines) + "\n", + ), + RetrievedMemoryFile( + filename="profile.md", + file_path=f"{root}/wiki/synthesis/profile.md", + mtime_ms=len(files) + 2, + content=self._render_profile(memories, synthesis), + ), + RetrievedMemoryFile( + filename="timeline.md", + file_path=f"{root}/wiki/synthesis/timeline.md", + mtime_ms=len(files) + 3, + content=synthesis.get("timeline", self._render_timeline(memories)), + ), + RetrievedMemoryFile( + filename="decisions.md", + file_path=f"{root}/wiki/synthesis/decisions.md", + mtime_ms=len(files) + 4, + content=synthesis.get("decisions", self._render_decisions(memories)), + ), + RetrievedMemoryFile( + filename=".wiki-schema.md", + file_path=f"{root}/.wiki-schema.md", + mtime_ms=0, + content=self._render_schema(wiki_mode), + ), + RetrievedMemoryFile( + filename="log.md", + file_path=f"{root}/log.md", + mtime_ms=len(files) + 5, + content=( + f"# Build log\n\nmode: {self.mode}\n" + f"sources: {len(sources)}\nmemories: {len(memories)}\n" + ), + ), + RetrievedMemoryFile( + filename="sources.json", + file_path=f"{root}/raw/sources.json", + mtime_ms=0, + content=json.dumps( + [source.__dict__ for source in sources], + ensure_ascii=False, + indent=2, + ), + ), + ] + ) + return files + + @staticmethod + def _render_profile(memories: list[SemanticMemory], synthesis: dict[str, str]) -> str: + if synthesis.get("profile"): + return "# Profile\n\n" + synthesis["profile"] + "\n" + lines = [ + "# Profile", + "", + "Stable facts, preferences, skills, and decisions consolidated from memory.", + "", + ] + for kind in ("preference", "skill", "decision", "constraint", "fact", "context"): + items = [memory for memory in memories if memory.kind == kind] + if items: + lines.extend([f"## {kind.title()}", *[f"- {item.content}" for item in items], ""]) + return "\n".join(lines) + + @staticmethod + def _render_timeline(memories: list[SemanticMemory]) -> str: + rows = sorted( + (memory for memory in memories if memory.timestamp), + key=lambda item: item.timestamp, + ) + lines = ["# Timeline", ""] + lines.extend(f"- {memory.timestamp}: {memory.content}" for memory in rows) + return "\n".join(lines) + "\n" + + @staticmethod + def _render_decisions(memories: list[SemanticMemory]) -> str: + lines = ["# Decisions", ""] + lines.extend( + f"- {memory.content}" + for memory in memories + if memory.kind in {"decision", "constraint", "preference"} + ) + return "\n".join(lines) + "\n" + + @staticmethod + def _render_schema(wiki_mode: str) -> str: + return ( + "# Canonical Memory Wiki Schema\n\n" + f"- memory kinds: {', '.join(MEMORY_KINDS)}\n" + "- every memory page carries evidence, provenance, entities, confidence, " + "and relations\n" + f"- source modality: {wiki_mode}\n" + ) + + @staticmethod + def _slug(value: str) -> str: + slug = re.sub(r"[^\w-]+", "-", value, flags=re.UNICODE).strip("-")[:100] + return slug or f"memory-{hashlib.sha1(value.encode('utf-8')).hexdigest()[:12]}" + + @staticmethod + def _session_number(value: str) -> int | None: + match = re.search(r"(?:session_|D)(\d+)", value) + return int(match.group(1)) if match else None diff --git a/src/construction/evolver_impl/__init__.py b/src/construction/evolver_impl/__init__.py index 082136a7..cd929bde 100644 --- a/src/construction/evolver_impl/__init__.py +++ b/src/construction/evolver_impl/__init__.py @@ -9,5 +9,6 @@ import_module(".orchestrating_evolver", __name__) import_module(".dynamic_evolver", __name__) +import_module(".wikimem_baseline_evolver", __name__) __all__ = ["EvolverProducer"] diff --git a/src/construction/evolver_impl/wikimem_baseline_evolver.py b/src/construction/evolver_impl/wikimem_baseline_evolver.py new file mode 100644 index 00000000..f5d046de --- /dev/null +++ b/src/construction/evolver_impl/wikimem_baseline_evolver.py @@ -0,0 +1,150 @@ +"""wikimem baseline consolidation for mem2.0 MemoryUnit records.""" + +from __future__ import annotations + +import re +import uuid +from copy import deepcopy +from dataclasses import dataclass, field + +from common.type_def import LifecycleState, MemoryTier, MemoryUnit, Segment +from construction.base import OperatorType +from construction.evolver import EvolveMode, EvolveResult, Evolver, EvolverProducer +from construction.extractor_impl.wikimem_baseline_extractor import ( + WIKIMEM_ACTION, + WIKIMEM_KEY, + WIKIMEM_MEMORY_TYPE, + WIKIMEM_OBSERVED_AT_MS, + WIKIMEM_SCORE, + WIKIMEM_SKIP_REASON, + WIKIMEM_SOURCE_MESSAGE_ID, + WIKIMEM_VALUE, +) + +WIKIMEM_DESCRIPTION = "wikimem.description" +WIKIMEM_KIND = "wikimem.kind" +WIKIMEM_RECORD_KIND = "record" + +_TOKEN_RE = re.compile(r"[A-Za-z0-9_-]+") + + +@dataclass +class WikimemBaselineOutcome: + """Consolidation output plus diagnostics for skipped candidates.""" + + units: list[MemoryUnit] = field(default_factory=list) + result: EvolveResult = field(default_factory=EvolveResult) + skipped_ids: list[str] = field(default_factory=list) + + +class WikimemBaselineEvolver(Evolver): + """Apply wikimem baseline upsert/forget candidates to MemoryUnit records.""" + + def operator_type(self) -> OperatorType: + return OperatorType.EVOLVER + + def health(self) -> None: + return None + + def evolve(self, units: list[MemoryUnit], mode: EvolveMode) -> EvolveResult: + if mode != EvolveMode.CONSOLIDATE: + return EvolveResult() + prior = [unit for unit in units if not is_wikimem_candidate(unit)] + candidates = [unit for unit in units if is_wikimem_candidate(unit)] + return self.consolidate(prior, candidates).result + + def consolidate( + self, prior: list[MemoryUnit], candidates: list[MemoryUnit] + ) -> WikimemBaselineOutcome: + outcome = WikimemBaselineOutcome(units=[deepcopy(unit) for unit in prior]) + for candidate in candidates: + metadata = candidate.metadata + if metadata.get(WIKIMEM_SKIP_REASON): + outcome.skipped_ids.append(candidate.id) + continue + action = metadata.get(WIKIMEM_ACTION) + if action == "upsert": + self._upsert(outcome, candidate) + elif action == "forget": + self._forget(outcome, candidate) + return outcome + + def _upsert(self, outcome: WikimemBaselineOutcome, candidate: MemoryUnit) -> None: + key = candidate.metadata[WIKIMEM_KEY] + supersedes = "" + for unit in outcome.units: + if ( + is_wikimem_record(unit) + and unit.lifecycle == LifecycleState.ACTIVE + and unit.metadata.get(WIKIMEM_KEY) == key + ): + unit.lifecycle = LifecycleState.SUPERSEDED + outcome.result.superseded_ids.append(unit.id) + supersedes = unit.id + + record = _record_from_candidate(candidate, supersedes=supersedes) + outcome.units.append(record) + outcome.result.created_ids.append(record.id) + + def _forget(self, outcome: WikimemBaselineOutcome, candidate: MemoryUnit) -> None: + key = candidate.metadata[WIKIMEM_KEY] + value = candidate.metadata[WIKIMEM_VALUE] + normalized_value = _normalize_memory_text(value) + for unit in outcome.units: + if not is_wikimem_record(unit) or unit.lifecycle == LifecycleState.FORGOTTEN: + continue + if ( + unit.metadata.get(WIKIMEM_KEY) == key + or _normalize_memory_text(unit.content) == normalized_value + ): + unit.lifecycle = LifecycleState.FORGOTTEN + outcome.result.forgotten_ids.append(unit.id) + + +def is_wikimem_candidate(unit: MemoryUnit) -> bool: + return "wikimem_candidate" in unit.tags and WIKIMEM_ACTION in unit.metadata + + +def is_wikimem_record(unit: MemoryUnit) -> bool: + return unit.metadata.get(WIKIMEM_KIND) == WIKIMEM_RECORD_KIND + + +def _record_from_candidate(candidate: MemoryUnit, *, supersedes: str) -> MemoryUnit: + metadata = dict(candidate.metadata) + memory_type = metadata.get(WIKIMEM_MEMORY_TYPE, "project") + metadata[WIKIMEM_KIND] = WIKIMEM_RECORD_KIND + metadata[WIKIMEM_DESCRIPTION] = _format_memory_description(metadata) + tags = [tag for tag in candidate.tags if tag != "wikimem_candidate"] + if "wikimem_record" not in tags: + tags.append("wikimem_record") + return MemoryUnit( + id=f"wikimem_record_{uuid.uuid4().hex}", + scope=candidate.scope, + tier=MemoryTier.SEMANTIC, + segments=[Segment(content=metadata[WIKIMEM_VALUE], source=candidate.source)], + source_ref=metadata.get(WIKIMEM_SOURCE_MESSAGE_ID, ""), + temporal=candidate.temporal, + provenance=list(candidate.provenance), + supersedes=supersedes, + tags=tags, + metadata=metadata, + lifecycle=LifecycleState.ACTIVE, + ) + + +def _format_memory_description(metadata: dict[str, str]) -> str: + memory_type = metadata.get(WIKIMEM_MEMORY_TYPE, "general") + source = metadata.get(WIKIMEM_SOURCE_MESSAGE_ID, "unknown-message") + return ( + f"remembered {memory_type} context from {source}: " + f"{metadata[WIKIMEM_KEY]} = {metadata[WIKIMEM_VALUE]}" + ) + + +def _normalize_memory_text(text: str) -> str: + return " ".join(match.group(0).lower() for match in _TOKEN_RE.finditer(text)) + + +@EvolverProducer.register("wikimem_baseline") +def _build(config): + return WikimemBaselineEvolver() diff --git a/src/construction/extractor_impl/__init__.py b/src/construction/extractor_impl/__init__.py index 3440365e..16c1a724 100644 --- a/src/construction/extractor_impl/__init__.py +++ b/src/construction/extractor_impl/__init__.py @@ -11,5 +11,6 @@ import_module(".keyword_extractor", __name__) import_module(".llm_extractor", __name__) import_module(".dynamic_llm_extractor", __name__) +import_module(".wikimem_baseline_extractor", __name__) __all__ = ["ExtractorProducer"] diff --git a/src/construction/extractor_impl/wikimem_baseline_extractor.py b/src/construction/extractor_impl/wikimem_baseline_extractor.py new file mode 100644 index 00000000..bf988ba6 --- /dev/null +++ b/src/construction/extractor_impl/wikimem_baseline_extractor.py @@ -0,0 +1,324 @@ +"""wikimem baseline candidate extraction for mem2.0 construction.""" + +from __future__ import annotations + +import re +import time +import uuid + +from common.type_def import LifecycleState, MemoryTier, MemoryUnit, Segment +from construction.base import OperatorType +from construction.extractor import Extractor, ExtractorProducer + +WIKIMEM_ACTION = "wikimem.action" +WIKIMEM_KEY = "wikimem.key" +WIKIMEM_VALUE = "wikimem.value" +WIKIMEM_MEMORY_TYPE = "wikimem.memory_type" +WIKIMEM_PREFERRED_SCOPE = "wikimem.preferred_scope" +WIKIMEM_SOURCE_MESSAGE_ID = "wikimem.source_message_id" +WIKIMEM_OBSERVED_AT_MS = "wikimem.observed_at_ms" +WIKIMEM_SCORE = "wikimem.score" +WIKIMEM_SKIP_REASON = "wikimem.skip_reason" + +_TOKEN_RE = re.compile(r"[A-Za-z0-9_-]+") +_INTERROGATIVE_KEYS = {"who", "what", "when", "where", "why", "how", "which"} +_SECRET_MARKERS = ( + "ghp_", + "gho_", + "ghu_", + "ghs_", + "ghr_", + "github_pat_", + "sk-ant-", + "sk-proj-", + "sk-svcacct-", + "sk-admin-", + "xoxb-", + "xoxp-", + "xoxe-", + "xapp-", + "akia", + "asia", + "abia", + "acca", +) +_SCOPE_PREFIXES = ( + ("this just for me", "auto"), + ("just for me", "auto"), + ("this for me only", "auto"), + ("for me only", "auto"), + ("this from private memory", "auto"), + ("from private memory", "auto"), + ("in private memory", "auto"), + ("to private memory", "auto"), + ("for private memory", "auto"), + ("this in auto memory", "auto"), + ("in auto memory", "auto"), + ("to auto memory", "auto"), + ("for auto memory", "auto"), + ("this for the team", "team"), + ("for the team", "team"), + ("this from team memory", "team"), + ("from team memory", "team"), + ("in team memory", "team"), + ("to team memory", "team"), + ("for team memory", "team"), + ("只给我", "auto"), + ("私有记忆里", "auto"), + ("私有记忆", "auto"), + ("私人记忆里", "auto"), + ("私人记忆", "auto"), + ("团队记忆里", "team"), + ("团队记忆", "team"), + ("团队共享", "team"), + ("给团队", "team"), +) + + +class WikimemBaselineExtractor(Extractor): + """Extract wikimem baseline candidates from raw MemoryUnit content.""" + + def operator_type(self) -> OperatorType: + return OperatorType.EXTRACTOR + + def health(self) -> None: + return None + + def extract(self, units: list[MemoryUnit]) -> list[MemoryUnit]: + candidates: list[MemoryUnit] = [] + for unit in units: + if unit.lifecycle != LifecycleState.ACTIVE or unit.provenance: + continue + parsed = parse_memory_candidate(unit.content, unit.source_ref or unit.id) + if parsed is None: + continue + candidates.append(_candidate_unit(unit, parsed)) + return candidates + + +def parse_memory_candidate(text: str, source_message_id: str) -> dict[str, str] | None: + """Parse one wikimem baseline candidate from text.""" + + stripped = text.strip() + if not stripped: + return None + return ( + _parse_explicit_forget_instruction(stripped, source_message_id) + or _parse_explicit_memory_instruction(stripped, source_message_id) + or _parse_key_value_candidate(stripped, source_message_id) + ) + + +def _parse_key_value_candidate(text: str, source_message_id: str) -> dict[str, str] | None: + for separator in (":", " is ", "="): + if separator not in text: + continue + key, value = text.split(separator, 1) + key = key.strip() + value = value.strip() + if not key or not value or key.lower() in _INTERROGATIVE_KEYS: + continue + memory_type = _infer_memory_type(key, value) + return _candidate_metadata( + action="upsert", + key=key, + value=value, + source_message_id=source_message_id, + memory_type=memory_type, + preferred_scope="", + ) + return None + + +def _parse_explicit_memory_instruction( + text: str, source_message_id: str +) -> dict[str, str] | None: + prefixes = ( + "please remember that ", + "please remember ", + "remember that ", + "remember ", + ) + lower = text.lower() + raw_value = "" + for prefix in prefixes: + if lower.startswith(prefix): + raw_value = text[len(prefix) :].strip() + break + if not raw_value: + if text.startswith("请记住"): + raw_value = text[len("请记住") :].strip() + elif text.startswith("记住"): + raw_value = text[len("记住") :].strip() + else: + return None + + preferred_scope, value = _extract_scope_hint(raw_value) + value = _clean_instruction_value(value) + if len(value) < 8: + return None + return _candidate_metadata( + action="upsert", + key=_derive_memory_note_key(value), + value=value, + source_message_id=source_message_id, + memory_type=_infer_memory_type("memory", value), + preferred_scope=preferred_scope, + ) + + +def _parse_explicit_forget_instruction( + text: str, source_message_id: str +) -> dict[str, str] | None: + prefixes = ( + "please forget that ", + "please forget ", + "forget that ", + "forget ", + ) + lower = text.lower() + raw_target = "" + for prefix in prefixes: + if lower.startswith(prefix): + raw_target = text[len(prefix) :].strip() + break + if not raw_target: + if text.startswith("请忘记"): + raw_target = text[len("请忘记") :].strip() + elif text.startswith("忘记"): + raw_target = text[len("忘记") :].strip() + else: + return None + + preferred_scope, target = _extract_scope_hint(raw_target) + target = _clean_instruction_value(target) + target = _strip_forget_article(target) + if not target: + return None + key = _derive_memory_note_key(target) if " " in target else target + return _candidate_metadata( + action="forget", + key=key, + value=target, + source_message_id=source_message_id, + memory_type="", + preferred_scope=preferred_scope, + ) + + +def _candidate_metadata( + *, + action: str, + key: str, + value: str, + source_message_id: str, + memory_type: str, + preferred_scope: str, +) -> dict[str, str]: + metadata = { + WIKIMEM_ACTION: action, + WIKIMEM_KEY: key, + WIKIMEM_VALUE: value, + WIKIMEM_SOURCE_MESSAGE_ID: source_message_id, + WIKIMEM_OBSERVED_AT_MS: str(int(time.time() * 1000)), + } + if preferred_scope: + metadata[WIKIMEM_PREFERRED_SCOPE] = preferred_scope + if memory_type: + metadata[WIKIMEM_MEMORY_TYPE] = memory_type + metadata[WIKIMEM_SCORE] = str(_candidate_score(memory_type)) + if preferred_scope == "team" and action == "upsert" and _has_potential_secret(value): + metadata[WIKIMEM_SKIP_REASON] = "potential_secret" + return metadata + + +def _candidate_unit(source: MemoryUnit, metadata: dict[str, str]) -> MemoryUnit: + return MemoryUnit( + id=f"wikimem_candidate_{uuid.uuid4().hex}", + scope=source.scope, + tier=MemoryTier.SEMANTIC, + segments=[Segment(content=metadata[WIKIMEM_VALUE], source=source.source)], + source_ref=metadata[WIKIMEM_SOURCE_MESSAGE_ID], + temporal=source.temporal, + provenance=[source.id], + tags=[*source.tags, "wikimem_candidate"], + metadata=metadata, + lifecycle=LifecycleState.ACTIVE, + ) + + +def _extract_scope_hint(text: str) -> tuple[str, str]: + trimmed = text.strip() + lower = trimmed.lower() + for prefix, scope in _SCOPE_PREFIXES: + if lower.startswith(prefix): + return scope, trimmed[len(prefix) :].lstrip(":, \t").strip() + return "", trimmed + + +def _clean_instruction_value(value: str) -> str: + return value.strip(":, \t").rstrip(".!?。！？").strip() + + +def _strip_forget_article(target: str) -> str: + lower = target.lower() + for prefix in ("the ", "this "): + if lower.startswith(prefix): + return target[len(prefix) :].strip() + return target + + +def _derive_memory_note_key(value: str) -> str: + tokens = _ordered_tokens(value)[:4] + if not tokens: + return "memory_note" + return f"memory_note_{'_'.join(tokens)}" + + +def _ordered_tokens(text: str) -> list[str]: + return [match.group(0).lower() for match in _TOKEN_RE.finditer(text) if len(match.group(0)) > 1] + + +def _infer_memory_type(key: str, value: str) -> str: + normalized_key = key.lower() + normalized_value = value.lower() + if normalized_key in {"user", "role", "experience", "knowledge"}: + return "user" + if ( + normalized_key in {"feedback", "preference", "preferences", "rule"} + or "prefer" in normalized_value + or "don't" in normalized_value + or "must" in normalized_value + ): + return "feedback" + if ( + normalized_key in {"reference", "dashboard", "linear", "slack", "grafana", "url"} + or "http" in normalized_value + or "grafana" in normalized_value + or "linear" in normalized_value + or "slack" in normalized_value + ): + return "reference" + return "project" + + +def _candidate_score(memory_type: str) -> float: + if memory_type == "feedback": + return 1.25 + if memory_type == "reference": + return 1.15 + if memory_type == "user": + return 1.1 + return 1.0 + + +def _has_potential_secret(value: str) -> bool: + lower = value.lower() + if "-----begin" in lower and "private key-----" in lower: + return True + return any(marker in lower for marker in _SECRET_MARKERS) + + +@ExtractorProducer.register("wikimem_baseline") +def _build(config): + return WikimemBaselineExtractor() diff --git a/src/retrieval/recaller_impl/__init__.py b/src/retrieval/recaller_impl/__init__.py index 227db754..94b3836e 100644 --- a/src/retrieval/recaller_impl/__init__.py +++ b/src/retrieval/recaller_impl/__init__.py @@ -10,5 +10,6 @@ import_module(".graph_recaller", __name__) import_module(".keyword_recaller", __name__) import_module(".vector_recaller", __name__) +import_module(".wikimem_memdir_recaller", __name__) __all__ = ["RecallerProducer"] diff --git a/src/retrieval/recaller_impl/wikimem_memdir_recaller.py b/src/retrieval/recaller_impl/wikimem_memdir_recaller.py new file mode 100644 index 00000000..7b2e3a46 --- /dev/null +++ b/src/retrieval/recaller_impl/wikimem_memdir_recaller.py @@ -0,0 +1,87 @@ +"""wikimem Markdown memory directory DOCUMENT recaller.""" + +from __future__ import annotations + +import hashlib +from pathlib import Path + +from common.type_def import Scope +from retrieval.base import RetrievalOperatorType +from retrieval.wikimem_memdir import ( + WikimemDirectory, + load_memory_entrypoints, + load_relevant_memory_files, +) +from retrieval.wikimem_options import parse_wikimem_options +from retrieval.recaller import Recaller, RecallerProducer +from retrieval.types import ParsedQuery, RecallChannel, ScoredUnit + + +def memory_file_unit_id(scope_name: str, file_path: str) -> str: + """Return the deterministic MemoryUnit id used for a projected memory file.""" + + normalized = str(Path(file_path).expanduser()).replace("\\", "/").lower() + digest = hashlib.sha256(f"{scope_name}:{normalized}".encode("utf-8")).hexdigest()[:24] + return f"wikimem:file:{scope_name}:{digest}" + + +def memory_entrypoint_unit_id(scope_name: str, file_path: str) -> str: + """Return the deterministic MemoryUnit id used for a projected MEMORY.md entrypoint.""" + + normalized = str(Path(file_path).expanduser()).replace("\\", "/").lower() + digest = hashlib.sha256(f"{scope_name}:{normalized}".encode("utf-8")).hexdigest()[:24] + return f"wikimem:entrypoint:{scope_name}:{digest}" + + +class WikimemMemdirRecaller(Recaller): + """Recall wikimem Markdown file candidates through the DOCUMENT channel.""" + + def operator_type(self) -> RetrievalOperatorType: + return RetrievalOperatorType.RECALLER + + def health(self) -> None: + return None + + def channel(self) -> RecallChannel: + return RecallChannel.DOCUMENT + + def recall(self, scope: Scope, query: ParsedQuery, top_k: int) -> list[ScoredUnit]: + options = parse_wikimem_options(query.extensions) + if not options.memory_dirs: + return [] + + directories = [ + WikimemDirectory(scope=directory.scope, path=directory.path) + for directory in options.memory_dirs + ] + text = query.rewritten or query.raw + files = load_relevant_memory_files( + text, + directories, + top_k=top_k, + recent_tools=options.recent_tools, + already_surfaced_file_paths=options.already_surfaced_file_paths, + ) + results = [ + ScoredUnit( + unit_id=memory_file_unit_id(file.scope, file.file_path), + score=float(max(top_k - index, 1)), + channel=RecallChannel.DOCUMENT, + ) + for index, file in enumerate(files) + ] + if options.include_entrypoints: + for entrypoint in load_memory_entrypoints(directories): + results.append( + ScoredUnit( + unit_id=memory_entrypoint_unit_id(entrypoint.scope, entrypoint.file_path), + score=1.0, + channel=RecallChannel.DOCUMENT, + ) + ) + return results[:top_k] + + +@RecallerProducer.register("wikimem_memdir") +def _build(config): + return WikimemMemdirRecaller() diff --git a/src/retrieval/wikimem_memdir.py b/src/retrieval/wikimem_memdir.py new file mode 100644 index 00000000..51a447d3 --- /dev/null +++ b/src/retrieval/wikimem_memdir.py @@ -0,0 +1,523 @@ +"""wikimem-compatible Markdown memory directory helpers.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path + +DEFAULT_MAX_MEMORY_FILES = 200 +DEFAULT_TOP_K = 5 +FRONTMATTER_MAX_LINES = 30 +MAX_MEMORY_LINES = 200 +MAX_MEMORY_BYTES = 4096 +MAX_ENTRYPOINT_LINES = 200 +MAX_ENTRYPOINT_BYTES = 25000 +MIN_HEADER_SELECTION_SCORE = 4.0 +MIN_BODY_SELECTION_SCORE = 4.5 +RELATIVE_SELECTION_SCORE_RATIO = 0.45 +ENTRYPOINT_FILENAME = "MEMORY.md" + +_TOKEN_RE = re.compile(r"[A-Za-z0-9_]+|[\u4e00-\u9fff]") +_DAILY_LOG_RE = re.compile(r"^logs/[0-9]{4}/[0-9]{2}/[0-9]{4}-[0-9]{2}-[0-9]{2}\.md$") +_DRIVE_RE = re.compile(r"^[A-Za-z]:/") +_SECRET_MARKERS = ( + "-----begin", + "private key-----", + "ghp_", + "gho_", + "ghu_", + "ghs_", + "ghr_", + "github_pat_", + "sk-ant-", + "sk-proj-", + "sk-svcacct-", + "sk-admin-", + "xoxb-", + "xoxp-", + "xoxe-", + "xapp-", + "akia", + "asia", + "abia", + "acca", +) +_WARNING_SIGNALS = { + "warning", + "warnings", + "gotcha", + "gotchas", + "issue", + "issues", + "bug", + "bugs", + "risk", + "risks", + "known", +} + + +@dataclass(frozen=True) +class WikimemDirectory: + """A wikimem memory directory.""" + + scope: str + path: str + + +@dataclass(frozen=True) +class MemoryFileHeader: + """Header-stage Markdown memory candidate.""" + + filename: str + file_path: str + mtime_ms: int + description: str | None = None + memory_type: str | None = None + + +@dataclass(frozen=True) +class RecalledMemoryFile: + """Selected memory file before body materialization.""" + + scope: str + filename: str + file_path: str + mtime_ms: int + description: str | None = None + memory_type: str | None = None + score: float = 0.0 + + +@dataclass(frozen=True) +class RetrievedMemoryFile: + """Materialized Markdown memory file.""" + + scope: str + filename: str + file_path: str + mtime_ms: int + content: str + description: str | None = None + memory_type: str | None = None + + +@dataclass(frozen=True) +class RetrievedMemoryEntrypoint: + """Materialized MEMORY.md entrypoint.""" + + scope: str + file_path: str + content: str + + +def normalize_memory_relative_path(relative_path: str) -> str: + """Normalize a memory relative path while preserving leading traversal.""" + + normalized: list[str] = [] + value = relative_path.replace("\\", "/") + for raw_part in value.split("/"): + if raw_part in ("", "."): + continue + if raw_part == "..": + if normalized and normalized[-1] != "..": + normalized.pop() + else: + normalized.append("..") + else: + normalized.append(raw_part) + return "/".join(normalized) + + +def should_include_memory_topic_file(relative_path: str) -> bool: + """Return whether a relative path is a wikimem topic Markdown file.""" + + value = relative_path.replace("\\", "/") + normalized = normalize_memory_relative_path(value) + if value.startswith("/") or _DRIVE_RE.match(value): + return False + if not normalized.endswith(".md"): + return False + if normalized == ENTRYPOINT_FILENAME: + return False + if ".." in normalized.split("/"): + return False + return _DAILY_LOG_RE.match(normalized) is None + + +def scan_memory_directory( + memory_dir: str | Path, max_files: int = DEFAULT_MAX_MEMORY_FILES +) -> list[MemoryFileHeader]: + """Scan Markdown topic files under a memory directory.""" + + root = Path(memory_dir) + if not root.is_dir(): + return [] + + headers: list[MemoryFileHeader] = [] + for path in root.rglob("*.md"): + relative = path.relative_to(root).as_posix() + if not should_include_memory_topic_file(relative): + continue + try: + headers.append(_read_memory_header(root, path)) + except OSError: + continue + + headers.sort(key=lambda item: (-item.mtime_ms, item.filename)) + return headers[: max(0, max_files)] + + +def format_memory_manifest(memories: list[MemoryFileHeader]) -> str: + """Render selector-facing memory manifest text.""" + + lines: list[str] = [] + for memory in memories: + type_tag = f"[{memory.memory_type}] " if memory.memory_type else "" + timestamp = _format_iso8601_ms(memory.mtime_ms) + if memory.description: + lines.append(f"- {type_tag}{memory.filename} ({timestamp}): {memory.description}") + else: + lines.append(f"- {type_tag}{memory.filename} ({timestamp})") + return "\n".join(lines) + + +def select_relevant_memory_files( + query: str, memories: list[MemoryFileHeader], top_k: int +) -> list[MemoryFileHeader]: + """Select high-confidence memories from headers.""" + + scored = [ + (score, memory.mtime_ms, memory) + for memory in memories + if (score := _score_header(query, memory)) is not None + ] + return [memory for _, _, memory in _select_confident(scored, top_k, MIN_HEADER_SELECTION_SCORE)] + + +def recall_relevant_memory_files_from_headers( + query: str, + memories: list[MemoryFileHeader], + top_k: int, + scope: str = "auto", +) -> list[RecalledMemoryFile]: + """Select preloaded headers without rescanning directories.""" + + scored = [ + (score, memory.mtime_ms, memory) + for memory in memories + if (score := _score_header(query, memory)) is not None + ] + return [ + _recalled_from_header(memory, scope, score) + for score, _, memory in _select_confident(scored, top_k, MIN_HEADER_SELECTION_SCORE) + ] + + +def load_relevant_memory_files( + query: str, + directories: list[WikimemDirectory], + top_k: int = DEFAULT_TOP_K, + recent_tools: list[str] | None = None, + already_surfaced_file_paths: list[str] | None = None, +) -> list[RetrievedMemoryFile]: + """Scan, select, and materialize relevant wikimem Markdown files.""" + + candidates = _collect_candidates( + directories, recent_tools or [], already_surfaced_file_paths or [] + ) + header_selected = [ + _recalled_from_header(memory, scope, _score_header(query, memory) or 0.0) + for memory, scope in candidates + if _score_header(query, memory) is not None + ] + selected = _select_confident_recalled(header_selected, top_k, MIN_HEADER_SELECTION_SCORE) + if not selected: + body_scored = [] + for memory, scope in candidates: + score = _score_body(query, memory.file_path) + if score is not None: + body_scored.append(_recalled_from_header(memory, scope, score)) + selected = _select_confident_recalled(body_scored, top_k, MIN_BODY_SELECTION_SCORE) + return materialize_recalled_memory_files(selected) + + +def materialize_recalled_memory_files(files: list[RecalledMemoryFile]) -> list[RetrievedMemoryFile]: + """Load selected file bodies with wikimem truncation limits.""" + + materialized: list[RetrievedMemoryFile] = [] + for file in files: + if file.scope == "team" and _content_has_potential_secrets(file.file_path): + continue + content = _read_limited_text(file.file_path, MAX_MEMORY_LINES, MAX_MEMORY_BYTES) + materialized.append( + RetrievedMemoryFile( + scope=file.scope, + filename=file.filename, + file_path=file.file_path, + mtime_ms=file.mtime_ms, + content=content, + description=file.description, + memory_type=file.memory_type, + ) + ) + return materialized + + +def load_memory_entrypoints( + directories: list[WikimemDirectory], +) -> list[RetrievedMemoryEntrypoint]: + """Load MEMORY.md entrypoints for directories that have one.""" + + seen: set[str] = set() + entrypoints: list[RetrievedMemoryEntrypoint] = [] + for directory in directories: + path = Path(directory.path) / ENTRYPOINT_FILENAME + normalized = str(path.resolve()) if path.exists() else str(path) + if normalized in seen or not path.is_file(): + continue + seen.add(normalized) + if directory.scope == "team" and _content_has_potential_secrets(path): + continue + entrypoints.append( + RetrievedMemoryEntrypoint( + scope=directory.scope, + file_path=str(path), + content=_read_limited_text(path, MAX_ENTRYPOINT_LINES, MAX_ENTRYPOINT_BYTES), + ) + ) + return entrypoints + + +def _read_memory_header(root: Path, path: Path) -> MemoryFileHeader: + stat = path.stat() + description, memory_type = _parse_frontmatter(path) + return MemoryFileHeader( + filename=path.relative_to(root).as_posix(), + file_path=str(path), + mtime_ms=stat.st_mtime_ns // 1_000_000, + description=description, + memory_type=memory_type, + ) + + +def _parse_frontmatter(path: Path) -> tuple[str | None, str | None]: + lines = path.read_text(encoding="utf-8", errors="replace").splitlines() + if not lines or lines[0].strip() != "---": + return None, None + + fields: dict[str, str] = {} + index = 1 + while index < min(len(lines), FRONTMATTER_MAX_LINES + 1): + line = lines[index] + if line.strip() == "---": + break + if ":" not in line: + index += 1 + continue + key, value = line.split(":", 1) + key = key.strip().lower() + value = value.strip() + if value in (">", "|"): + folded: list[str] = [] + index += 1 + while index < len(lines) and lines[index].startswith((" ", "\t")): + folded.append(lines[index].strip()) + index += 1 + fields[key] = " ".join(part for part in folded if part) + continue + fields[key] = _strip_quotes(value) + index += 1 + return fields.get("description") or None, fields.get("type") or None + + +def _strip_quotes(value: str) -> str: + if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}: + return value[1:-1] + return value + + +def _collect_candidates( + directories: list[WikimemDirectory], + recent_tools: list[str], + already_surfaced_file_paths: list[str], +) -> list[tuple[MemoryFileHeader, str]]: + explicit_team_roots = [ + Path(directory.path).resolve() + for directory in directories + if directory.scope == "team" and Path(directory.path).exists() + ] + surfaced = {_normalize_file_path(path) for path in already_surfaced_file_paths} + candidates: list[tuple[MemoryFileHeader, str]] = [] + for directory in directories: + root = Path(directory.path) + headers = scan_memory_directory(root) + if directory.scope == "auto" and explicit_team_roots: + headers = [ + header + for header in headers + if not any( + _path_starts_with(header.file_path, team_root) + for team_root in explicit_team_roots + ) + ] + for header in headers: + if _normalize_file_path(header.file_path) in surfaced: + continue + if _is_recent_tool_reference(header, recent_tools): + continue + candidates.append((header, directory.scope)) + return candidates + + +def _select_confident( + scored: list[tuple[float, int, MemoryFileHeader]], top_k: int, threshold: float +) -> list[tuple[float, int, MemoryFileHeader]]: + if not scored: + return [] + limit = DEFAULT_TOP_K if top_k == 0 else max(1, top_k) + scored.sort(key=lambda item: (-item[0], -item[1], item[2].filename)) + best = scored[0][0] + min_allowed = max(threshold, best * RELATIVE_SELECTION_SCORE_RATIO) + return [item for item in scored if item[0] >= min_allowed][:limit] + + +def _select_confident_recalled( + files: list[RecalledMemoryFile], top_k: int, threshold: float +) -> list[RecalledMemoryFile]: + if not files: + return [] + limit = DEFAULT_TOP_K if top_k == 0 else max(1, top_k) + ordered = sorted(files, key=lambda item: (-item.score, -item.mtime_ms, item.filename)) + best = ordered[0].score + min_allowed = max(threshold, best * RELATIVE_SELECTION_SCORE_RATIO) + return [item for item in ordered if item.score >= min_allowed][:limit] + + +def _recalled_from_header( + memory: MemoryFileHeader, scope: str, score: float +) -> RecalledMemoryFile: + return RecalledMemoryFile( + scope=scope, + filename=memory.filename, + file_path=memory.file_path, + mtime_ms=memory.mtime_ms, + description=memory.description, + memory_type=memory.memory_type, + score=score, + ) + + +def _score_header(query: str, memory: MemoryFileHeader) -> float | None: + query_tokens = _tokenize(query) + if not query_tokens: + return None + filename_tokens = _tokenize(memory.filename) + description_tokens = _tokenize(memory.description or "") + type_tokens = _tokenize(memory.memory_type or "") + score = 0.0 + for token in query_tokens: + if token in filename_tokens: + score += 1.2 + if token in description_tokens: + score += 1.5 + if token in type_tokens: + score += 0.5 + if memory.memory_type == "feedback" and {"test", "safe", "safety"} & set(query_tokens): + score += 0.8 + return score if score > 0 else None + + +def _score_body(query: str, file_path: str | Path) -> float | None: + query_tokens = _tokenize(query) + if not query_tokens: + return None + try: + content = Path(file_path).read_text(encoding="utf-8", errors="replace") + except OSError: + return None + body_tokens = _tokenize(content) + score = sum(1.5 for token in query_tokens if token in body_tokens) + return score if score > 0 else None + + +def _tokenize(text: str) -> list[str]: + result: list[str] = [] + seen: set[str] = set() + for token in _TOKEN_RE.findall(text.lower()): + normalized = _normalize_token(token) + if normalized and normalized not in seen: + result.append(normalized) + seen.add(normalized) + return result + + +def _normalize_token(token: str) -> str: + if len(token) > 5 and token.endswith("ing"): + token = token[:-3] + if len(token) > 4 and token.endswith("ly"): + token = token[:-2] + if len(token) > 3 and token.endswith("ies"): + token = f"{token[:-3]}y" + if len(token) > 3 and token.endswith("s"): + token = token[:-1] + if token == "safety": + return "safe" + return token + + +def _is_recent_tool_reference(header: MemoryFileHeader, recent_tools: list[str]) -> bool: + if header.memory_type != "reference" or not recent_tools: + return False + haystack = " ".join([header.filename, header.description or ""]).lower() + tokens = set(_tokenize(haystack)) + if tokens & _WARNING_SIGNALS: + return False + return any(_normalize_token(tool.lower()) in tokens for tool in recent_tools) + + +def _content_has_potential_secrets(path: str | Path) -> bool: + try: + content = Path(path).read_text(encoding="utf-8", errors="replace").lower() + except OSError: + return False + if "-----begin" in content and "private key-----" in content: + return True + return any(marker in content for marker in _SECRET_MARKERS if not marker.startswith("-----")) + + +def _read_limited_text(path: str | Path, max_lines: int, max_bytes: int) -> str: + text = Path(path).read_text(encoding="utf-8", errors="replace") + lines = text.splitlines() + truncated = False + if len(lines) > max_lines: + lines = lines[:max_lines] + truncated = True + content = "\n".join(lines) + encoded = content.encode("utf-8") + if len(encoded) > max_bytes: + content = encoded[:max_bytes].decode("utf-8", errors="ignore") + truncated = True + if truncated: + content = f"{content}\n[truncated]" + return content + + +def _format_iso8601_ms(mtime_ms: int) -> str: + return datetime.fromtimestamp(mtime_ms / 1000, tz=timezone.utc).isoformat().replace( + "+00:00", "Z" + ) + + +def _normalize_file_path(path: str) -> str: + return str(Path(path).expanduser()).replace("\\", "/").lower() + + +def _path_starts_with(path: str, root: Path) -> bool: + try: + Path(path).resolve().relative_to(root) + return True + except ValueError: + return False diff --git a/src/retrieval/wikimem_options.py b/src/retrieval/wikimem_options.py new file mode 100644 index 00000000..252a9874 --- /dev/null +++ b/src/retrieval/wikimem_options.py @@ -0,0 +1,129 @@ +"""wikimem compatibility option parsing for retrieval adapters.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from typing import Any + +from common.errors import ValidationError + +KEY_RECENT_TOOLS = "wikimem.recent_tools" +KEY_ALREADY_SURFACED_FILE_PATHS = "wikimem.already_surfaced_file_paths" +KEY_INCLUDE_ENTRYPOINTS = "wikimem.include_entrypoints" +KEY_MEMORY_DIRS = "wikimem.memory_dirs" +KEY_PROFILE = "wikimem.profile" +KEY_SELECTOR_MODEL = "wikimem.selector_model" +KEY_SELECTOR_FALLBACK_MODEL = "wikimem.selector_fallback_model" +KEY_MEMORY_PARALLELISM = "wikimem.memory_parallelism" + +_VALID_MEMORY_DIR_SCOPES = {"auto", "team"} + + +@dataclass(frozen=True) +class WikimemDirectory: + """A wikimem Markdown memory directory declaration.""" + + scope: str + path: str + + +@dataclass(frozen=True) +class WikimemRetrievalOptions: + """Parsed wikimem retrieval options carried by ``Context.extensions``.""" + + recent_tools: list[str] = field(default_factory=list) + already_surfaced_file_paths: list[str] = field(default_factory=list) + include_entrypoints: bool = False + memory_dirs: list[WikimemDirectory] = field(default_factory=list) + profile: str = "" + selector_model: str = "" + selector_fallback_model: str = "" + memory_parallelism: int | None = None + + +def parse_wikimem_options(extensions: dict[str, str] | None) -> WikimemRetrievalOptions: + """Parse transport-safe ``wikimem.*`` extension values. + + Core mem2.0 ignores these keys. wikimem-compatible adapters call this helper + at their boundary and receive typed values plus explicit configuration errors. + """ + + values = extensions or {} + return WikimemRetrievalOptions( + recent_tools=_parse_string_array(values, KEY_RECENT_TOOLS), + already_surfaced_file_paths=_parse_string_array( + values, KEY_ALREADY_SURFACED_FILE_PATHS + ), + include_entrypoints=_parse_bool(values, KEY_INCLUDE_ENTRYPOINTS, default=False), + memory_dirs=_parse_memory_dirs(values), + profile=values.get(KEY_PROFILE, "").strip(), + selector_model=values.get(KEY_SELECTOR_MODEL, "").strip(), + selector_fallback_model=values.get(KEY_SELECTOR_FALLBACK_MODEL, "").strip(), + memory_parallelism=_parse_positive_floor_int(values, KEY_MEMORY_PARALLELISM), + ) + + +def _parse_json_value(values: dict[str, str], key: str) -> Any: + raw = values.get(key) + if raw is None or raw.strip() == "": + return None + try: + return json.loads(raw) + except json.JSONDecodeError as exc: + raise ValidationError(f"{key} must be valid JSON") from exc + + +def _parse_string_array(values: dict[str, str], key: str) -> list[str]: + parsed = _parse_json_value(values, key) + if parsed is None: + return [] + if not isinstance(parsed, list) or not all(isinstance(item, str) for item in parsed): + raise ValidationError(f"{key} must be a JSON string array") + return list(parsed) + + +def _parse_bool(values: dict[str, str], key: str, *, default: bool) -> bool: + raw = values.get(key) + if raw is None or raw.strip() == "": + return default + normalized = raw.strip().lower() + if normalized == "true": + return True + if normalized == "false": + return False + raise ValidationError(f"{key} must be \"true\" or \"false\"") + + +def _parse_positive_floor_int(values: dict[str, str], key: str) -> int | None: + raw = values.get(key) + if raw is None or raw.strip() == "": + return None + try: + value = int(raw.strip()) + except ValueError as exc: + raise ValidationError(f"{key} must be an integer string") from exc + return max(1, value) + + +def _parse_memory_dirs(values: dict[str, str]) -> list[WikimemDirectory]: + parsed = _parse_json_value(values, KEY_MEMORY_DIRS) + if parsed is None: + return [] + if not isinstance(parsed, list): + raise ValidationError(f"{KEY_MEMORY_DIRS} must be a JSON array") + + dirs: list[WikimemDirectory] = [] + for index, item in enumerate(parsed): + if not isinstance(item, dict): + raise ValidationError(f"{KEY_MEMORY_DIRS}[{index}] must be an object") + scope = item.get("scope") + path = item.get("path") + if not isinstance(scope, str) or scope not in _VALID_MEMORY_DIR_SCOPES: + raise ValidationError( + f"{KEY_MEMORY_DIRS}[{index}].scope must be one of auto, team" + ) + if not isinstance(path, str) or not path.strip(): + raise ValidationError(f"{KEY_MEMORY_DIRS}[{index}].path must be non-empty") + dirs.append(WikimemDirectory(scope=scope, path=path)) + return dirs diff --git a/tests/unit/agent_plugin/test_wikimem_team_sync.py b/tests/unit/agent_plugin/test_wikimem_team_sync.py new file mode 100644 index 00000000..6d235043 --- /dev/null +++ b/tests/unit/agent_plugin/test_wikimem_team_sync.py @@ -0,0 +1,227 @@ +"""wikimem team memory sync compatibility tests.""" + +from __future__ import annotations + +import asyncio +from collections import deque + +import pytest + +from agent_plugin.wikimem.team_sync import ( + FetchOutcome, + HashesProbe, + PutOutcome, + RemoteTeamMemoryData, + SyncState, + TeamMemoryFailureKind, + batch_delta_by_bytes, + hash_content, + pull_team_memory, + push_team_memory, + read_local_team_memory, + validate_relative_team_memory_key, +) + +pytestmark = pytest.mark.unit + + +class FakeRemote: + def __init__(self, *, fetch_results=None, hash_results=None, put_results=None) -> None: + self.fetch_results = deque(fetch_results or []) + self.hash_results = deque(hash_results or []) + self.put_results = deque(put_results or []) + self.put_requests = [] + + async def fetch(self, repo_slug: str, if_none_match: str | None): + assert repo_slug == "owner/repo" + return self.fetch_results.popleft() + + async def fetch_hashes(self, repo_slug: str): + assert repo_slug == "owner/repo" + return self.hash_results.popleft() + + async def put_entries(self, repo_slug: str, if_match: str | None, entries: dict[str, str]): + assert repo_slug == "owner/repo" + self.put_requests.append((if_match, dict(entries))) + return self.put_results.popleft() + + +def test_hashes_content_with_sha256_prefix() -> None: + assert ( + hash_content("hello") + == "sha256:2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824" + ) + + +def test_batches_delta_by_serialized_body_size() -> None: + delta = { + "a.md": "alpha" * 20, + "b.md": "beta" * 20, + "c.md": "gamma" * 20, + } + + batches = batch_delta_by_bytes(delta, 120) + + assert [list(batch) for batch in batches] == [["a.md"], ["b.md"], ["c.md"]] + + +def test_validates_safe_relative_team_memory_keys() -> None: + assert validate_relative_team_memory_key("nested/release.md") == "nested/release.md" + + for key in [ + "%2e%2e%2fsecret.md", + "..\\secret.md", + "/secret.md", + "../secret.md", + "C:/secret.md", + "nested/%2fsecret.md", + ]: + with pytest.raises(ValueError): + validate_relative_team_memory_key(key) + + +def test_pull_team_memory_handles_not_modified_empty_and_data(tmp_path) -> None: + remote = FakeRemote( + fetch_results=[ + FetchOutcome.not_modified("etag-1"), + FetchOutcome.empty(), + FetchOutcome.data( + RemoteTeamMemoryData( + checksum="etag-4", + entries={ + "MEMORY.md": "server memory", + "nested/patterns.md": "nested pattern", + }, + entry_checksums={}, + ) + ), + ] + ) + state = SyncState( + last_known_checksum="etag-1", + server_checksums={"MEMORY.md": "sha256:1"}, + ) + + first = asyncio.run(pull_team_memory(remote, state, tmp_path, "owner/repo")) + assert first.success + assert first.not_modified + assert first.checksum == "etag-1" + assert remote.fetch_results + + second = asyncio.run( + pull_team_memory(remote, state, tmp_path, "owner/repo", skip_etag_cache=True) + ) + assert second.success + assert second.is_empty + assert state.last_known_checksum is None + assert state.server_checksums == {} + + third = asyncio.run(pull_team_memory(remote, state, tmp_path, "owner/repo")) + assert third.success + assert third.files_written == 2 + assert (tmp_path / "MEMORY.md").read_text(encoding="utf-8") == "server memory" + assert (tmp_path / "nested" / "patterns.md").read_text(encoding="utf-8") == "nested pattern" + assert state.server_checksums["MEMORY.md"] == hash_content("server memory") + + +def test_pull_team_memory_rejects_remote_path_traversal(tmp_path) -> None: + remote = FakeRemote( + fetch_results=[ + FetchOutcome.data( + RemoteTeamMemoryData( + checksum="etag-bad", + entries={"../secret.md": "bad"}, + entry_checksums={}, + ) + ) + ] + ) + + result = asyncio.run(pull_team_memory(remote, SyncState(), tmp_path, "owner/repo")) + + assert not result.success + assert result.failure is not None + assert result.failure.kind is TeamMemoryFailureKind.UNKNOWN + assert not (tmp_path.parent / "secret.md").exists() + + +def test_push_team_memory_retries_conflicts_after_hash_probe(tmp_path) -> None: + (tmp_path / "MEMORY.md").write_text("local memory", encoding="utf-8") + (tmp_path / "patterns.md").write_text("shared pattern", encoding="utf-8") + remote = FakeRemote( + hash_results=[ + HashesProbe( + checksum="etag-2", + entry_checksums={ + "MEMORY.md": hash_content("remote memory"), + "patterns.md": hash_content("shared pattern"), + }, + ) + ], + put_results=[ + PutOutcome.conflict(), + PutOutcome.success("etag-3"), + ], + ) + state = SyncState(last_known_checksum="etag-1") + + result = asyncio.run(push_team_memory(remote, state, tmp_path, "owner/repo")) + + assert result.success + assert result.files_uploaded == 1 + assert result.checksum == "etag-3" + assert len(remote.put_requests) == 2 + assert len(remote.put_requests[0][1]) == 2 + assert list(remote.put_requests[1][1]) == ["MEMORY.md"] + assert state.server_checksums["MEMORY.md"] == hash_content("local memory") + assert state.server_checksums["patterns.md"] == hash_content("shared pattern") + + +def test_push_team_memory_learns_server_max_entries_from_413(tmp_path) -> None: + for name, content in [("a.md", "alpha"), ("b.md", "beta"), ("c.md", "gamma")]: + (tmp_path / name).write_text(content, encoding="utf-8") + remote = FakeRemote( + put_results=[PutOutcome.too_many_entries(max_entries=2, received_entries=3)] + ) + state = SyncState() + + result = asyncio.run(push_team_memory(remote, state, tmp_path, "owner/repo")) + + assert not result.success + assert result.server_max_entries == 2 + assert state.server_max_entries == 2 + assert result.failure is not None + assert result.failure.http_status == 413 + + +def test_push_team_memory_updates_checksums_across_multiple_batches(tmp_path, monkeypatch) -> None: + (tmp_path / "a.md").write_text("alpha " * 36_000, encoding="utf-8") + (tmp_path / "b.md").write_text("beta " * 36_000, encoding="utf-8") + monkeypatch.setattr("agent_plugin.wikimem.team_sync.MAX_PUT_BODY_BYTES", 200_000) + remote = FakeRemote( + put_results=[ + PutOutcome.success("etag-1"), + PutOutcome.success("etag-2"), + ] + ) + state = SyncState() + + result = asyncio.run(push_team_memory(remote, state, tmp_path, "owner/repo")) + + assert result.success + assert result.files_uploaded == 2 + assert len(remote.put_requests) == 2 + assert state.last_known_checksum == "etag-2" + assert set(state.server_checksums) == {"a.md", "b.md"} + + +def test_read_local_team_memory_skips_potential_secret(tmp_path) -> None: + (tmp_path / "safe.md").write_text("safe shared note", encoding="utf-8") + (tmp_path / "secret.md").write_text("token: github_pat_secret", encoding="utf-8") + + snapshot = read_local_team_memory(tmp_path) + + assert snapshot.entries == {"safe.md": "safe shared note"} + assert [(item.path, item.reason) for item in snapshot.skipped_secrets] == [ + ("secret.md", "potential_secret") + ] diff --git a/tests/unit/construction/test_wikimem_baseline_evolver.py b/tests/unit/construction/test_wikimem_baseline_evolver.py new file mode 100644 index 00000000..8678ec44 --- /dev/null +++ b/tests/unit/construction/test_wikimem_baseline_evolver.py @@ -0,0 +1,92 @@ +"""wikimem baseline consolidation compatibility tests.""" + +from __future__ import annotations + +import pytest + +from common.type_def import LifecycleState +from construction.evolver_impl.wikimem_baseline_evolver import ( + WIKIMEM_DESCRIPTION, + WikimemBaselineEvolver, + is_wikimem_record, +) +from construction.extractor_impl.wikimem_baseline_extractor import ( + WIKIMEM_KEY, + WIKIMEM_MEMORY_TYPE, + WIKIMEM_SCORE, + WIKIMEM_SKIP_REASON, + WIKIMEM_SOURCE_MESSAGE_ID, + WikimemBaselineExtractor, +) +from tests.unit.construction.fixtures import create_test_unit + +pytestmark = pytest.mark.unit + + +def _candidate(text: str, source_id: str = "msg-1"): + [candidate] = WikimemBaselineExtractor().extract([create_test_unit(source_id, text)]) + return candidate + + +def _record_by_key(units, key: str): + return next(unit for unit in units if unit.metadata.get(WIKIMEM_KEY) == key) + + +def test_consolidate_upsert_creates_wikimem_record() -> None: + candidate = _candidate("preference: must run staging checks", "msg-1") + + outcome = WikimemBaselineEvolver().consolidate([], [candidate]) + + assert outcome.result.created_ids + [record] = outcome.units + assert is_wikimem_record(record) + assert record.lifecycle == LifecycleState.ACTIVE + assert record.content == "must run staging checks" + assert record.metadata[WIKIMEM_KEY] == "preference" + assert record.metadata[WIKIMEM_MEMORY_TYPE] == "feedback" + assert record.metadata[WIKIMEM_SCORE] == "1.25" + assert record.metadata[WIKIMEM_SOURCE_MESSAGE_ID] == "msg-1" + assert "remembered feedback context" in record.metadata[WIKIMEM_DESCRIPTION] + + +def test_consolidate_upsert_supersedes_existing_key() -> None: + old = WikimemBaselineEvolver().consolidate( + [], + [_candidate("preference: use old staging gate", "msg-1")], + ).units + new_candidate = _candidate("preference: use new staging gate", "msg-2") + + outcome = WikimemBaselineEvolver().consolidate(old, [new_candidate]) + + old_record = next(unit for unit in outcome.units if unit.id == old[0].id) + new_record = _record_by_key( + [unit for unit in outcome.units if unit.lifecycle == LifecycleState.ACTIVE], + "preference", + ) + assert old_record.lifecycle == LifecycleState.SUPERSEDED + assert new_record.supersedes == old_record.id + assert new_record.content == "use new staging gate" + assert old_record.id in outcome.result.superseded_ids + + +def test_consolidate_forget_marks_key_and_normalized_value_forgotten() -> None: + prior = WikimemBaselineEvolver().consolidate( + [], + [_candidate("remember Release gate must run staging smoke tests", "msg-1")], + ).units + forget = _candidate("forget the Release gate must run staging smoke tests", "msg-2") + + outcome = WikimemBaselineEvolver().consolidate(prior, [forget]) + + assert outcome.units[0].lifecycle == LifecycleState.FORGOTTEN + assert outcome.units[0].id in outcome.result.forgotten_ids + + +def test_consolidate_skips_team_secret_candidate_with_diagnostic() -> None: + candidate = _candidate("remember for team memory token is github_pat_secret", "msg-1") + assert candidate.metadata[WIKIMEM_SKIP_REASON] == "potential_secret" + + outcome = WikimemBaselineEvolver().consolidate([], [candidate]) + + assert outcome.units == [] + assert outcome.skipped_ids == [candidate.id] diff --git a/tests/unit/construction/test_wikimem_baseline_extractor.py b/tests/unit/construction/test_wikimem_baseline_extractor.py new file mode 100644 index 00000000..da217f75 --- /dev/null +++ b/tests/unit/construction/test_wikimem_baseline_extractor.py @@ -0,0 +1,92 @@ +"""wikimem baseline candidate extractor compatibility tests.""" + +from __future__ import annotations + +import pytest + +from common.type_def import MemoryTier +from construction.extractor_impl.wikimem_baseline_extractor import ( + WIKIMEM_ACTION, + WIKIMEM_KEY, + WIKIMEM_MEMORY_TYPE, + WIKIMEM_PREFERRED_SCOPE, + WIKIMEM_SCORE, + WIKIMEM_SKIP_REASON, + WIKIMEM_SOURCE_MESSAGE_ID, + WIKIMEM_VALUE, + WikimemBaselineExtractor, +) + +from tests.unit.construction.fixtures import create_test_unit + +pytestmark = pytest.mark.unit + + +def _candidate_by_key(candidates, key: str): + return next(item for item in candidates if item.metadata[WIKIMEM_KEY] == key) + + +def test_extracts_key_value_candidates_with_type_score_and_source() -> None: + source = create_test_unit("msg-1", "preference: must run staging checks") + extractor = WikimemBaselineExtractor() + + candidates = extractor.extract([source]) + + candidate = _candidate_by_key(candidates, "preference") + assert candidate.tier == MemoryTier.SEMANTIC + assert candidate.provenance == ["msg-1"] + assert candidate.metadata[WIKIMEM_ACTION] == "upsert" + assert candidate.metadata[WIKIMEM_VALUE] == "must run staging checks" + assert candidate.metadata[WIKIMEM_MEMORY_TYPE] == "feedback" + assert candidate.metadata[WIKIMEM_SCORE] == "1.25" + assert candidate.metadata[WIKIMEM_SOURCE_MESSAGE_ID] == "msg-1" + + +def test_ignores_interrogative_key_value_text() -> None: + source = create_test_unit("msg-1", "what: should I remember") + + assert WikimemBaselineExtractor().extract([source]) == [] + + +def test_extracts_explicit_remember_with_scope_hint_and_stable_note_key() -> None: + source = create_test_unit( + "msg-2", + "remember for team memory Release gate must run staging smoke tests.", + ) + + [candidate] = WikimemBaselineExtractor().extract([source]) + + assert candidate.metadata[WIKIMEM_ACTION] == "upsert" + assert candidate.metadata[WIKIMEM_KEY] == "memory_note_release_gate_must_run" + assert candidate.metadata[WIKIMEM_VALUE] == "Release gate must run staging smoke tests" + assert candidate.metadata[WIKIMEM_PREFERRED_SCOPE] == "team" + assert candidate.metadata[WIKIMEM_MEMORY_TYPE] == "feedback" + + +def test_extracts_chinese_scope_hint() -> None: + source = create_test_unit("msg-3", "请记住团队共享 release gate must run") + + [candidate] = WikimemBaselineExtractor().extract([source]) + + assert candidate.metadata[WIKIMEM_PREFERRED_SCOPE] == "team" + assert candidate.metadata[WIKIMEM_KEY] == "memory_note_release_gate_must_run" + + +def test_extracts_explicit_forget_by_multiword_target() -> None: + source = create_test_unit("msg-4", "forget the Release gate must run staging smoke tests") + + [candidate] = WikimemBaselineExtractor().extract([source]) + + assert candidate.metadata[WIKIMEM_ACTION] == "forget" + assert candidate.metadata[WIKIMEM_KEY] == "memory_note_release_gate_must_run" + assert candidate.metadata[WIKIMEM_VALUE] == "Release gate must run staging smoke tests" + assert WIKIMEM_MEMORY_TYPE not in candidate.metadata + + +def test_team_secret_candidate_keeps_diagnostic_skip_reason() -> None: + source = create_test_unit("msg-5", "remember for team memory token is github_pat_secret") + + [candidate] = WikimemBaselineExtractor().extract([source]) + + assert candidate.metadata[WIKIMEM_PREFERRED_SCOPE] == "team" + assert candidate.metadata[WIKIMEM_SKIP_REASON] == "potential_secret" diff --git a/tests/unit/evaluation/test_wikimem_ama_bench.py b/tests/unit/evaluation/test_wikimem_ama_bench.py new file mode 100644 index 00000000..ffc70310 --- /dev/null +++ b/tests/unit/evaluation/test_wikimem_ama_bench.py @@ -0,0 +1,378 @@ +"""AMA-Bench retrieval-only harness compatibility tests.""" + +from __future__ import annotations + +import json + +import pytest + +from evaluation.wikimem.ama_bench import ( + AmaEpisode, + AmaQuestion, + AmaTurn, + aggregate_ama_method_summaries, + build_ama_precision_comparison, + load_ama_dataset_split, + load_ama_episodes_from_jsonl, + run_python_ama_retrieval_eval, + run_python_wikimem_qmd_retrieval, + score_ama_retrieval_proxy, + select_ama_lexical_turn_ids, +) + +pytestmark = pytest.mark.unit + + +def test_load_ama_episodes_and_score_proxy_metrics(tmp_path) -> None: + dataset = tmp_path / "open_end_qa_set.jsonl" + dataset.write_text( + json.dumps( + { + "episode_id": 7, + "domain": "household", + "task_type": "open", + "trajectory": [ + { + "turn_idx": 1, + "action": "open drawer", + "observation": "The drawer contains brass tools.", + }, + { + "turn_idx": 2, + "action": "inspect shelf", + "observation": "The shelf contains paper.", + }, + ], + "qa_pairs": [ + { + "question": "Where were the brass tools?", + "answer": "drawer", + "type": "open_end", + } + ], + } + ) + + "\n", + encoding="utf-8", + ) + + episodes = load_ama_episodes_from_jsonl(dataset) + metrics = score_ama_retrieval_proxy( + episodes[0], + episodes[0].qa_pairs[0], + retrieved_turn_ids=[1, 2, 1], + ) + + assert episodes[0].episode_id == "7" + assert metrics.proxy_gold_turn_ids == [1] + assert metrics.retrieved_turn_ids == [1, 2] + assert metrics.hit_turn_ids == [1] + assert metrics.proxy_recall_at_k == 1.0 + assert metrics.proxy_precision_at_k == 0.5 + assert metrics.proxy_hit_at_k == 1.0 + + +def test_load_ama_episodes_tolerates_multiline_json_records(tmp_path) -> None: + dataset = tmp_path / "open_end_qa_set.jsonl" + dataset.write_text( + '{"episode_id":"ep-1","trajectory":[{"turn_idx":1,' + '"observation":"first line\n' + 'second line"}],"qa_pairs":[{"question":"q","answer":"line"}]}\n', + encoding="utf-8", + ) + + episodes = load_ama_episodes_from_jsonl(dataset) + + assert episodes[0].trajectory[0].observation == "first line\nsecond line" + + +def test_load_ama_dataset_split_stops_after_sample_limit(tmp_path) -> None: + dataset_root = tmp_path / "dataset" + dataset = dataset_root / "test" / "open_end_qa_set.jsonl" + dataset.parent.mkdir(parents=True) + dataset.write_text( + json.dumps({"episode_id": "ep-1", "trajectory": [], "qa_pairs": []}) + + "\n" + + '{"episode_id": "broken"\n', + encoding="utf-8", + ) + + episodes = load_ama_dataset_split(dataset_root, "open_end", sample_limit=1) + + assert [episode.episode_id for episode in episodes] == ["ep-1"] + + +def test_python_wikimem_qmd_retrieval_produces_rust_summary_shape() -> None: + episode = AmaEpisode( + episode_id="ep-1", + task="find tools", + task_type="open", + domain="household", + trajectory=[ + AmaTurn( + turn_idx=1, + action="open drawer", + observation="Alice found brass tools in the drawer.", + ), + AmaTurn( + turn_idx=2, + action="check shelf", + observation="Alice found paper labels on the shelf.", + ), + ], + qa_pairs=[ + AmaQuestion( + question="What did Alice find in the drawer?", + answer="brass tools", + qa_type="open_end", + ) + ], + ) + + result = run_python_wikimem_qmd_retrieval( + episode, + episode.qa_pairs[0], + top_k=3, + ) + summaries = aggregate_ama_method_summaries([result]) + + assert result.method_name == "wikimem_qmd" + assert result.proxy_metrics.proxy_gold_turn_ids == [1] + assert 1 in result.retrieved_turn_ids + assert result.proxy_metrics.proxy_recall_at_k == 1.0 + assert summaries[0].method_name == "wikimem_qmd" + assert summaries[0].proxy_recall_at_k == 1.0 + assert summaries[0].proxy_precision_at_k >= 0.5 + + +def test_select_ama_lexical_turn_ids_matches_rust_turn_token_ordering() -> None: + turns = [ + AmaTurn( + turn_idx=0, + action="inspect file", + observation="saw the same repeated placeholder", + ), + AmaTurn( + turn_idx=7, + action="inspect file", + observation="saw the same repeated placeholder", + ), + ] + + turn_ids = select_ama_lexical_turn_ids( + turns, + "At step 7, what exactly did the agent inspect?", + top_k=1, + ) + + assert turn_ids == [7] + + +def test_wikimem_qmd_ama_retrieval_uses_ama_lexical_primary() -> None: + episode = AmaEpisode( + episode_id="ep-ama", + task="inspect release notes", + task_type="open", + domain="software", + trajectory=[ + AmaTurn( + turn_idx=0, + action="open browser", + observation="search page loaded", + ), + AmaTurn( + turn_idx=1, + action="search changelog", + observation="release note mentions the delta sync fix", + ), + AmaTurn( + turn_idx=2, + action="open ticket", + observation="ticket summary is unrelated", + ), + ], + qa_pairs=[ + AmaQuestion( + question="Which turn mentions the delta sync fix?", + answer="turn 1", + qa_type="open_end", + ) + ], + ) + + result = run_python_wikimem_qmd_retrieval( + episode, + episode.qa_pairs[0], + top_k=2, + method_name="wikimem_qmd_ama", + ) + + assert result.method_name == "wikimem_qmd_ama" + assert result.retrieved_turn_ids == [1, 0] + assert result.retrieved_file_paths == ["ama_turns/T1.md", "ama_turns/T0.md"] + assert "proxy_turn_source=ama_lexical_primary" in result.retrieval_notes + + +def test_ama_golden_answer_does_not_affect_retrieval_ranking() -> None: + episode = AmaEpisode( + episode_id="ep-no-leak", + task="inspect notes", + task_type="open", + domain="software", + trajectory=[ + AmaTurn(turn_idx=0, action="open note", observation="alpha release"), + AmaTurn(turn_idx=1, action="open note", observation="beta release"), + ], + qa_pairs=[], + ) + first = run_python_wikimem_qmd_retrieval( + episode, + AmaQuestion(question="Which note mentions beta?", answer="turn 1", qa_type="open"), + top_k=2, + method_name="wikimem_qmd_ama", + ) + second = run_python_wikimem_qmd_retrieval( + episode, + AmaQuestion( + question="Which note mentions beta?", + answer="deliberately different gold answer", + qa_type="open", + ), + top_k=2, + method_name="wikimem_qmd_ama", + ) + + assert first.retrieved_turn_ids == second.retrieved_turn_ids + assert first.retrieved_file_paths == second.retrieved_file_paths + + +def test_ama_precision_comparison_marks_regressions_against_rust_baseline() -> None: + comparison = build_ama_precision_comparison( + python_summary={ + "method_name": "wikimem_qmd", + "questions": 2, + "proxy_recall_at_k": 0.7, + "proxy_precision_at_k": 0.8, + "proxy_hit_at_k": 1.0, + }, + rust_baseline={ + "method_name": "wikimem_qmd", + "questions": 2, + "proxy_recall_at_k": 0.75, + "proxy_precision_at_k": 0.79, + "proxy_hit_at_k": 1.0, + }, + tolerance=0.001, + ) + + assert comparison["status"] == "regression" + assert comparison["method_name"] == "wikimem_qmd" + assert comparison["metric_deltas"]["proxy_recall_at_k"] == -0.05 + assert comparison["metric_status"]["proxy_recall_at_k"] == "regression" + assert comparison["metric_status"]["proxy_precision_at_k"] == "pass" + + +def test_run_python_ama_retrieval_eval_writes_report_summary_and_comparison(tmp_path) -> None: + dataset_root = tmp_path / "dataset" + dataset_file = dataset_root / "test" / "open_end_qa_set.jsonl" + dataset_file.parent.mkdir(parents=True) + dataset_file.write_text( + json.dumps( + { + "episode_id": "ep-1", + "domain": "household", + "task_type": "open", + "trajectory": [ + { + "turn_idx": 1, + "action": "open drawer", + "observation": "Alice found brass tools in the drawer.", + } + ], + "qa_pairs": [ + { + "question": "What did Alice find in the drawer?", + "answer": "brass tools", + "type": "open_end", + } + ], + } + ) + + "\n", + encoding="utf-8", + ) + output_dir = tmp_path / "output" + + result = run_python_ama_retrieval_eval( + dataset_root=dataset_root, + output_dir=output_dir, + subset="open_end", + top_k=3, + rust_baseline={ + "method_name": "wikimem_qmd", + "questions": 1, + "proxy_recall_at_k": 1.0, + "proxy_precision_at_k": 1.0, + "proxy_hit_at_k": 1.0, + "answer_support_coverage": 1.0, + }, + ) + + report = json.loads((output_dir / "report_openend.json").read_text(encoding="utf-8")) + summary = json.loads((output_dir / "summary_openend.json").read_text(encoding="utf-8")) + comparison = json.loads( + (output_dir / "comparison_openend.json").read_text(encoding="utf-8") + ) + assert result["comparison"]["status"] == "pass" + assert report["metric_label_source"] == "answer_derived_proxy" + assert summary["metric_label_source"] == "answer_derived_proxy" + assert report["total_episodes"] == 1 + assert summary["method_summaries"][0]["method_name"] == "wikimem_qmd" + assert comparison["metric_status"]["proxy_recall_at_k"] == "pass" + + +def test_run_python_ama_retrieval_eval_accepts_qmd_ama_method(tmp_path) -> None: + dataset_root = tmp_path / "dataset" + dataset_file = dataset_root / "test" / "open_end_qa_set.jsonl" + dataset_file.parent.mkdir(parents=True) + dataset_file.write_text( + json.dumps( + { + "episode_id": "ep-ama", + "domain": "software", + "task_type": "open", + "trajectory": [ + { + "turn_idx": 0, + "action": "open browser", + "observation": "search page loaded", + }, + { + "turn_idx": 1, + "action": "search changelog", + "observation": "release note mentions the delta sync fix", + }, + ], + "qa_pairs": [ + { + "question": "Which turn mentions the delta sync fix?", + "answer": "turn 1", + "type": "open_end", + } + ], + } + ) + + "\n", + encoding="utf-8", + ) + + result = run_python_ama_retrieval_eval( + dataset_root=dataset_root, + output_dir=tmp_path / "output", + subset="open_end", + top_k=2, + methods=["wikimem_qmd_ama"], + ) + + assert result["summary"]["config"]["methods"] == ["wikimem_qmd_ama"] + assert result["summary"]["method_summaries"][0]["method_name"] == "wikimem_qmd_ama" diff --git a/tests/unit/evaluation/test_wikimem_example_eval.py b/tests/unit/evaluation/test_wikimem_example_eval.py new file mode 100644 index 00000000..8ed47930 --- /dev/null +++ b/tests/unit/evaluation/test_wikimem_example_eval.py @@ -0,0 +1,1345 @@ +"""Example dataset runner tests for wikimem migration.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest + +from evaluation.wikimem import example_eval +from evaluation.wikimem.example_eval import ( + build_mem_gallery_python_workspaces, + build_meta_crag_python_workspaces, + run_evermembench_python_eval, + run_filesystem_proxy_eval, + run_locomo_refined_offline_eval, +) +from evaluation.wikimem.qmd_consensus import RetrievedMemoryFile + +pytestmark = pytest.mark.unit + + +def _mark_python_workspace(root, dataset_name: str) -> None: + (root / ".wikimem-workspace.json").write_text( + json.dumps( + { + "producer": "mem2.0.wikimem.python", + "dataset_name": dataset_name, + "schema_version": 1, + } + ), + encoding="utf-8", + ) + + +def test_build_mem_gallery_python_workspaces_writes_clean_multimodal_wiki(tmp_path) -> None: + dialog_root = tmp_path / "dialog" + dialog_root.mkdir() + (dialog_root / "Robot_Art.json").write_text( + json.dumps( + { + "character_profile": {"name": "Ada"}, + "multi_session_dialogues": [ + { + "session_id": "D6", + "date": "2026-01-02", + "dialogues": [ + { + "round": "D6:1", + "user": "I set up a robot artist easel.", + "assistant": "The easel has a cobalt palette.", + "input_image": ["../image/Robot_Art/d6.png"], + "image_caption": ["robot artist with cobalt palette"], + "image_id": ["IMG_7"], + } + ], + } + ], + "human-annotated QAs": [ + { + "question": "Which palette was on the robot artist easel?", + "answer": "cobalt", + "point": "visual", + "clue": ["D6:1"], + "session_id": ["D6"], + "question_image": "../image/Robot_Art/q.png", + "image_caption": "robot artist easel", + } + ], + } + ), + encoding="utf-8", + ) + + cases = build_mem_gallery_python_workspaces( + dialog_root=dialog_root, + workspace_root=tmp_path / "workspace", + ) + + assert cases == [ + { + "case_id": "Robot_Art::q1", + "question": "Which palette was on the robot artist easel?", + "silver_evidence_ids": ["D6:1"], + "knowledge_base_root": str(tmp_path / "workspace" / "Robot_Art"), + "question_image_caption": "robot artist easel", + "source_session_ids": ["D6"], + } + ] + kb = tmp_path / "workspace" / "Robot_Art" + assert "Image caption: robot artist with cobalt palette" in ( + kb / "raw" / "turns" / "DD6_1.md" + ).read_text(encoding="utf-8") + assert "Evidence: D6:1" in ( + kb / "wiki" / "turns" / "D6_1.md" + ).read_text(encoding="utf-8") + assert json.loads((kb / ".mem-gallery" / "artifact_support_map.json").read_text(encoding="utf-8")) == [ + { + "memory_path": (kb / "wiki" / "memories" / "clue-summary-D6_1.md").as_posix(), + "linked_clue_ids": ["D6:1"], + } + ] + retrieval = json.loads( + (kb / ".kb-research" / "retrieval" / "D6_1.json").read_text(encoding="utf-8") + ) + assert retrieval["evidence_id"] == "D6:1" + assert "robot artist with cobalt palette" in retrieval["searchable_text"] + provenance = json.loads((kb / ".wikimem-workspace.json").read_text(encoding="utf-8")) + assert provenance["producer"] == "mem2.0.wikimem.python" + + +def test_build_meta_crag_python_workspaces_labels_without_rust_cases(tmp_path) -> None: + data_root = tmp_path / "meta" + data_root.mkdir() + (data_root / "validation.json").write_text( + json.dumps( + [ + { + "session_id": "s1", + "image_url": "https://example.test/city.png", + "turns": { + "query": ["I visited Paris.", "Which city did I visit?"], + "search_query": ["visited Paris", "visited city"], + "answer": ["Paris", ""], + }, + "answers": {"ans_full": ["Paris", "Paris"]}, + } + ] + ), + encoding="utf-8", + ) + + cases = build_meta_crag_python_workspaces( + data_root=data_root, + workspace_root=tmp_path / "workspace", + dataset_variant="single_turn", + ) + + q2 = cases[1] + assert q2["case_id"] == "s1::q2" + assert q2["silver_evidence_ids"] == ["turn-0"] + kb = tmp_path / "workspace" / "s1__q2" + assert "Evidence: turn-0" in (kb / "wiki" / "memories" / "turn-0.md").read_text(encoding="utf-8") + retrieval = json.loads((kb / ".kb-research" / "retrieval" / "turn-0.json").read_text(encoding="utf-8")) + assert retrieval["evidence_id"] == "turn-0" + assert "Paris" in retrieval["searchable_text"] + + +def test_run_locomo_refined_offline_eval_reuses_retained_runner(tmp_path) -> None: + dataset = tmp_path / "locomo_refined.json" + dataset.write_text( + json.dumps( + [ + { + "sample_id": "conv-x", + "conversation": { + "session_1": [ + { + "speaker": "Ana", + "dia_id": "D1:1", + "text": "I painted a blue lighthouse.", + "blip_caption": "blue lighthouse painting", + } + ] + }, + "qa": [ + { + "question": "What did Ana paint?", + "answer": "blue lighthouse", + "evidence": ["D1:1"], + "category": 4, + } + ], + "session_summary": {"session_1": "Ana painted a blue lighthouse."}, + "event_summary": {}, + "observation": {}, + } + ] + ), + encoding="utf-8", + ) + + result = run_locomo_refined_offline_eval( + dataset_path=dataset, + output_dir=tmp_path / "out", + workspace_root=tmp_path / "workspace", + top_k=4, + ) + + assert result["summary"]["dataset_name"] == "locomo_refined" + assert result["summary"]["evidence_recall_macro"] == 1.0 + + +def test_run_filesystem_proxy_eval_retrieves_expected_file_stem(tmp_path) -> None: + kb = tmp_path / "kb" + (kb / "wiki" / "memories").mkdir(parents=True) + (kb / "MEMORY.md").write_text("- [image](wiki/memories/image-main.md)", encoding="utf-8") + (kb / "wiki" / "memories" / "image-main.md").write_text( + "koshary egypt national dish", + encoding="utf-8", + ) + cases = tmp_path / "cases.json" + cases.write_text( + json.dumps( + [ + { + "case_id": "c1", + "question": "which country has koshary as national dish?", + "silver_evidence_ids": ["image-main"], + "knowledge_base_root": str(kb), + } + ] + ), + encoding="utf-8", + ) + + result = run_filesystem_proxy_eval( + dataset_name="meta_crag", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=2, + ) + + assert result["summary"]["recall_at_k"] == 1.0 + assert result["summary"]["metric_label_source"] == "answer_derived_proxy" + assert result["cases"][0]["hit_evidence_ids"] == ["image-main"] + + +def test_run_filesystem_proxy_eval_rejects_stale_retrieval_fields(tmp_path) -> None: + kb = tmp_path / "kb" + kb.mkdir() + (kb / "turn.md").write_text("Evidence: hit\nneedle", encoding="utf-8") + cases = tmp_path / "cases.json" + cases.write_text( + json.dumps( + [ + { + "case_id": "polluted", + "question": "needle", + "silver_evidence_ids": ["hit"], + "final_retrieved_ids": ["hit"], + "knowledge_base_root": str(kb), + } + ] + ), + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="stale retrieval"): + run_filesystem_proxy_eval( + dataset_name="meta_crag", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=1, + ) + + +def test_run_filesystem_proxy_eval_reads_yaml_evidence_id(tmp_path) -> None: + kb = tmp_path / "kb" + (kb / "wiki" / "memories").mkdir(parents=True) + (kb / "wiki" / "memories" / "artifact-1.md").write_text( + "\n".join( + [ + "---", + 'artifact_id: "artifact-1"', + 'evidence_id: "silver-1"', + "---", + "", + "# Artifact", + "blue lighthouse dish marker", + ] + ), + encoding="utf-8", + ) + cases = tmp_path / "cases.json" + cases.write_text( + json.dumps( + [ + { + "case_id": "c-yaml", + "question": "Which artifact mentions the blue lighthouse dish marker?", + "silver_evidence_ids": ["silver-1"], + "knowledge_base_root": str(kb), + } + ] + ), + encoding="utf-8", + ) + + result = run_filesystem_proxy_eval( + dataset_name="meta_crag", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=1, + ) + + assert result["summary"]["recall_at_k"] == 1.0 + assert result["cases"][0]["hit_evidence_ids"] == ["silver-1"] + + +def test_run_filesystem_proxy_eval_indexes_meta_crag_retrieval_json(tmp_path) -> None: + kb = tmp_path / "kb" + (kb / "wiki" / "memories").mkdir(parents=True) + (kb / ".kb-research" / "retrieval").mkdir(parents=True) + _mark_python_workspace(kb, "meta_crag") + memory_path = kb / "wiki" / "memories" / "artifact-2.md" + memory_path.write_text( + "\n".join( + [ + "---", + 'artifact_id: "artifact-2"', + 'evidence_id: "silver-2"', + "---", + "", + "# Artifact", + "generic page text", + ] + ), + encoding="utf-8", + ) + (kb / ".kb-research" / "retrieval" / "artifact-2.json").write_text( + json.dumps( + { + "artifact_id": "artifact-2", + "evidence_id": "silver-2", + "title": "Rare orchid propagation", + "searchable_text": "rare orchid should be propagated in april", + "memory_path": str(memory_path), + } + ), + encoding="utf-8", + ) + cases = tmp_path / "cases.json" + cases.write_text( + json.dumps( + [ + { + "case_id": "c-json", + "question": "Should the rare orchid be propagated in april?", + "silver_evidence_ids": ["silver-2"], + "knowledge_base_root": str(kb), + } + ] + ), + encoding="utf-8", + ) + + result = run_filesystem_proxy_eval( + dataset_name="meta_crag", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=1, + ) + + assert result["summary"]["recall_at_k"] == 1.0 + assert result["cases"][0]["hit_evidence_ids"] == ["silver-2"] + + +def test_run_filesystem_proxy_eval_ignores_unprovenanced_retrieval_json(tmp_path) -> None: + kb = tmp_path / "kb" + (kb / ".kb-research" / "retrieval").mkdir(parents=True) + (kb / ".kb-research" / "retrieval" / "borrowed.json").write_text( + json.dumps( + { + "evidence_id": "silver-2", + "title": "Borrowed ranking input", + "searchable_text": "rare orchid should be propagated in april", + } + ), + encoding="utf-8", + ) + cases = tmp_path / "cases.json" + cases.write_text( + json.dumps( + [ + { + "case_id": "unprovenanced-json", + "question": "Should the rare orchid be propagated in april?", + "silver_evidence_ids": ["silver-2"], + "knowledge_base_root": str(kb), + } + ] + ), + encoding="utf-8", + ) + + result = run_filesystem_proxy_eval( + dataset_name="meta_crag", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=1, + ) + + assert result["cases"][0]["retrieved_evidence_ids"] == [] + + +def test_run_filesystem_proxy_eval_summary_ignores_unlabeled_cases(tmp_path) -> None: + kb = tmp_path / "kb" + kb.mkdir() + (kb / "hit.md").write_text("Evidence: hit\nlabeled marker", encoding="utf-8") + cases = tmp_path / "cases.json" + cases.write_text( + json.dumps( + [ + { + "case_id": "unlabeled", + "question": "missing marker", + "silver_evidence_ids": [], + "knowledge_base_root": str(kb), + }, + { + "case_id": "labeled", + "question": "labeled marker", + "silver_evidence_ids": ["hit"], + "knowledge_base_root": str(kb), + }, + ] + ), + encoding="utf-8", + ) + + result = run_filesystem_proxy_eval( + dataset_name="meta_crag", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=1, + ) + + assert result["summary"]["total_cases"] == 2 + assert result["summary"]["evidence_labeled_cases"] == 1 + assert result["summary"]["recall_at_k"] == 1.0 + + +def test_run_filesystem_proxy_eval_recomputes_meta_crag_candidates(tmp_path) -> None: + kb = tmp_path / "kb" + kb.mkdir() + (kb / "turn.md").write_text("Evidence: gold::abc123\nquestion text", encoding="utf-8") + cases = tmp_path / "cases.json" + cases.write_text( + json.dumps( + [ + { + "case_id": "baseline-hit", + "question": "question text", + "silver_evidence_ids": ["gold::abc123"], + "baseline_retrieved_ids": ["baseline::stale"], + "baseline_retrieved_file_paths": ["/tmp/retrieved.txt"], + "knowledge_base_root": str(kb), + } + ] + ), + encoding="utf-8", + ) + + result = run_filesystem_proxy_eval( + dataset_name="meta_crag", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=1, + allow_stale_retrieval_fields=True, + ) + + assert result["summary"]["recall_at_k"] == 1.0 + assert result["cases"][0]["retrieved_evidence_ids"] == ["gold::abc123"] + + +def test_run_filesystem_proxy_eval_ignores_meta_crag_path_ids(tmp_path) -> None: + kb = tmp_path / "kb" + (kb / "wiki" / "turns").mkdir(parents=True) + path = kb / "wiki" / "turns" / "t1_assistant.md" + path.write_text("Evidence: assistant\nrare answer token", encoding="utf-8") + cases = tmp_path / "cases.json" + cases.write_text( + json.dumps( + [ + { + "case_id": "meta-path-id", + "question": "rare answer token", + "silver_evidence_ids": ["assistant"], + "final_retrieved_ids": ["baseline::abc123"], + "final_retrieved_file_paths": [path.as_posix()], + "knowledge_base_root": str(kb), + } + ] + ), + encoding="utf-8", + ) + + result = run_filesystem_proxy_eval( + dataset_name="meta_crag", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=1, + allow_stale_retrieval_fields=True, + ) + + assert result["summary"]["recall_at_k"] == 1.0 + assert result["cases"][0]["retrieved_evidence_ids"] == ["assistant"] + + +def test_run_filesystem_proxy_eval_ignores_meta_crag_final_paths(tmp_path, monkeypatch) -> None: + kb = tmp_path / "kb" + (kb / "wiki" / "turns").mkdir(parents=True) + path = kb / "wiki" / "turns" / "hit.md" + path.write_text("Evidence: hit\nrare answer token", encoding="utf-8") + stale_path = kb / "wiki" / "turns" / "stale.md" + stale_path.write_text("Evidence: stale\nunrelated", encoding="utf-8") + cases = tmp_path / "cases.json" + cases.write_text( + json.dumps( + [ + { + "case_id": "meta-final-paths", + "question": "rare answer token", + "silver_evidence_ids": ["hit"], + "final_retrieved_ids": ["stale"], + "final_retrieved_file_paths": [stale_path.as_posix()], + "knowledge_base_root": str(kb), + } + ] + ), + encoding="utf-8", + ) + monkeypatch.setattr( + example_eval, + "retrieve_qmd_consensus_files", + lambda **_: SimpleNamespace( + files=[ + RetrievedMemoryFile( + filename=path.name, + file_path=path.as_posix(), + mtime_ms=1, + content=path.read_text(encoding="utf-8"), + ) + ] + ), + ) + + result = run_filesystem_proxy_eval( + dataset_name="meta_crag", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=12, + allow_stale_retrieval_fields=True, + ) + + assert result["cases"][0]["retrieved_evidence_ids"] == ["hit"] + + +def test_run_filesystem_proxy_eval_accepts_cases_envelope(tmp_path) -> None: + kb = tmp_path / "kb" + (kb / "wiki" / "observations").mkdir(parents=True) + (kb / "wiki" / "observations" / "D1_4_obs_4.md").write_text( + "artificial intelligence education image", + encoding="utf-8", + ) + cases = tmp_path / "mem_gallery_eval.json" + cases.write_text( + json.dumps( + { + "summary": {}, + "cases": [ + { + "case_id": "gallery::q1", + "question": "Which image is about artificial intelligence education?", + "gold_clue_ids": ["D1:4"], + "knowledge_base_root": str(kb), + } + ], + } + ), + encoding="utf-8", + ) + + result = run_filesystem_proxy_eval( + dataset_name="mem_gallery", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=1, + ) + + assert result["summary"]["recall_at_k"] == 1.0 + + +def test_run_filesystem_proxy_eval_recomputes_mem_gallery_clues(tmp_path) -> None: + kb = tmp_path / "kb" + kb.mkdir() + (kb / "D1_4.md").write_text("Evidence: D1:4\nartificial intelligence education image", encoding="utf-8") + cases = tmp_path / "mem_gallery_eval.json" + cases.write_text( + json.dumps( + { + "cases": [ + { + "case_id": "gallery::q", + "question": "Which image is about artificial intelligence education?", + "gold_clue_ids": ["D1:4"], + "retrieved_clue_ids": ["D9:9"], + "ranked_clues": [{"clue_id": "D9:9", "source_path": "/tmp/D9.md"}], + "knowledge_base_root": str(kb), + } + ] + } + ), + encoding="utf-8", + ) + + result = run_filesystem_proxy_eval( + dataset_name="mem_gallery", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=1, + allow_stale_retrieval_fields=True, + ) + + assert result["summary"]["recall_at_k"] == 1.0 + assert result["cases"][0]["retrieved_evidence_ids"] == ["D1:4"] + + +def test_run_filesystem_proxy_eval_uses_mem_gallery_question_image_caption(tmp_path, monkeypatch) -> None: + kb = tmp_path / "kb" + (kb / ".kb-research" / "retrieval").mkdir(parents=True) + _mark_python_workspace(kb, "mem_gallery") + artifact = kb / ".kb-research" / "retrieval" / "artifact.json" + artifact.write_text( + json.dumps( + { + "evidence_id": "D6:1", + "artifact_id": "artifact-d6-1", + "title": "artifact", + "searchable_text": "cartoon robot artist easel palette", + } + ), + encoding="utf-8", + ) + (kb / ".kb-research" / "retrieval" / "wrong.json").write_text( + json.dumps( + { + "evidence_id": "D9:9", + "artifact_id": "artifact-d9-9", + "title": "artifact", + "searchable_text": "option matches provided picture", + } + ), + encoding="utf-8", + ) + monkeypatch.setattr( + example_eval, + "retrieve_qmd_consensus_files", + lambda **_: SimpleNamespace( + files=[ + RetrievedMemoryFile( + filename="wrong.json", + file_path=(kb / ".kb-research" / "retrieval" / "wrong.json").as_posix(), + mtime_ms=1, + content="Evidence: D9:9\noption matches provided picture", + ) + ] + ), + ) + cases = tmp_path / "mem_gallery_eval.json" + cases.write_text( + json.dumps( + { + "cases": [ + { + "case_id": "gallery::q", + "question": "Which option matches the provided picture?", + "question_image_caption": "cartoon robot artist easel palette", + "gold_clue_ids": ["D6:1"], + "knowledge_base_root": str(kb), + } + ] + } + ), + encoding="utf-8", + ) + + result = run_filesystem_proxy_eval( + dataset_name="mem_gallery", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=2, + ) + + assert result["summary"]["recall_at_k"] == 1.0 + + +def test_run_filesystem_proxy_eval_loads_mem_gallery_caption_from_dialog_root(tmp_path, monkeypatch) -> None: + dialog_root = tmp_path / "dialog" + dialog_root.mkdir() + (dialog_root / "Topic.json").write_text( + json.dumps( + { + "human-annotated QAs": [ + { + "question": "Which option matches the provided picture?", + "image_caption": "cartoon robot artist easel palette", + "session_id": ["D6"], + } + ] + } + ), + encoding="utf-8", + ) + kb = tmp_path / "kb" + (kb / ".kb-research" / "retrieval").mkdir(parents=True) + _mark_python_workspace(kb, "mem_gallery") + (kb / ".kb-research" / "retrieval" / "artifact.json").write_text( + json.dumps( + { + "evidence_id": "D6:1", + "artifact_id": "artifact-d6-1", + "title": "artifact", + "searchable_text": "cartoon robot artist easel palette", + } + ), + encoding="utf-8", + ) + monkeypatch.setattr( + example_eval, + "retrieve_qmd_consensus_files", + lambda **_: SimpleNamespace(files=[]), + ) + cases = tmp_path / "mem_gallery_eval.json" + cases.write_text( + json.dumps( + { + "cases": [ + { + "case_id": "Topic::q1", + "question": "Which option matches the provided picture?", + "gold_clue_ids": ["D6:1"], + "knowledge_base_root": str(kb), + } + ] + } + ), + encoding="utf-8", + ) + + result = run_filesystem_proxy_eval( + dataset_name="mem_gallery", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=1, + mem_gallery_dialog_root=dialog_root, + ) + + assert result["summary"]["recall_at_k"] == 1.0 + + +def test_run_filesystem_proxy_eval_ignores_mem_gallery_retrieved_paths(tmp_path, monkeypatch) -> None: + kb = tmp_path / "kb" + (kb / "wiki" / "observations").mkdir(parents=True) + (kb / "wiki" / "observations" / "D1_4_obs_4.md").write_text("Evidence: D1:4\nplain image", encoding="utf-8") + (kb / "wiki" / "observations" / "D9_9_obs_9.md").write_text( + "Evidence: D9:9\nmatching query terms", + encoding="utf-8", + ) + cases = tmp_path / "mem_gallery_eval.json" + cases.write_text( + json.dumps( + { + "cases": [ + { + "case_id": "gallery::q", + "question": "matching query terms", + "gold_clue_ids": ["D9:9"], + "retrieved_clue_ids": ["D9:9"], + "retrieved_file_paths": [ + (kb / "wiki" / "observations" / "D1_4_obs_4.md").as_posix() + ], + "knowledge_base_root": str(kb), + } + ] + } + ), + encoding="utf-8", + ) + hit_path = kb / "wiki" / "observations" / "D9_9_obs_9.md" + monkeypatch.setattr( + example_eval, + "retrieve_qmd_consensus_files", + lambda **_: SimpleNamespace( + files=[ + RetrievedMemoryFile( + filename=hit_path.name, + file_path=hit_path.as_posix(), + mtime_ms=1, + content=hit_path.read_text(encoding="utf-8"), + ) + ] + ), + ) + + result = run_filesystem_proxy_eval( + dataset_name="mem_gallery", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=1, + allow_stale_retrieval_fields=True, + ) + + assert result["summary"]["recall_at_k"] == 1.0 + assert result["cases"][0]["retrieved_evidence_ids"] == ["D9:9"] + + +def test_run_filesystem_proxy_eval_reads_bulleted_evidence_marker(tmp_path) -> None: + kb = tmp_path / "kb" + (kb / "wiki" / "observations").mkdir(parents=True) + path = kb / "wiki" / "observations" / "D1_4_obs_4.md" + path.write_text("- Session: D1\n- Evidence: D1:4\nplain image", encoding="utf-8") + cases = tmp_path / "mem_gallery_eval.json" + cases.write_text( + json.dumps( + { + "cases": [ + { + "case_id": "gallery::q", + "question": "plain image", + "gold_clue_ids": ["D1:4"], + "knowledge_base_root": str(kb), + } + ] + } + ), + encoding="utf-8", + ) + + result = run_filesystem_proxy_eval( + dataset_name="mem_gallery", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=1, + ) + + assert result["cases"][0]["retrieved_evidence_ids"] == ["D1:4"] + + +def test_run_filesystem_proxy_eval_maps_source_relative_turn_links(tmp_path) -> None: + kb = tmp_path / "kb" + (kb / "wiki" / "sources").mkdir(parents=True) + (kb / "wiki" / "turns").mkdir(parents=True) + source = kb / "wiki" / "sources" / "session_6.md" + turn = kb / "wiki" / "turns" / "D6_1.md" + turn2 = kb / "wiki" / "turns" / "D6_2.md" + source.write_text( + "- [turn D6:1](../turns/D6_1.md) matching query\n" + "- [turn D6:2](../turns/D6_2.md) matching query", + encoding="utf-8", + ) + turn.write_text("Evidence: D6:1\nmatching query", encoding="utf-8") + turn2.write_text("Evidence: D6:2\nmatching query", encoding="utf-8") + cases = tmp_path / "mem_gallery_eval.json" + cases.write_text( + json.dumps( + { + "cases": [ + { + "case_id": "gallery::q", + "question": "matching query", + "gold_clue_ids": ["D6:1", "D6:2"], + "knowledge_base_root": str(kb), + } + ] + } + ), + encoding="utf-8", + ) + + result = run_filesystem_proxy_eval( + dataset_name="mem_gallery", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=1, + ) + + assert result["cases"][0]["retrieved_evidence_ids"] == ["D6:1"] + + +def test_run_filesystem_proxy_eval_maps_mem_gallery_clue_summary_source(tmp_path) -> None: + kb = tmp_path / "kb" + (kb / "wiki" / "memories").mkdir(parents=True) + path = kb / "wiki" / "memories" / "clue-summary-d2-24-9c29c32f.md" + path.write_text( + 'sources: ["clue:D2:D2:24"]\n\n# Clue summary D2:24\n- source ref: clue:D2:D2:24', + encoding="utf-8", + ) + cases = tmp_path / "mem_gallery_eval.json" + cases.write_text( + json.dumps( + { + "cases": [ + { + "case_id": "gallery::q", + "question": "summary", + "gold_clue_ids": ["D2:24"], + "knowledge_base_root": str(kb), + } + ] + } + ), + encoding="utf-8", + ) + + result = run_filesystem_proxy_eval( + dataset_name="mem_gallery", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=1, + ) + + assert result["cases"][0]["retrieved_evidence_ids"] == ["D2:24"] + + +def test_run_filesystem_proxy_eval_prefers_mem_gallery_support_map(tmp_path) -> None: + kb = tmp_path / "kb" + (kb / "wiki" / "memories").mkdir(parents=True) + (kb / ".mem-gallery").mkdir(parents=True) + path = kb / "wiki" / "memories" / "clue-summary-d6-1.md" + path.write_text('summary\n- source ref: clue:D6:D6:1', encoding="utf-8") + (kb / ".mem-gallery" / "artifact_support_map.json").write_text( + json.dumps( + [ + { + "memory_path": path.as_posix(), + "linked_clue_ids": ["D6:IMG_001"], + } + ] + ), + encoding="utf-8", + ) + cases = tmp_path / "mem_gallery_eval.json" + cases.write_text( + json.dumps( + { + "cases": [ + { + "case_id": "gallery::q", + "question": "summary", + "gold_clue_ids": ["D6:IMG_001"], + "knowledge_base_root": str(kb), + } + ] + } + ), + encoding="utf-8", + ) + + result = run_filesystem_proxy_eval( + dataset_name="mem_gallery", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=1, + ) + + assert result["cases"][0]["retrieved_evidence_ids"] == ["D6:IMG_001"] + + +def test_run_filesystem_proxy_eval_reads_mem_gallery_retrieval_json_support(tmp_path) -> None: + kb = tmp_path / "kb" + (kb / "wiki" / "memories").mkdir(parents=True) + (kb / ".mem-gallery").mkdir(parents=True) + (kb / ".kb-research" / "retrieval").mkdir(parents=True) + _mark_python_workspace(kb, "mem_gallery") + memory = kb / "wiki" / "memories" / "turn-image-d6-3.md" + memory.write_text("image placeholder", encoding="utf-8") + (kb / ".mem-gallery" / "artifact_support_map.json").write_text( + json.dumps( + [ + { + "memory_path": memory.as_posix(), + "linked_clue_ids": ["D6:IMG_001"], + } + ] + ), + encoding="utf-8", + ) + (kb / ".kb-research" / "retrieval" / "turn-image-d6-3.json").write_text( + json.dumps( + { + "artifact_id": "turn-image-d6-3", + "memory_path": memory.as_posix(), + "title": "Education robot image", + "searchable_text": "bright classroom robot tutoring students", + } + ), + encoding="utf-8", + ) + cases = tmp_path / "mem_gallery_eval.json" + cases.write_text( + json.dumps( + { + "cases": [ + { + "case_id": "gallery::q", + "question": "Which image shows a robot tutoring students?", + "gold_clue_ids": ["D6:IMG_001"], + "knowledge_base_root": str(kb), + } + ] + } + ), + encoding="utf-8", + ) + + result = run_filesystem_proxy_eval( + dataset_name="mem_gallery", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=1, + ) + + assert result["cases"][0]["retrieved_evidence_ids"] == ["D6:IMG_001"] + + +def test_mem_gallery_source_session_labels_do_not_affect_retrieval( + tmp_path, + monkeypatch, +) -> None: + kb = tmp_path / "kb" + retrieval = kb / ".kb-research" / "retrieval" + retrieval.mkdir(parents=True) + _mark_python_workspace(kb, "mem_gallery") + (retrieval / "labeled-session.json").write_text( + json.dumps( + { + "evidence_id": "D6:1", + "title": "unrelated artifact", + "searchable_text": "plain unrelated note", + } + ), + encoding="utf-8", + ) + (retrieval / "matching.json").write_text( + json.dumps( + { + "evidence_id": "D9:9", + "title": "telescope artifact", + "searchable_text": "blue telescope on balcony", + } + ), + encoding="utf-8", + ) + monkeypatch.setattr( + example_eval, + "retrieve_qmd_consensus_files", + lambda **_: SimpleNamespace(files=[]), + ) + cases = tmp_path / "mem_gallery_eval.json" + cases.write_text( + json.dumps( + { + "cases": [ + { + "case_id": "gallery::q", + "question": "Which clue mentions a blue telescope?", + "gold_clue_ids": ["D9:9"], + "source_session_ids": ["D6"], + "knowledge_base_root": str(kb), + } + ] + } + ), + encoding="utf-8", + ) + + result = run_filesystem_proxy_eval( + dataset_name="mem_gallery", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=1, + ) + + assert result["cases"][0]["retrieved_evidence_ids"] == ["D9:9"] + + +def test_run_filesystem_proxy_eval_reuses_workspace_read_for_same_root(tmp_path, monkeypatch) -> None: + kb = tmp_path / "kb" + kb.mkdir() + (kb / "D1_4.md").write_text("artificial intelligence education image", encoding="utf-8") + cases = tmp_path / "cases.json" + cases.write_text( + json.dumps( + [ + { + "case_id": "gallery::q1", + "question": "artificial intelligence education", + "gold_clue_ids": ["D1:4"], + "knowledge_base_root": str(kb), + }, + { + "case_id": "gallery::q2", + "question": "education image", + "gold_clue_ids": ["D1:4"], + "knowledge_base_root": str(kb), + }, + ] + ), + encoding="utf-8", + ) + calls = [] + original = example_eval._read_workspace_files + + def counted(root, *, include_retrieval_json=False): + calls.append(root) + return original(root, include_retrieval_json=include_retrieval_json) + + monkeypatch.setattr(example_eval, "_read_workspace_files", counted) + + result = run_filesystem_proxy_eval( + dataset_name="mem_gallery", + cases_path=cases, + output_dir=tmp_path / "out", + top_k=1, + ) + + assert len(calls) == 1 + assert result["summary"]["recall_at_k"] == 1.0 + + +def test_run_evermembench_python_eval_reads_dialogue_references(tmp_path) -> None: + topic = tmp_path / "01" + topic.mkdir() + (topic / "dialogue.json").write_text( + json.dumps( + [ + { + "topic_id": "01", + "date": "2025-01-01", + "dialogues": { + "Empty Group": None, + "Group 1": [ + { + "speaker": "Lin", + "time": "2025-01-01 09:00:00", + "dialogue": "The peak CPU usage was 65 percent.", + "message_index": 4, + } + ] + }, + } + ] + ), + encoding="utf-8", + ) + (topic / "qa_01.json").write_text( + json.dumps( + [ + { + "topic_id": "01", + "id": "q1", + "Q": "What was the peak CPU usage?", + "A": "65 percent", + "R": [{"date": "2025-01-01", "group": "Group 1", "message_index": "4"}], + } + ] + ), + encoding="utf-8", + ) + + result = run_evermembench_python_eval( + data_root=tmp_path, + output_dir=tmp_path / "out", + top_k=1, + ) + + assert result["summary"]["recall_at_k"] == 1.0 + assert result["cases"][0]["hit_evidence_ids"] == ["D1:4"] + + +def test_run_evermembench_python_eval_supports_topic_and_question_shards(tmp_path) -> None: + for topic_id in ("01", "02"): + topic = tmp_path / topic_id + topic.mkdir() + (topic / "dialogue.json").write_text( + json.dumps( + [ + { + "topic_id": topic_id, + "date": "2025-01-01", + "dialogues": { + "Group 1": [ + { + "speaker": "Lin", + "time": "2025-01-01 09:00:00", + "dialogue": f"Topic {topic_id} alpha marker.", + "message_index": 1, + }, + { + "speaker": "Lin", + "time": "2025-01-01 09:01:00", + "dialogue": f"Topic {topic_id} beta marker.", + "message_index": 2, + }, + ] + }, + } + ] + ), + encoding="utf-8", + ) + (topic / f"qa_{topic_id}.json").write_text( + json.dumps( + [ + { + "topic_id": topic_id, + "id": "q1", + "Q": f"What alpha marker belongs to topic {topic_id}?", + "A": "alpha", + "R": [{"date": "2025-01-01", "group": "Group 1", "message_index": "1"}], + }, + { + "topic_id": topic_id, + "id": "q2", + "Q": f"What beta marker belongs to topic {topic_id}?", + "A": "beta", + "R": [{"date": "2025-01-01", "group": "Group 1", "message_index": "2"}], + }, + ] + ), + encoding="utf-8", + ) + + result = run_evermembench_python_eval( + data_root=tmp_path, + output_dir=tmp_path / "out", + top_k=2, + topic_names=["02"], + question_offset=1, + question_limit=1, + ) + + assert [case["case_id"] for case in result["cases"]] == ["02::q2"] + assert result["cases"][0]["hit_evidence_ids"] == ["D1:2"] + + +def test_run_evermembench_python_eval_expands_dialogue_neighbors(tmp_path) -> None: + topic = tmp_path / "01" + topic.mkdir() + (topic / "dialogue.json").write_text( + json.dumps( + [ + { + "topic_id": "01", + "date": "2025-01-01", + "dialogues": { + "Group 1": [ + { + "speaker": "Lin", + "time": "2025-01-01 09:00:00", + "dialogue": "CPU setup marker was discussed first.", + "message_index": 1, + }, + { + "speaker": "Lin", + "time": "2025-01-01 09:01:00", + "dialogue": "The value was 65 percent.", + "message_index": 2, + }, + { + "speaker": "Lin", + "time": "2025-01-01 09:02:00", + "dialogue": "Unrelated distractor.", + "message_index": 3, + }, + ] + }, + } + ] + ), + encoding="utf-8", + ) + (topic / "qa_01.json").write_text( + json.dumps( + [ + { + "topic_id": "01", + "id": "q-neighbor", + "Q": "What followed the CPU setup marker?", + "A": "65 percent", + "R": [{"date": "2025-01-01", "group": "Group 1", "message_index": "2"}], + } + ] + ), + encoding="utf-8", + ) + + result = run_evermembench_python_eval( + data_root=tmp_path, + output_dir=tmp_path / "out", + top_k=2, + ) + + assert result["summary"]["recall_at_k"] == 1.0 + assert result["cases"][0]["hit_evidence_ids"] == ["D1:2"] + + +def test_evermem_retrieval_adds_same_group_neighbors() -> None: + rows = [ + { + "evidence_id": "D1:1", + "session_id": "D1", + "date": "2025-01-01", + "group": "Group 1", + }, + { + "evidence_id": "D1:2", + "session_id": "D1", + "date": "2025-01-01", + "group": "Group 1", + }, + { + "evidence_id": "D1:3", + "session_id": "D1", + "date": "2025-01-01", + "group": "Group 1", + }, + ] + files = [ + RetrievedMemoryFile( + filename=f"{index}.md", + file_path=f"/{index}.md", + mtime_ms=index, + content=content, + ) + for index, content in enumerate( + ["CPU setup marker.", "The value was 65 percent.", "Unrelated distractor."], + start=1, + ) + ] + evidence_by_path = { + file.file_path: row["evidence_id"] for file, row in zip(files, rows) + } + + retrieved, _ = example_eval._retrieve_evermem_evidence( + "What followed the CPU setup marker?", + files=files, + evidence_by_path=evidence_by_path, + rows=rows, + top_k=2, + ) + + assert retrieved == ["D1:1", "D1:2"] diff --git a/tests/unit/evaluation/test_wikimem_longmemeval.py b/tests/unit/evaluation/test_wikimem_longmemeval.py new file mode 100644 index 00000000..ac36da91 --- /dev/null +++ b/tests/unit/evaluation/test_wikimem_longmemeval.py @@ -0,0 +1,216 @@ +"""LongMemEval example harness compatibility tests.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest + +from evaluation.wikimem import ( + ConversationRecord, + LoCoMoQuestion, + PreparedSample, + SessionEvents, + adapt_longmemeval_to_locomo_samples, + load_longmemeval_samples, + run_python_longmemeval_retrieval_eval, + run_retained_qmd_eval, +) +from evaluation.wikimem.longmemeval import convert_retained_output_to_longmemeval +from evaluation.wikimem.retained_eval import build_retained_memory_files + +pytestmark = pytest.mark.unit + + +def _longmemeval_rows() -> list[dict]: + return [ + { + "question_id": "q-1", + "question_type": "single-session-user", + "question": "What hobby do I enjoy?", + "answer": "Chess", + "question_date": "2024-01-03", + "haystack_session_ids": ["sess-a", "sess-b"], + "haystack_dates": ["2024-01-01", "2024-01-02"], + "haystack_sessions": [ + [ + { + "role": "user", + "content": "I really enjoy playing chess.", + "has_answer": True, + }, + {"role": "assistant", "content": "That's a great hobby."}, + ], + [ + {"role": "user", "content": "I also went grocery shopping."}, + {"role": "assistant", "content": "Nice."}, + ], + ], + "answer_session_ids": ["sess-a"], + } + ] + + +def test_load_longmemeval_samples_maps_session_and_turn_targets(tmp_path) -> None: + dataset = tmp_path / "longmemeval.json" + dataset.write_text(json.dumps(_longmemeval_rows()), encoding="utf-8") + + samples = load_longmemeval_samples(dataset) + + assert len(samples) == 1 + assert samples[0]["question_id"] == "q-1" + assert samples[0]["answer_session_ids"] == ["sess-a"] + assert samples[0]["answer_turn_ids"] == ["sess-a_1"] + assert samples[0]["session_documents"][0]["corpus_id"] == "sess-a" + assert samples[0]["turn_documents"][0]["corpus_id"] == "sess-a_1" + + +def test_load_longmemeval_samples_uses_raw_turn_index_for_answer_turn_ids(tmp_path) -> None: + rows = _longmemeval_rows() + rows[0]["haystack_sessions"][0] = [ + {"role": "assistant", "content": "Preface."}, + {"role": "user", "content": "I enjoy chess.", "has_answer": True}, + ] + dataset = tmp_path / "longmemeval.json" + dataset.write_text(json.dumps(rows), encoding="utf-8") + + samples = load_longmemeval_samples(dataset) + + assert samples[0]["answer_turn_ids"] == ["sess-a_2"] + assert samples[0]["_locomo_evidence_by_turn_id"]["sess-a_2"] == "D1:1" + + +def test_run_python_longmemeval_retrieval_eval_writes_summary(tmp_path) -> None: + dataset = tmp_path / "longmemeval.json" + dataset.write_text(json.dumps(_longmemeval_rows()), encoding="utf-8") + output_dir = tmp_path / "output" + + result = run_python_longmemeval_retrieval_eval( + dataset_path=dataset, + output_dir=output_dir, + workspace_root=tmp_path / "workspace", + top_k=4, + granularity="turn", + ) + + assert result["summary"]["dataset_name"] == "longmemeval" + assert result["summary"]["total_cases"] == 1 + assert result["summary"]["averaged_metrics"]["turn"]["recall_any@1"] == 1.0 + assert result["summary"]["averaged_metrics"]["turn"]["recall_all@1"] == 1.0 + assert result["summary"]["averaged_metrics"]["turn"]["ndcg_any@1"] == 1.0 + assert result["cases"][0]["retrieval_results"]["metrics"]["turn"]["recall_any@1"] == 1.0 + assert (output_dir / "longmemeval_retrieval_eval.json").exists() + assert (output_dir / "harness" / "run_manifest.json").exists() + + +def test_run_python_longmemeval_retrieval_eval_skips_turn_cases_without_answer_turns( + tmp_path, +) -> None: + rows = _longmemeval_rows() + rows[0]["question_id"] = "q-no-turn-target" + rows[0]["haystack_sessions"][0][0]["has_answer"] = False + dataset = tmp_path / "longmemeval.json" + dataset.write_text(json.dumps(rows), encoding="utf-8") + output_dir = tmp_path / "output" + + result = run_python_longmemeval_retrieval_eval( + dataset_path=dataset, + output_dir=output_dir, + workspace_root=tmp_path / "workspace", + top_k=4, + granularity="turn", + ) + + assert result["summary"]["evaluated_cases"] == 0 + assert result["summary"]["skipped_no_target_cases"] == 1 + assert result["summary"]["averaged_metrics"] == {"turn": {}, "session": {}} + assert result["cases"][0]["evaluated"] is False + assert result["cases"][0]["retrieval_results"]["metrics"] == {"turn": {}, "session": {}} + + +def test_longmemeval_adapter_builds_rust_fallback_multiview_files(tmp_path) -> None: + samples = adapt_longmemeval_to_locomo_samples(_longmemeval_rows()) + sample = samples[0] + + assert sample.observations[1][0].evidence_id == "D1:1" + assert sample.event_summaries[1].items_by_speaker["User"] == [ + "Evidence D1:1: I really enjoy playing chess." + ] + paths = [ + file.file_path + for file in build_retained_memory_files(sample, tmp_path / "workspace") + ] + assert any("/wiki/observations/" in path for path in paths) + assert any("/wiki/events/" in path for path in paths) + assert any("/wiki/entities/" in path for path in paths) + + +def test_longmemeval_metrics_count_ranked_page_support_like_rust(tmp_path) -> None: + rows = _longmemeval_rows() + rows[0]["haystack_sessions"][0].insert( + 1, + {"role": "user", "content": "I also enjoy go.", "has_answer": True}, + ) + dataset = tmp_path / "longmemeval.json" + dataset.write_text(json.dumps(rows), encoding="utf-8") + samples = load_longmemeval_samples(dataset) + adapt_longmemeval_to_locomo_samples(samples) + score = SimpleNamespace( + sample_id="q-1", + question=rows[0]["question"], + knowledge_base_root="workspace/q-1", + retrieved_file_paths=["workspace/q-1/wiki/sources/session_1.md"], + retrieved_evidence=[], + ) + + result = convert_retained_output_to_longmemeval( + samples=samples, + retained_cases=[score], + granularity="turn", + workspace_root="workspace", + ) + + case = result["cases"][0] + assert case.retrieval_results.metrics["turn"]["recall_all@1"] == 1.0 + + +def test_run_retained_qmd_eval_uses_multiview_profile_for_longmemeval(tmp_path) -> None: + sample = PreparedSample( + sample_id="long-1", + raw_sample={}, + records=[ + ConversationRecord( + dia_id="D1:1", + session_id="D1", + speaker="User", + text="The answer token is azure-harbor.", + ) + ], + questions=[ + LoCoMoQuestion( + question="Which token did I mention?", + answer="azure-harbor", + evidence=["D1:1"], + ) + ], + session_datetimes={1: "2024-01-01"}, + session_summaries={1: "The answer token is azure-harbor."}, + event_summaries={ + 1: SessionEvents( + date="2024-01-01", + items_by_speaker={"User": ["The answer token is azure-harbor."]}, + ) + }, + observations={}, + ) + + output = run_retained_qmd_eval( + dataset_name="longmemeval", + samples=[sample], + workspace_root=tmp_path / "workspace", + top_k=8, + ) + + paths = output.cases[0].retrieved_file_paths + assert any("/wiki/events/" in path or "/wiki/entities/" in path for path in paths) diff --git a/tests/unit/evaluation/test_wikimem_qmd_consensus.py b/tests/unit/evaluation/test_wikimem_qmd_consensus.py new file mode 100644 index 00000000..b3aae35a --- /dev/null +++ b/tests/unit/evaluation/test_wikimem_qmd_consensus.py @@ -0,0 +1,189 @@ +"""qmd_consensus retrieval plugin compatibility tests.""" + +from __future__ import annotations + +from evaluation.wikimem.qmd_consensus import ( + CandidateFile, + RetrievedMemoryFile, + apply_query_augmentation, + build_cached_file_lexical_features, + build_qmd_consensus_augmentation, + build_qmd_consensus_candidate_proposals, + build_qmd_consensus_late_bridge_proposals, + build_qmd_consensus_rerank_proposals, + build_question_profile, + candidate_query_hits_with_features, + normalize_memory_path, + qmd_consensus_is_conservative, + score_line, +) + + +def test_build_question_profile_matches_rust_flags_and_tokens() -> None: + profile = build_question_profile( + "Where would Alice pursue art planning with Bob?", + ["Alice", "Bob", "Charlie"], + ) + + assert profile.query_tokens == ["alice", "pursue", "art", "planning", "with", "bob"] + assert "Alice" in profile.named_entities + assert "Bob" in profile.named_entities + assert profile.location is True + assert profile.hypothetical is True + assert qmd_consensus_is_conservative(profile) is False + + +def test_qmd_conservative_profile_skips_plain_temporal_questions() -> None: + profile = build_question_profile("When did Alice visit the museum?", ["Alice"]) + + assert profile.temporal is True + assert qmd_consensus_is_conservative(profile) is True + + +def test_cached_features_collect_normalized_bridge_targets() -> None: + source = _file( + "/tmp/kb/wiki/sources/session_1.md", + "[turn](../turns/T7.md)\n[event](../events/session_2.md)\n[web](https://example.com)", + ) + features = build_cached_file_lexical_features(source) + + assert normalize_memory_path(r"C:\tmp\KB\wiki\turns\T7.md") == "c:/tmp/kb/wiki/turns/t7.md" + assert features.bridge_target_paths == [ + "/tmp/kb/wiki/turns/t7.md", + "/tmp/kb/wiki/events/session_2.md", + ] + + +def test_cached_features_reuse_equal_files_without_stale_content() -> None: + build_cached_file_lexical_features.cache_clear() + file = _file("/tmp/kb/wiki/turns/T7.md", "Alice inspected the drawer.") + + first = build_cached_file_lexical_features(file) + second = build_cached_file_lexical_features(file) + changed = build_cached_file_lexical_features( + _file("/tmp/kb/wiki/turns/T7.md", "Bob inspected the drawer.") + ) + + assert second is first + assert changed is not first + + +def test_score_line_uses_exact_fuzzy_phrase_and_soft_overlap() -> None: + profile = build_question_profile("What musical activities did Alice pursue?", ["Alice"]) + + assert score_line("Alice pursued musicals at the academy.", profile, profile.question) > 6.0 + assert score_line("Unrelated weather note.", profile, profile.question) == 0.0 + + +def test_candidate_query_hits_counts_token_and_fuzzy_overlap() -> None: + profile = build_question_profile("What did Alice inspect in the drawer?", ["Alice"]) + features = build_cached_file_lexical_features( + _file("/tmp/kb/wiki/turns/T7.md", "Alice inspected the drawer and found tools.") + ) + + assert candidate_query_hits_with_features(features, profile) >= 4 + + +def test_qmd_augmentation_uses_supported_seed_lines() -> None: + profile = build_question_profile("What musical activities would Alice pursue?", ["Alice"]) + root_files = [ + _file( + "/tmp/kb/wiki/sources/session_1.md", + "Alice discussed ceramics and musical theatre with Bob.\n" + "Alice discussed ceramics and musical theatre again.", + ), + _file( + "/tmp/kb/wiki/entities/alice.md", + "Alice would pursue ceramics and musical theatre after classes.", + ), + ] + + augmentation = build_qmd_consensus_augmentation(profile.question, profile, root_files) + + assert "ceramics" in augmentation.tokens + assert any(phrase == "musical theatre" for phrase in augmentation.phrases) + expanded = apply_query_augmentation(profile, augmentation) + assert "ceramics" in expanded.expansion_tokens + + +def test_qmd_candidate_proposals_fuse_source_anchor_and_linked_views() -> None: + files = [ + _file( + "/tmp/kb/wiki/sources/session_1.md", + "## Summary\nAlice inspected the drawer and found brass tools.\n" + "## Turn Index\n- [turn 7](../turns/T7.md)", + ), + _file( + "/tmp/kb/wiki/observations/D1_obs.md", + "- Session: D1\n- Evidence: D1:7\nAlice found brass tools in the drawer.", + ), + _file( + "/tmp/kb/wiki/turns/T7.md", + "- Session: D1\n- Evidence: D1:7\nAlice inspected the drawer for brass tools.", + ), + ] + profile = build_question_profile("What did Alice find in the drawer?", ["Alice"]) + + proposals = build_qmd_consensus_candidate_proposals(profile.question, profile, files) + + paths = [proposal.file_path for proposal in proposals] + assert "/tmp/kb/wiki/sources/session_1.md" in paths + assert "/tmp/kb/wiki/observations/d1_obs.md" in paths + assert "/tmp/kb/wiki/turns/t7.md" in paths + + +def test_qmd_rerank_proposals_add_linked_targets_from_seed_sources() -> None: + files = [ + _file( + "/tmp/kb/wiki/sources/session_1.md", + "## Summary\nAlice inspected the drawer and found brass tools.\n" + "## Turn Index\n- [turn 7](../turns/T7.md)", + ), + _file( + "/tmp/kb/wiki/turns/T7.md", + "- Session: D1\n- Evidence: D1:7\nAlice inspected the drawer and found brass tools.", + ), + ] + profile = build_question_profile("What did Alice find in the drawer?", ["Alice"]) + ranked = [CandidateFile(file=files[0], query_hits=3, seed_boost=7.0)] + + proposals = build_qmd_consensus_rerank_proposals(profile.question, profile, ranked, files) + + assert any(proposal.file_path == "/tmp/kb/wiki/turns/t7.md" for proposal in proposals) + + +def test_qmd_late_bridge_proposals_follow_seed_links_to_cross_session_targets() -> None: + files = [ + _file( + "/tmp/kb/wiki/sources/session_1.md", + "## Summary\nAlice mentioned a linked workshop note.\n" + "## Turn Index\n- [turn 8](../turns/T8.md)", + ), + _file( + "/tmp/kb/wiki/turns/T8.md", + "- Session: D8\n- Evidence: D8:2\n" + "The brass tools were displayed during the workshop.", + ), + ] + profile = build_question_profile( + "What activities did Alice connect with brass tools?", + ["Alice"], + ) + + proposals = build_qmd_consensus_late_bridge_proposals( + profile.question, + profile, + seed_files=[files[0]], + files=files, + ) + + assert [proposal.file_path for proposal in proposals] == ["/tmp/kb/wiki/turns/t8.md"] + + +def _file(path: str, content: str) -> RetrievedMemoryFile: + return RetrievedMemoryFile( + filename=path.rsplit("/", maxsplit=1)[-1], + file_path=path, + mtime_ms=0, + content=content, + ) diff --git a/tests/unit/evaluation/test_wikimem_retained_eval.py b/tests/unit/evaluation/test_wikimem_retained_eval.py new file mode 100644 index 00000000..3f43ea1b --- /dev/null +++ b/tests/unit/evaluation/test_wikimem_retained_eval.py @@ -0,0 +1,568 @@ +"""wikimem retained_eval compatibility tests.""" + +from __future__ import annotations + +import json +from copy import deepcopy + +import pytest + +from evaluation.wikimem import ( + CaseScore, + EvalHarnessConfig, + EvalOutput, + ProgressUpdate, + RetrievalCoverageSummary, + RetrievedMemoryFile, + StageProfileArtifact, + StageProfiler, + StageTimingRecord, + build_eval_cases, + build_question_profile, + format_progress_message, + parse_retrieval_plugin_list, + parse_sample_filter, + prepare_locomo_samples, + run_python_locomo_retrieval_eval, + run_retained_qmd_eval, + summarize_scores, + summarize_scores_by_locomo_category, + write_harness_artifacts, +) +from evaluation.wikimem.retained_eval import ( + _extract_retrieved_evidence_ids, + build_retained_memory_files, +) + +pytestmark = pytest.mark.unit + + +def _sample_payload() -> dict: + return { + "sample_id": "conv-1", + "conversation": { + "session_2": [ + { + "speaker": "Alice", + "dia_id": "D2:1", + "text": "I adopted a dog.", + "blip_caption": "photo of dog", + } + ], + "session_1": [ + {"speaker": "Bob", "dia_id": "D1:1", "text": "I like chess."} + ], + "session_1_date_time": "2026-01-01", + }, + "qa": [ + { + "question": "What game does Bob like?", + "answer": "chess", + "evidence": ["D1:1"], + "category": 4, + }, + { + "question": "What animal did Alice adopt?", + "answer": ["dog"], + "evidence": ["D2:1"], + "category": 2, + }, + ], + "session_summary": {"session_1": "Bob likes chess."}, + "event_summary": { + "session_2": {"date": "2026-01-02", "Alice": ["Adopted a dog"]} + }, + "observation": {"session_2": {"Alice": [{"evidence_id": "D2:1", "text": "dog"}]}}, + } + + +def test_prepare_locomo_samples_preserves_records_questions_and_notes() -> None: + prepared = prepare_locomo_samples([_sample_payload()]) + + assert len(prepared) == 1 + sample = prepared[0] + assert sample.sample_id == "conv-1" + assert [record.session_id for record in sample.records] == ["D1", "D2"] + assert sample.records[1].text == "I adopted a dog. [caption: photo of dog]" + assert sample.session_datetimes == {1: "2026-01-01"} + assert sample.session_summaries == {1: "Bob likes chess."} + assert sample.event_summaries[2].items_by_speaker == {"Alice": ["Adopted a dog"]} + assert sample.observations[2][0].evidence_id == "D2:1" + + +def test_retained_builder_adds_rust_style_multimodal_memory_pages() -> None: + sample = prepare_locomo_samples([_sample_payload()])[0] + + memory_pages = [ + file + for file in build_retained_memory_files(sample, "/tmp/kb") + if "/wiki/memories/" in file.file_path + ] + + assert len(memory_pages) == 1 + assert memory_pages[0].file_path.endswith("D2_1_multimodal.md") + assert "Evidence: D2:1" in memory_pages[0].content + assert "photo of dog" in memory_pages[0].content + + +def test_prepare_locomo_samples_parses_rust_session_summary_keys() -> None: + payload = _sample_payload() + payload["session_summary"] = {"session_1_summary": "Bob likes chess."} + + sample = prepare_locomo_samples([payload])[0] + + assert sample.session_summaries == {1: "Bob likes chess."} + + +def test_prepare_locomo_samples_keeps_events_session_summaries_in_kb(tmp_path) -> None: + payload = _sample_payload() + payload["event_summary"] = { + "events_session_2": { + "date": "2026-01-02", + "Alice": ["Alice adopted a dog after lunch."], + } + } + + sample = prepare_locomo_samples([payload])[0] + files = build_retained_memory_files(sample, tmp_path / "workspace" / sample.sample_id) + + assert sample.event_summaries[2].date == "2026-01-02" + event_file = next( + file for file in files if file.file_path.endswith("/wiki/events/session_2_event_1.md") + ) + assert "Alice adopted a dog after lunch." in event_file.content + assert ( + "- [session_2 events](wiki/topics/session_2_events.md) - dated event index" + in files[0].content + ) + + +def test_prepare_locomo_samples_keeps_session_observation_pairs_in_kb(tmp_path) -> None: + payload = _sample_payload() + payload["observation"] = { + "session_2_observation": { + "Alice": [["Alice adopted a dog after lunch.", "D2:1"]] + } + } + + sample = prepare_locomo_samples([payload])[0] + files = build_retained_memory_files(sample, tmp_path / "workspace" / sample.sample_id) + + assert sample.observations[2][0].evidence_id == "D2:1" + observations_file = next( + file for file in files if file.file_path.endswith("/wiki/observations/D2_1_obs_1.md") + ) + assert "- Evidence: D2:1" in observations_file.content + assert "Alice adopted a dog after lunch." in observations_file.content + + +def test_build_retained_memory_files_writes_entity_pages_with_turn_links(tmp_path) -> None: + sample = prepare_locomo_samples([_sample_payload()])[0] + + files = build_retained_memory_files(sample, tmp_path / "workspace" / sample.sample_id) + + entity_file = next(file for file in files if file.file_path.endswith("/wiki/entities/Alice.md")) + assert "# Alice" in entity_file.content + assert "[turn D2:1](../turns/D2_1.md)" in entity_file.content + assert "- [profile](wiki/synthesis/profile.md)" in files[0].content + + +def test_build_retained_memory_files_caps_entity_turn_links_like_rust(tmp_path) -> None: + payload = _sample_payload() + for session_number in range(3, 15): + payload["conversation"][f"session_{session_number}"] = [ + { + "speaker": "Alice", + "dia_id": f"D{session_number}:1", + "text": f"Alice note {session_number} first.", + }, + { + "speaker": "Alice", + "dia_id": f"D{session_number}:2", + "text": f"Alice note {session_number} last.", + }, + ] + sample = prepare_locomo_samples([payload])[0] + + files = build_retained_memory_files(sample, tmp_path / "workspace" / sample.sample_id) + + entity = next(file for file in files if file.file_path.endswith("/wiki/entities/Alice.md")) + assert entity.content.count("(../turns/") == 10 + + +def test_prepare_locomo_samples_infers_evidence_id_when_turn_dia_id_is_missing() -> None: + payload = _sample_payload() + del payload["conversation"]["session_1"][0]["dia_id"] + + sample = prepare_locomo_samples([payload])[0] + + assert sample.records[0].dia_id == "D1:1" + + +def test_parse_sample_filter_accepts_numbers_commas_and_all() -> None: + assert parse_sample_filter(None) is None + assert parse_sample_filter("all") is None + assert parse_sample_filter("1, conv-2 / custom; 3") == {"conv-1", "conv-2", "custom", "conv-3"} + + +def test_build_eval_cases_uses_one_based_case_ids_and_stringified_answer() -> None: + sample = prepare_locomo_samples([_sample_payload()])[0] + + cases = build_eval_cases([sample], question_limit=1) + + assert len(cases) == 1 + assert cases[0].case_id == "conv-1::q1" + assert cases[0].question_index == 0 + assert cases[0].answer == "chess" + assert cases[0].expected_evidence == ["D1:1"] + + +def test_retained_retrieval_ranking_ignores_answer_and_evidence_labels(tmp_path) -> None: + baseline_payload = _sample_payload() + changed_payload = deepcopy(baseline_payload) + changed_payload["qa"][0]["answer"] = "unrelated answer" + changed_payload["qa"][0]["evidence"] = ["D2:1"] + + baseline = run_retained_qmd_eval( + dataset_name="locomo", + samples=prepare_locomo_samples([baseline_payload]), + workspace_root=tmp_path / "baseline", + top_k=4, + question_limit=1, + ).cases[0] + changed = run_retained_qmd_eval( + dataset_name="locomo", + samples=prepare_locomo_samples([changed_payload]), + workspace_root=tmp_path / "changed", + top_k=4, + question_limit=1, + ).cases[0] + + assert changed.retrieved_evidence == baseline.retrieved_evidence + assert [path.split("/conv-1/", 1)[-1] for path in changed.retrieved_file_paths] == [ + path.split("/conv-1/", 1)[-1] for path in baseline.retrieved_file_paths + ] + + +def test_parse_retrieval_plugin_list_normalizes_and_drops_empty_tokens() -> None: + assert parse_retrieval_plugin_list(" qmd-consensus, sparse_query ,, ama_anchor_consensus ") == [ + "qmd_consensus", + "sparse_query", + "ama_anchor_consensus", + ] + + +def test_summarize_scores_matches_retained_eval_metric_rounding() -> None: + scores = [ + _score("c1", expected=["e1", "e2"], retrieved=["e1"], hits=["e1"], category=4), + _score("c2", expected=["e3"], retrieved=["e3", "x"], hits=["e3"], category=2), + ] + + summary = summarize_scores("locomo", scores) + categories = summarize_scores_by_locomo_category(scores) + + assert summary.total_cases == 2 + assert summary.cases_with_evidence == 2 + assert summary.evidence_precision_macro == 0.75 + assert summary.evidence_precision_micro == 0.6667 + assert summary.evidence_recall_macro == 0.75 + assert summary.evidence_recall_micro == 0.6667 + assert summary.full_evidence_hit_rate == 0.5 + assert [(item.category, item.label, item.evidence_recall_macro) for item in categories] == [ + (2, "2 Temporal", 1.0), + (4, "4 Single Hop", 0.5), + ] + + +def test_write_harness_artifacts_outputs_manifest_cases_and_stage_profile(tmp_path) -> None: + score = _score( + "conv-1::q1", + expected=["D1:1"], + retrieved=[], + hits=[], + category=4, + coverage=RetrievalCoverageSummary(miss_stage="root"), + ) + output = EvalOutput( + summary=summarize_scores("locomo", [score]), + cases=[score], + stage_profile=StageProfileArtifact( + created_at_ms=1, + total_samples=1, + total_cases=1, + stages=[StageTimingRecord(stage="root", calls=1, total_ms=2.0, avg_ms=2.0, max_ms=2.0)], + ), + ) + config = EvalHarnessConfig( + dataset_name="locomo", + samples="conv-1", + question_limit=1, + top_k=24, + workspace_root="/tmp/kb", + llm_provider=None, + retrieval_plugins=["qmd_consensus"], + ) + + write_harness_artifacts(tmp_path, output, config) + + assert (tmp_path / "run_manifest.json").exists() + assert (tmp_path / "failure_report.md").exists() + assert (tmp_path / "stage_profile.json").exists() + assert (tmp_path / "cases" / "conv-1__q1.json").exists() + manifest = json.loads((tmp_path / "run_manifest.json").read_text(encoding="utf-8")) + assert manifest["config"]["retrieval_plugins"] == ["qmd_consensus"] + assert manifest["total_failures"] == 1 + + +def test_stage_profiler_aggregates_and_sorts_stage_timings() -> None: + profiler = StageProfiler() + + profiler.record_duration_ms("candidate_ranking", 2.34567) + profiler.record_duration_ms("root_selector_retrieve", 6.0) + profiler.record_duration_ms("candidate_ranking", 3.0) + + artifact = profiler.snapshot(total_samples=2, total_cases=5, created_at_ms=123) + + assert artifact.created_at_ms == 123 + assert artifact.total_samples == 2 + assert artifact.total_cases == 5 + assert artifact.stages == [ + StageTimingRecord( + stage="root_selector_retrieve", + calls=1, + total_ms=6.0, + avg_ms=6.0, + max_ms=6.0, + ), + StageTimingRecord( + stage="candidate_ranking", + calls=2, + total_ms=5.3457, + avg_ms=2.6728, + max_ms=3.0, + ), + ] + + +def test_run_retained_qmd_eval_retrieves_locomo_evidence_and_writes_artifacts(tmp_path) -> None: + sample = prepare_locomo_samples([_sample_payload()])[0] + + output = run_retained_qmd_eval( + dataset_name="locomo", + samples=[sample], + workspace_root=tmp_path / "workspace", + top_k=4, + question_limit=1, + harness_root=tmp_path / "harness", + ) + + assert output.summary.total_cases == 1 + assert output.summary.evidence_recall_macro == 1.0 + assert output.summary.full_evidence_hit_rate == 1.0 + assert output.cases[0].hit_evidence == ["D1:1"] + assert output.cases[0].retrieval_coverage is not None + assert output.cases[0].retrieval_coverage.final_hit_evidence == ["D1:1"] + assert (tmp_path / "harness" / "run_manifest.json").exists() + assert (tmp_path / "harness" / "stage_profile.json").exists() + + +def test_extract_retrieved_evidence_ids_filters_unmatched_entity_turns() -> None: + profile = build_question_profile("What game did Bob play after dinner?", ["Bob"]) + file = RetrievedMemoryFile( + filename="bob.md", + file_path="/sample/wiki/entities/bob.md", + mtime_ms=1, + content="\n".join( + [ + "- [D1:1](../turns/d1_1.md) After dinner Bob played chess with Nora.", + "- [D1:2](../turns/d1_2.md) Bob mentioned tennis during breakfast.", + ] + ), + ) + + assert _extract_retrieved_evidence_ids(file, "What game did Bob play after dinner?", profile) == [ + "D1:1" + ] + + +def test_extract_retrieved_evidence_ids_keeps_family_entity_context() -> None: + question = "What do Melanie's kids like?" + profile = build_question_profile(question, ["Melanie"]) + file = RetrievedMemoryFile( + filename="melanie.md", + file_path="/sample/wiki/entities/melanie.md", + mtime_ms=1, + content="\n".join( + [ + "- [D4:8](../turns/d4_8.md) Her kids like painting.", + "- [D6:6](../turns/d6_6.md) Her children like piano.", + ] + ), + ) + + assert _extract_retrieved_evidence_ids(file, question, profile) == ["D4:8", "D6:6"] + + +def test_run_retained_qmd_eval_does_not_fill_top_k_with_unranked_source_roots(tmp_path) -> None: + target_session = 52 + payload = { + "sample_id": "conv-many", + "conversation": { + f"session_{index}": [ + { + "speaker": "User", + "dia_id": f"D{index}:1", + "text": ( + "I graduated with a degree in Business Administration." + if index == target_session + else f"I talked about unrelated topic {index}." + ), + } + ] + for index in range(1, 61) + }, + "qa": [ + { + "question": "What degree did I graduate with?", + "answer": "Business Administration", + "evidence": [f"D{target_session}:1"], + "category": 4, + } + ], + "session_summary": { + f"session_{index}": ( + "The user graduated with a Business Administration degree." + if index == target_session + else f"Unrelated session {index}." + ) + for index in range(1, 61) + }, + "event_summary": {}, + "observation": {}, + } + sample = prepare_locomo_samples([payload])[0] + + output = run_retained_qmd_eval( + dataset_name="locomo", + samples=[sample], + workspace_root=tmp_path / "workspace", + top_k=100, + question_limit=1, + ) + + target_path = f"/wiki/sources/session_{target_session}.md" + target_rank = next( + index + for index, path in enumerate(output.cases[0].retrieved_file_paths, start=1) + if path.endswith(target_path) + ) + assert target_rank <= 10 + assert output.cases[0].hit_evidence == [f"D{target_session}:1"] + assert output.cases[0].retrieval_coverage is not None + assert output.cases[0].retrieval_coverage.final_hit_evidence == [f"D{target_session}:1"] + + +def test_run_python_locomo_retrieval_eval_loads_dataset_and_writes_summary(tmp_path) -> None: + dataset = tmp_path / "locomo.json" + dataset.write_text(json.dumps([_sample_payload()]), encoding="utf-8") + output_dir = tmp_path / "output" + + result = run_python_locomo_retrieval_eval( + dataset_path=dataset, + output_dir=output_dir, + workspace_root=tmp_path / "workspace", + top_k=4, + question_limit=1, + ) + + assert result["summary"]["dataset_name"] == "locomo" + assert result["summary"]["total_cases"] == 1 + assert result["summary"]["evidence_recall_macro"] == 1.0 + assert result["cases"][0]["hit_evidence"] == ["D1:1"] + assert (output_dir / "locomo_retrieval_eval.json").exists() + assert (output_dir / "harness" / "run_manifest.json").exists() + + +def test_stage_profile_artifact_merges_additional_stage_and_keeps_rust_sort() -> None: + artifact = StageProfileArtifact( + created_at_ms=123, + total_samples=1, + total_cases=2, + stages=[ + StageTimingRecord( + stage="candidate_ranking", + calls=1, + total_ms=5.0, + avg_ms=5.0, + max_ms=5.0, + ) + ], + ) + + updated = artifact.with_additional_stage("root_selector_retrieve", 5.0) + updated = updated.with_additional_stage("candidate_ranking", 2.25) + + assert updated.stages == [ + StageTimingRecord( + stage="candidate_ranking", + calls=2, + total_ms=7.25, + avg_ms=3.625, + max_ms=5.0, + ), + StageTimingRecord( + stage="root_selector_retrieve", + calls=1, + total_ms=5.0, + avg_ms=5.0, + max_ms=5.0, + ), + ] + + +def test_format_progress_message_matches_rust_shape() -> None: + message = format_progress_message( + ProgressUpdate( + completed_cases=3, + total_cases=9, + sample_index=1, + total_samples=4, + sample_id="conv-1", + question_index=2, + sample_question_total=5, + ) + ) + + assert message == "sample 1/4 conv-1 q2/5 overall 3/9" + + +def _score( + case_id: str, + *, + expected: list[str], + retrieved: list[str], + hits: list[str], + category: int | None = None, + coverage: RetrievalCoverageSummary | None = None, +) -> CaseScore: + return CaseScore( + case_id=case_id, + sample_id=case_id.split("::", maxsplit=1)[0], + question_index=0, + question="question", + category=category, + expected_evidence=expected, + retrieved_evidence=retrieved, + hit_evidence=hits, + retrieved_record_ids=[], + retrieved_file_paths=[], + retrieved_entrypoint_paths=[], + knowledge_base_root="/tmp/kb", + retrieved_record_count=0, + retrieved_file_count=0, + retrieved_entrypoint_count=0, + evidence_precision=0.0 if not retrieved else len(hits) / len(retrieved), + evidence_recall=0.0 if not expected else len(hits) / len(expected), + full_evidence_hit=set(expected) <= set(hits), + retrieval_coverage=coverage, + ) diff --git a/tests/unit/evaluation/test_wikimem_retrieval_profile.py b/tests/unit/evaluation/test_wikimem_retrieval_profile.py new file mode 100644 index 00000000..c21ec0aa --- /dev/null +++ b/tests/unit/evaluation/test_wikimem_retrieval_profile.py @@ -0,0 +1,490 @@ +"""wikimem retained retrieval profile tests.""" + +from __future__ import annotations + +from evaluation.wikimem.qmd_consensus import ( + CandidateFile, + RetrievedMemoryFile, + build_question_profile, +) +from evaluation.wikimem.retrieval_profile import ( + build_corpus_consensus_augmentation, + build_session_source_files, + collect_session_source_companions, + infer_session_number_from_path, + rank_global_session_sources, + retrieve_qmd_consensus_files, + scoped_budgets, + select_scoped_candidate_files, + source_companion_budget, + source_injection_budget, +) + + +def test_retrieve_qmd_consensus_files_preserves_root_then_adds_plugin_candidates() -> None: + files = [ + _file( + "/tmp/kb/wiki/sources/session_1.md", + "## Summary\nAlice inspected the drawer and found brass tools.\n" + "## Turn Index\n- [turn 7](../turns/T7.md)", + ), + _file( + "/tmp/kb/wiki/observations/D1_obs.md", + "- Session: D1\n- Evidence: D1:7\nAlice found brass tools in the drawer.", + ), + _file( + "/tmp/kb/wiki/turns/T7.md", + "- Session: D1\n- Evidence: D1:7\nAlice inspected the drawer for brass tools.", + ), + ] + + result = retrieve_qmd_consensus_files( + question="What did Alice find in the drawer?", + files=files, + root_files=[files[0]], + entity_names=["Alice"], + top_k=3, + ) + + assert result.profile.query_tokens == ["alice", "find", "drawer"] + assert result.files[0].file_path == "/tmp/kb/wiki/sources/session_1.md" + assert {file.file_path for file in result.files} == { + "/tmp/kb/wiki/sources/session_1.md", + "/tmp/kb/wiki/observations/D1_obs.md", + "/tmp/kb/wiki/turns/T7.md", + } + assert result.coverage.root_file_paths == ["/tmp/kb/wiki/sources/session_1.md"] + assert "/tmp/kb/wiki/turns/t7.md" in result.coverage.late_bridge_file_paths + + +def test_auto_root_selection_uses_bounded_projection_before_late_atomic_pages() -> None: + files = [ + _file( + "/tmp/kb/wiki/sources/session_1.md", + "## Summary\nThe session covered a general project update.\n", + ), + _file( + "/tmp/kb/wiki/turns/T1.md", + "- Session: D1\n- Evidence: D1:1\nA general project update.", + ), + ] + files.extend( + _file( + f"/tmp/kb/wiki/observations/D1_obs_{index}.md", + "A neutral observation with no query terms.", + mtime_ms=index, + ) + for index in range(1, 205) + ) + files.append( + _file( + "/tmp/kb/wiki/entities/Caroline.md", + "Caroline identity details that should stay in the scoped pool.", + ) + ) + + result = retrieve_qmd_consensus_files( + question="What is Caroline's identity?", + files=files, + entity_names=["Caroline"], + top_k=3, + ) + + assert all("/wiki/observations/" not in path for path in result.coverage.root_file_paths) + + +def test_auto_root_selection_uses_body_phrase_fallback_for_plain_markdown() -> None: + file = _file( + "/tmp/kb/D1_4.md", + "Evidence: D1:4\nartificial intelligence education image", + ) + + result = retrieve_qmd_consensus_files( + question="education image", + files=[file], + top_k=1, + ) + + assert result.files[0].file_path == "/tmp/kb/D1_4.md" + + +def test_retrieve_qmd_consensus_files_respects_top_k_and_deduplicates_paths() -> None: + files = [ + _file( + "/tmp/kb/wiki/sources/session_1.md", + "## Summary\nAlice inspected the drawer and found brass tools.\n" + "## Turn Index\n- [turn 7](../turns/T7.md)", + ), + _file( + "/tmp/kb/wiki/turns/T7.md", + "- Session: D1\n- Evidence: D1:7\nAlice inspected the drawer for brass tools.", + ), + ] + + result = retrieve_qmd_consensus_files( + question="What did Alice inspect?", + files=files, + root_files=[files[0], files[0]], + entity_names=["Alice"], + top_k=1, + ) + + assert [file.file_path for file in result.files] == ["/tmp/kb/wiki/sources/session_1.md"] + assert result.coverage.final_file_paths == ["/tmp/kb/wiki/sources/session_1.md"] + + +def test_corpus_consensus_augmentation_corrects_query_typo_from_sources() -> None: + files = [ + _file( + "/tmp/kb/wiki/sources/session_1.md", + "## Summary\nAlice repaired the drawer latch with brass tools.\n" + "## Turn Index\n- [turn 1](../turns/T1.md)", + ), + _file( + "/tmp/kb/wiki/sources/session_2.md", + "## Summary\nAlice cleaned the drawer after the repair.\n" + "## Turn Index\n- [turn 2](../turns/T2.md)", + ), + ] + profile = build_question_profile("What did Alice put in the drawre?", ["Alice"]) + + augmentation = build_corpus_consensus_augmentation( + "What did Alice put in the drawre?", + profile, + build_session_source_files(files), + ) + + assert augmentation.tokens == ["drawer"] + assert "drawe" in augmentation.fuzzy_tokens + assert augmentation.phrases == [] + + +def test_corpus_consensus_augmentation_adds_supported_anchor_tokens_and_phrases() -> None: + files = [ + _file( + "/tmp/kb/wiki/sources/session_1.md", + "## Summary\nAlice organized the drawer brass tools for pottery studio work.\n" + "## Turn Index\n- [turn 1](../turns/T1.md)", + ), + _file( + "/tmp/kb/wiki/sources/session_2.md", + "## Summary\nAlice photographed drawer brass tools beside ceramic glaze.\n" + "## Turn Index\n- [turn 2](../turns/T2.md)", + ), + ] + profile = build_question_profile("What activities did Alice do with the drawre?", ["Alice"]) + + augmentation = build_corpus_consensus_augmentation( + "What activities did Alice do with the drawre?", + profile, + build_session_source_files(files), + ) + + assert augmentation.tokens[:2] == ["drawer", "brass"] + assert "tools" in augmentation.tokens + assert "brass tools" in augmentation.phrases + + +def test_retrieve_qmd_consensus_files_uses_corpus_correction_for_candidates() -> None: + files = [ + _file( + "/tmp/kb/wiki/entities/alice.md", + "Alice keeps household repair notes.", + ), + _file( + "/tmp/kb/wiki/sources/session_1.md", + "## Summary\nAlice repaired the drawer latch with brass tools.\n" + "## Turn Index\n- [turn 1](../turns/T1.md)", + ), + _file( + "/tmp/kb/wiki/observations/D1_obs.md", + "- Session: D1\n- Evidence: D1:2\nAlice placed brass tools in the drawer.", + ), + ] + + result = retrieve_qmd_consensus_files( + question="What did Alice put in the drawre?", + files=files, + root_files=[files[0]], + entity_names=["Alice"], + top_k=3, + ) + + assert "drawer" in result.profile.expansion_tokens + assert "/tmp/kb/wiki/observations/D1_obs.md" in result.coverage.scoped_file_paths + assert "/tmp/kb/wiki/observations/D1_obs.md" in result.coverage.final_file_paths + + +def test_rank_global_session_sources_prefers_query_phrase_entity_and_mtime() -> None: + files = [ + _file( + "/tmp/kb/wiki/sources/session_1.md", + "## Summary\nAlice repaired the brass drawer latch.\n" + "## Turn Index\n- [turn 1](../turns/T1.md)", + mtime_ms=10, + ), + _file( + "/tmp/kb/wiki/sources/session_2.md", + "## Summary\nAlice repaired the brass drawer latch.\n" + "## Turn Index\n- [turn 2](../turns/T2.md)", + mtime_ms=30, + ), + _file( + "/tmp/kb/wiki/sources/session_3.md", + "## Summary\nBlake reviewed kitchen plans.\n" + "## Turn Index\n- [turn 3](../turns/T3.md)", + mtime_ms=99, + ), + ] + profile = build_question_profile("What did Alice repair in the brass drawer?", ["Alice"]) + + ranked = rank_global_session_sources( + "What did Alice repair in the brass drawer?", + profile, + build_session_source_files(files), + ) + + assert [source.session_number for source in ranked] == [2, 1, 3] + + +def test_source_budget_helpers_match_retained_profile_flags() -> None: + aggregate = build_question_profile("What activities did Alice do?", ["Alice"]) + temporal = build_question_profile("When did Alice visit the gallery?", ["Alice"]) + location = build_question_profile("Where did Alice meet Blake?", ["Alice"]) + neutral = build_question_profile("What did Alice repair?", ["Alice"]) + + assert source_injection_budget(aggregate) == 8 + assert source_injection_budget(temporal) == 4 + assert source_injection_budget(location) == 3 + assert source_injection_budget(neutral) == 3 + + assert source_companion_budget(aggregate, has_plugin_retrieval=False, top_k=24) == 2 + assert source_companion_budget(aggregate, has_plugin_retrieval=True, top_k=24) == 4 + assert source_companion_budget(temporal, has_plugin_retrieval=True, top_k=24) == 3 + assert source_companion_budget(aggregate, has_plugin_retrieval=True, top_k=2) == 2 + + +def test_scoped_budgets_match_retained_profile_flags() -> None: + temporal = build_question_profile("When did Alice visit the gallery?", ["Alice"]) + identity = build_question_profile("Who is Alice's mentor?", ["Alice"]) + location = build_question_profile("Where did Alice meet Blake?", ["Alice"]) + neutral = build_question_profile("What did Alice repair?", ["Alice"]) + + assert scoped_budgets(temporal) == [ + ("wiki/sources", 3, 5.0), + ("wiki/observations", 3, 4.0), + ("wiki/events", 2, 3.0), + ("wiki/turns", 1, 1.0), + ("wiki/entities", 1, 2.0), + ] + assert scoped_budgets(identity) == [ + ("wiki/entities", 2, 4.0), + ("wiki/sources", 3, 5.0), + ("wiki/observations", 3, 4.0), + ("wiki/turns", 1, 2.0), + ("wiki/events", 1, 1.0), + ] + assert scoped_budgets(location) == [ + ("wiki/sources", 2, 4.0), + ("wiki/observations", 3, 4.0), + ("wiki/events", 2, 3.0), + ("wiki/turns", 1, 1.0), + ("wiki/entities", 1, 1.0), + ] + assert scoped_budgets(neutral) == [ + ("wiki/sources", 2, 3.0), + ("wiki/observations", 3, 3.0), + ("wiki/events", 2, 2.0), + ("wiki/entities", 1, 2.0), + ("wiki/turns", 1, 1.0), + ] + + +def test_select_scoped_candidate_files_prefers_matching_files_per_scope() -> None: + files = [ + _file( + "/tmp/kb/wiki/observations/D7_obs.md", + "- Session: D7\n- Evidence: D7:2\nAlice repaired the drawer on Monday.", + ), + _file( + "/tmp/kb/wiki/observations/D2_obs.md", + "- Session: D2\n- Evidence: D2:1\nBlake planned groceries.", + ), + _file( + "/tmp/kb/wiki/events/session_7_event.md", + "Alice repair event: the drawer latch was fixed on Monday.", + ), + _file( + "/tmp/kb/wiki/turns/T7.md", + "- Session: D7\nAlice repaired the drawer during the evening.", + ), + ] + profile = build_question_profile("When did Alice repair the drawer?", ["Alice"]) + + candidates = select_scoped_candidate_files( + "When did Alice repair the drawer?", + profile, + files, + ) + + assert [candidate.file.file_path for candidate in candidates][:3] == [ + "/tmp/kb/wiki/observations/D7_obs.md", + "/tmp/kb/wiki/events/session_7_event.md", + "/tmp/kb/wiki/turns/T7.md", + ] + assert all(candidate.query_hits >= 1 for candidate in candidates[:3]) + + +def test_retrieve_qmd_consensus_files_adds_scoped_candidates_for_temporal_query() -> None: + files = [ + _file( + "/tmp/kb/wiki/sources/session_7.md", + "## Summary\nAlice repaired the drawer after lunch.\n" + "## Turn Index\n- [turn 7](../turns/T7.md)", + ), + _file( + "/tmp/kb/wiki/observations/D7_obs.md", + "- Session: D7\n- Evidence: D7:2\nAlice repaired the drawer after lunch.", + ), + _file( + "/tmp/kb/wiki/events/session_7_event.md", + "Alice drawer repair event happened after lunch.", + ), + _file( + "/tmp/kb/wiki/entities/blake.md", + "Blake discussed unrelated travel plans.", + ), + ] + + result = retrieve_qmd_consensus_files( + question="When did Alice repair the drawer?", + files=files, + root_files=[files[0]], + entity_names=["Alice"], + top_k=4, + ) + + assert "/tmp/kb/wiki/observations/D7_obs.md" in result.coverage.scoped_file_paths + assert "/tmp/kb/wiki/events/session_7_event.md" in result.coverage.scoped_file_paths + assert "/tmp/kb/wiki/observations/D7_obs.md" in result.coverage.final_file_paths + + +def test_retrieve_qmd_consensus_files_records_late_bridge_targets_from_seed_links() -> None: + files = [ + _file( + "/tmp/kb/wiki/sources/session_1.md", + "## Summary\nAlice mentioned a linked workshop note.\n" + "## Turn Index\n- [turn 8](../turns/T8.md)", + ), + _file( + "/tmp/kb/wiki/turns/T8.md", + "- Session: D8\n- Evidence: D8:2\n" + "The brass tools were displayed during the workshop.", + ), + ] + + result = retrieve_qmd_consensus_files( + question="What activities did Alice connect with brass tools?", + files=files, + root_files=[files[0]], + entity_names=["Alice"], + top_k=2, + ) + + assert result.coverage.late_bridge_file_paths == ["/tmp/kb/wiki/turns/t8.md"] + assert "/tmp/kb/wiki/turns/T8.md" in result.coverage.final_file_paths + + +def test_retrieve_qmd_consensus_files_injects_ranked_sources_before_candidates() -> None: + files = [ + _file( + "/tmp/kb/wiki/entities/alice.md", + "Alice keeps notes about household repairs.", + ), + _file( + "/tmp/kb/wiki/sources/session_7.md", + "## Summary\nAlice repaired the brass drawer latch.\n" + "## Turn Index\n- [turn 7](../turns/T7.md)", + mtime_ms=70, + ), + _file( + "/tmp/kb/wiki/turns/T7.md", + "- Session: D7\n- Evidence: D7:1\nAlice repaired the brass drawer latch.", + ), + _file( + "/tmp/kb/wiki/observations/D7_obs.md", + "- Session: D7\n- Evidence: D7:1\nThe brass drawer latch repair succeeded.", + ), + ] + + result = retrieve_qmd_consensus_files( + question="What did Alice repair in the brass drawer?", + files=files, + root_files=[files[0]], + entity_names=["Alice"], + top_k=3, + ) + + assert [file.file_path for file in result.files][:2] == [ + "/tmp/kb/wiki/entities/alice.md", + "/tmp/kb/wiki/sources/session_7.md", + ] + assert len(result.files) == 3 + assert "/tmp/kb/wiki/sources/session_7.md" in result.coverage.source_file_paths + + +def test_collect_session_source_companions_uses_path_and_content_session_signals() -> None: + source_7 = _file( + "/tmp/kb/wiki/sources/session_7.md", + "## Summary\nAlice repaired the brass drawer latch.\n" + "## Turn Index\n- [turn 7](../turns/T7.md)", + ) + source_8 = _file( + "/tmp/kb/wiki/sources/session_8.md", + "## Summary\nBlake documented D8 gallery logistics.\n" + "## Turn Index\n- [turn 8](../turns/T8.md)", + ) + sources = build_session_source_files([source_7, source_8]) + sources_by_session = {source.session_number: source for source in sources} + profile = build_question_profile("What did Alice repair in the brass drawer?", ["Alice"]) + candidates = [ + CandidateFile( + file=_file( + "/tmp/kb/wiki/turns/T7.md", + "- Session: D7\n- Evidence: D7:1\nAlice repaired the brass drawer latch.", + ), + query_hits=2, + seed_boost=1.5, + ), + CandidateFile( + file=_file( + "/tmp/kb/wiki/entities/alice.md", + "Alice asked Blake to compare session_8 and D7 logistics.", + ), + query_hits=1, + seed_boost=0.5, + ), + ] + + companions = collect_session_source_companions( + "What did Alice repair in the brass drawer?", + profile, + candidates, + sources_by_session, + ) + + assert [companion.file.file_path for companion in companions] == [ + "/tmp/kb/wiki/sources/session_7.md", + "/tmp/kb/wiki/sources/session_8.md", + ] + assert infer_session_number_from_path("/tmp/kb/wiki/turns/T7.md") == 7 + assert infer_session_number_from_path("/tmp/kb/wiki/observations/D8_obs.md") == 8 + + +def _file(path: str, content: str, mtime_ms: int = 0) -> RetrievedMemoryFile: + return RetrievedMemoryFile( + filename=path.rsplit("/", maxsplit=1)[-1], + file_path=path, + mtime_ms=mtime_ms, + content=content, + ) diff --git a/tests/unit/retrieval/test_wikimem_memdir.py b/tests/unit/retrieval/test_wikimem_memdir.py new file mode 100644 index 00000000..7462ae17 --- /dev/null +++ b/tests/unit/retrieval/test_wikimem_memdir.py @@ -0,0 +1,147 @@ +"""wikimem Markdown memory directory compatibility tests.""" + +from __future__ import annotations + +import time + +import pytest + +from retrieval.wikimem_memdir import ( + WikimemDirectory, + MemoryFileHeader, + format_memory_manifest, + load_memory_entrypoints, + load_relevant_memory_files, + normalize_memory_relative_path, + scan_memory_directory, + select_relevant_memory_files, + should_include_memory_topic_file, +) + +pytestmark = pytest.mark.unit + + +def _write(path, content: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(content, encoding="utf-8") + + +def test_path_rules_match_wikimem_topic_file_boundaries() -> None: + assert normalize_memory_relative_path("project/../feedback/testing.md") == "feedback/testing.md" + assert normalize_memory_relative_path("../../outside.md") == "../../outside.md" + + assert should_include_memory_topic_file("feedback/testing.md") + assert not should_include_memory_topic_file("MEMORY.md") + assert not should_include_memory_topic_file("logs/2026/04/2026-04-02.md") + assert not should_include_memory_topic_file("../../outside.md") + assert not should_include_memory_topic_file("/absolute/outside.md") + + +def test_scan_memory_directory_parses_frontmatter_and_excludes_entrypoint(tmp_path) -> None: + _write( + tmp_path / "feedback" / "testing.md", + "---\ndescription: \"Deploy: run staging gate before prod\"\ntype: feedback\n---\nBody.\n", + ) + time.sleep(0.005) + _write( + tmp_path / "project" / "multiline.md", + "---\n" + "description: >\n" + " Release verification:\n" + " run smoke tests\n\n" + "type: project\n" + "---\n" + "Body.\n", + ) + _write(tmp_path / "MEMORY.md", "# entrypoint\n") + _write(tmp_path / "logs" / "2026" / "04" / "2026-04-02.md", "# daily\n") + + headers = scan_memory_directory(tmp_path) + + assert [header.filename for header in headers] == [ + "project/multiline.md", + "feedback/testing.md", + ] + assert headers[0].description == "Release verification: run smoke tests" + assert headers[0].memory_type == "project" + assert headers[1].description == "Deploy: run staging gate before prod" + assert headers[1].memory_type == "feedback" + + +def test_manifest_and_header_selection_are_high_confidence() -> None: + headers = [ + MemoryFileHeader( + filename="feedback/testing.md", + file_path="/mem/feedback/testing.md", + mtime_ms=2000, + description="Integration tests for database migration safety", + memory_type="feedback", + ), + MemoryFileHeader( + filename="reference/database-glossary.md", + file_path="/mem/reference/database-glossary.md", + mtime_ms=1000, + description="Database glossary and terminology reference", + memory_type="reference", + ), + ] + + manifest = format_memory_manifest(headers) + selected = select_relevant_memory_files( + "how should I test this database migration safely before release", + headers, + 5, + ) + + assert "[feedback] feedback/testing.md" in manifest + assert "1970-" in manifest + assert selected == [headers[0]] + + +def test_load_relevant_memory_files_uses_body_fallback_and_team_dedupe(tmp_path) -> None: + _write( + tmp_path / "feedback" / "testing.md", + "---\ndescription: Private testing preference\ntype: feedback\n---\nPrefer local checks.\n", + ) + _write( + tmp_path / "team" / "feedback" / "release.md", + "---\n" + "description: Team release policy\n" + "type: feedback\n" + "---\n" + "Run shared staging database verification before release.\n", + ) + + files = load_relevant_memory_files( + "how should we verify shared staging database behavior", + [ + WikimemDirectory(scope="auto", path=str(tmp_path)), + WikimemDirectory(scope="team", path=str(tmp_path / "team")), + ], + top_k=5, + recent_tools=[], + already_surfaced_file_paths=[], + ) + + assert len(files) == 1 + assert files[0].scope == "team" + assert files[0].filename == "feedback/release.md" + assert "shared staging database" in files[0].content + + +def test_load_memory_entrypoints_is_opt_in_truncated_and_secret_safe(tmp_path) -> None: + _write(tmp_path / "MEMORY.md", "\n".join(f"line {i}" for i in range(250))) + _write(tmp_path / "team" / "MEMORY.md", "github_pat_secret") + + entrypoints = load_memory_entrypoints( + [ + WikimemDirectory(scope="auto", path=str(tmp_path)), + WikimemDirectory(scope="team", path=str(tmp_path / "team")), + ] + ) + + assert len(entrypoints) == 1 + assert entrypoints[0].scope == "auto" + assert "line 199" in entrypoints[0].content + assert "line 249" not in entrypoints[0].content + assert "[truncated]" in entrypoints[0].content diff --git a/tests/unit/retrieval/test_wikimem_memdir_recaller.py b/tests/unit/retrieval/test_wikimem_memdir_recaller.py new file mode 100644 index 00000000..f1f1be46 --- /dev/null +++ b/tests/unit/retrieval/test_wikimem_memdir_recaller.py @@ -0,0 +1,60 @@ +"""wikimem memdir Recaller integration tests.""" + +from __future__ import annotations + +import json + +import pytest + +from retrieval.recaller_impl.wikimem_memdir_recaller import ( + WikimemMemdirRecaller, + memory_file_unit_id, +) +from retrieval.types import ParsedQuery, RecallChannel + +pytestmark = pytest.mark.unit + + +def _write(path, content: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(content, encoding="utf-8") + + +def test_wikimem_memdir_recaller_reads_extensions_and_returns_document_units( + tmp_path, scope +) -> None: + memory_dir = tmp_path / "mem" + _write( + memory_dir / "feedback" / "release.md", + "---\n" + "description: Team release database policy\n" + "type: feedback\n" + "---\n" + "Use shared staging DB.\n", + ) + parsed = ParsedQuery( + raw="how should release database checks run", + rewritten="how should release database checks run", + extensions={ + "wikimem.memory_dirs": json.dumps( + [{"scope": "team", "path": str(memory_dir)}] + ) + }, + ) + recaller = WikimemMemdirRecaller() + + results = recaller.recall(scope, parsed, 5) + + assert recaller.channel() == RecallChannel.DOCUMENT + assert len(results) == 1 + assert results[0].channel == RecallChannel.DOCUMENT + assert results[0].unit_id == memory_file_unit_id( + "team", str(memory_dir / "feedback" / "release.md") + ) + assert results[0].score > 0 + + +def test_wikimem_memdir_recaller_without_memory_dirs_returns_empty(scope) -> None: + recaller = WikimemMemdirRecaller() + + assert recaller.recall(scope, ParsedQuery(raw="anything"), 5) == [] diff --git a/tests/unit/retrieval/test_wikimem_options.py b/tests/unit/retrieval/test_wikimem_options.py new file mode 100644 index 00000000..cf5f8d88 --- /dev/null +++ b/tests/unit/retrieval/test_wikimem_options.py @@ -0,0 +1,69 @@ +"""wikimem retrieval option compatibility tests.""" + +from __future__ import annotations + +import pytest + +from common.errors import ValidationError +from retrieval.wikimem_options import ( + WikimemDirectory, + WikimemRetrievalOptions, + parse_wikimem_options, +) + +pytestmark = pytest.mark.unit + + +def test_parse_wikimem_options_reads_json_lists_and_defaults() -> None: + options = parse_wikimem_options( + { + "wikimem.recent_tools": '["Read", "Edit"]', + "wikimem.already_surfaced_file_paths": '["docs/old.md"]', + "wikimem.memory_dirs": '[{"scope": "auto", "path": "C:/repo/.memory"}]', + } + ) + + assert options == WikimemRetrievalOptions( + recent_tools=["Read", "Edit"], + already_surfaced_file_paths=["docs/old.md"], + include_entrypoints=False, + memory_dirs=[WikimemDirectory(scope="auto", path="C:/repo/.memory")], + memory_parallelism=None, + ) + + +def test_parse_wikimem_options_reads_bool_int_and_profile_fields() -> None: + options = parse_wikimem_options( + { + "wikimem.include_entrypoints": "true", + "wikimem.memory_parallelism": "0", + "wikimem.profile": "memdir", + "wikimem.selector_model": "primary", + "wikimem.selector_fallback_model": "fallback", + "max_tokens": "128", + } + ) + + assert options.include_entrypoints is True + assert options.memory_parallelism == 1 + assert options.profile == "memdir" + assert options.selector_model == "primary" + assert options.selector_fallback_model == "fallback" + + +@pytest.mark.parametrize( + ("extensions", "message"), + [ + ({"wikimem.recent_tools": "Read,Edit"}, "wikimem.recent_tools"), + ({"wikimem.already_surfaced_file_paths": "[1]"}, "string array"), + ({"wikimem.memory_dirs": '[{"scope": "global", "path": "x"}]'}, "scope"), + ({"wikimem.memory_dirs": '[{"scope": "auto", "path": ""}]'}, "path"), + ({"wikimem.include_entrypoints": "yes"}, "include_entrypoints"), + ({"wikimem.memory_parallelism": "many"}, "memory_parallelism"), + ], +) +def test_parse_wikimem_options_rejects_invalid_transport_values( + extensions: dict[str, str], message: str +) -> None: + with pytest.raises(ValidationError, match=message): + parse_wikimem_options(extensions) From b8a5f4432052ccbd01fd86d9f92b4c66c3626b2e Mon Sep 17 00:00:00 2001 From: xinyao1994 Date: Wed, 5 Aug 2026 20:52:09 +0800 Subject: [PATCH 2/3] docs(memory): document wikimem integration --- agent_plugin/wikimem/AGENTS.md | 16 + docs/Roadmap.md | 2 - docs/design/architecture.md | 2 +- docs/features/F02-wikimem-compat.md | 346 ++++++++++++++++++ ...memory-compat.md => F04-wikimem-compat.md} | 60 +-- docs/specs/S02-memory-api.md | 12 +- docs/specs/S03-control.md | 5 +- docs/specs/S04-retrieval.md | 2 +- docs/specs/S05-construction.md | 23 +- docs/specs/S07-common.md | 2 +- evaluation/wikimem/AGENTS.md | 17 + ...-vs-llm-wiki.md => wikimem-vs-llm-wiki.md} | 8 +- ...s-llm-wiki.mmd => wikimem-vs-llm-wiki.mmd} | 4 +- ...s-llm-wiki.png => wikimem-vs-llm-wiki.png} | Bin 14 files changed, 441 insertions(+), 58 deletions(-) create mode 100644 agent_plugin/wikimem/AGENTS.md create mode 100644 docs/features/F02-wikimem-compat.md rename docs/features/construction/{F04-cc-memory-compat.md => F04-wikimem-compat.md} (80%) create mode 100644 evaluation/wikimem/AGENTS.md rename figures/{cc-memory-vs-llm-wiki.md => wikimem-vs-llm-wiki.md} (92%) rename figures/{cc-memory-vs-llm-wiki.mmd => wikimem-vs-llm-wiki.mmd} (96%) rename figures/{cc-memory-vs-llm-wiki.png => wikimem-vs-llm-wiki.png} (100%) diff --git a/agent_plugin/wikimem/AGENTS.md b/agent_plugin/wikimem/AGENTS.md new file mode 100644 index 00000000..c04b9e1a --- /dev/null +++ b/agent_plugin/wikimem/AGENTS.md @@ -0,0 +1,16 @@ +# wikimem agent adapter + +本目录承接 `rust/wikimem` 中不属于 mem2.0 core API 的 agent 侧能力。 + +## 文件职责 + +- `team_sync.py`：团队记忆目录同步 adapter，保留 `SyncState`、pull / push、ETag、checksum、delta、412/413、路径校验和 secret skip 语义。 +- `__init__.py`：导出 wikimem agent adapter 的稳定入口。 + +## 边界 + +- `team_sync.py` 可以读写本地 team memory 目录，也可以通过 `TeamMemoryRemote` 与远端同步。 +- `team_sync.py` 不得被 `MemoryAPI.recall`、`Retriever.retrieve` 或默认 Recaller 自动调用。 +- 远端副作用必须显式由 agent adapter、CLI 或评测准备脚本触发。 +- 远端 key 写入本地前必须经过相对路径校验，拒绝绝对路径、父级跳转、反斜杠和百分号编码跳转。 +- push 前必须跳过潜在 secret 文件，并在 summary 中返回 `skipped_secrets`，不能静默丢弃。 diff --git a/docs/Roadmap.md b/docs/Roadmap.md index bb4631c1..0484ea4f 100644 --- a/docs/Roadmap.md +++ b/docs/Roadmap.md @@ -155,5 +155,3 @@ | 日期 | 说明 | | ---------- | ---------------------- | | 2026-08-03 | 按照当日代码及与各团队对其的情况进行初稿规划 | - - diff --git a/docs/design/architecture.md b/docs/design/architecture.md index 22333e4d..58f3e0c0 100644 --- a/docs/design/architecture.md +++ b/docs/design/architecture.md @@ -660,7 +660,7 @@ agent-memory/ - **三类统一契约**：接口代码落地为「算子 + 插件 + 存储」三类自描述契约——各层算子（`IngestOperator`/`ConstructionOperator`/`RetrievalOperator`/`ControlOperator`）、共享能力插件（`Plugin`）、存储后端（`BaseStore`），均为「类型枚举方法（`operatorType`/`pluginType`/`storeType`）+ `health()` 探活」，路由按类型不按实现名。 -- **兼容报告单独归档**：跨层 legacy 兼容（例如 `rust/cc_memory` 的 `MemoryIngestor`/`MemoryRetriever`、`memdir`、`retained_eval`）不塞进单层接口；统一归 `docs/features/construction/F04-cc-memory-compat.md`，再映射回 `src/api` / `src/retrieval` / `evaluation` / `agent_plugin`。 +- **兼容报告单独归档**：跨层 legacy 兼容（例如 `rust/wikimem` 的 `MemoryIngestor`/`MemoryRetriever`、`memdir`、`retained_eval`）不塞进单层接口；统一归 `docs/features/F02-wikimem-compat.md` 与 `docs/features/construction/F04-wikimem-compat.md`，再映射回 `src/api` / `src/retrieval` / `evaluation` / `agent_plugin`。 - **写入边界**：`ingest` 只做规约与转换（RawPayload → MemoryUnit），**不落盘**；`construction` 负责把 MemoryUnit 写入真源、在其上挖掘分层记忆并构建索引。构建层**没有编排 service**，六个算子（extractor/abstractor/associator/classifier/index_builder/evolver）由上层/控制层驱动。 - **索引「构建」与「持久化」分离**：`src/construction/index_builder` 负责构建/更新索引（逻辑），`src/storage` 负责持久化驱动（后端），经 Store 抽象解耦；统一 CRUD 动词为 `insert`（增，冲突抛 `ConflictError`）/ `delete`（删，幂等）/ `update`（改，缺失抛 `NotFoundError`）/ `get`（查，按 id 点读）。检索型 Store 的记录/查询带一等 `scope` 字段，按 scope **原生隔离**（kv/fs 为通用键值/二进制原语，不引入 scope）。 - **共享插件保证两侧一致**：分词/切分/向量化/特征抽取/LLM/规约/重排抽到 `src/common`，构建侧与检索侧（以及重建/演进路径）注入**同一实现**——同词表、同向量空间、同切分规则、同规约器，是「派生可重建」与召回对齐的前提。 diff --git a/docs/features/F02-wikimem-compat.md b/docs/features/F02-wikimem-compat.md new file mode 100644 index 00000000..2c5e20e7 --- /dev/null +++ b/docs/features/F02-wikimem-compat.md @@ -0,0 +1,346 @@ +# wikimem 系统架构、接口与评测报告 + +## LLM-native Wiki extension + +The existing retained Wiki remains the deterministic benchmark baseline. The +same evaluation surface now also supports an injected LLM path for open-domain +conversation, agent traces, and document source blocks: + +`raw sources -> semantic extraction -> entity resolution -> memory +consolidation -> profile/timeline/decision synthesis -> wiki/memory retrieval`. + +The LLM path uses the shared memory ontology, preserves evidence/provenance and +multimodal source metadata, and adds `wiki/memory//` pages to qmd +candidates. `wiki_builder_mode="deterministic"` or an absent/failed LLM keeps +the historical deterministic renderer, so existing recall baselines remain +reproducible. `query_llm` can be enabled independently for conservative query +intent/entity/time expansion. No answer labels or Rust intermediate outputs +are construction inputs. + +## 元信息 + +| 项 | 值 | +|---|---| +| 系统范围 | 记忆写入、基线记忆、Markdown 目录检索、团队记忆同步、retained evaluation | +| 日期 | 2026-07-29 | +| 关联模块 | `agent_plugin/wikimem`、`src/construction`、`src/retrieval`、`evaluation/wikimem` | +| 适用调用面 | `MemoryAPI`、构建算子、Recaller、评测运行器、团队同步 adapter | + +## 1. 系统目标和边界 + +wikimem 是建立在统一 `MemoryAPI` 之上的结构化记忆系统。它支持把对话中的显式记忆指令转成可演进记录，把本地 Markdown 记忆目录接入 DOCUMENT 召回通道，并提供可追踪的多阶段检索评测。 + +系统不把数据集专用字段放进通用 `RetrievalQuery`，不在 `recall` 中执行网络同步，也不改变现有 `MemoryAPI`、`Retriever`、`Recaller` 的必填参数。 + +## 2. 总体架构 + +```mermaid +flowchart TB + subgraph Client[调用方] + C1[对话 / 应用消息] + C2[检索问题] + C3[评测样本] + end + + subgraph API[统一接口层] + A1[MemoryAPI.write] + A2[MemoryAPI.recall] + A3[MemoryAPI.evolve] + CTX[Context
scope + extensions] + end + + subgraph Baseline[基线记忆构建] + B1[WikimemBaselineExtractor
remember / forget / key-value] + B2[WikimemBaselineEvolver
upsert / supersede / forget] + B3[MemoryUnit records] + end + + subgraph Retrieval[Markdown 目录检索] + R1[parse_wikimem_options] + R2[WikimemMemdirRecaller
DOCUMENT channel] + R3[header selection] + R4[body fallback] + R5[entrypoint loading] + end + + subgraph Retained[多阶段 retained evaluation] + E1[样本归一化 / 多视图 workspace] + E2[QuestionProfile
tokens / fuzzy / entities / time / relation] + E3[Root Retrieval] + E4[Scoped Retrieval + Source Ranking] + E5[Candidate + Rerank] + E6[Late Bridge + Final Assembly] + E7[指标、case trace、failure report、stage profile] + end + + subgraph Team[团队记忆同步] + T1[本地安全扫描] + T2[ETag / checksum / delta] + T3[TeamMemoryRemote] + T4[远端 team memory] + end + + C1 --> A1 --> B1 --> B2 --> B3 + C2 --> CTX --> A2 --> R1 --> R2 --> R3 + R3 --> R4 --> R5 + C3 --> E1 --> E2 --> E3 --> E4 --> E5 --> E6 --> E7 + B3 -.可检索记录.-> A2 + R5 -.ScoredUnit.-> A2 + T1 --> T2 --> T3 <--> T4 + T1 -.本地目录.-> R2 +``` + +### 模块职责 + +| 层 | 组件 | 职责 | 不负责 | +|---|---|---|---| +| 入口 | `MemoryAPI` | 权限检查、审计、write/recall/evolve 统一入口 | 解释 wikimem 专用配置 | +| 构建 | `WikimemBaselineExtractor` | 识别 remember、forget、`key: value`、范围提示、潜在 secret | 直接覆盖既有记录 | +| 构建 | `WikimemBaselineEvolver` | 将候选转为记录；last-write-wins；标记 superseded / forgotten | 数据集评测 | +| 检索 | `WikimemMemdirRecaller` | 将 Markdown 目录映射到 DOCUMENT 通道的 `ScoredUnit` | 修改底层存储或调用网络 | +| 评测 | `evaluation/wikimem` | 构建 workspace、阶段检索、独立标签计分、诊断产物 | 修改通用 API 类型 | +| 同步 | `team_sync` | team memory 的 pull/push、checksum、冲突重试、路径与 secret 防护 | 普通 recall 的同步副作用 | + +## 3. 公共接口 + +### 3.1 统一记忆接口 + +```python +api.write( + content: str, + scope: Scope, + *, + identity: Scope, + assets: list[str] | None = None, + tags: list[str] | None = None, + metadata: dict[str, str] | None = None, +) -> list[MemoryUnit] + +api.recall( + query: str, + context: Context, + *, + identity: Scope, + filters: list[FilterClause] | None = None, + as_of: datetime | None = None, + top_k: int = 10, + disclosure: DisclosureLevel = DisclosureLevel.L0, + with_trajectory: bool = False, +) -> RetrievalResult + +api.evolve( + scope: Scope, + mode: EvolveMode, + *, + identity: Scope, +) -> str +``` + +`scope` 是目标记忆范围，`identity` 是调用方身份；二者必须显式传入。`Context.extensions` 只能携带字符串配置，经过 API 边界后进入 `RetrievalQuery.extensions` 与 `ParsedQuery.extensions`。 + +### 3.2 wikimem 检索配置 + +| extension key | 值 | 默认 | 作用 | +|---|---|---:|---| +| `wikimem.memory_dirs` | JSON array | `[]` | Markdown 目录声明，每项含 `scope=auto/team` 与 `path` | +| `wikimem.recent_tools` | JSON string array | `[]` | 降低纯工具文档的优先级 | +| `wikimem.already_surfaced_file_paths` | JSON string array | `[]` | 避免本轮重复展示已返回文件 | +| `wikimem.include_entrypoints` | `"true"` / `"false"` | `false` | 是否追加目录 `MEMORY.md` | +| `wikimem.profile` | string | `""` | 调用方选择的兼容 profile 名称 | +| `wikimem.selector_model` | string | `""` | header selector 名称 | +| `wikimem.selector_fallback_model` | string | `""` | selector 不可用时的替代名称 | +| `wikimem.memory_parallelism` | integer string | unset | 目录扫描并行度，最小为 1 | + +示例： + +```python +result = api.recall( + query="上次团队约定的发布回滚步骤是什么？", + context=Context( + scope=Scope(namespace="project"), + extensions={ + "wikimem.memory_dirs": ( + '[{"scope":"team","path":"./team-memory"}]' + ), + "wikimem.include_entrypoints": "true", + "wikimem.already_surfaced_file_paths": "[]", + }, + ), + identity=Scope(namespace="project"), + top_k=5, + with_trajectory=True, +) +``` + +无效 JSON、未知目录 scope、空路径或非法布尔值必须在 `parse_wikimem_options` 抛出 `ValidationError`，不得静默忽略。 + +## 4. 基线记忆生命周期 + +### 4.1 写入和演进 + +```mermaid +sequenceDiagram + participant U as 调用方 + participant A as MemoryAPI + participant X as BaselineExtractor + participant E as BaselineEvolver + participant M as MemoryUnit + + U->>A: write(content, scope, identity) + A->>X: 提取显式记忆候选 + X-->>E: upsert / forget candidate + E->>M: 创建 record 或标记旧 record + E-->>A: created / superseded / forgotten ids + A-->>U: MemoryUnit 列表或 evolve job id +``` + +| 输入模式 | 产生的动作 | 处理规则 | +|---|---|---| +| `remember ...` / `请记住 ...` | `upsert` | 提取 note key、value、记忆类型和可选 scope hint | +| `forget ...` / `请忘记 ...` | `forget` | 以 key 或归一化内容匹配，标记为 `FORGOTTEN` | +| `key: value`、`key = value`、`key is value` | `upsert` | 排除疑问词 key，识别记忆类型 | +| 同 key 新值 | `upsert` | 旧 active record 标记 `SUPERSEDED`，新记录成为 active | +| 潜在 secret | skip | 写入 skip 原因，不生成可同步记录 | + +记录使用 `MemoryUnit` 表示，关键 metadata 包括：`wikimem.key`、`wikimem.value`、`wikimem.action`、`wikimem.memory_type`、`wikimem.preferred_scope`、`wikimem.observed_at_ms`、`wikimem.score`。 + +## 5. Markdown memory directory 检索 + +### 5.1 文件模型 + +| 类型 | 内容 | 是否可作为 topic 候选 | +|---|---|---| +| topic Markdown | frontmatter、标题、正文 | 是 | +| `MEMORY.md` | 目录入口和摘要 | 默认否；`include_entrypoints=true` 时追加 | +| `logs/YYYY/MM/YYYY-MM-DD.md` | 日志 | 否 | +| 非 `.md`、绝对路径、含 `..` 的路径 | 不安全或不适用 | 否 | + +`MemoryFileHeader` 保存 filename、file path、mtime、frontmatter `description` 和 `type`。`RetrievedMemoryFile` 在正文 materialize 后附加内容；`RetrievedMemoryEntrypoint` 表示入口文件。 + +### 5.2 检索流程 + +```mermaid +flowchart LR + D[wikimem.memory_dirs] --> S[扫描安全 Markdown 文件] + S --> H[解析 header / frontmatter] + H --> F[过滤 daily log、entrypoint、secret、recent tool、已展示路径] + F --> HS{header score 达标?} + HS -- 是 --> M[materialize 文件正文] + HS -- 否 --> BS[body fallback score] + BS --> M + M --> E{include_entrypoints?} + E -- 是 --> EP[加载 MEMORY.md] + E -- 否 --> O[DOCUMENT ScoredUnit] + EP --> O +``` + +| 参数 | 值 | 含义 | +|---|---:|---| +| `DEFAULT_MAX_MEMORY_FILES` | 200 | 单目录最多扫描 topic 文件数 | +| `DEFAULT_TOP_K` | 5 | memory directory 默认返回数 | +| `MIN_HEADER_SELECTION_SCORE` | 4.0 | header 选择绝对阈值 | +| `MIN_BODY_SELECTION_SCORE` | 4.5 | body fallback 绝对阈值 | +| `RELATIVE_SELECTION_SCORE_RATIO` | 0.45 | 保留相对最佳分数足够高的候选 | +| `MAX_MEMORY_LINES` / bytes | 200 / 4096 | topic 正文截断预算 | +| `MAX_ENTRYPOINT_LINES` / bytes | 200 / 25000 | `MEMORY.md` 截断预算 | + +`WikimemMemdirRecaller` 使用稳定 SHA-256 派生 unit id，并通过 DOCUMENT channel 返回 `ScoredUnit`；后续仍由既有的 Fuser、UnitReader、Discloser 处理。 + +## 6. 团队记忆同步 + +### 6.1 端口与结果类型 + +```python +class TeamMemoryRemote(Protocol): + async def fetch(repo_slug, if_none_match) -> FetchOutcome: ... + async def fetch_hashes(repo_slug) -> HashesProbe: ... + async def put_entries(repo_slug, if_match, entries) -> PutOutcome: ... + +async def pull_team_memory(remote, state, team_memory_root, repo_slug) -> PullOutcome: ... +async def push_team_memory(remote, state, team_memory_root, repo_slug) -> PushOutcomeSummary: ... +``` + +| 操作 | 正常路径 | 异常与边界 | +|---|---|---| +| pull | 以 `If-None-Match` 拉取；304 不写文件；200 写入 entries 并更新 checksum | 404 视为空远端；失败返回结构化 `TeamMemorySyncFailure` | +| push | 扫描本地文件，过滤 secret，计算 checksum delta，按 body bytes 分批 PUT | 412 拉取 hashes 后重试；413 记录 server max entries；失败返回结构化原因 | +| key 校验 | 只接受安全相对路径 | 拒绝 NUL、反斜杠、绝对路径、`..`、百分号编码逃逸和 Windows drive prefix | +| secret 检测 | 跳过并记录 `SkippedSecretFile` | 不把潜在凭据上传到远端 | + +同步是独立 adapter：它把数据写到 team memory 目录，但 `MemoryAPI.recall` 本身不触发 pull 或 push。 + +## 7. retained evaluation 引擎 + +### 7.1 输入与内部对象 + +| 对象 | 作用 | +|---|---| +| `PreparedSample` | 统一样本：records、questions、sessions、events、observations、原始 payload | +| `LoCoMoQuestion` | 问题、答案、evidence、类别；答案不进入检索输入 | +| `RetrievedMemoryFile` | 检索阶段消费的内存文件投影 | +| `QuestionProfile` | tokens、fuzzy tokens、实体、时间、地点、关系和聚合信号 | +| `CaseScore` | 单题 expected / retrieved / hits、precision、recall、full hit、文件路径和 coverage | +| `EvalOutput` | summary、cases、stage profile | + +### 7.2 多阶段检索 + +```mermaid +flowchart LR + Q[问题] --> P[QuestionProfile] + P --> R1[Phase 1: Root Retrieval] + R1 --> R2[Phase 2: Scoped Retrieval
Source Ranking / Companion] + R2 --> R3[Phase 3: Candidate + Rerank] + R3 --> R4[Phase 4: Late Bridge] + R4 --> F[Final Assembly] + F --> C[CaseScore + Coverage] + C --> A[summary / category breakdown / failure report / stage profile] + + P -.query expansion.-> X[qmd_consensus] + X -.candidate proposals.-> R2 + X -.rerank proposals.-> R3 + X -.bridge proposals.-> R4 +``` + +| 阶段 | 输入 | 输出 | 诊断 | +|---|---|---|---| +| Root Retrieval | question、memory root | 根级候选 | root hit / miss | +| Scoped Retrieval | source、turn、observation、event、entity 等视图 | 分区候选池 | candidate pool 覆盖 | +| Candidate + Rerank | lexical features、profile、linked pages | 有序候选 | source ranking、rerank proposal | +| Late Bridge | 已选 page 中的链接和会话线索 | 补充 atomic evidence | late bridge hit / miss | +| Final Assembly | 去重后的文件 | evidence id、路径、最终统计 | full hit、unexpected evidence | + +评测必须写出：`summary`、每题 case trace、`failure_report`、`failure_buckets`、`category_breakdown`、`stage_profile` 和 run manifest。 + +## 8. 数据集运行接口与标签规则 + +| 运行器 | 主入口 | 标签来源 | 指标状态 | +|---|---|---|---| +| LoCoMo | `run_python_locomo_retrieval_eval` | 数据集 `evidence` | 确认的 evidence retrieval 指标 | +| LongMemEval | `run_python_longmemeval_retrieval_eval` | 原始 turn `has_answer`、session id | 确认的 turn/session retrieval 指标 | +| MemGallery | `build_mem_gallery_python_workspaces` + `run_filesystem_proxy_eval` | 人工 `clue` | clue retrieval 指标 | +| EverMemBench | `run_evermembench_python_eval` | dataset reference | reference retrieval 指标 | +| AMA-Bench | `run_python_ama_retrieval_eval` | 答案与轨迹的词重叠 | `answer_derived_proxy`，不能作为独立精度 | +| Meta-CRAG | `build_meta_crag_python_workspaces` + `run_filesystem_proxy_eval` | 答案与 artifact 的支持关系 | `answer_derived_proxy`，不能作为独立精度 | + +### 8.1 评测数据完整性 + +1. LoCoMo 和 LongMemEval 在检索完成后才读取标签计分。 +2. MemGallery 的 `source_session_ids` 仅为元数据，不参与排序。 +3. case 中的 `baseline_retrieved_*`、`final_retrieved_*`、`ranked_clues`、`retrieved_file_paths` 默认拒绝。 +4. 文件系统评测只读取有 `producer=mem2.0.wikimem.python`、匹配 dataset name 和 schema version 的 workspace retrieval JSON。 +5. provenance 是输入来源检查，不是密码学签名；面对主动伪造输入时，调用方仍必须控制 workspace 的写入权限。 + +## 9. 全量数据集 Recall 测试结果 + +| 数据集 | 检索指标 | 数值 | 最优方案 | 关键配置 | +|---|---|---|---|---| +| `LoCoMo`
纯文本长程对话情景记忆检索基准 | `recall_macro` | `0.9413` | Python retained `qmd_consensus` | `locomo10.json`；`1986 cases`；`top_k=24`；`wiki_mode=text`；远程全量运行（2026-07-16） | +| `LoCoMo_refined`
带图像证据的多模态长程对话情景记忆检索基准 | `recall_macro` | `0.9556` | Python retained + multimodal adjunct | `1382 cases`；`top_k=24`；adjunct `top_k=6`；`wiki_mode=multimodal`；远程全量运行（2026-07-16） | +| `LongMemEval`
长程对话记忆检索基准 | `turn/session recall_any/all@5,@10,@30` | `turn @5 any 0.9212 / all 0.8115；@10 any 0.9642 / all 0.8926；@30 any 0.9833 / all 0.9379`
`session @5 any 0.9308 / all 0.7947；@10 any 0.9690 / all 0.8807；@30 any 0.9857 / all 0.9332` | Python retained `qmd_consensus` | `longmemeval_s_cleaned.json`；`500 cases`；`419 evaluated`；`top_k=100`；`threads=4`；`wiki_mode=text` | +| `AMA-Bench`
自主 agent 轨迹记忆检索基准 | `proxy_recall@112` | `0.9142` | Python `wikimem_qmd_ama` retrieval-only | `208 episodes`；`2496 questions`；`top_k=112`；`methods=["wikimem_qmd_ama"]`；`wiki_mode=text`；远程全量运行（2026-07-16） | + +## 10. 验收要求 + +所有测试和评测必须在项目指定的远程环境执行。每次精度记录至少包含数据集路径与版本、样本范围、`top_k`、方法名、代码版本、summary、case trace、failure report 和 stage profile。 + +以下情况不能作为精度对比：本地运行、只跑 smoke、不同数据版本、不同 `top_k`、不同标签来源，以及未带 workspace provenance 的文件系统评测。 diff --git a/docs/features/construction/F04-cc-memory-compat.md b/docs/features/construction/F04-wikimem-compat.md similarity index 80% rename from docs/features/construction/F04-cc-memory-compat.md rename to docs/features/construction/F04-wikimem-compat.md index 8d30fd16..1ddb3982 100644 --- a/docs/features/construction/F04-cc-memory-compat.md +++ b/docs/features/construction/F04-wikimem-compat.md @@ -1,4 +1,4 @@ -# cc_memory 系统架构、接口与评测报告 +# wikimem 系统架构、接口与评测报告 ## 元信息 @@ -6,12 +6,12 @@ |---|---| | 系统范围 | 记忆写入、基线记忆、Markdown 目录检索、团队记忆同步、retained evaluation | | 日期 | 2026-07-22 | -| 关联模块 | `agent_plugin/cc_memory`、`src/construction`、`src/retrieval`、`evaluation/cc_memory` | +| 关联模块 | `agent_plugin/wikimem`、`src/construction`、`src/retrieval`、`evaluation/wikimem` | | 适用调用面 | `MemoryAPI`、构建算子、Recaller、评测运行器、团队同步 adapter | ## 1. 系统目标和边界 -cc_memory 是建立在统一 `MemoryAPI` 之上的结构化记忆系统。它支持把对话中的显式记忆指令转成可演进记录，把本地 Markdown 记忆目录接入 DOCUMENT 召回通道，并提供可追踪的多阶段检索评测。 +wikimem 是建立在统一 `MemoryAPI` 之上的结构化记忆系统。它支持把对话中的显式记忆指令转成可演进记录，把本地 Markdown 记忆目录接入 DOCUMENT 召回通道，并提供可追踪的多阶段检索评测。 系统不把数据集专用字段放进通用 `RetrievalQuery`，不在 `recall` 中执行网络同步，也不改变现有 `MemoryAPI`、`Retriever`、`Recaller` 的必填参数。 @@ -33,14 +33,14 @@ flowchart TB end subgraph Baseline[基线记忆构建] - B1[CcMemoryBaselineExtractor
remember / forget / key-value] - B2[CcMemoryBaselineEvolver
upsert / supersede / forget] + B1[WikimemBaselineExtractor
remember / forget / key-value] + B2[WikimemBaselineEvolver
upsert / supersede / forget] B3[MemoryUnit records] end subgraph Retrieval[Markdown 目录检索] - R1[parse_cc_memory_options] - R2[CcMemoryMemdirRecaller
DOCUMENT channel] + R1[parse_wikimem_options] + R2[WikimemMemdirRecaller
DOCUMENT channel] R3[header selection] R4[body fallback] R5[entrypoint loading] @@ -77,11 +77,11 @@ flowchart TB | 层 | 组件 | 职责 | 不负责 | |---|---|---|---| -| 入口 | `MemoryAPI` | 权限检查、审计、write/recall/evolve 统一入口 | 解释 cc_memory 专用配置 | -| 构建 | `CcMemoryBaselineExtractor` | 识别 remember、forget、`key: value`、范围提示、潜在 secret | 直接覆盖既有记录 | -| 构建 | `CcMemoryBaselineEvolver` | 将候选转为记录；last-write-wins；标记 superseded / forgotten | 数据集评测 | -| 检索 | `CcMemoryMemdirRecaller` | 将 Markdown 目录映射到 DOCUMENT 通道的 `ScoredUnit` | 修改底层存储或调用网络 | -| 评测 | `evaluation/cc_memory` | 构建 workspace、阶段检索、独立标签计分、诊断产物 | 修改通用 API 类型 | +| 入口 | `MemoryAPI` | 权限检查、审计、write/recall/evolve 统一入口 | 解释 wikimem 专用配置 | +| 构建 | `WikimemBaselineExtractor` | 识别 remember、forget、`key: value`、范围提示、潜在 secret | 直接覆盖既有记录 | +| 构建 | `WikimemBaselineEvolver` | 将候选转为记录；last-write-wins；标记 superseded / forgotten | 数据集评测 | +| 检索 | `WikimemMemdirRecaller` | 将 Markdown 目录映射到 DOCUMENT 通道的 `ScoredUnit` | 修改底层存储或调用网络 | +| 评测 | `evaluation/wikimem` | 构建 workspace、阶段检索、独立标签计分、诊断产物 | 修改通用 API 类型 | | 同步 | `team_sync` | team memory 的 pull/push、checksum、冲突重试、路径与 secret 防护 | 普通 recall 的同步副作用 | ## 3. 公共接口 @@ -121,18 +121,18 @@ api.evolve( `scope` 是目标记忆范围，`identity` 是调用方身份；二者必须显式传入。`Context.extensions` 只能携带字符串配置，经过 API 边界后进入 `RetrievalQuery.extensions` 与 `ParsedQuery.extensions`。 -### 3.2 cc_memory 检索配置 +### 3.2 wikimem 检索配置 | extension key | 值 | 默认 | 作用 | |---|---|---:|---| -| `cc_memory.memory_dirs` | JSON array | `[]` | Markdown 目录声明，每项含 `scope=auto/team` 与 `path` | -| `cc_memory.recent_tools` | JSON string array | `[]` | 降低纯工具文档的优先级 | -| `cc_memory.already_surfaced_file_paths` | JSON string array | `[]` | 避免本轮重复展示已返回文件 | -| `cc_memory.include_entrypoints` | `"true"` / `"false"` | `false` | 是否追加目录 `MEMORY.md` | -| `cc_memory.profile` | string | `""` | 调用方选择的兼容 profile 名称 | -| `cc_memory.selector_model` | string | `""` | header selector 名称 | -| `cc_memory.selector_fallback_model` | string | `""` | selector 不可用时的替代名称 | -| `cc_memory.memory_parallelism` | integer string | unset | 目录扫描并行度，最小为 1 | +| `wikimem.memory_dirs` | JSON array | `[]` | Markdown 目录声明，每项含 `scope=auto/team` 与 `path` | +| `wikimem.recent_tools` | JSON string array | `[]` | 降低纯工具文档的优先级 | +| `wikimem.already_surfaced_file_paths` | JSON string array | `[]` | 避免本轮重复展示已返回文件 | +| `wikimem.include_entrypoints` | `"true"` / `"false"` | `false` | 是否追加目录 `MEMORY.md` | +| `wikimem.profile` | string | `""` | 调用方选择的兼容 profile 名称 | +| `wikimem.selector_model` | string | `""` | header selector 名称 | +| `wikimem.selector_fallback_model` | string | `""` | selector 不可用时的替代名称 | +| `wikimem.memory_parallelism` | integer string | unset | 目录扫描并行度，最小为 1 | 示例： @@ -142,11 +142,11 @@ result = api.recall( context=Context( scope=Scope(namespace="project"), extensions={ - "cc_memory.memory_dirs": ( + "wikimem.memory_dirs": ( '[{"scope":"team","path":"./team-memory"}]' ), - "cc_memory.include_entrypoints": "true", - "cc_memory.already_surfaced_file_paths": "[]", + "wikimem.include_entrypoints": "true", + "wikimem.already_surfaced_file_paths": "[]", }, ), identity=Scope(namespace="project"), @@ -155,7 +155,7 @@ result = api.recall( ) ``` -无效 JSON、未知目录 scope、空路径或非法布尔值必须在 `parse_cc_memory_options` 抛出 `ValidationError`，不得静默忽略。 +无效 JSON、未知目录 scope、空路径或非法布尔值必须在 `parse_wikimem_options` 抛出 `ValidationError`，不得静默忽略。 ## 4. 基线记忆生命周期 @@ -185,7 +185,7 @@ sequenceDiagram | 同 key 新值 | `upsert` | 旧 active record 标记 `SUPERSEDED`，新记录成为 active | | 潜在 secret | skip | 写入 skip 原因，不生成可同步记录 | -记录使用 `MemoryUnit` 表示，关键 metadata 包括：`cc_memory.key`、`cc_memory.value`、`cc_memory.action`、`cc_memory.memory_type`、`cc_memory.preferred_scope`、`cc_memory.observed_at_ms`、`cc_memory.score`。 +记录使用 `MemoryUnit` 表示，关键 metadata 包括：`wikimem.key`、`wikimem.value`、`wikimem.action`、`wikimem.memory_type`、`wikimem.preferred_scope`、`wikimem.observed_at_ms`、`wikimem.score`。 ## 5. Markdown memory directory 检索 @@ -204,7 +204,7 @@ sequenceDiagram ```mermaid flowchart LR - D[cc_memory.memory_dirs] --> S[扫描安全 Markdown 文件] + D[wikimem.memory_dirs] --> S[扫描安全 Markdown 文件] S --> H[解析 header / frontmatter] H --> F[过滤 daily log、entrypoint、secret、recent tool、已展示路径] F --> HS{header score 达标?} @@ -227,7 +227,7 @@ flowchart LR | `MAX_MEMORY_LINES` / bytes | 200 / 4096 | topic 正文截断预算 | | `MAX_ENTRYPOINT_LINES` / bytes | 200 / 25000 | `MEMORY.md` 截断预算 | -`CcMemoryMemdirRecaller` 使用稳定 SHA-256 派生 unit id，并通过 DOCUMENT channel 返回 `ScoredUnit`；后续仍由既有的 Fuser、UnitReader、Discloser 处理。 +`WikimemMemdirRecaller` 使用稳定 SHA-256 派生 unit id，并通过 DOCUMENT channel 返回 `ScoredUnit`；后续仍由既有的 Fuser、UnitReader、Discloser 处理。 ## 6. 团队记忆同步 @@ -310,7 +310,7 @@ flowchart LR 1. LoCoMo 和 LongMemEval 在检索完成后才读取标签计分。 2. MemGallery 的 `source_session_ids` 仅为元数据，不参与排序。 3. case 中的 `baseline_retrieved_*`、`final_retrieved_*`、`ranked_clues`、`retrieved_file_paths` 默认拒绝。 -4. 文件系统评测只读取有 `producer=mem2.0.cc_memory.python`、匹配 dataset name 和 schema version 的 workspace retrieval JSON。 +4. 文件系统评测只读取有 `producer=mem2.0.wikimem.python`、匹配 dataset name 和 schema version 的 workspace retrieval JSON。 5. provenance 是输入来源检查，不是密码学签名；面对主动伪造输入时，调用方仍必须控制 workspace 的写入权限。 ## 9. 全量数据集 Recall 测试结果 @@ -320,7 +320,7 @@ flowchart LR | `LoCoMo`
纯文本长程对话情景记忆检索基准 | `recall_macro` | `0.9413` | Python retained `qmd_consensus` | `locomo10.json`；`1986 cases`；`top_k=24`；`wiki_mode=text`；远程全量运行（2026-07-16） | | `LoCoMo_refined`
带图像证据的多模态长程对话情景记忆检索基准 | `recall_macro` | `0.9556` | Python retained + multimodal adjunct | `1382 cases`；`top_k=24`；adjunct `top_k=6`；`wiki_mode=multimodal`；远程全量运行（2026-07-16） | | `LongMemEval`
长程对话记忆检索基准 | `turn/session recall_any/all@5,@10,@30` | `turn @5 any 0.9212 / all 0.8115；@10 any 0.9642 / all 0.8926；@30 any 0.9833 / all 0.9379`
`session @5 any 0.9308 / all 0.7947；@10 any 0.9690 / all 0.8807；@30 any 0.9857 / all 0.9332` | Python retained `qmd_consensus` | `longmemeval_s_cleaned.json`；`500 cases`；`419 evaluated`；`top_k=100`；`threads=4`；`wiki_mode=text` | -| `AMA-Bench`
自主 agent 轨迹记忆检索基准 | `proxy_recall@112` | `0.9142` | Python `cc_memory_qmd_ama` retrieval-only | `208 episodes`；`2496 questions`；`top_k=112`；`methods=["cc_memory_qmd_ama"]`；`wiki_mode=text`；远程全量运行（2026-07-16） | +| `AMA-Bench`
自主 agent 轨迹记忆检索基准 | `proxy_recall@112` | `0.9142` | Python `wikimem_qmd_ama` retrieval-only | `208 episodes`；`2496 questions`；`top_k=112`；`methods=["wikimem_qmd_ama"]`；`wiki_mode=text`；远程全量运行（2026-07-16） | ## 10. 验收要求 diff --git a/docs/specs/S02-memory-api.md b/docs/specs/S02-memory-api.md index c327af0d..091ce482 100644 --- a/docs/specs/S02-memory-api.md +++ b/docs/specs/S02-memory-api.md @@ -5,8 +5,8 @@ | 项 | 值 | |---|---| | 关联模块 | src/api/ | -| 最近一次修订日期 | 2026-08-05 | -| 关联特性文档 | docs/features/F01-system-spec-design.md，docs/features/api/F01-memory-api-impl-design.md，docs/features/api/F02-write-infer-extract.md，docs/features/api/F03-batch-write-api.md，docs/features/construction/F02-dynamic-extraction-consolidation.md，docs/features/construction/F04-cc-memory-compat.md，docs/features/common/F03-scope-space-isolation.md，docs/features/retrieval/F03-metadata-filtering.md，docs/features/control/F04-permission-context-routing.md，docs/features/control/F05-cloud-engine-design.md | +| 最近一次修订日期 | 2026-07-30 | +| 关联特性文档 | docs/features/F01-system-spec-design.md，docs/features/api/F01-memory-api-impl-design.md，docs/features/api/F02-write-infer-extract.md，docs/features/construction/F02-dynamic-extraction-consolidation.md，docs/features/F02-wikimem-compat.md，docs/features/construction/F04-wikimem-compat.md，docs/features/common/F03-scope-space-isolation.md，docs/features/retrieval/F03-metadata-filtering.md，docs/features/control/F04-permission-context-routing.md，docs/features/control/F05-cloud-engine-design.md | ## 范围 / 边界 **管什么**： @@ -50,8 +50,6 @@ |------|------|------| | `write` | `(content, scope, source=TEXT, *, identity, assets, tags, metadata, occurred_at) -> list[MemoryUnit]` | 同步写入：鉴权 WRITE→委托 Engine→阻塞至 hot path 完成。infer/procedural 触发时返回 `created_ids` 对应的派生单元（可空），否则返回原始单元 | | `write_async` | `async (同签名) -> list[MemoryUnit]` | 异步写入：直通 Engine 协程，供事件循环形态使用 | -| `batch_write` | `(items: list[BatchWriteItem], scope=None, source=TEXT, *, identity, tags, metadata, occurred_at, stream_id="", continue_on_error=True) -> BatchWriteResult` | 同步桥接批量写入；逐项归一化、WRITE 鉴权、space 校验与审计，结果始终按输入索引对齐 | -| `batch_write_async` | `async (同签名) -> BatchWriteResult` | 串行保序批量写入；默认归集单项错误，`continue_on_error=False` 时后续项为 `Skipped` | | `recall` | `(query, context: Context, *, identity, filters, as_of, top_k, disclosure, with_trajectory) -> RetrievalResult` | 混合检索：鉴权 READ→拆 Context→装配 RetrievalQuery→委托 Engine | | `list` | `(scope, *, identity, offset=0, limit=100, memory_types=None, extensions=None, filters=None) -> MemoryListResult` | 列出已建索引记忆：支持类型/FilterExpr 过滤、自定义参数透传和分页前精确总数；只返回 `/memory/` 真源记录 | | `get` | `(unit_id, scope, *, identity, as_of=None) -> MemoryUnit` | 真源点读：鉴权 READ→委托 Engine | @@ -244,12 +242,6 @@ scope 不走 filters。metadata 比较严格保留类型：number、string、boo - `items: list[MemoryUnit]`：当前分页结果。 - `count: int`：同一 Scope 和过滤条件下的分页前精确总数，不受 offset/limit 影响。 -### BatchWriteItem / BatchWriteOutcome / BatchWriteResult（batch_write，`control/types.py`） - -- `BatchWriteItem` 表达单项内容与可选 scope/source/tags/metadata/occurred_at 覆盖；`stream_id`、`sequence`、`idempotency_key` 首版仅用于调度和回显，不写入真源。 -- `BatchWriteOutcome` 包含输入索引、归一化 item、该项产生的 `units` 与可归集的 `error` / `error_type`；成功且 units 为空仍是成功。Engine 的非领域异常也必须归集为 `InternalError`，不能使整批 HTTP 请求退化为 500。 -- `BatchWriteResult.outcomes` 与输入严格一一对应。相同 `(Scope, stream_id)` 的非空 `sequence` 不得重复；接口不自动重排。 - ### DisclosureLevel / RetrievalResult（recall 返回，`retrieval/types.py`） `DisclosureLevel`：`L0`（摘要）/ `L1`（片段）/ `L2`（全文）/ `ADAPTIVE`（按 `max_tokens` 预算自动选层级）。 diff --git a/docs/specs/S03-control.md b/docs/specs/S03-control.md index a65cb90c..03bd041e 100644 --- a/docs/specs/S03-control.md +++ b/docs/specs/S03-control.md @@ -5,8 +5,8 @@ | 项 | 值 | |---|---| | 关联模块 | src/control/ | -| 最近一次修订日期 | 2026-08-05 | -| 关联特性文档 | docs/features/F01-system-spec-design.md，docs/features/api/F01-memory-api-impl-design.md，docs/features/api/F02-write-infer-extract.md，docs/features/api/F03-batch-write-api.md，docs/features/construction/F02-dynamic-extraction-consolidation.md，docs/features/construction/F04-cc-memory-compat.md，docs/features/control/F02-control-isolation-and-audit.md，docs/features/control/F03-control-pipeline-routing.md，docs/features/control/F04-permission-context-routing.md，docs/features/control/F05-cloud-engine-design.md，docs/features/common/F03-scope-space-isolation.md，docs/features/retrieval/F03-metadata-filtering.md | +| 最近一次修订日期 | 2026-07-30 | +| 关联特性文档 | docs/features/F01-system-spec-design.md，docs/features/api/F01-memory-api-impl-design.md，docs/features/api/F02-write-infer-extract.md，docs/features/construction/F02-dynamic-extraction-consolidation.md，docs/features/F02-wikimem-compat.md，docs/features/construction/F04-wikimem-compat.md，docs/features/control/F02-control-isolation-and-audit.md，docs/features/control/F03-control-pipeline-routing.md，docs/features/control/F04-permission-context-routing.md，docs/features/control/F05-cloud-engine-design.md，docs/features/common/F03-scope-space-isolation.md，docs/features/retrieval/F03-metadata-filtering.md | ## 范围 / 边界 **管什么**： @@ -70,7 +70,6 @@ class ControlOperator(ABC): | 方法 | 签名 | 语义 | |------|------|------| | `write` | `async (content, scope, source, *, assets, tags, metadata: dict[str, Any] \| None, occurred_at) -> list[MemoryUnit]` | 规约→可选抽取/分类→落盘+建索引；`infer=true` 时返回 `created_ids` 对应的派生结果，否则处理原始单元（直写不去重） | -| `batch_write` | `async (items: list[BatchWriteItem], *, continue_on_error=True) -> BatchWriteResult` | 只接收 API 已归一化并完成鉴权/space 前置校验的项；按输入顺序复用 `write`，归集领域异常及非领域异常（后者为 `InternalError`）；fail-fast 时填充 `Skipped` outcomes | | `recall` | `async (scope, query: RetrievalQuery) -> RetrievalResult` | 委托 Retriever 完整检索链路 | | `list` | `async (scope, *, offset=0, limit=100, memory_types=None, extensions=None, filters=None) -> MemoryListResult` | 校验分页参数并完整委托 `KVStore.list`；返回当前页和分页前匹配总数 | | `permission_context_for_unit` | `async (unit_id, scope) -> PermissionContext` | 读取已有记忆的权限上下文，只返回 memory_type/tags/metadata 等鉴权元数据，不返回 content/assets | diff --git a/docs/specs/S04-retrieval.md b/docs/specs/S04-retrieval.md index c9aaa616..c3874469 100644 --- a/docs/specs/S04-retrieval.md +++ b/docs/specs/S04-retrieval.md @@ -6,7 +6,7 @@ |---|---| | 关联模块 | src/retrieval/ | | 最近一次修订日期 | 2026-07-30 | -| 关联特性文档 | docs/features/F01-system-spec-design.md、docs/features/construction/F04-cc-memory-compat.md、docs/features/retrieval/F02-retrieval-threshold-topk-design.md、docs/features/retrieval/F03-metadata-filtering.md、docs/features/retrieval/F04-score-max-fusion.md | +| 关联特性文档 | docs/features/F01-system-spec-design.md、docs/features/retrieval/F02-retrieval-threshold-topk-design.md、docs/features/retrieval/F03-metadata-filtering.md、docs/features/retrieval/F04-score-max-fusion.md、docs/features/F02-wikimem-compat.md、docs/features/construction/F04-wikimem-compat.md | ## 范围 / 边界 diff --git a/docs/specs/S05-construction.md b/docs/specs/S05-construction.md index a5bcef35..3576be1c 100644 --- a/docs/specs/S05-construction.md +++ b/docs/specs/S05-construction.md @@ -1,12 +1,28 @@ # S05 — 构建层（Construction Layer） +## LLM-native Wiki construction contract + +The evaluation-side open-domain Wiki path accepts `SemanticSource` records and +produces provenance-preserving `SemanticMemory` records. The supported memory +ontology is `entity`, `fact`, `event`, `preference`, `skill`, `relationship`, +`decision`, `constraint`, `context`, and `artifact`. + +`WikiBuilder(llm=None, mode="llm")` uses the LLM extractor, entity resolver, +consolidator, and profile/timeline/decision synthesis when an LLM is injected. +With no model, or after an LLM failure, it falls back to the deterministic +`TemplateExtractor`; retained benchmark adapters may continue using their +historical deterministic renderer for exact regression compatibility. Every +generated page carries evidence, provenance, metadata, stable IDs, and source +links. Query understanding is optional and must not replace the deterministic +qmd retrieval path. + ## 元信息 | 项 | 值 | |---|---| | 关联模块 | src/construction/ | -| 最近一次修订日期 | 2026-08-04 | -| 关联特性文档 | docs/features/F01-system-spec-design.md, docs/features/construction/F01-construction-spec-design.md, docs/features/construction/F02-dynamic-extraction-consolidation.md, docs/features/construction/F03-extraction-layer-integrity.md, docs/features/construction/F04-cc-memory-compat.md, docs/features/common/F01-memory-layer.md, docs/features/common/F03-scope-space-isolation.md, docs/features/retrieval/F03-metadata-filtering.md | +| 最近一次修订日期 | 2026-08-05 | +| 关联特性文档 | docs/features/F01-system-spec-design.md, docs/features/construction/F01-construction-spec-design.md, docs/features/construction/F02-dynamic-extraction-consolidation.md, docs/features/construction/F03-extraction-layer-integrity.md, docs/features/F02-wikimem-compat.md, docs/features/construction/F04-wikimem-compat.md, docs/features/common/F01-memory-layer.md, docs/features/common/F03-scope-space-isolation.md, docs/features/retrieval/F03-metadata-filtering.md | ## 范围 / 边界 @@ -76,8 +92,7 @@ class ConstructionOperator(ABC): |------|------|------| | `extract` | `(units: list[MemoryUnit], *, context: ExtractContext \| None = None) -> list[MemoryUnit]` | 从一批原始记忆单元中提取零或多条低抽象粒度的派生单元；context 只作 prompt 参考 | -派生单元的 `tier` 由 LLM 在抽取时产出；`tags` 为源 unit 的 write tags ∪ LLM 主题 -tags ∪ 系统标记（`extracted` / `procedural`）。`layers`（L0/L1 分层标注）不由 Extractor +派生单元的 `tier`/`tags` 由 LLM 在抽取时产出。`layers`（L0/L1 分层标注）不由 Extractor 产出——由 Evolver 抽取后委托 `LayerAnnotator` 生成（见下文 LayerAnnotator 节 + F01-memory-layer）。 LLM 抽取只合并同一实体同一关系或同一事件。派生单元的 L2 只保存紧凑抽取陈述， diff --git a/docs/specs/S07-common.md b/docs/specs/S07-common.md index 1b89cdb7..58be7956 100644 --- a/docs/specs/S07-common.md +++ b/docs/specs/S07-common.md @@ -6,7 +6,7 @@ |---|-------------| | 关联模块 | src/common/ | | 最近一次修订日期 | 2026-08-03 | -| 关联特性文档 | docs/features/F01-system-spec-design.md，docs/features/api/F01-memory-api-impl-design.md，docs/features/construction/F04-cc-memory-compat.md，docs/features/common/F01-memory-layer.md，docs/features/common/F02-dashscope-llm-provider.md，docs/features/common/F03-scope-space-isolation.md，docs/features/common/F04-security-interfaces-and-encryption.md，docs/features/control/F02-control-isolation-and-audit.md，docs/features/retrieval/F03-metadata-filtering.md，docs/features/common/F05-model-service-ssl.md，docs/features/common/F06-distributed-lock.md | +| 关联特性文档 | docs/features/F01-system-spec-design.md，docs/features/api/F01-memory-api-impl-design.md，docs/features/construction/F04-wikimem-compat.md，docs/features/common/F01-memory-layer.md，docs/features/common/F02-dashscope-llm-provider.md，docs/features/common/F03-scope-space-isolation.md，docs/features/common/F04-security-interfaces-and-encryption.md，docs/features/control/F02-control-isolation-and-audit.md，docs/features/retrieval/F03-metadata-filtering.md，docs/features/common/F05-model-service-ssl.md，docs/features/common/F06-distributed-lock.md | ## 范围 / 边界 diff --git a/evaluation/wikimem/AGENTS.md b/evaluation/wikimem/AGENTS.md new file mode 100644 index 00000000..bfdc93b5 --- /dev/null +++ b/evaluation/wikimem/AGENTS.md @@ -0,0 +1,17 @@ +# wikimem retained evaluation + +本目录承接 `rust/wikimem/src/retained_eval` 在 mem2.0 中的 evaluation/profile 能力。 + +## 文件职责 + +- `__init__.py`：导出 retained_eval 兼容公共面。 +- `retained_eval.py`：LoCoMo 样本归一化、EvalCase / CaseScore DTO、指标摘要、plugin list 解析、progress 文案和 harness artifact 写出。 +- `example_eval.py` / `ama_bench.py`：文件系统评测和 AMA retrieval-only 运行器；必须标明标签来源，答案派生指标只能标为 proxy。 + +## 边界 + +- 本目录属于 evaluation，不得修改 `MemoryAPI`、`RetrievalQuery` 或 `Retriever.retrieve` 签名。 +- LoCoMo / LongMemEval 字段只能在 evaluation adapter 内消费，不得进入通用 core 类型。 +- retained 检索 profile 后续可调用现有 mem2.0 API / Recaller，但不得把数据集字段塞进 `src/retrieval/types.py`。 +- artifact 写出必须保留 summary、case trace、failure report、category breakdown 和 stage profile，便于精度回归定位 miss stage。 +- `.kb-research/retrieval` 记录仅可由带 Python provenance 的 workspace 消费；外部或 Rust 记录不得作为检索输入。 diff --git a/figures/cc-memory-vs-llm-wiki.md b/figures/wikimem-vs-llm-wiki.md similarity index 92% rename from figures/cc-memory-vs-llm-wiki.md rename to figures/wikimem-vs-llm-wiki.md index 22890845..8f450055 100644 --- a/figures/cc-memory-vs-llm-wiki.md +++ b/figures/wikimem-vs-llm-wiki.md @@ -1,6 +1,6 @@ -# cc-memory 与 LLM wiki 的三项关键差异 +# wikimem 与 LLM wiki 的三项关键差异 -下面的流程图将文本 wiki 基线与 cc-memory 的多模态、qmd 多视图融合和 agent 轨迹链路并置展示。黄色节点为差异增强点。 +下面的流程图将文本 wiki 基线与 wikimem 的多模态、qmd 多视图融合和 agent 轨迹链路并置展示。黄色节点为差异增强点。 ```mermaid flowchart TB @@ -16,7 +16,7 @@ flowchart TB input --> capture - subgraph build["cc-memory 构建：保留 wiki，并增加三条增强链路"] + subgraph build["wikimem 构建：保留 wiki，并增加三条增强链路"] normalize["PreparedSample 归一化
records / sessions / questions / raw"] text_views["多视图文本 wiki
source / turn / observation / event / entity / topic"] snapshot["MemorySnapshot
轻量 lexical root recall"] @@ -38,7 +38,7 @@ flowchart TB mm_download -."未启用或下载失败：保留 URL 与状态".-> mm_artifact mm_vision -."未配置 vision：text / caption / query 仍可检索".-> mm_artifact - subgraph query["cc-memory 查询：多阶段、可诊断"] + subgraph query["wikimem 查询：多阶段、可诊断"] question["问题 / agent 子任务"] profile["QuestionProfile
tokens / fuzzy / entities / temporal / relation"] root["Root retrieval
snapshot keyword + root header"] diff --git a/figures/cc-memory-vs-llm-wiki.mmd b/figures/wikimem-vs-llm-wiki.mmd similarity index 96% rename from figures/cc-memory-vs-llm-wiki.mmd rename to figures/wikimem-vs-llm-wiki.mmd index 0b4da138..660d3dba 100644 --- a/figures/cc-memory-vs-llm-wiki.mmd +++ b/figures/wikimem-vs-llm-wiki.mmd @@ -11,7 +11,7 @@ flowchart TB input --> capture - subgraph build["cc-memory 构建：保留 wiki，并增加三条增强链路"] + subgraph build["wikimem 构建：保留 wiki，并增加三条增强链路"] normalize["PreparedSample 归一化
records / sessions / questions / raw"] text_views["多视图文本 wiki
source / turn / observation / event / entity / topic"] snapshot["MemorySnapshot
轻量 lexical root recall"] @@ -33,7 +33,7 @@ flowchart TB mm_download -."未启用或下载失败：保留 URL 与状态".-> mm_artifact mm_vision -."未配置 vision：text / caption / query 仍可检索".-> mm_artifact - subgraph query["cc-memory 查询：多阶段、可诊断"] + subgraph query["wikimem 查询：多阶段、可诊断"] question["问题 / agent 子任务"] profile["QuestionProfile
tokens / fuzzy / entities / temporal / relation"] root["Root retrieval
snapshot keyword + root header"] diff --git a/figures/cc-memory-vs-llm-wiki.png b/figures/wikimem-vs-llm-wiki.png similarity index 100% rename from figures/cc-memory-vs-llm-wiki.png rename to figures/wikimem-vs-llm-wiki.png From 0a0ad002e23d336d87ec59772da5e4e0a376ba98 Mon Sep 17 00:00:00 2001 From: xinyao1994 Date: Thu, 6 Aug 2026 17:29:12 +0800 Subject: [PATCH 3/3] fix(memory): consolidate exported CI repairs --- agent_plugin/wikimem/team_sync.py | 7 +- evaluation/wikimem/ama_bench.py | 10 +- evaluation/wikimem/example_eval.py | 54 ++- evaluation/wikimem/llm_semantics.py | 5 +- evaluation/wikimem/locomo_refined.py | 119 ++++-- evaluation/wikimem/longmemeval.py | 40 +- evaluation/wikimem/qmd_consensus.py | 277 ++++++------- evaluation/wikimem/retained_eval.py | 171 ++++---- evaluation/wikimem/retrieval_profile.py | 388 +++++++++--------- evaluation/wikimem/wiki_builder.py | 14 +- .../wikimem_baseline_extractor.py | 46 ++- src/retrieval/wikimem_memdir.py | 38 +- .../evaluation/test_wikimem_example_eval.py | 26 +- .../evaluation/test_wikimem_qmd_consensus.py | 7 +- 14 files changed, 640 insertions(+), 562 deletions(-) diff --git a/agent_plugin/wikimem/team_sync.py b/agent_plugin/wikimem/team_sync.py index 088aad06..7a1c7b42 100644 --- a/agent_plugin/wikimem/team_sync.py +++ b/agent_plugin/wikimem/team_sync.py @@ -414,11 +414,14 @@ def hash_content(content: str) -> str: def validate_relative_team_memory_key(key: str) -> str: - if "\0" in key or "\\" in key or key.startswith("/") or _has_windows_prefix(key): + has_unsafe_character = "\0" in key or "\\" in key + has_absolute_prefix = key.startswith("/") or _has_windows_prefix(key) + if has_unsafe_character or has_absolute_prefix: raise ValueError(f"Invalid team memory key: {key}") decoded = urllib.parse.unquote(key) - if decoded != key and (".." in decoded or "/" in decoded or "\\" in decoded): + has_encoded_traversal = ".." in decoded or "/" in decoded or "\\" in decoded + if decoded != key and has_encoded_traversal: raise ValueError(f"Invalid team memory key: {key}") parts = key.split("/") diff --git a/evaluation/wikimem/ama_bench.py b/evaluation/wikimem/ama_bench.py index 6d8cc5b8..a3a38152 100644 --- a/evaluation/wikimem/ama_bench.py +++ b/evaluation/wikimem/ama_bench.py @@ -227,11 +227,11 @@ def run_python_wikimem_qmd_retrieval( entity_names=_candidate_entity_names(episode, question), top_k=top_k, ) - retrieved_turn_ids = [ - turn_id - for file in result.files - if (turn_id := _extract_turn_id_from_path(file.file_path)) is not None - ] + retrieved_turn_ids = [] + for file in result.files: + turn_id = _extract_turn_id_from_path(file.file_path) + if turn_id is not None: + retrieved_turn_ids.append(turn_id) metrics = score_ama_retrieval_proxy(episode, question, retrieved_turn_ids) return AmaMethodQuestionResult( method_name=method_name, diff --git a/evaluation/wikimem/example_eval.py b/evaluation/wikimem/example_eval.py index e0582641..31a87be1 100644 --- a/evaluation/wikimem/example_eval.py +++ b/evaluation/wikimem/example_eval.py @@ -319,7 +319,10 @@ def _meta_crag_samples_from_row( "turn_index": index, "query": query, "search_query": searches[index] if index < len(searches) else query, - "answer": (answers[index] if index < len(answers) else "") or (full_answers[index] if index < len(full_answers) else ""), + "answer": ( + (answers[index] if index < len(answers) else "") + or (full_answers[index] if index < len(full_answers) else "") + ), } for index, query in enumerate(queries) ] @@ -467,7 +470,9 @@ def _meta_crag_answer_supported(answer: str, artifact: dict[str, Any]) -> bool: normalized_answer = _normalize_support_text(answer) text = _normalize_support_text(f"{artifact['title']}\n{artifact['text']}") answer_tokens = _support_tokens(answer) - return bool(normalized_answer and normalized_answer in text) or _support_tokens_match(answer_tokens, _support_tokens(text)) + return bool(normalized_answer and normalized_answer in text) or _support_tokens_match( + answer_tokens, _support_tokens(text) + ) def _meta_crag_support_proxy_score(sample: dict[str, Any], artifact: dict[str, Any]) -> int: @@ -585,12 +590,20 @@ def _write_mem_gallery_python_workspace( / f"{evidence_slug}_obs_{_observation_suffix(turn['clue_id'])}.md" ) observation_path.write_text( - f"# Observation {turn['clue_id']}\nEvidence: {turn['clue_id']}\nSession: {turn['session_id']}\n\n{searchable}\n", + ( + f"# Observation {turn['clue_id']}\n" + f"Evidence: {turn['clue_id']}\nSession: {turn['session_id']}\n\n" + f"{searchable}\n" + ), encoding="utf-8", ) memory_path = root / "wiki" / "memories" / f"clue-summary-{evidence_slug}.md" memory_path.write_text( - f"# Clue summary {turn['clue_id']}\nEvidence: {turn['clue_id']}\nSession: {turn['session_id']}\n\n{clue_summary}\n", + ( + f"# Clue summary {turn['clue_id']}\n" + f"Evidence: {turn['clue_id']}\nSession: {turn['session_id']}\n\n" + f"{clue_summary}\n" + ), encoding="utf-8", ) retrieval_path = root / ".kb-research" / "retrieval" / f"{evidence_slug}.json" @@ -1066,7 +1079,10 @@ def _meta_crag_ids_for_paths( for path in paths: normalized = path.replace("\\", "/") evidence_id = evidence_by_path.get(path) or evidence_by_path.get(normalized) or "" - if "example/meta-crag/" in normalized and "/wiki/memories/" not in normalized and "/.kb-research/" not in normalized: + is_meta_crag = "example/meta-crag/" in normalized + is_memory_page = "/wiki/memories/" in normalized + is_retrieval_artifact = "/.kb-research/" in normalized + if is_meta_crag and not is_memory_page and not is_retrieval_artifact: evidence_id = _rust_stable_path_id(_rust_relative_meta_crag_path(normalized)) selected.append(evidence_id) return _ordered_unique(selected) @@ -1089,12 +1105,13 @@ def _retrieve_mem_gallery_artifact_ids( } if not query_tokens: return [] - candidates = [ - file - for file in files - if "/.kb-research/retrieval/" in file.file_path.replace("\\", "/") - and _is_mem_gallery_clue_id(evidence_by_path.get(file.file_path, "")) - ] + candidates = [] + for file in files: + normalized_path = file.file_path.replace("\\", "/") + is_retrieval_artifact = "/.kb-research/retrieval/" in normalized_path + is_clue = _is_mem_gallery_clue_id(evidence_by_path.get(file.file_path, "")) + if is_retrieval_artifact and is_clue: + candidates.append(file) token_by_path = { file.file_path: set(_cached_content_keywords(file.content)) for file in candidates @@ -1164,7 +1181,7 @@ def sip_round() -> None: end = len(data) - (len(data) % 8) for offset in range(0, end, 8): - chunk = int.from_bytes(data[offset : offset + 8], "little") + chunk = int.from_bytes(data[offset:offset + 8], "little") v3 ^= chunk sip_round() v0 ^= chunk @@ -1259,13 +1276,12 @@ def _retrieve_evermem_evidence( candidate_scores.get(row["evidence_id"], 0.0), score * 0.45, ) - selected = [ - evidence_id - for evidence_id, _ in sorted( - candidate_scores.items(), - key=lambda item: (-item[1], item[0]), - )[: max(top_k, 1)] - ] + selected = [] + ranked_candidates = sorted( + candidate_scores.items(), key=lambda item: (-item[1], item[0]) + ) + for evidence_id, _ in ranked_candidates[: max(top_k, 1)]: + selected.append(evidence_id) paths = [ file_by_id[evidence_id].file_path for evidence_id in selected diff --git a/evaluation/wikimem/llm_semantics.py b/evaluation/wikimem/llm_semantics.py index b5184b39..a3cb407a 100644 --- a/evaluation/wikimem/llm_semantics.py +++ b/evaluation/wikimem/llm_semantics.py @@ -17,7 +17,6 @@ from common.llm.base import LLM from common.type_def.chat import ChatMessage - MEMORY_KINDS = ( "entity", "fact", @@ -127,7 +126,7 @@ def extract_semantic_memories( materialized = [source for source in sources if source.text.strip()] result: list[SemanticMemory] = [] for start in range(0, len(materialized), max(1, batch_size)): - batch = materialized[start : start + max(1, batch_size)] + batch = materialized[start:start + max(1, batch_size)] source_text = "\n".join( "---\n" f"source_id: {source.source_id}\n" @@ -213,7 +212,7 @@ def _parse_memory_item(item: Any, valid_ids: set[str]) -> SemanticMemory | None: relations = tuple(value for value in relations if value is not None) stable_key = f"{source_id}:{kind}:{content}".encode("utf-8") memory_id = _clean_scalar(item.get("memory_id")) or ( - f"{source_id}:{kind}:{hashlib.sha1(stable_key).hexdigest()[:16]}" + f"{source_id}:{kind}:{hashlib.sha256(stable_key).hexdigest()[:16]}" ) return SemanticMemory( memory_id=memory_id, diff --git a/evaluation/wikimem/locomo_refined.py b/evaluation/wikimem/locomo_refined.py index b2e23fdf..fb10a49a 100644 --- a/evaluation/wikimem/locomo_refined.py +++ b/evaluation/wikimem/locomo_refined.py @@ -104,6 +104,18 @@ class _VisionSummary: error: str | None = None +@dataclass(frozen=True) +class _DownloadAttempt: + opener: Any + url: str + strategy: str + browser_headers: bool + force_identity: bool + local_path: Path + source_url: str + timeout_secs: int + + def build_multimodal_memory_artifacts( sample: PreparedSample, kb_root: str | Path, @@ -457,16 +469,15 @@ def _build_download_source_report(workspace_root: Path) -> dict[str, Any]: error = str(item.get("error") or "unknown_error") if error not in entry["example_errors"] and len(entry["example_errors"]) < 3: entry["example_errors"].append(error) - successful_domains = [ - {"domain": domain, "downloaded_count": count} - for domain, count in sorted(successful.items(), key=lambda pair: (-pair[1], pair[0])) - ] - failed_domains = [ - {"domain": domain, **value} - for domain, value in sorted( - failed.items(), key=lambda pair: (-pair[1]["failed_count"], pair[0]) - ) - ] + successful_domains = [] + for domain, count in sorted(successful.items(), key=lambda pair: (-pair[1], pair[0])): + successful_domains.append({"domain": domain, "downloaded_count": count}) + failed_domains = [] + failed_items = sorted( + failed.items(), key=lambda pair: (-pair[1]["failed_count"], pair[0]) + ) + for domain, value in failed_items: + failed_domains.append({"domain": domain, **value}) return { "total_downloaded_images": sum(successful.values()), "total_failed_images": sum(item["failed_count"] for item in failed.values()), @@ -683,15 +694,12 @@ def _collect_multimodal_turns(raw_sample: dict[str, Any]) -> list[dict[str, Any] conversation = raw_sample.get("conversation") if not isinstance(conversation, dict): raise ValueError("raw_sample.conversation must be an object") - sessions = sorted( - ( - int(key.removeprefix("session_")), - value, - ) - for key, value in conversation.items() - if key.startswith("session_") - and key.removeprefix("session_").isdigit() - ) + sessions = [] + for key, value in conversation.items(): + session_suffix = key.removeprefix("session_") + if key.startswith("session_") and session_suffix.isdigit(): + sessions.append((int(session_suffix), value)) + sessions.sort() turns: list[dict[str, Any]] = [] for session_number, messages in sessions: if not isinstance(messages, list): @@ -803,14 +811,16 @@ def _materialize_image_assets( last_strategy, last_url = strategy, url try: mime, resolved = _execute_download_attempt( - opener, - url, - strategy, - browser_headers, - identity, - local_path, - source_url, - options.request_timeout_secs, + _DownloadAttempt( + opener=opener, + url=url, + strategy=strategy, + browser_headers=browser_headers, + force_identity=identity, + local_path=local_path, + source_url=source_url, + timeout_secs=options.request_timeout_secs, + ) ) assets.append( _DownloadedImageAsset( @@ -859,11 +869,33 @@ def _build_download_attempt_plans(source_url: str) -> list[tuple[str, str, bool, ] parsed = urllib.parse.urlparse(source_url) if parsed.scheme == "http": - candidates.append((urllib.parse.urlunparse(parsed._replace(scheme="https")), "scheme_swap_browser_identity", True, True)) + candidates.append( + ( + urllib.parse.urlunparse(parsed._replace(scheme="https")), + "scheme_swap_browser_identity", + True, + True, + ) + ) if parsed.hostname and parsed.hostname.lower() == "imgur.com": - candidates.append((urllib.parse.urlunparse(parsed._replace(netloc="i.imgur.com")), "normalized_direct_browser_identity", True, True)) + candidates.append( + ( + urllib.parse.urlunparse(parsed._replace(netloc="i.imgur.com")), + "normalized_direct_browser_identity", + True, + True, + ) + ) if parsed.hostname and parsed.hostname.lower() == "i.redd.it": - candidates.append(("https://www.reddit.com/media?url=" + urllib.parse.quote(source_url, safe=""), "reddit_media_browser_identity", True, True)) + candidates.append( + ( + "https://www.reddit.com/media?url=" + + urllib.parse.quote(source_url, safe=""), + "reddit_media_browser_identity", + True, + True, + ) + ) seen: set[tuple[str, bool, bool]] = set() for item in candidates: key = (item[0], item[2], item[3]) @@ -873,29 +905,23 @@ def _build_download_attempt_plans(source_url: str) -> list[tuple[str, str, bool, return plans -def _execute_download_attempt( - opener: Any, - url: str, - strategy: str, - browser_headers: bool, - force_identity: bool, - local_path: Path, - source_url: str, - timeout_secs: int, -) -> tuple[str | None, str]: +def _execute_download_attempt(attempt: _DownloadAttempt) -> tuple[str | None, str]: + opener = attempt.opener + url = attempt.url + local_path = attempt.local_path headers = {} - if browser_headers: + if attempt.browser_headers: headers.update( { "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 Chrome/135.0 Safari/537.36", "Accept": "image/avif,image/webp,image/apng,image/*,*/*;q=0.8", - "Referer": _derive_referer(source_url), + "Referer": _derive_referer(attempt.source_url), } ) - if force_identity: + if attempt.force_identity: headers["Accept-Encoding"] = "identity" request = urllib.request.Request(url, headers=headers) - with opener.open(request, timeout=max(timeout_secs, 1)) as response: + with opener.open(request, timeout=max(attempt.timeout_secs, 1)) as response: status = getattr(response, "status", response.getcode()) if status < 200 or status >= 300: raise RuntimeError(f"http {status}") @@ -1117,7 +1143,10 @@ def _render_memory_page( f"attributes: {', '.join(vision.attributes) or '(none)'}\n" f"keywords: {', '.join(vision.keywords) or '(none)'}\n" ) - quote = lambda value: str(value).replace('"', '\\"') + + def quote(value: object) -> str: + return str(value).replace('"', '\\"') + sources = ", ".join(f'"{quote(item)}"' for item in manifest["sources"]) topics = ", ".join(f'"{quote(item)}"' for item in manifest["linked_topics"]) return ( diff --git a/evaluation/wikimem/longmemeval.py b/evaluation/wikimem/longmemeval.py index 177ecf2d..298bb205 100644 --- a/evaluation/wikimem/longmemeval.py +++ b/evaluation/wikimem/longmemeval.py @@ -181,13 +181,11 @@ def adapt_longmemeval_to_locomo_samples( ] }, ) - session_turn_ids = [ - f"{session_id}_{index + 1}" - for index, turn in enumerate( - _get_index(sample.get("haystack_sessions", []), session_index) or [] - ) - if str(turn.get("role", "")).lower() == "user" - ] + session_turn_ids = [] + session_turns = _get_index(sample.get("haystack_sessions", []), session_index) or [] + for index, turn in enumerate(session_turns): + if str(turn.get("role", "")).lower() == "user": + session_turn_ids.append(f"{session_id}_{index + 1}") session_support = {"turn_ids": session_turn_ids, "session_ids": [session_id]} support_by_path[f"wiki/sources/session_{session_number}.md"] = session_support support_by_path[ @@ -196,7 +194,7 @@ def adapt_longmemeval_to_locomo_samples( support_by_path[f"wiki/topics/session_{session_number}_events.md"] = session_support for ordinal, note in enumerate(observations.get(session_number, []), start=1): turn_support = { - "turn_ids": session_turn_ids[ordinal - 1 : ordinal], + "turn_ids": session_turn_ids[ordinal - 1:ordinal], "session_ids": [session_id], } support_by_path[ @@ -233,14 +231,9 @@ def adapt_longmemeval_to_locomo_samples( category=None, ) ], - session_datetimes={ - session_number_by_id[session_id]: str(date) - for session_id, date in zip( - sample.get("haystack_session_ids", []), - sample.get("haystack_dates", []), - ) - if session_id in session_number_by_id - }, + session_datetimes=_session_datetimes_from_longmemeval_sample( + sample, session_number_by_id + ), session_summaries=session_summaries, event_summaries=session_events, observations=observations, @@ -249,6 +242,19 @@ def adapt_longmemeval_to_locomo_samples( return prepared +def _session_datetimes_from_longmemeval_sample( + sample: dict[str, Any], + session_number_by_id: dict[str, int], +) -> dict[int, str]: + values: dict[int, str] = {} + session_ids = sample.get("haystack_session_ids", []) + dates = sample.get("haystack_dates", []) + for session_id, date in zip(session_ids, dates): + if session_id in session_number_by_id: + values[session_number_by_id[session_id]] = str(date) + return values + + def convert_retained_output_to_longmemeval( *, samples: list[dict[str, Any]], @@ -471,7 +477,7 @@ def _relative_retrieved_path(path: str, knowledge_root: str) -> str: normalized = path.replace("\\", "/") root = knowledge_root.replace("\\", "/").rstrip("/") prefix = f"{root}/" - return normalized[len(prefix) :] if normalized.startswith(prefix) else normalized.lstrip("/") + return normalized[len(prefix):] if normalized.startswith(prefix) else normalized.lstrip("/") def _page_support_metrics( diff --git a/evaluation/wikimem/qmd_consensus.py b/evaluation/wikimem/qmd_consensus.py index 39fcc342..f2e4878a 100644 --- a/evaluation/wikimem/qmd_consensus.py +++ b/evaluation/wikimem/qmd_consensus.py @@ -143,6 +143,27 @@ class QmdConsensusFileMetrics: support_density: float +@dataclass(frozen=True) +class _CachedQmdFiles: + files: dict[str, RetrievedMemoryFile] + features: dict[str, CachedFileLexicalFeatures] + + +@dataclass(frozen=True) +class _BridgeProposalConfig: + significant: list[str] + max_candidates: int + min_target_score: float + + +@dataclass(frozen=True) +class _QmdScoreContext: + profile: QuestionProfile + significant: list[str] + ngrams: list[str] + named_entities_lower: list[str] + + @dataclass(frozen=True) class CandidateProposal: file_path: str @@ -290,37 +311,32 @@ def build_qmd_consensus_augmentation( phrase_scores[phrase] = phrase_scores.get(phrase, 0.0) + weighted_score phrase_support.setdefault(phrase, set()).add(normalized_path) - tokens = [ - token - for token, _ in sorted( - ( - (token, score) - for token, score in token_scores.items() - if len(token_support.get(token, set())) >= 2 or score >= 7.0 - ), - key=lambda item: ( - -len(token_support.get(item[0], set())), - -item[1], - item[0], - ), - )[:6] - ] - phrases = [ - phrase - for phrase, _ in sorted( - ( - (phrase, score) - for phrase, score in phrase_scores.items() - if (len(phrase_support.get(phrase, set())) >= 1 and score >= 6.0) - or len(phrase_support.get(phrase, set())) >= 2 - ), - key=lambda item: ( - -len(phrase_support.get(item[0], set())), - -item[1], - item[0], - ), - )[:4] - ] + ranked_tokens = [] + for token, score in token_scores.items(): + support_count = len(token_support.get(token, set())) + if support_count >= 2 or score >= 7.0: + ranked_tokens.append((token, score)) + ranked_tokens.sort( + key=lambda item: ( + -len(token_support.get(item[0], set())), + -item[1], + item[0], + ) + ) + tokens = [token for token, _ in ranked_tokens[:6]] + ranked_phrases = [] + for phrase, score in phrase_scores.items(): + support_count = len(phrase_support.get(phrase, set())) + if (support_count >= 1 and score >= 6.0) or support_count >= 2: + ranked_phrases.append((phrase, score)) + ranked_phrases.sort( + key=lambda item: ( + -len(phrase_support.get(item[0], set())), + -item[1], + item[0], + ) + ) + phrases = [phrase for phrase, _ in ranked_phrases[:4]] return QueryAugmentation( tokens=tokens, fuzzy_tokens=tokenize_fuzzy_query(" ".join(tokens)), @@ -365,21 +381,11 @@ def build_qmd_consensus_candidate_proposals( if kind is None: continue features = cached_features[normalized_path] - metrics = _metrics_for( - metrics_cache, - normalized_path, - file, - features, - focused, - significant, - ) + metrics = _metrics_for(metrics_cache, normalized_path, features, focused, significant) score = _score_cached_file_with_metrics( features, normalized_path, - focused, - significant, - ngrams, - named_lower, + _QmdScoreContext(focused, significant, ngrams, named_lower), metrics, ) if kind == "source": @@ -398,20 +404,22 @@ def build_qmd_consensus_candidate_proposals( _sort_qmd_ranked_items(source_view) _sort_qmd_ranked_items(anchor_view) linked_view = _build_qmd_linked_candidate_view( - question, focused, - cached_files, - cached_features, + _CachedQmdFiles(cached_files, cached_features), source_view, anchor_view, + significant, ) - return [ - CandidateProposal(file_path=path, query_hits=query_hits, seed_boost=boost) - for path, query_hits, boost in _fuse_ranked_views( - [(2.0, source_view), (1.5, anchor_view), (1.0, linked_view)], - 12, + proposals = [] + fused = _fuse_ranked_views( + [(2.0, source_view), (1.5, anchor_view), (1.0, linked_view)], + 12, + ) + for path, query_hits, boost in fused: + proposals.append( + CandidateProposal(file_path=path, query_hits=query_hits, seed_boost=boost) ) - ] + return proposals def build_qmd_consensus_rerank_proposals( @@ -435,14 +443,7 @@ def build_qmd_consensus_rerank_proposals( features = cached_features.get(normalized_path) if features is None: continue - metrics = _metrics_for( - metrics_cache, - normalized_path, - candidate.file, - features, - focused, - significant, - ) + metrics = _metrics_for(metrics_cache, normalized_path, features, focused, significant) item = ( normalized_path, max(candidate.query_hits, 1), @@ -458,25 +459,27 @@ def build_qmd_consensus_rerank_proposals( source_view = source_view[:6] anchor_view = anchor_view[:6] seed_files = [candidate.file for candidate in ranked_candidates[:6]] - linked_view = [ - (proposal.file_path, max(proposal.query_hits, 1), proposal.seed_boost) - for proposal in _build_bridge_proposals( - question, - focused, - seed_files, - cached_files, - cached_features, - 8, - 4.2, - ) - ] - return [ - RerankProposal(file_path=path, query_hits=query_hits, seed_boost=boost) - for path, query_hits, boost in _fuse_ranked_views( - [(2.0, source_view), (1.5, anchor_view), (1.0, linked_view)], - 12, - ) - ] + linked_view = [] + bridge_config = _BridgeProposalConfig( + significant=significant_phrases(question), + max_candidates=8, + min_target_score=4.2, + ) + for proposal in _build_bridge_proposals( + focused, + seed_files, + _CachedQmdFiles(cached_files, cached_features), + bridge_config, + ): + linked_view.append((proposal.file_path, max(proposal.query_hits, 1), proposal.seed_boost)) + proposals = [] + fused = _fuse_ranked_views( + [(2.0, source_view), (1.5, anchor_view), (1.0, linked_view)], + 12, + ) + for path, query_hits, boost in fused: + proposals.append(RerankProposal(file_path=path, query_hits=query_hits, seed_boost=boost)) + return proposals def build_qmd_consensus_late_bridge_proposals( @@ -491,13 +494,14 @@ def build_qmd_consensus_late_bridge_proposals( cached_files, cached_features = _cache_files(files) focused = _qmd_focus_profile(profile) return _build_bridge_proposals( - question, focused, seed_files, - cached_files, - cached_features, - 8, - 4.2, + _CachedQmdFiles(cached_files, cached_features), + _BridgeProposalConfig( + significant=significant_phrases(question), + max_candidates=8, + min_target_score=4.2, + ), ) @@ -549,7 +553,7 @@ def keyword_ngrams(question: str) -> list[str]: ngrams = [] for size in (2, 3): for index in range(0, max(0, len(tokens) - size + 1)): - phrase = " ".join(tokens[index : index + size]) + phrase = " ".join(tokens[index:index + size]) if phrase not in seen: ngrams.append(phrase) seen.add(phrase) @@ -616,7 +620,6 @@ def _qmd_keeps_focus_token(token: str, named_entities: list[str]) -> bool: def _metrics_for( metrics_cache: dict[str, QmdConsensusFileMetrics], normalized_path: str, - _file: RetrievedMemoryFile, features: CachedFileLexicalFeatures, profile: QuestionProfile, significant: list[str], @@ -638,18 +641,18 @@ def _metrics_for( def _score_cached_file_with_metrics( features: CachedFileLexicalFeatures, normalized_path: str, - profile: QuestionProfile, - significant: list[str], - ngrams: list[str], - named_entities_lower: list[str], + context: _QmdScoreContext, metrics: QmdConsensusFileMetrics, ) -> float: + profile = context.profile score = metrics.best_line_score score += _candidate_path_weight(normalized_path) * 1.8 score += _token_overlap_with_set(profile.query_tokens, features.path_token_set) * 1.8 - score += sum(2.3 for phrase in significant if phrase in features.lower_content) - score += sum(1.2 for phrase in ngrams if phrase in features.lower_content) - score += sum(1.5 for entity in named_entities_lower if entity in features.lower_content) + score += sum(2.3 for phrase in context.significant if phrase in features.lower_content) + score += sum(1.2 for phrase in context.ngrams if phrase in features.lower_content) + score += sum( + 1.5 for entity in context.named_entities_lower if entity in features.lower_content + ) if features.has_session_marker: score += 1.0 if features.has_evidence_marker: @@ -658,14 +661,14 @@ def _score_cached_file_with_metrics( def _build_qmd_linked_candidate_view( - question: str, profile: QuestionProfile, - cached_files: dict[str, RetrievedMemoryFile], - cached_features: dict[str, CachedFileLexicalFeatures], + corpus: _CachedQmdFiles, source_view: list[tuple[str, int, float]], anchor_view: list[tuple[str, int, float]], + significant: list[str], ) -> list[tuple[str, int, float]]: - significant = significant_phrases(question) + cached_files = corpus.files + cached_features = corpus.features metrics_cache: dict[str, QmdConsensusFileMetrics] = {} direct_atomic: list[tuple[str, int, float]] = [] for normalized_path, file in cached_files.items(): @@ -676,14 +679,7 @@ def _build_qmd_linked_candidate_view( ): continue features = cached_features[normalized_path] - metrics = _metrics_for( - metrics_cache, - normalized_path, - file, - features, - profile, - significant, - ) + metrics = _metrics_for(metrics_cache, normalized_path, features, profile, significant) direct_score = ( metrics.best_line_score * 1.15 + metrics.support_density * 0.2 @@ -720,7 +716,6 @@ def _build_qmd_linked_candidate_view( metrics = _metrics_for( metrics_cache, normalized_target, - target, target_features, profile, significant, @@ -771,15 +766,14 @@ def _build_qmd_linked_candidate_view( def _build_bridge_proposals( - question: str, profile: QuestionProfile, seed_files: list[RetrievedMemoryFile], - cached_files: dict[str, RetrievedMemoryFile], - cached_features: dict[str, CachedFileLexicalFeatures], - max_candidates: int, - min_target_score: float, + corpus: _CachedQmdFiles, + config: _BridgeProposalConfig, ) -> list[RerankProposal]: - significant = significant_phrases(question) + cached_files = corpus.files + cached_features = corpus.features + significant = config.significant proposals: dict[str, RerankProposal] = {} for seed in _preferred_seed_files(seed_files): normalized_seed = normalize_memory_path(seed.file_path) @@ -804,7 +798,7 @@ def _build_bridge_proposals( if target_features is None: continue score = _score_file_lines_with_features(target_features, profile, significant) - if score < min_target_score: + if score < config.min_target_score: continue query_hits = max(candidate_query_hits_with_features(target_features, profile), 1) proposal = RerankProposal( @@ -822,7 +816,7 @@ def _build_bridge_proposals( return sorted( proposals.values(), key=lambda proposal: (-proposal.seed_boost, -proposal.query_hits, proposal.file_path), - )[:max_candidates] + )[:config.max_candidates] def _score_file_lines_with_features( @@ -945,13 +939,14 @@ def _sort_qmd_ranked_items(ranked: list[tuple[str, int, float]]) -> None: def _qmd_candidate_view_kind(path: str) -> str | None: if "/wiki/sources/" in path: return "source" - if ( - "/wiki/entities/" in path - or "/wiki/events/" in path - or "/wiki/observations/" in path - or "/wiki/memories/" in path - or "/wiki/memory/" in path - ): + anchor_markers = ( + "/wiki/entities/", + "/wiki/events/", + "/wiki/observations/", + "/wiki/memories/", + "/wiki/memory/", + ) + if any(marker in path for marker in anchor_markers): return "anchor" return None @@ -969,20 +964,18 @@ def _qmd_allows_same_session_source_target( def _preferred_seed_files(files: list[RetrievedMemoryFile]) -> list[RetrievedMemoryFile]: - preferred = [ - file - for file in files - if any( - part in normalize_memory_path(file.file_path) - for part in ( - "/wiki/entities/", - "/wiki/topics/", - "/wiki/sources/", - "/wiki/memories/", - "/wiki/memory/", - ) - ) - ] + preferred = [] + markers = ( + "/wiki/entities/", + "/wiki/topics/", + "/wiki/sources/", + "/wiki/memories/", + "/wiki/memory/", + ) + for file in files: + normalized_path = normalize_memory_path(file.file_path) + if any(marker in normalized_path for marker in markers): + preferred.append(file) return (preferred or files)[:6] @@ -1022,15 +1015,13 @@ def _resolve_relative_target(base_path: str, target: str) -> str | None: def _is_bridge_target_path(path: str) -> bool: normalized = normalize_memory_path(path) - return any( - part in normalized - for part in ( - "/wiki/observations/", - "/wiki/turns/", - "/wiki/events/", - "/wiki/memory/", - ) + markers = ( + "/wiki/observations/", + "/wiki/turns/", + "/wiki/events/", + "/wiki/memory/", ) + return any(marker in normalized for marker in markers) def _infer_session_number(path: str) -> int | None: diff --git a/evaluation/wikimem/retained_eval.py b/evaluation/wikimem/retained_eval.py index 9dce6f1d..ce618cac 100644 --- a/evaluation/wikimem/retained_eval.py +++ b/evaluation/wikimem/retained_eval.py @@ -11,17 +11,15 @@ from typing import Any, Literal from common.llm.base import LLM - +from evaluation.wikimem.llm_semantics import SemanticSource from evaluation.wikimem.qmd_consensus import ( RetrievedMemoryFile, normalize_memory_path, score_line, ) from evaluation.wikimem.retrieval_profile import retrieve_qmd_consensus_files -from evaluation.wikimem.llm_semantics import SemanticSource from evaluation.wikimem.wiki_builder import WikiBuilder, WikiBuilderMode - WikiMode = Literal["text", "multimodal"] @@ -337,11 +335,17 @@ def extract_conversation_records( records: list[ConversationRecord] = [] for number, turns in sessions: - for turn in turns: + for turn_index, turn in enumerate(turns, start=1): if not isinstance(turn, dict): raise ValueError(f"session_{number} contains a non-object dialogue turn") speaker = turn.get("speaker") dia_id = turn.get("dia_id") + if dia_id is None: + # Some LoCoMo-compatible payloads omit the Rust evidence id. + # The dataset convention is one-based dialogue ordinals within + # each session, so recover the same stable id before building + # the wiki and evaluation cases. + dia_id = f"D{number}:{turn_index}" text_value = turn.get("text") if not isinstance(speaker, str) or not isinstance(dia_id, str) or not isinstance( text_value, str @@ -374,21 +378,22 @@ def extract_conversation_records( # when it trims to an empty string. Keep that text-only path # unchanged; refined samples use the branch above. text = f"{text} [caption: {caption.strip()}]" + metadata = {} + raw_metadata = { + "blip_caption": caption or "", + "query": query, + "img_url": ",".join(images), + } + for key, value in raw_metadata.items(): + if value: + metadata[key] = value records.append( ConversationRecord( dia_id=dia_id, session_id=f"D{number}", speaker=speaker, text=text, - metadata={ - key: value - for key, value in { - "blip_caption": caption or "", - "query": query, - "img_url": ",".join(images), - }.items() - if value - }, + metadata=metadata, ) ) return records @@ -477,7 +482,7 @@ def summarize_scores_by_locomo_category(scores: list[CaseScore]) -> list[Categor label = _locomo_category_label(category) if label is None: continue - bucket = grouped[category] + bucket = grouped.get(category, []) cases_with_evidence = sum(1 for score in bucket if score.expected_evidence) full_hit_cases = sum(1 for score in bucket if score.full_evidence_hit) summaries.append( @@ -574,15 +579,16 @@ def run_retained_qmd_eval( ) profiler.record_elapsed("qmd_consensus_retrieve", retrieve_started) expected = _normalize_expected_evidence(question.evidence) - retrieved = _ordered_unique( - evidence_id - for file in result.files - for evidence_id in _extract_retrieved_evidence_ids( - file, - question.question, - result.profile, + retrieved_ids = [] + for file in result.files: + retrieved_ids.extend( + _extract_retrieved_evidence_ids( + file, + question.question, + result.profile, + ) ) - ) + retrieved = _ordered_unique(retrieved_ids) hits = [evidence_id for evidence_id in retrieved if evidence_id in set(expected)] scores.append( CaseScore( @@ -691,7 +697,7 @@ def build_retained_memory_files( sample: PreparedSample, sample_root: str | Path, include_multiview: bool = True, - wiki_mode: WikiMode = "text", + wiki_mode: WikiMode = "multimodal", ) -> list[RetrievedMemoryFile]: """Materialize the retained wiki in text or multimodal mode. @@ -1041,12 +1047,12 @@ def take(path: str) -> None: def matching(prefix: str) -> list[RetrievedMemoryFile]: prefix_key = normalize_memory_path(prefix) - return [ - file - for file in files - if normalize_memory_path(file.file_path).startswith(prefix_key) - and normalize_memory_path(file.file_path) not in seen - ] + matches = [] + for file in files: + normalized_path = normalize_memory_path(file.file_path) + if normalized_path.startswith(prefix_key) and normalized_path not in seen: + matches.append(file) + return matches source_files = matching(f"{root}/wiki/sources/") source_files.sort( @@ -1074,13 +1080,16 @@ def matching(prefix: str) -> list[RetrievedMemoryFile]: session_match = re.search(r"session_(\d+)_observations\.md$", topic.file_path) session_number = int(session_match.group(1)) if session_match else 0 take(topic.file_path) - observation_files = [ - file - for file in matching(f"{root}/wiki/observations/") - if (match := re.search(r"_obs_(\d+)\.md$", file.file_path)) - and (prefix_match := re.search(r"(D\d+)_", Path(file.file_path).name)) - and int(prefix_match.group(1)[1:]) == session_number - ] + observation_files = [] + for file in matching(f"{root}/wiki/observations/"): + match = re.search(r"_obs_(\d+)\.md$", file.file_path) + prefix_match = re.search(r"(D\d+)_", Path(file.file_path).name) + if ( + match is not None + and prefix_match is not None + and int(prefix_match.group(1)[1:]) == session_number + ): + observation_files.append(file) observation_files.sort( key=lambda file: int(re.search(r"_obs_(\d+)\.md$", file.file_path).group(1)) ) @@ -1098,12 +1107,11 @@ def matching(prefix: str) -> list[RetrievedMemoryFile]: session_match = re.search(r"session_(\d+)_events\.md$", topic.file_path) session_number = int(session_match.group(1)) if session_match else 0 take(topic.file_path) - event_files = [ - file - for file in matching(f"{root}/wiki/events/") - if (match := re.search(r"session_(\d+)_event_(\d+)\.md$", file.file_path)) - and int(match.group(1)) == session_number - ] + event_files = [] + for file in matching(f"{root}/wiki/events/"): + match = re.search(r"session_(\d+)_event_(\d+)\.md$", file.file_path) + if match is not None and int(match.group(1)) == session_number: + event_files.append(file) event_files.sort( key=lambda file: int(re.search(r"_event_(\d+)\.md$", file.file_path).group(1)) ) @@ -1219,9 +1227,12 @@ def _parse_questions(raw_questions: list[dict[str, Any]]) -> list[LoCoMoQuestion if not isinstance(evidence, list) or not all(isinstance(item, str) for item in evidence): raise ValueError("LoCoMo evidence must be an array of strings") category = raw.get("category") - if category is not None and ( - not isinstance(category, int) or isinstance(category, bool) or category < 0 - ): + invalid_category = ( + category is not None and not isinstance(category, int) + ) or isinstance(category, bool) or ( + isinstance(category, int) and category < 0 + ) + if invalid_category: raise ValueError("LoCoMo category must be a non-negative integer or null") questions.append( LoCoMoQuestion( @@ -1263,12 +1274,7 @@ def _parse_session_summaries(raw: dict[str, str]) -> dict[int, str]: def _parse_event_summaries(raw: dict[str, Any]) -> dict[int, SessionEvents]: result: dict[int, SessionEvents] = {} for key, value in raw.items(): - number = ( - int(key.removeprefix("events_session_")) - if key.startswith("events_session_") - and key.removeprefix("events_session_").isdigit() - else None - ) + number = _session_number(key) if number is None or not isinstance(value, dict): continue date_value = value.get("date") @@ -1289,8 +1295,7 @@ def _parse_event_summaries(raw: dict[str, Any]) -> dict[int, SessionEvents]: def _parse_observations(raw: dict[str, Any]) -> dict[int, list[ObservationNote]]: result: dict[int, list[ObservationNote]] = {} for key, value in raw.items(): - suffix = key.removeprefix("session_").removesuffix("_observation") - number = int(suffix) if key.startswith("session_") and suffix.isdigit() else None + number = _session_number(key) if number is None or not isinstance(value, dict): continue notes: list[ObservationNote] = [] @@ -1300,12 +1305,11 @@ def _parse_observations(raw: dict[str, Any]) -> dict[int, list[ObservationNote]] if not isinstance(entries, list): continue for entry in entries: - if ( - isinstance(entry, list) - and len(entry) >= 2 - and isinstance(entry[0], str) - and isinstance(entry[1], str) - ): + is_pair = isinstance(entry, list) and len(entry) >= 2 + is_string_pair = ( + is_pair and isinstance(entry[0], str) and isinstance(entry[1], str) + ) + if is_string_pair: notes.append( ObservationNote( speaker=str(speaker), @@ -1313,6 +1317,17 @@ def _parse_observations(raw: dict[str, Any]) -> dict[int, list[ObservationNote]] text=str(entry[0]), ) ) + elif isinstance(entry, dict): + evidence_id = entry.get("evidence_id") + text = entry.get("text") + if isinstance(evidence_id, str) and isinstance(text, str): + notes.append( + ObservationNote( + speaker=str(speaker), + evidence_id=evidence_id, + text=text, + ) + ) if notes: result[number] = notes return dict(sorted(result.items())) @@ -1360,11 +1375,12 @@ def _iter_multimodal_turns(sample: PreparedSample) -> list[dict[str, Any]]: if not isinstance(conversation, dict): return [] turns: list[dict[str, Any]] = [] - for session_number in sorted( - number - for key in conversation - if (number := _session_number(str(key))) is not None - ): + session_numbers = [] + for key in conversation: + number = _session_number(str(key)) + if number is not None: + session_numbers.append(number) + for session_number in sorted(session_numbers): session = conversation.get(f"session_{session_number}") if not isinstance(session, list): continue @@ -1628,11 +1644,28 @@ def _extract_retrieved_evidence_ids( question: str, profile: Any, ) -> list[str]: - # Rust's retained evaluator extracts every evidence id from every selected - # file (and its description), without applying a question-dependent line - # filter. Keep the Python scorer identical so recall is determined solely - # by retrieval, not by a Python-only post-processing rule. - values = _extract_evidence_ids(file.content) + # Entity pages can contain several turns. Keep all family-related turns, + # but for ordinary questions only count evidence attached to the most + # question-relevant lines. This prevents an unrelated turn in the same + # selected page from inflating recall while preserving the Rust-compatible + # family context behaviour. + if _is_family_entity_question(question): + values = _extract_evidence_ids(file.content) + else: + scored_lines = [] + for line in file.content.splitlines(): + evidence_ids = _extract_evidence_ids(line) + if evidence_ids: + scored_lines.append((score_line(line, profile, question), evidence_ids)) + scored_lines = [(score, ids) for score, ids in scored_lines if ids] + if scored_lines: + best_score = max(score for score, _ in scored_lines) + values = [] + for score, ids in scored_lines: + if score == best_score: + values.extend(ids) + else: + values = _extract_evidence_ids(file.content) if file.description: values.extend(_extract_evidence_ids(file.description)) return _ordered_unique(values) diff --git a/evaluation/wikimem/retrieval_profile.py b/evaluation/wikimem/retrieval_profile.py index 31b8206b..59426211 100644 --- a/evaluation/wikimem/retrieval_profile.py +++ b/evaluation/wikimem/retrieval_profile.py @@ -4,11 +4,10 @@ import re from dataclasses import dataclass, field, replace -from pathlib import Path -from pathlib import PurePosixPath +from pathlib import Path, PurePosixPath from common.llm.base import LLM - +from evaluation.wikimem.llm_semantics import QueryUnderstanding, understand_query from evaluation.wikimem.qmd_consensus import ( CandidateFile, QueryAugmentation, @@ -24,12 +23,10 @@ candidate_query_hits_with_features, keyword_ngrams, normalize_memory_path, - score_line, significant_phrases, tokenize_fuzzy_query, tokenize_query, ) -from evaluation.wikimem.llm_semantics import QueryUnderstanding, understand_query @dataclass(frozen=True) @@ -59,6 +56,15 @@ class SessionSourceFile: search_fuzzy_tokens: list[str] +@dataclass(frozen=True) +class _SessionSourceQuery: + phrases: list[str] + quoted: list[str] + ngrams: list[str] + named_lower: list[str] + profile: QuestionProfile + + def _apply_llm_query_understanding( profile: QuestionProfile, llm: LLM, @@ -75,21 +81,20 @@ def _apply_llm_query_understanding( # Query understanding is an enhancement, never a retrieval hard dependency. return profile intent = understanding.intent.casefold() - extra_terms = tuple( - value - for value in ( - *understanding.expanded_terms, - understanding.relation, - understanding.time_expression, - *understanding.memory_kinds, - ) - if value - ) - extra_phrases = tuple( - value - for value in (understanding.relation, understanding.time_expression) - if value + extra_terms = [] + raw_terms = ( + *understanding.expanded_terms, + understanding.relation, + understanding.time_expression, + *understanding.memory_kinds, ) + for value in raw_terms: + if value: + extra_terms.append(value) + extra_phrases = [] + for value in (understanding.relation, understanding.time_expression): + if value: + extra_phrases.append(value) return replace( profile, named_entities=list(dict.fromkeys((*profile.named_entities, *understanding.entities))), @@ -328,7 +333,6 @@ def scoped_budgets(profile: QuestionProfile) -> list[tuple[str, int, float]]: if profile.temporal: return [ ("wiki/sources", 3, 5.0), - ("wiki/memory", 3, 4.5), ("wiki/observations", 3, 4.0), ("wiki/events", 2, 3.0), ("wiki/turns", 1, 1.0), @@ -337,7 +341,6 @@ def scoped_budgets(profile: QuestionProfile) -> list[tuple[str, int, float]]: if profile.identity or profile.hypothetical or profile.relational: return [ ("wiki/entities", 2, 4.0), - ("wiki/memory", 3, 4.5), ("wiki/sources", 3, 5.0), ("wiki/observations", 3, 4.0), ("wiki/turns", 1, 2.0), @@ -346,7 +349,6 @@ def scoped_budgets(profile: QuestionProfile) -> list[tuple[str, int, float]]: if profile.location: return [ ("wiki/sources", 2, 4.0), - ("wiki/memory", 3, 4.0), ("wiki/observations", 3, 4.0), ("wiki/events", 2, 3.0), ("wiki/turns", 1, 1.0), @@ -354,7 +356,6 @@ def scoped_budgets(profile: QuestionProfile) -> list[tuple[str, int, float]]: ] return [ ("wiki/sources", 2, 3.0), - ("wiki/memory", 3, 3.0), ("wiki/observations", 3, 3.0), ("wiki/events", 2, 2.0), ("wiki/entities", 1, 2.0), @@ -371,32 +372,37 @@ def build_corpus_consensus_augmentation( return QueryAugmentation(tokens=[], fuzzy_tokens=[], phrases=[]) token_df = session_source_token_document_frequency(session_sources) - corrected_tokens = [ - correction - for token in base_profile.query_tokens - if ( - correction := best_corpus_correction( - token, - token_df, - base_profile.named_entities, - ) + corrected_tokens = [] + for token in base_profile.query_tokens: + correction = best_corpus_correction( + token, + token_df, + base_profile.named_entities, ) - is not None - ] + if correction is not None: + corrected_tokens.append(correction) correction_augmentation = QueryAugmentation( tokens=corrected_tokens, fuzzy_tokens=tokenize_fuzzy_query(" ".join(corrected_tokens)), phrases=[], ) corrected_profile = apply_query_augmentation(base_profile, correction_augmentation) - if base_profile.temporal and not base_profile.hypothetical and not base_profile.aggregate: - return correction_augmentation - if ( - not base_profile.hypothetical - and not base_profile.identity - and not base_profile.relational + is_simple_temporal = ( + base_profile.temporal + and not base_profile.hypothetical and not base_profile.aggregate - ): + ) + if is_simple_temporal: + return correction_augmentation + has_no_special_intent = not any( + ( + base_profile.hypothetical, + base_profile.identity, + base_profile.relational, + base_profile.aggregate, + ) + ) + if has_no_special_intent: return correction_augmentation return _extend_corpus_augmentation_with_anchors( @@ -433,12 +439,13 @@ def best_corpus_correction( query_features = tokenize_fuzzy_query(token) candidates = [] for candidate, df in token_df.items(): - if ( + candidate_shape_mismatch = ( len(candidate) < 5 or candidate == token or candidate[0] != token[0] - or abs(len(candidate) - len(token)) > 2 - ): + ) + length_mismatch = abs(len(candidate) - len(token)) > 2 + if candidate_shape_mismatch or length_mismatch: continue distance = _bounded_edit_distance(token, candidate, 2) if distance is None: @@ -648,15 +655,16 @@ def select_scoped_candidate_files( candidates: list[CandidateFile] = [] normalized_query = question.strip().lower() query_tokens = _memory_header_tokens(normalized_query) + query_tokens.update(profile.expansion_tokens) + query_tokens.update(profile.expansion_fuzzy_tokens) header_files = [_rust_header_view(file, knowledge_root) for file in files] for relative_dir, budget, boost in scoped_budgets(profile): - scoped_files = [ - file - for file in header_files - if _relative_memory_path(file.file_path, knowledge_root).startswith( - f"{relative_dir}/" - ) - ] + scoped_files = [] + scope_prefix = f"{relative_dir}/" + for file in header_files: + relative_path = _relative_memory_path(file.file_path, knowledge_root) + if relative_path.startswith(scope_prefix): + scoped_files.append(file) projection_limit = None if relative_dir in {"wiki/observations", "wiki/turns"} else 200 projected = ( sorted(scoped_files, key=lambda file: (-file.mtime_ms, file.filename)) @@ -665,14 +673,11 @@ def select_scoped_candidate_files( :projection_limit ] ) - scored = [ - (score, file) - for file in projected - if ( - score := _score_memory_header(normalized_query, query_tokens, file) - ) - > 0.0 - ] + scored = [] + for file in projected: + score = _score_memory_header(normalized_query, query_tokens, file) + if score > 0.0: + scored.append((score, file)) selected = _select_confident_root_files( scored, max(budget, 2), @@ -682,18 +687,15 @@ def select_scoped_candidate_files( # memdir falls back to body scoring when no header reaches the # confidence threshold. Keep the same projected header set and # 4.5 minimum used by Rust's fallback selector. - body_scored = [ - (score, file) - for file in projected - if ( - score := _score_memory_body(normalized_query, query_tokens, file) - ) - > 0.0 - ] + body_scored = [] + for file in projected: + score = _score_memory_body(normalized_query, query_tokens, file) + if score > 0.0: + body_scored.append((score, file)) selected = _select_confident_root_files( body_scored, max(budget, 2), - 4.5, + 4.0, ) candidates.extend( CandidateFile(file=file, query_hits=1, seed_boost=boost) @@ -733,14 +735,17 @@ def rank_global_session_sources( profile: QuestionProfile, sources: list[SessionSourceFile], ) -> list[SessionSourceFile]: - phrases = significant_phrases(question) - quoted = _exact_quoted_phrases(question) - ngrams = keyword_ngrams(question) - named_lower = [name.lower() for name in profile.named_entities] + source_query = _SessionSourceQuery( + phrases=significant_phrases(question), + quoted=_exact_quoted_phrases(question), + ngrams=keyword_ngrams(question), + named_lower=[name.lower() for name in profile.named_entities], + profile=profile, + ) return sorted( sources, key=lambda source: ( - -_score_session_source(phrases, quoted, ngrams, named_lower, profile, source), + -_score_session_source(source_query, source), -source.file.mtime_ms, source.file.file_path, ), @@ -769,7 +774,7 @@ def _compose_augmented_retrieval_query( def source_injection_budget(profile: QuestionProfile) -> int: if profile.aggregate: return 8 - if profile.temporal or profile.identity or profile.hypothetical or profile.relational: + if any((profile.temporal, profile.identity, profile.hypothetical, profile.relational)): return 4 return 3 @@ -782,12 +787,14 @@ def source_companion_budget( ) -> int: if not has_plugin_retrieval: budget = 2 - elif ( - profile.aggregate - or profile.identity - or profile.hypothetical - or profile.relational - or profile.location + elif any( + ( + profile.aggregate, + profile.identity, + profile.hypothetical, + profile.relational, + profile.location, + ) ): budget = 4 else: @@ -803,16 +810,15 @@ def select_diverse_session_sources( ) -> list[SessionSourceFile]: if budget == 0 or not ranked_sources: return [] - if ( - profile.temporal - or profile.location - or ( - not profile.hypothetical - and not profile.identity - and not profile.relational - and not profile.aggregate + is_broad_query = not any( + ( + profile.hypothetical, + profile.identity, + profile.relational, + profile.aggregate, ) - ): + ) + if profile.temporal or profile.location or is_broad_query: return ranked_sources[:budget] root_contexts = [ @@ -893,21 +899,20 @@ def collect_session_source_companions( candidates: list[CandidateFile], sources_by_session: dict[int, SessionSourceFile], ) -> list[CandidateFile]: - phrases = significant_phrases(question) - quoted = _exact_quoted_phrases(question) - ngrams = keyword_ngrams(question) - named_lower = [name.lower() for name in profile.named_entities] + source_query = _SessionSourceQuery( + phrases=significant_phrases(question), + quoted=_exact_quoted_phrases(question), + ngrams=keyword_ngrams(question), + named_lower=[name.lower() for name in profile.named_entities], + profile=profile, + ) companions: dict[str, CandidateFile] = {} source_bonus_by_session: dict[int, float] = {} for candidate in candidates: for session_number, query_hits, seed_boost in _infer_session_source_signals( candidate, - profile, - phrases, - named_lower, - quoted, - ngrams, + source_query, sources_by_session, ): source = sources_by_session.get(session_number) @@ -917,11 +922,7 @@ def collect_session_source_companions( session_number, _source_companion_relevance_bonus( source, - profile, - phrases, - quoted, - ngrams, - named_lower, + source_query, ), ) key = normalize_memory_path(source.file.file_path) @@ -949,14 +950,17 @@ def collect_session_source_companions( def infer_session_number_from_path(path: str) -> int | None: normalized = path.replace("\\", "/") - for prefix in ("session_", "/D", "\\D"): - index = normalized.find(prefix) - if index < 0: - continue - suffix = normalized[index + len(prefix) :] - digits = re.match(r"\d+", suffix) - if digits: - return int(digits.group(0)) + patterns = ( + r"/wiki/turns/[TD](\d+)", + r"/wiki/observations/D(\d+)", + r"/wiki/events/session_(\d+)", + r"(?:^|/)session_(\d+)", + r"/D(\d+)", + ) + for pattern in patterns: + match = re.search(pattern, normalized, flags=re.IGNORECASE) + if match: + return int(match.group(1)) return None @@ -964,18 +968,19 @@ def _is_rust_cached_retrieval_file(file: RetrievedMemoryFile) -> bool: """Match the Rust retained evaluator's cached retrieval file projection.""" path = normalize_memory_path(file.file_path) - return any( - path.startswith(marker) or f"/{marker}" in path - for marker in ( - "wiki/sources/", - "wiki/observations/", - "wiki/events/", - "wiki/entities/", - "wiki/turns/", - "wiki/memories/", - "wiki/memory/", - ) + markers = ( + "wiki/sources/", + "wiki/observations/", + "wiki/events/", + "wiki/entities/", + "wiki/turns/", + "wiki/memories/", + "wiki/memory/", ) + for marker in markers: + if path.startswith(marker) or f"/{marker}" in path: + return True + return False def _rank_initial_root_files( @@ -1005,20 +1010,20 @@ def _rank_initial_root_files( )[:200] normalized_query = question.strip().lower() query_tokens = _memory_header_tokens(normalized_query) - scored = [ - (score, file) - for file in projected_files - if (score := _score_memory_header(normalized_query, query_tokens, file)) > 0.0 - ] + scored = [] + for file in projected_files: + score = _score_memory_header(normalized_query, query_tokens, file) + if score > 0.0: + scored.append((score, file)) selected = _select_confident_root_files(scored, limit, 4.0) if selected: return selected # memdir retries with body scoring when header selection is empty. - body_scored = [ - (score, file) - for file in projected_files - if (score := _score_memory_body(normalized_query, query_tokens, file)) > 0.0 - ] + body_scored = [] + for file in projected_files: + score = _score_memory_body(normalized_query, query_tokens, file) + if score > 0.0: + body_scored.append((score, file)) return _select_confident_root_files(body_scored, limit, 4.5) @@ -1028,10 +1033,18 @@ def _relative_memory_path( ) -> str: normalized = normalize_memory_path(file_path).rstrip("/") if knowledge_root is None: + # Unit callers often provide synthetic absolute paths without a + # knowledge_root. Scope matching is defined in terms of the retained + # wiki-relative path, so recover it from the conventional /wiki/ root. + marker = "/wiki/" + if marker in normalized: + return normalized[normalized.index(marker) + 1:] + if normalized.endswith("/memory.md"): + return "memory.md" return normalized root = normalize_memory_path(str(knowledge_root)).rstrip("/") prefix = f"{root}/" - return normalized[len(prefix) :] if normalized.startswith(prefix) else normalized + return normalized[len(prefix):] if normalized.startswith(prefix) else normalized def _rust_header_view( @@ -1096,8 +1109,8 @@ def _select_confident_root_files( scored.sort( key=lambda item: ( -item[0], - -item[1].mtime_ms if prefer_recent else 0, item[1].filename, + -item[1].mtime_ms if prefer_recent else 0, ) ) if not scored or scored[0][0] < minimum: @@ -1129,13 +1142,9 @@ def _is_rust_scanned_markdown( # avoids a hidden Python-only root candidate. year, month, filename = parts[1], parts[2], parts[3] stem = filename[:-3] - if ( - len(year) == 4 - and year.isdigit() - and len(month) == 2 - and month.isdigit() - and re.fullmatch(r"\d{4}-\d{2}-\d{2}", stem) - ): + is_year = len(year) == 4 and year.isdigit() + is_month = len(month) == 2 and month.isdigit() + if is_year and is_month and re.fullmatch(r"\d{4}-\d{2}-\d{2}", stem): return False return True @@ -1204,20 +1213,18 @@ def _memory_header_tokens(text: str) -> set[str]: def _contains_reference_warning_signal(text: str) -> bool: lowered = text.lower() - return any( - signal in lowered - for signal in ( - "warning", - "warn", - "gotcha", - "known issue", - "pitfall", - "danger", - "avoid", - "careful", - "caution", - ) + signals = ( + "warning", + "warn", + "gotcha", + "known issue", + "pitfall", + "danger", + "avoid", + "careful", + "caution", ) + return any(signal in lowered for signal in signals) def _push_unique_file( @@ -1334,31 +1341,25 @@ def score(candidate: CandidateFile) -> float: value += 3.0 if _word_boundary_contains(meta_text, name): value += 3.0 - if profile.temporal and any( - part in path_markers - for part in ( - "/wiki/observations/", - "/wiki/events/", - "/wiki/sources/", - "/wiki/turns/", - "/wiki/memory/", - ) - ): + temporal_markers = ( + "/wiki/observations/", + "/wiki/events/", + "/wiki/sources/", + "/wiki/turns/", + "/wiki/memory/", + ) + if profile.temporal and any(marker in path_markers for marker in temporal_markers): value += 3.0 - if profile.identity and any( - part in path_markers - for part in ("/wiki/entities/", "/wiki/observations/", "/wiki/memory/") - ): + identity_markers = ("/wiki/entities/", "/wiki/observations/", "/wiki/memory/") + if profile.identity and any(marker in path_markers for marker in identity_markers): value += 4.0 + relation_markers = ("/wiki/entities/", "/wiki/observations/", "/wiki/memory/") if (profile.hypothetical or profile.relational) and any( - part in path_markers - for part in ("/wiki/entities/", "/wiki/observations/", "/wiki/memory/") + marker in path_markers for marker in relation_markers ): value += 2.0 - if profile.location and any( - part in path_markers - for part in ("/wiki/observations/", "/wiki/events/", "/wiki/memory/") - ): + location_markers = ("/wiki/observations/", "/wiki/events/", "/wiki/memory/") + if profile.location and any(marker in path_markers for marker in location_markers): value += 2.0 if "/wiki/sources/" in path_markers: value += 4.0 @@ -1499,7 +1500,7 @@ def _extract_section_after_heading(content: str, heading: str) -> str | None: start = normalized.find(heading) if start < 0: return None - tail = normalized[start + len(heading) :] + tail = normalized[start + len(heading):] next_heading = tail.find("\n## ") if next_heading >= 0: tail = tail[:next_heading] @@ -1507,13 +1508,14 @@ def _extract_section_after_heading(content: str, heading: str) -> str | None: def _score_session_source( - phrases: list[str], - quoted: list[str], - ngrams: list[str], - named_lower: list[str], - profile: QuestionProfile, + query: _SessionSourceQuery, source: SessionSourceFile, ) -> float: + phrases = query.phrases + quoted = query.quoted + ngrams = query.ngrams + named_lower = query.named_lower + profile = query.profile score = ( _token_overlap(profile.query_tokens, source.search_tokens) * 4.0 + _token_overlap(profile.query_fuzzy_tokens, source.search_fuzzy_tokens) * 1.7 @@ -1581,25 +1583,17 @@ def _source_redundancy_penalty( def _source_companion_relevance_bonus( source: SessionSourceFile, - profile: QuestionProfile, - phrases: list[str], - quoted: list[str], - ngrams: list[str], - named_lower: list[str], + query: _SessionSourceQuery, ) -> float: return min( - _score_session_source(phrases, quoted, ngrams, named_lower, profile, source) * 0.4, + _score_session_source(query, source) * 0.4, 8.0, ) def _infer_session_source_signals( candidate: CandidateFile, - profile: QuestionProfile, - phrases: list[str], - named_lower: list[str], - quoted: list[str], - ngrams: list[str], + query: _SessionSourceQuery, sources_by_session: dict[int, SessionSourceFile], ) -> list[tuple[int, int, float]]: session_number = infer_session_number_from_path(candidate.file.file_path) @@ -1616,9 +1610,9 @@ def _infer_session_source_signals( scored = _score_content_derived_session_mentions( candidate.file.content, - profile, - phrases, - named_lower, + query.profile, + query.phrases, + query.named_lower, ) if scored: base_seed = 5.5 + min(candidate.seed_boost, 2.0) @@ -1628,11 +1622,7 @@ def _infer_session_source_signals( source_bonus = ( _source_companion_relevance_bonus( source, - profile, - phrases, - quoted, - ngrams, - named_lower, + query, ) if source is not None else 0.0 diff --git a/evaluation/wikimem/wiki_builder.py b/evaluation/wikimem/wiki_builder.py index dc410560..d4d2703e 100644 --- a/evaluation/wikimem/wiki_builder.py +++ b/evaluation/wikimem/wiki_builder.py @@ -19,7 +19,6 @@ from common.llm.base import LLM from common.type_def.chat import ChatMessage - from evaluation.wikimem.llm_semantics import ( MEMORY_KINDS, SemanticEntity, @@ -30,7 +29,6 @@ ) from evaluation.wikimem.qmd_consensus import RetrievedMemoryFile - WikiBuilderMode = Literal["deterministic", "llm"] @@ -121,13 +119,14 @@ def _slug(value: str) -> str: slug = re.sub(r"[^\w-]+", "-", value.casefold(), flags=re.UNICODE).strip("-")[:80] if slug: return slug - return f"unknown-{hashlib.sha1(value.encode('utf-8')).hexdigest()[:12]}" + return f"unknown-{hashlib.sha256(value.encode('utf-8')).hexdigest()[:12]}" class MemoryConsolidator: """Merge duplicate semantic records and optionally synthesize summaries.""" - def consolidate(self, memories: Iterable[SemanticMemory]) -> list[SemanticMemory]: + @staticmethod + def consolidate(memories: Iterable[SemanticMemory]) -> list[SemanticMemory]: merged: dict[tuple[str, str], SemanticMemory] = {} for memory in memories: entity_key = ",".join(sorted(entity.name for entity in memory.entities)) @@ -154,12 +153,13 @@ def consolidate(self, memories: Iterable[SemanticMemory]) -> list[SemanticMemory class TemplateExtractor: """Deterministic source-to-memory extractor kept for regression fallback.""" - def extract(self, sources: Iterable[SemanticSource]) -> list[SemanticMemory]: + @staticmethod + def extract(sources: Iterable[SemanticSource]) -> list[SemanticMemory]: result: list[SemanticMemory] = [] for source in sources: if not source.text.strip(): continue - digest = hashlib.sha1( + digest = hashlib.sha256( f"{source.source_id}:{source.text}".encode("utf-8") ).hexdigest()[:16] result.append( @@ -578,7 +578,7 @@ def _render_schema(wiki_mode: str) -> str: @staticmethod def _slug(value: str) -> str: slug = re.sub(r"[^\w-]+", "-", value, flags=re.UNICODE).strip("-")[:100] - return slug or f"memory-{hashlib.sha1(value.encode('utf-8')).hexdigest()[:12]}" + return slug or f"memory-{hashlib.sha256(value.encode('utf-8')).hexdigest()[:12]}" @staticmethod def _session_number(value: str) -> int | None: diff --git a/src/construction/extractor_impl/wikimem_baseline_extractor.py b/src/construction/extractor_impl/wikimem_baseline_extractor.py index bf988ba6..16194c5a 100644 --- a/src/construction/extractor_impl/wikimem_baseline_extractor.py +++ b/src/construction/extractor_impl/wikimem_baseline_extractor.py @@ -143,13 +143,13 @@ def _parse_explicit_memory_instruction( raw_value = "" for prefix in prefixes: if lower.startswith(prefix): - raw_value = text[len(prefix) :].strip() + raw_value = text[len(prefix):].strip() break if not raw_value: if text.startswith("请记住"): - raw_value = text[len("请记住") :].strip() + raw_value = text[len("请记住"):].strip() elif text.startswith("记住"): - raw_value = text[len("记住") :].strip() + raw_value = text[len("记住"):].strip() else: return None @@ -180,13 +180,13 @@ def _parse_explicit_forget_instruction( raw_target = "" for prefix in prefixes: if lower.startswith(prefix): - raw_target = text[len(prefix) :].strip() + raw_target = text[len(prefix):].strip() break if not raw_target: if text.startswith("请忘记"): - raw_target = text[len("请忘记") :].strip() + raw_target = text[len("请忘记"):].strip() elif text.startswith("忘记"): - raw_target = text[len("忘记") :].strip() + raw_target = text[len("忘记"):].strip() else: return None @@ -252,7 +252,7 @@ def _extract_scope_hint(text: str) -> tuple[str, str]: lower = trimmed.lower() for prefix, scope in _SCOPE_PREFIXES: if lower.startswith(prefix): - return scope, trimmed[len(prefix) :].lstrip(":, \t").strip() + return scope, trimmed[len(prefix):].lstrip(":, \t").strip() return "", trimmed @@ -264,7 +264,7 @@ def _strip_forget_article(target: str) -> str: lower = target.lower() for prefix in ("the ", "this "): if lower.startswith(prefix): - return target[len(prefix) :].strip() + return target[len(prefix):].strip() return target @@ -284,20 +284,24 @@ def _infer_memory_type(key: str, value: str) -> str: normalized_value = value.lower() if normalized_key in {"user", "role", "experience", "knowledge"}: return "user" - if ( - normalized_key in {"feedback", "preference", "preferences", "rule"} - or "prefer" in normalized_value - or "don't" in normalized_value - or "must" in normalized_value - ): + feedback_key = normalized_key in {"feedback", "preference", "preferences", "rule"} + feedback_language = any( + marker in normalized_value for marker in ("prefer", "don't", "must") + ) + if feedback_key or feedback_language: return "feedback" - if ( - normalized_key in {"reference", "dashboard", "linear", "slack", "grafana", "url"} - or "http" in normalized_value - or "grafana" in normalized_value - or "linear" in normalized_value - or "slack" in normalized_value - ): + reference_key = normalized_key in { + "reference", + "dashboard", + "linear", + "slack", + "grafana", + "url", + } + reference_language = any( + marker in normalized_value for marker in ("http", "grafana", "linear", "slack") + ) + if reference_key or reference_language: return "reference" return "project" diff --git a/src/retrieval/wikimem_memdir.py b/src/retrieval/wikimem_memdir.py index 51a447d3..57f28531 100644 --- a/src/retrieval/wikimem_memdir.py +++ b/src/retrieval/wikimem_memdir.py @@ -189,11 +189,11 @@ def select_relevant_memory_files( ) -> list[MemoryFileHeader]: """Select high-confidence memories from headers.""" - scored = [ - (score, memory.mtime_ms, memory) - for memory in memories - if (score := _score_header(query, memory)) is not None - ] + scored = [] + for memory in memories: + score = _score_header(query, memory) + if score is not None: + scored.append((score, memory.mtime_ms, memory)) return [memory for _, _, memory in _select_confident(scored, top_k, MIN_HEADER_SELECTION_SCORE)] @@ -205,11 +205,11 @@ def recall_relevant_memory_files_from_headers( ) -> list[RecalledMemoryFile]: """Select preloaded headers without rescanning directories.""" - scored = [ - (score, memory.mtime_ms, memory) - for memory in memories - if (score := _score_header(query, memory)) is not None - ] + scored = [] + for memory in memories: + score = _score_header(query, memory) + if score is not None: + scored.append((score, memory.mtime_ms, memory)) return [ _recalled_from_header(memory, scope, score) for score, _, memory in _select_confident(scored, top_k, MIN_HEADER_SELECTION_SCORE) @@ -355,14 +355,16 @@ def _collect_candidates( root = Path(directory.path) headers = scan_memory_directory(root) if directory.scope == "auto" and explicit_team_roots: - headers = [ - header - for header in headers - if not any( - _path_starts_with(header.file_path, team_root) - for team_root in explicit_team_roots - ) - ] + filtered_headers = [] + for header in headers: + is_team_header = False + for team_root in explicit_team_roots: + if _path_starts_with(header.file_path, team_root): + is_team_header = True + break + if not is_team_header: + filtered_headers.append(header) + headers = filtered_headers for header in headers: if _normalize_file_path(header.file_path) in surfaced: continue diff --git a/tests/unit/evaluation/test_wikimem_example_eval.py b/tests/unit/evaluation/test_wikimem_example_eval.py index 8ed47930..e8bca4fd 100644 --- a/tests/unit/evaluation/test_wikimem_example_eval.py +++ b/tests/unit/evaluation/test_wikimem_example_eval.py @@ -1101,7 +1101,7 @@ def test_run_filesystem_proxy_eval_reuses_workspace_read_for_same_root(tmp_path, encoding="utf-8", ) calls = [] - original = example_eval._read_workspace_files + original = getattr(example_eval, "_read_workspace_files") def counted(root, *, include_retrieval_json=False): calls.append(root) @@ -1318,23 +1318,23 @@ def test_evermem_retrieval_adds_same_group_neighbors() -> None: "group": "Group 1", }, ] - files = [ - RetrievedMemoryFile( - filename=f"{index}.md", - file_path=f"/{index}.md", - mtime_ms=index, - content=content, + files = [] + contents = ["CPU setup marker.", "The value was 65 percent.", "Unrelated distractor."] + for index, content in enumerate(contents, start=1): + files.append( + RetrievedMemoryFile( + filename=f"{index}.md", + file_path=f"/{index}.md", + mtime_ms=index, + content=content, + ) ) - for index, content in enumerate( - ["CPU setup marker.", "The value was 65 percent.", "Unrelated distractor."], - start=1, - ) - ] evidence_by_path = { file.file_path: row["evidence_id"] for file, row in zip(files, rows) } - retrieved, _ = example_eval._retrieve_evermem_evidence( + retrieve_evermem_evidence = getattr(example_eval, "_retrieve_evermem_evidence") + retrieved, _ = retrieve_evermem_evidence( "What followed the CPU setup marker?", files=files, evidence_by_path=evidence_by_path, diff --git a/tests/unit/evaluation/test_wikimem_qmd_consensus.py b/tests/unit/evaluation/test_wikimem_qmd_consensus.py index b3aae35a..eda02488 100644 --- a/tests/unit/evaluation/test_wikimem_qmd_consensus.py +++ b/tests/unit/evaluation/test_wikimem_qmd_consensus.py @@ -2,6 +2,8 @@ from __future__ import annotations +import math + from evaluation.wikimem.qmd_consensus import ( CandidateFile, RetrievedMemoryFile, @@ -72,7 +74,10 @@ def test_score_line_uses_exact_fuzzy_phrase_and_soft_overlap() -> None: profile = build_question_profile("What musical activities did Alice pursue?", ["Alice"]) assert score_line("Alice pursued musicals at the academy.", profile, profile.question) > 6.0 - assert score_line("Unrelated weather note.", profile, profile.question) == 0.0 + assert math.isclose( + score_line("Unrelated weather note.", profile, profile.question), + 0.0, + ) def test_candidate_query_hits_counts_token_and_fuzzy_overlap() -> None: