Update megumin_agent/retrieval.py

megumin_agent/retrieval.py

@@ -9,6 +9,7 @@ from dataclasses import dataclass
 from functools import lru_cache
 from pathlib import Path
 from typing import Any
+from typing import Iterable
 
 import faiss
 import numpy as np
@@ -42,6 +43,8 @@ FAISS_METADATA_FILENAME = os.getenv(
     "MEGUMIN_FAISS_METADATA_FILENAME",
     "megumin_questions_meta.json",
 )
+PERSONA_DATASET_PATTERNS = ("megumin_qa_dataset.json",)
+FACT_DATASET_PATTERNS = ("namuwiki*.json",)
 
 
 def _normalize_text(value: Any) -> str:
@@ -58,6 +61,11 @@ def _safe_excerpt(text: str, limit: int = 220) -> str:
     return compact[: limit - 3].rstrip() + "..."
 
 
+def _normalize_patterns(patterns: Iterable[str] | None) -> tuple[str, ...]:
+    normalized = tuple(pattern.strip() for pattern in (patterns or ()) if pattern.strip())
+    return normalized
+
+
 @dataclass(frozen=True)
 class QaRecord:
     question: str
@@ -165,14 +173,29 @@ def _load_metadata_records(path: Path) -> tuple[QaRecord, ...]:
     return tuple(records)
 
 
-@lru_cache(maxsize=8)
-def _load_records(dataset_dir: str) -> tuple[QaRecord, ...]:
+def _iter_matching_paths(root: Path, include_patterns: tuple[str, ...]) -> list[Path]:
+    if not include_patterns:
+        return sorted(root.glob("*.json"))
+
+    seen: set[Path] = set()
+    paths: list[Path] = []
+    for pattern in include_patterns:
+        for path in sorted(root.glob(pattern)):
+            if path in seen or path.suffix.lower() != ".json":
+                continue
+            seen.add(path)
+            paths.append(path)
+    return paths
+
+
+@lru_cache(maxsize=16)
+def _load_records(dataset_dir: str, include_patterns: tuple[str, ...] = ()) -> tuple[QaRecord, ...]:
     root = Path(dataset_dir)
     if not root.exists():
         return tuple()
 
     all_records: list[QaRecord] = []
-    for path in sorted(root.glob("*.json")):
+    for path in _iter_matching_paths(root, include_patterns):
         try:
             all_records.extend(_load_json_records(path))
         except OSError:
@@ -238,9 +261,10 @@ def build_and_save_faiss_index(
     output_dimensionality: int = EMBEDDING_DIMENSION,
     index_filename: str = FAISS_INDEX_FILENAME,
     metadata_filename: str = FAISS_METADATA_FILENAME,
+    include_patterns: Iterable[str] | None = None,
 ) -> tuple[Path, Path]:
     root = Path(dataset_dir)
-    records = _load_records(str(root.resolve()))
+    records = _load_records(str(root.resolve()), _normalize_patterns(include_patterns))
     if not records:
         raise FileNotFoundError(f"No JSON records found under {root}")
 
@@ -283,9 +307,17 @@ def _load_vector_store(
     dataset_dir: str,
     embedding_model: str,
     output_dimensionality: int,
+    include_patterns: tuple[str, ...] = (),
+    index_filename: str | None = FAISS_INDEX_FILENAME,
+    metadata_filename: str | None = FAISS_METADATA_FILENAME,
 ) -> VectorStore:
-    index_path, metadata_path = _index_artifact_paths(dataset_dir)
-    if index_path.exists() and metadata_path.exists():
+    if index_filename and metadata_filename:
+        index_path = Path(dataset_dir) / index_filename
+        metadata_path = Path(dataset_dir) / metadata_filename
+    else:
+        index_path = metadata_path = None
+
+    if index_path and metadata_path and index_path.exists() and metadata_path.exists():
         index = faiss.read_index(str(index_path))
         records = _load_metadata_records(metadata_path)
         if index.ntotal != len(records):
@@ -299,7 +331,7 @@ def _load_vector_store(
             dimension=index.d,
         )
 
-    records = _load_records(dataset_dir)
+    records = _load_records(dataset_dir, include_patterns)
     if not records:
         empty_index = faiss.IndexFlatIP(output_dimensionality)
         return VectorStore(
@@ -334,16 +366,25 @@ class JsonQaRetriever:
         *,
         embedding_model: str = EMBEDDING_MODEL_NAME,
         output_dimensionality: int = EMBEDDING_DIMENSION,
+        include_patterns: Iterable[str] | None = None,
+        index_filename: str | None = FAISS_INDEX_FILENAME,
+        metadata_filename: str | None = FAISS_METADATA_FILENAME,
     ):
         self.dataset_dir = Path(dataset_dir)
         self.embedding_model = embedding_model
         self.output_dimensionality = output_dimensionality
+        self.include_patterns = _normalize_patterns(include_patterns)
+        self.index_filename = index_filename
+        self.metadata_filename = metadata_filename
 
     def warmup(self) -> None:
         _load_vector_store(
             str(self.dataset_dir.resolve()),
             self.embedding_model,
             self.output_dimensionality,
+            self.include_patterns,
+            self.index_filename,
+            self.metadata_filename,
         )
 
     def _style_notes(self, matches: list[dict[str, Any]]) -> list[str]:
@@ -376,6 +417,9 @@ class JsonQaRetriever:
             str(self.dataset_dir.resolve()),
             self.embedding_model,
             self.output_dimensionality,
+            self.include_patterns,
+            self.index_filename,
+            self.metadata_filename,
         )
         if not store.records:
             return {