search.py
python
sha256:cf6265cea8c21d9228d90dec13ef6ec2dab5103d466db9cc4590681832de4bf8
docs(KD-STAGING): sync governance after KD-6b DONE
Human
13 days ago
| 1 | """Lexical + vector search helpers for ``knowtation/search`` (§KD-5a.A(a)).""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import json |
| 6 | import logging |
| 7 | import os |
| 8 | import pathlib |
| 9 | import re |
| 10 | import urllib.error |
| 11 | import urllib.parse |
| 12 | import urllib.request |
| 13 | from typing import Any, TypedDict |
| 14 | |
| 15 | from muse.plugins.knowtation.hooks import DEFAULT_PORT, ENV_PORT, _resolve_port |
| 16 | from muse.plugins.knowtation.link_index import _MAX_NOTE_BYTES, build_link_index |
| 17 | from muse.plugins.knowtation.mcp_helpers import _note_body_and_frontmatter |
| 18 | from muse.plugins.knowtation.parser import parse_frontmatter |
| 19 | from muse.plugins.knowtation.plugin import _ALWAYS_IGNORE_DIRS |
| 20 | from muse.plugins.knowtation.stats import collect_vault_stats |
| 21 | |
| 22 | logger = logging.getLogger(__name__) |
| 23 | |
| 24 | W_LEX: float = 0.4 |
| 25 | W_VEC: float = 0.6 |
| 26 | VECTOR_TIMEOUT_SECONDS: float = 2.0 |
| 27 | _MAX_HUB_RESPONSE_BYTES: int = 256 * 1024 |
| 28 | |
| 29 | _MARKDOWN_EXTS: frozenset[str] = frozenset({".md", ".markdown", ".mdx"}) |
| 30 | _IGNORE_DIRS: frozenset[str] = _ALWAYS_IGNORE_DIRS | frozenset({"templates", "meta"}) |
| 31 | |
| 32 | |
| 33 | class LexicalHit(TypedDict, total=False): |
| 34 | path: str |
| 35 | score: float |
| 36 | section_address: str |
| 37 | snippet: str |
| 38 | |
| 39 | |
| 40 | def _note_ext(path: str) -> bool: |
| 41 | lower = path.lower() |
| 42 | return any(lower.endswith(ext) for ext in _MARKDOWN_EXTS) |
| 43 | |
| 44 | |
| 45 | def _query_tokens(query: str) -> list[str]: |
| 46 | return [t for t in re.split(r"\s+", query.strip().lower()) if t] |
| 47 | |
| 48 | |
| 49 | def _score_text(text: str, tokens: list[str]) -> float: |
| 50 | if not tokens: |
| 51 | return 0.0 |
| 52 | hay = text.lower() |
| 53 | hits = sum(1 for t in tokens if t in hay) |
| 54 | return hits / len(tokens) |
| 55 | |
| 56 | |
| 57 | def lexical_search( |
| 58 | root: pathlib.Path, |
| 59 | query: str, |
| 60 | *, |
| 61 | top_k: int = 10, |
| 62 | ) -> list[LexicalHit]: |
| 63 | """Grep vault notes for *query* — title, headings, frontmatter, body.""" |
| 64 | tokens = _query_tokens(query) |
| 65 | if not tokens: |
| 66 | return [] |
| 67 | |
| 68 | results: list[LexicalHit] = [] |
| 69 | for dirpath, dirnames, filenames in os.walk(root): |
| 70 | dirnames[:] = [d for d in dirnames if d not in _IGNORE_DIRS and not d.startswith(".")] |
| 71 | for name in filenames: |
| 72 | if not _note_ext(name): |
| 73 | continue |
| 74 | full = pathlib.Path(dirpath) / name |
| 75 | try: |
| 76 | rel = full.relative_to(root).as_posix() |
| 77 | except ValueError: |
| 78 | continue |
| 79 | try: |
| 80 | size = full.stat().st_size |
| 81 | except OSError: |
| 82 | continue |
| 83 | if size > _MAX_NOTE_BYTES: |
| 84 | continue |
| 85 | try: |
| 86 | raw = full.read_bytes() |
| 87 | except OSError: |
| 88 | continue |
| 89 | try: |
| 90 | text = raw.decode("utf-8") |
| 91 | except UnicodeDecodeError: |
| 92 | continue |
| 93 | |
| 94 | fm_parsed = parse_frontmatter(raw) |
| 95 | title = pathlib.Path(rel).stem |
| 96 | project = None |
| 97 | source_type = None |
| 98 | tags = None |
| 99 | if fm_parsed is not None: |
| 100 | title = str(fm_parsed.title or title) |
| 101 | project = fm_parsed.project |
| 102 | tags = fm_parsed.tags |
| 103 | source_type = fm_parsed.effective_source |
| 104 | if text.startswith("---"): |
| 105 | parts = text.split("---", 2) |
| 106 | body = parts[2] if len(parts) >= 3 else text |
| 107 | else: |
| 108 | body = text |
| 109 | fm_text = text.split("---", 2)[1] if text.startswith("---") and len(text.split("---", 2)) >= 2 else "" |
| 110 | heading_hits = re.findall(r"^#{1,6}\s+(.+)$", body, re.MULTILINE) |
| 111 | combined = " ".join([title, fm_text, " ".join(heading_hits), body[:8000]]) |
| 112 | score = _score_text(combined, tokens) |
| 113 | if score <= 0: |
| 114 | continue |
| 115 | snippet_src = body[:300] if body else title |
| 116 | results.append( |
| 117 | LexicalHit( |
| 118 | path=rel, |
| 119 | score=min(1.0, score), |
| 120 | snippet=snippet_src.replace("\n", " ")[:240], |
| 121 | ) |
| 122 | ) |
| 123 | |
| 124 | results.sort(key=lambda h: (-float(h.get("score", 0)), h.get("path", ""))) |
| 125 | return results[:top_k] |
| 126 | |
| 127 | |
| 128 | def vector_search( |
| 129 | query: str, |
| 130 | *, |
| 131 | top_k: int = 10, |
| 132 | port: int | None = None, |
| 133 | ) -> tuple[list[dict[str, Any]], str]: |
| 134 | """Query Knowtation hub vector index via GET (§KD-5a.A(a)). |
| 135 | |
| 136 | Returns: |
| 137 | ``(hits, status)`` where *status* is ``ok`` or ``unavailable``. |
| 138 | """ |
| 139 | resolved_port = port |
| 140 | if resolved_port is None: |
| 141 | p, err = _resolve_port() |
| 142 | if err or p is None: |
| 143 | return [], "unavailable" |
| 144 | resolved_port = p |
| 145 | |
| 146 | params = urllib.parse.urlencode({"q": query, "k": str(top_k)}) |
| 147 | url = f"http://localhost:{resolved_port}/api/v1/search?{params}" |
| 148 | request = urllib.request.Request( |
| 149 | url=url, |
| 150 | method="GET", |
| 151 | headers={"User-Agent": "muse-mcp-search/1.0", "Accept": "application/json"}, |
| 152 | ) |
| 153 | try: |
| 154 | with urllib.request.urlopen(request, timeout=VECTOR_TIMEOUT_SECONDS) as resp: # noqa: S310 |
| 155 | raw = resp.read(_MAX_HUB_RESPONSE_BYTES + 1) |
| 156 | if len(raw) > _MAX_HUB_RESPONSE_BYTES: |
| 157 | return [], "unavailable" |
| 158 | data = json.loads(raw.decode("utf-8")) |
| 159 | except (urllib.error.URLError, TimeoutError, json.JSONDecodeError, OSError) as exc: |
| 160 | logger.debug("vector search unavailable: %s", type(exc).__name__) |
| 161 | return [], "unavailable" |
| 162 | |
| 163 | hits: list[dict[str, Any]] = [] |
| 164 | items = data.get("results") or data.get("hits") or data |
| 165 | if isinstance(items, list): |
| 166 | for item in items: |
| 167 | if not isinstance(item, dict): |
| 168 | continue |
| 169 | path = item.get("path") or item.get("note_path") or item.get("id") |
| 170 | if not isinstance(path, str) or not path: |
| 171 | continue |
| 172 | score_raw = item.get("score") or item.get("similarity") or 0.0 |
| 173 | try: |
| 174 | score = float(score_raw) |
| 175 | except (TypeError, ValueError): |
| 176 | score = 0.0 |
| 177 | snippet = str(item.get("snippet") or item.get("excerpt") or "")[:240] |
| 178 | hits.append({"path": path.replace("\\", "/"), "score": score, "snippet": snippet}) |
| 179 | return hits, "ok" |
| 180 | |
| 181 | |
| 182 | def _head_mtime(root: pathlib.Path, path: str) -> float: |
| 183 | try: |
| 184 | return (root / path).stat().st_mtime |
| 185 | except OSError: |
| 186 | return 0.0 |
| 187 | |
| 188 | |
| 189 | def hybrid_search( |
| 190 | root: pathlib.Path, |
| 191 | query: str, |
| 192 | *, |
| 193 | top_k: int = 10, |
| 194 | mode: str = "hybrid", |
| 195 | project: str | None = None, |
| 196 | ) -> dict[str, Any]: |
| 197 | """Fan-out lexical + vector backends and merge (§KD-5a.A(a)).""" |
| 198 | stats = collect_vault_stats(root) |
| 199 | index_freshness = stats.get("index_freshness") |
| 200 | |
| 201 | lexical_status = "ok" |
| 202 | vector_status = "ok" |
| 203 | degraded = False |
| 204 | deduped = 0 |
| 205 | |
| 206 | lex_hits: list[LexicalHit] = [] |
| 207 | vec_hits: list[dict[str, Any]] = [] |
| 208 | |
| 209 | if mode in ("hybrid", "lexical"): |
| 210 | try: |
| 211 | lex_hits = lexical_search(root, query, top_k=top_k * 2) |
| 212 | except Exception: |
| 213 | lexical_status = "error" |
| 214 | lex_hits = [] |
| 215 | |
| 216 | if mode in ("hybrid", "semantic"): |
| 217 | vec_hits, vs = vector_search(query, top_k=top_k * 2) |
| 218 | vector_status = vs |
| 219 | if vs != "ok": |
| 220 | degraded = True |
| 221 | |
| 222 | if mode == "semantic" and vector_status != "ok": |
| 223 | degraded = True |
| 224 | if mode == "hybrid" and vector_status != "ok": |
| 225 | degraded = True |
| 226 | |
| 227 | lex_max = max((float(h.get("score", 0)) for h in lex_hits), default=0.0) or 1.0 |
| 228 | vec_max = max((float(h.get("score", 0)) for h in vec_hits), default=0.0) or 1.0 |
| 229 | |
| 230 | merged: dict[str, dict[str, Any]] = {} |
| 231 | |
| 232 | for hit in lex_hits: |
| 233 | path = hit["path"] |
| 234 | if project and not path.startswith(project): |
| 235 | continue |
| 236 | norm = float(hit["score"]) / lex_max if lex_max else 0.0 |
| 237 | merged[path.lower()] = { |
| 238 | "path": path, |
| 239 | "lex": norm, |
| 240 | "vec": 0.0, |
| 241 | "snippet": hit.get("snippet", ""), |
| 242 | "section_address": hit.get("section_address"), |
| 243 | "source": "muse-lexical", |
| 244 | } |
| 245 | |
| 246 | for hit in vec_hits: |
| 247 | path = hit["path"] |
| 248 | if project and not path.startswith(project): |
| 249 | continue |
| 250 | key = path.lower() |
| 251 | norm = float(hit["score"]) / vec_max if vec_max else 0.0 |
| 252 | if key in merged: |
| 253 | merged[key]["vec"] = norm |
| 254 | merged[key]["source"] = "both" |
| 255 | if hit.get("snippet") and not merged[key].get("snippet"): |
| 256 | merged[key]["snippet"] = hit["snippet"] |
| 257 | else: |
| 258 | merged[key] = { |
| 259 | "path": path, |
| 260 | "lex": 0.0, |
| 261 | "vec": norm, |
| 262 | "snippet": hit.get("snippet", ""), |
| 263 | "section_address": None, |
| 264 | "source": "knowtation-vector", |
| 265 | } |
| 266 | |
| 267 | # Case-insensitive basename dedupe — first walk-order path wins (lex order first). |
| 268 | seen_basenames: dict[str, str] = {} |
| 269 | dedupe_keys: list[str] = [] |
| 270 | for key in sorted(merged.keys(), key=lambda k: merged[k]["path"]): |
| 271 | entry = merged[key] |
| 272 | base = pathlib.PurePosixPath(entry["path"]).name.lower() |
| 273 | if base in seen_basenames and seen_basenames[base] != entry["path"]: |
| 274 | deduped += 1 |
| 275 | continue |
| 276 | seen_basenames.setdefault(base, entry["path"]) |
| 277 | dedupe_keys.append(key) |
| 278 | |
| 279 | ranked: list[dict[str, Any]] = [] |
| 280 | for key in dedupe_keys: |
| 281 | entry = merged[key] |
| 282 | combined = W_LEX * entry["lex"] + W_VEC * entry["vec"] |
| 283 | ranked.append({**entry, "combined": combined, "key": key}) |
| 284 | |
| 285 | ranked.sort( |
| 286 | key=lambda e: ( |
| 287 | -e["combined"], |
| 288 | -_head_mtime(root, e["path"]), |
| 289 | e["path"], |
| 290 | ) |
| 291 | ) |
| 292 | |
| 293 | results: list[dict[str, Any]] = [] |
| 294 | for entry in ranked[:top_k]: |
| 295 | path = entry["path"] |
| 296 | try: |
| 297 | raw = (root / path).read_bytes() |
| 298 | fm_dict, _ = _note_body_and_frontmatter(raw) |
| 299 | frontmatter = { |
| 300 | "project": fm_dict.get("project"), |
| 301 | "tags": fm_dict.get("tags"), |
| 302 | "source_type": fm_dict.get("source_type") or fm_dict.get("source"), |
| 303 | } |
| 304 | except OSError: |
| 305 | frontmatter = {"project": None, "tags": None, "source_type": None} |
| 306 | |
| 307 | results.append( |
| 308 | { |
| 309 | "path": path, |
| 310 | "score": round(entry["combined"], 6), |
| 311 | "source": entry["source"], |
| 312 | "snippet": entry.get("snippet", ""), |
| 313 | "section_address": entry.get("section_address"), |
| 314 | "frontmatter": frontmatter, |
| 315 | } |
| 316 | ) |
| 317 | |
| 318 | if not results and lexical_status == "error" and vector_status != "ok": |
| 319 | degraded = True |
| 320 | |
| 321 | return { |
| 322 | "query": query, |
| 323 | "mode": mode, |
| 324 | "results": results, |
| 325 | "index_freshness": index_freshness, |
| 326 | "degraded": degraded, |
| 327 | "deduped": deduped, |
| 328 | "backends": {"lexical": lexical_status, "vector": vector_status}, |
| 329 | } |
| 330 | |
| 331 | |
| 332 | __all__ = [ |
| 333 | "W_LEX", |
| 334 | "W_VEC", |
| 335 | "VECTOR_TIMEOUT_SECONDS", |
| 336 | "hybrid_search", |
| 337 | "lexical_search", |
| 338 | "vector_search", |
| 339 | ] |
File History
1 commit
sha256:cf6265cea8c21d9228d90dec13ef6ec2dab5103d466db9cc4590681832de4bf8
docs(KD-STAGING): sync governance after KD-6b DONE
Human
13 days ago