thread-archive 0.0.2__py3-none-any.whl

thread_archive/_retrieval/_extract.py

@@ -0,0 +1,184 @@
+"""Searchable-content extraction: an event's payload → FTS shadow rows.
+
+Reads the payload dict keys directly — the importer's ``DefaultEventBuilder``
+produces exactly these payload shapes, so there's no dependency on a dataclass
+event-type hierarchy.
+
+Each extractor returns ``(content, content_type, tool_name)`` tuples; an event
+with no searchable content returns ``[]``.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Optional
+
+# Event types that carry searchable content (must match the dispatch below).
+#
+# NB: ``api_request_completed`` is deliberately *not* indexed. thread-archive never
+# live-streams — every assistant turn is *imported* via DefaultEventBuilder, which
+# emits the granular text_complete / thinking_complete events alongside an
+# api_request_completed summary that duplicates them. Indexing both double-counts
+# the same text, so we index the granular events and drop the summary. (A
+# live-streaming path would need the summary, since its stream emits text_delta
+# rather than text_complete — but that path doesn't apply here.)
+INDEXABLE_EVENT_TYPES = [
+    "user_message_sent",
+    "text_complete",
+    "thinking_complete",
+    "tool_use_started",
+    "tool_use_complete",
+    "tool_execution_completed",
+    "context_summary",
+    "tool_execution_error",
+    "thread_message_sent",
+    "ide_context",
+    "content_block",
+    "message",
+]
+
+
+def _to_str(value) -> str:
+    return value if isinstance(value, str) else json.dumps(value, default=str)
+
+
+def _block_search_text(block) -> str:
+    """Best-effort human-readable text from an arbitrary content block, skipping
+    binary/base64 payloads (the ``source`` blob on image/document blocks). Used to
+    make preserved-but-unmodeled blocks (``content_block``) searchable without
+    indexing megabytes of base64."""
+    if isinstance(block, str):
+        return block
+    if not isinstance(block, dict):
+        return _to_str(block)
+    parts: list[str] = []
+    for key, value in block.items():
+        if key in ("type", "source"):  # ``source`` carries base64 image/doc data
+            continue
+        if isinstance(value, str):
+            if len(value) > 1000 and " " not in value[:100]:
+                continue  # looks like an opaque/base64 blob
+            parts.append(value)
+        elif isinstance(value, (dict, list)):
+            parts.append(_to_str(value))
+    return " ".join(parts).strip()
+
+
+def _strip_frontmatter(text: str) -> str:
+    """Strip YAML frontmatter (---\\n...\\n---) from markdown content."""
+    if text.startswith("---"):
+        end = text.find("---", 3)
+        if end != -1:
+            return text[end + 3:].lstrip()
+    return text
+
+
+def _fts_tool_use(payload: dict) -> list[tuple[str, str, Optional[str]]]:
+    tool_name = payload.get("tool_name") or ""
+    tool_input = payload.get("input")
+    if tool_name.startswith("mcp__") and isinstance(tool_input, dict) and "command" in tool_input:
+        # Heredoc bodies span many lines; index the first line, truncated.
+        tool_name = tool_input["command"].split("\n", 1)[0][:200]
+
+    # Prioritize the real content keys over stringified metadata.
+    content_parts = [payload.get("tool_name") or ""]
+    if isinstance(tool_input, dict):
+        for key in ("content", "text", "body", "message", "query", "prompt", "description", "entry"):
+            val = tool_input.get(key)
+            if val and isinstance(val, str):
+                content_parts.append(_strip_frontmatter(val))
+                break
+        else:
+            content_parts.append(" ".join(f"{k}={v}" for k, v in tool_input.items()))
+    elif tool_input is not None:
+        content_parts.append(str(tool_input))
+
+    content = " ".join(content_parts).strip()
+    return [(content[:2000], "tool", tool_name or None)] if content else []
+
+
+def _fts_tool_completed(payload: dict) -> list[tuple[str, str, Optional[str]]]:
+    output = payload.get("output")
+    tool_name = payload.get("tool_name")
+    if payload.get("is_error"):
+        error_text = output or payload.get("error", "")
+        return [(_to_str(error_text), "tool_error", tool_name)] if error_text else []
+    if output:
+        return [(_to_str(output), "tool_result", tool_name)]
+    return []
+
+
+def extract_fts_content(event_type: str, payload: dict) -> list[tuple[str, str, Optional[str]]]:
+    """Extract searchable ``(content, content_type, tool_name)`` tuples from an event."""
+    if not payload:
+        return []
+
+    if event_type == "user_message_sent":
+        content = payload.get("content", "")
+        if not content:
+            return []
+        if content.lstrip().startswith("This session is being continued from a previous conversation"):
+            return [(content, "continuation_summary", None)]
+        return [(content, "user", None)]
+
+    if event_type == "api_request_completed":
+        results: list[tuple[str, str, Optional[str]]] = []
+        for block in payload.get("content_blocks", []) or []:
+            bt = block.get("type")
+            if bt == "thinking" and block.get("thinking"):
+                results.append((block["thinking"], "thinking", None))
+            elif bt == "text" and block.get("text"):
+                results.append((block["text"], "text", None))
+        return results
+
+    if event_type == "text_complete":
+        text = payload.get("text", "")
+        return [(text, "text", None)] if text else []
+
+    if event_type == "thinking_complete":
+        text = payload.get("text", "")
+        return [(text, "thinking", None)] if text else []
+
+    if event_type in ("tool_use_started", "tool_use_complete"):
+        return _fts_tool_use(payload)
+
+    if event_type == "tool_execution_completed":
+        return _fts_tool_completed(payload)
+
+    if event_type == "tool_execution_error":
+        error = payload.get("error", "")
+        return [(_to_str(error), "tool_error", payload.get("tool_name"))] if error else []
+
+    if event_type == "context_summary":
+        content = payload.get("content", "")
+        return [(content[:2000], "context_summary", None)] if content else []
+
+    if event_type == "thread_message_sent":
+        content = payload.get("content", "")
+        return [(content, "user", None)] if content else []
+
+    if event_type == "ide_context":
+        # Index the opened-file path / selection body so "what was I looking at"
+        # is searchable. file_path (when present) leads so a path query matches.
+        content = payload.get("content", "")
+        file_path = payload.get("file_path")
+        text = f"{file_path}\n{content}" if file_path else content
+        return [(text[:2000], "ide_context", None)] if text.strip() else []
+
+    if event_type == "content_block":
+        # An unmodeled block preserved verbatim; index its human-readable text.
+        text = _block_search_text(payload.get("data"))
+        content_type = payload.get("block_type") or "content_block"
+        return [(text[:2000], content_type, None)] if text.strip() else []
+
+    if event_type == "message":
+        # A preserved non-standard-role turn. Prefer its text; fall back to blocks.
+        content = payload.get("content", "")
+        content_type = payload.get("role") or "message"
+        if content.strip():
+            return [(content, content_type, None)]
+        blocks = payload.get("content_blocks") or []
+        text = " ".join(t for t in (_block_search_text(b) for b in blocks) if t)
+        return [(text[:2000], content_type, None)] if text.strip() else []
+
+    return []

thread_archive/_retrieval/embed.py

@@ -0,0 +1,186 @@
+"""In-process nomic embeddings (sentence-transformers) — the `[embeddings]` extra.
+
+The in-process embedding server: nomic model, ``search_query:`` /
+``search_document:`` prefixes, 768-dim ``list[float]`` out, ``None`` on failure.
+Heavy (torch); gated to the extra. The base install is lexical-only:
+``is_available()`` is False and the vector arm degrades out cleanly.
+
+Self-consistency is the contract that matters — vectors indexed with this provider
+must be *queried* with it (llama.cpp-nomic ≠ torch-nomic are different spaces). The
+``space_key`` tags a vector set by model so a cached set is only reused in-space.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import logging
+import os
+import sys
+import threading
+
+logger = logging.getLogger(__name__)
+
+# Cap input length before embedding. On the in-process torch path a batch of many
+# 6000-char docs (≈1500 tokens each) spikes MPS memory — it hangs and crawls
+# (~9 docs/s). 2048 chars keeps the gist of all but the longest code/tool dumps
+# (mean doc is ~300 chars) and runs ~7× faster with no hang.
+EMBEDDING_CHAR_CAP = 2048
+_MODEL_NAME = os.environ.get("THREAD_ARCHIVE_EMBED_MODEL", "nomic-ai/nomic-embed-text-v1.5")
+
+_model = None
+_load_failed = False
+# Serializes _load() so a background warm (mcp.server) and a concurrent first query can't
+# both construct the heavy model at once (a transient double-load / memory spike).
+_load_lock = threading.Lock()
+
+
+def model_name() -> str:
+    return _MODEL_NAME
+
+
+def space_key() -> str:
+    """Identity of the embedding space — vectors are only comparable within one."""
+    return "local:" + _MODEL_NAME
+
+
+def _device() -> str:
+    """``$THREAD_ARCHIVE_EMBED_DEVICE`` if set, else the best accelerator (Apple
+    ``mps`` / CUDA), falling back to ``cpu``. On Apple silicon the GPU is ~20× the
+    CPU — the difference between an hour and a day for the corpus embed."""
+    dev = os.environ.get("THREAD_ARCHIVE_EMBED_DEVICE")
+    if dev:
+        return dev
+    try:
+        import torch
+
+        if torch.backends.mps.is_available():
+            return "mps"
+        if torch.cuda.is_available():
+            return "cuda"
+    except Exception:
+        pass
+    return "cpu"
+
+
+def is_available() -> bool:
+    """True when sentence-transformers is importable (the extra is installed). Cheap —
+    does not load the model; a load failure degrades at call time."""
+    try:
+        if importlib.util.find_spec("sentence_transformers") is None:
+            return False
+    except ImportError:
+        return False
+    return not _load_failed
+
+
+def _hub_cache_dir() -> str:
+    """The HF hub cache directory, resolved the way huggingface_hub does — but as a pure
+    filesystem path, with no `huggingface_hub` import. Called *before* we set the offline
+    env, so it must not pull in the hub (which freezes ``HF_HUB_OFFLINE`` at import)."""
+    if os.environ.get("HUGGINGFACE_HUB_CACHE"):
+        return os.environ["HUGGINGFACE_HUB_CACHE"]
+    if os.environ.get("HF_HOME"):
+        return os.path.join(os.environ["HF_HOME"], "hub")
+    return os.path.expanduser("~/.cache/huggingface/hub")
+
+
+def _model_cached() -> bool:
+    """True when ``_MODEL_NAME`` is already in the HF hub cache (a non-empty ``snapshots/``).
+    Pure filesystem probe — no hub import — so it's safe to consult before going offline."""
+    folder = "models--" + _MODEL_NAME.replace("/", "--")
+    snaps = os.path.join(_hub_cache_dir(), folder, "snapshots")
+    try:
+        return os.path.isdir(snaps) and any(os.scandir(snaps))
+    except OSError:
+        return False
+
+
+def _pin_offline_if_cached() -> None:
+    """Pin the model load to the local cache once it's downloaded, so the hot path makes
+    NO Hub request: no revision re-resolve, no ``trust_remote_code`` re-fetch, no
+    unauthenticated-HF-Hub warning, and it still works with the network down. The model is
+    pinned (one fixed ``_MODEL_NAME``) — a Hub round-trip on every search buys nothing.
+
+    Skipped when the model isn't cached yet (so a first install still downloads it) or when
+    ``THREAD_ARCHIVE_EMBED_ONLINE=1`` forces an online load (first download / deliberate
+    refresh). ``huggingface_hub`` freezes ``HF_HUB_OFFLINE`` into a module constant at
+    import, so set the env *before* it's imported and, if it already is, flip the live
+    constant too."""
+    if os.environ.get("THREAD_ARCHIVE_EMBED_ONLINE") == "1" or not _model_cached():
+        return
+    os.environ.setdefault("HF_HUB_OFFLINE", "1")
+    os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
+    const = sys.modules.get("huggingface_hub.constants")
+    if const is not None:
+        const.HF_HUB_OFFLINE = True
+
+
+def _load():
+    """Lazily construct the cached SentenceTransformer (nomic needs trust_remote_code).
+    Double-checked under ``_load_lock`` so a background warm and a first query race to a
+    single construction, not two concurrent heavy loads."""
+    global _model, _load_failed
+    if _model is not None:
+        return _model
+    if _load_failed:
+        return None
+    with _load_lock:
+        if _model is not None:
+            return _model
+        if _load_failed:
+            return None
+        try:
+            _pin_offline_if_cached()
+            from sentence_transformers import SentenceTransformer
+
+            device = _device()
+            _model = SentenceTransformer(_MODEL_NAME, trust_remote_code=True, device=device)
+            dim = getattr(_model, "get_embedding_dimension", _model.get_sentence_embedding_dimension)()
+            logger.info("embed: loaded %s (dim=%s, device=%s)", _MODEL_NAME, dim, device)
+            return _model
+        except Exception as e:  # noqa: BLE001
+            _load_failed = True
+            logger.warning("embed: model load failed (%s) — vector arm degrades", e)
+            return None
+
+
+def warm() -> bool:
+    """Eagerly load the embedding model so it isn't cold-loaded inside the first query.
+    Fail-soft and idempotent: returns False when the ``[embeddings]`` extra is absent or the
+    load fails (search then stays lexical, exactly as it does without warming)."""
+    if not is_available():
+        return False
+    return _load() is not None
+
+
+def _cap(text: str) -> str:
+    return text[:EMBEDDING_CHAR_CAP]
+
+
+def _encode(prefixed: list[str]):
+    model = _load()
+    if model is None:
+        return None
+    try:
+        # Un-normalized to match the contract — the vector store normalizes on write.
+        vecs = model.encode(prefixed, normalize_embeddings=False, convert_to_numpy=True)
+        return [v.astype("float32").tolist() for v in vecs]
+    except Exception as e:  # noqa: BLE001
+        logger.warning("embed: encode failed (%s)", e)
+        return None
+
+
+def embed_query(text: str):
+    """Embed a search query (nomic ``search_query:`` prefix). None on any failure."""
+    if not text or not text.strip():
+        return None
+    vecs = _encode([f"search_query: {_cap(text)}"])
+    return vecs[0] if vecs else None
+
+
+def embed_documents(texts: list[str]):
+    """Embed indexed content (nomic ``search_document:`` prefix), in input order. None
+    on any failure (the caller writes nothing — never a partial/wrong-space batch)."""
+    if not texts:
+        return None
+    return _encode([f"search_document: {_cap(t)}" for t in texts])

thread_archive/_retrieval/format.py

@@ -0,0 +1,149 @@
+"""Render search hits for the CLI / MCP.
+
+Default render is a compact text list, but it carries the **match-quality signal**:
+a top-line ``quality=`` verdict (strong / partial / weak / semantic) plus a ``K/N``
+per hit — how many of the N query terms actually landed — so a reader can tell a
+solid keyword hit from a nearest-neighbour guess before trusting it. ``output``
+switches to ``count`` (a per-thread tally over the whole match pool) or ``linkable``
+(JSON of event/thread ids for batch linking)."""
+
+from __future__ import annotations
+
+import json
+import re
+from collections import Counter
+
+from sqlalchemy import text as sa_text
+
+from .._store import use_session
+from . import rank as _rank
+
+# output='count' wants a true tally, so the pipeline over-fetches to this cap; a
+# pool that reaches it was truncated and the tally renders as a floor ("N+").
+COUNT_FETCH_CAP = 1000
+
+
+def _hit_text(h: dict) -> str:
+    return h.get("full_content") or h.get("snippet") or ""
+
+
+def _term_hit_count(content: str, terms: list[str]) -> int:
+    """How many of ``terms`` literally appear in ``content``. Terms ≥4 chars match
+    by substring; shorter terms must hit a word boundary (so 'go' doesn't match
+    'good'). Each term counts at most once."""
+    if not terms or not content:
+        return 0
+    c = content.lower()
+    n = 0
+    for t in terms:
+        if len(t) >= 4:
+            if t in c:
+                n += 1
+        elif re.search(r"\b" + re.escape(t) + r"\b", c):
+            n += 1
+    return n
+
+
+def _search_quality(top_hit_count: int, n_terms: int, did_rerank: bool):
+    """Verdict for the top hit → ``(quality, note)`` or None. Rerank wins (the order
+    is by-meaning, not keyword overlap); else zero overlap is ``weak``, ≥⌈2/3·N⌉
+    terms is ``strong``, in-between is ``partial``."""
+    if n_terms <= 0:
+        return None
+    if did_rerank:
+        return ("semantic", "ranked by meaning, not keyword overlap — confirm the top hit "
+                            "actually answers the query before trusting it")
+    if top_hit_count == 0:
+        return ("weak", "no query term appears in the top hit — these are nearest-neighbour "
+                        "guesses and the log may simply not contain this. Rephrase the concept "
+                        "or switch data store; piling on more synonyms won't help")
+    strong_at = max(1, -(-2 * n_terms // 3))  # ceil(2/3 · n_terms)
+    if top_hit_count >= strong_at:
+        return ("strong", None)
+    return ("partial", "only some query terms matched the top hit — scan before trusting")
+
+
+def _query_terms(query: str) -> list[str]:
+    """Ranking terms minus the pipe-OR token (which isn't a content term)."""
+    return [t for t in _rank.search_terms(query) if t and t != "|"]
+
+
+def _format_count(hits: list[dict], query: str) -> str:
+    thread_counts = Counter(h["thread_id"] for h in hits)
+    with use_session() as s:
+        corpus = s.execute(
+            sa_text("SELECT count(*), count(DISTINCT thread_id) FROM event_search")
+        ).one()
+    capped = len(hits) >= COUNT_FETCH_CAP
+    total = f"{len(hits)}+ (tally capped)" if capped else str(len(hits))
+    lines = [
+        f"Total: {total} results across {len(thread_counts)} threads",
+        f"  (corpus: {corpus[0]:,} indexed events, {corpus[1]:,} threads)",
+    ]
+    for tid, count in thread_counts.most_common():
+        title = next((h.get("thread_title") for h in hits if h["thread_id"] == tid), None)
+        lines.append(f"  [{tid}] {title or '(untitled)'}: {count}")
+    return "\n".join(lines)
+
+
+def _format_linkable(hits: list[dict]) -> str:
+    out = []
+    for h in hits:
+        entry = {
+            "event_id": h["event_id"],
+            "thread_id": h["thread_id"],
+            "preview": (h.get("snippet") or "")[:80],
+        }
+        if "context_events" in h:
+            entry["context_events"] = h["context_events"]
+        out.append(entry)
+    return json.dumps(out, indent=2)
+
+
+def format_results(hits: list[dict], query: str, *, output: str | None = None) -> str:
+    if output == "count":
+        return _format_count(hits, query)
+    if output == "linkable":
+        return _format_linkable(hits)
+    if not hits:
+        return f'No results for "{query}".'
+
+    terms = _query_terms(query)
+    n_terms = len(terms)
+    did_rerank = bool(hits[0].get("_did_rerank"))
+    verdict = _search_quality(_term_hit_count(_hit_text(hits[0]), terms), n_terms, did_rerank) if n_terms else None
+
+    header = f'{len(hits)} result(s) for "{query}"'
+    if verdict:
+        header += f" · quality={verdict[0]}"
+    lines = [header]
+    if verdict and verdict[1]:
+        lines.append(f"  note: {verdict[1]}")
+    lines.append("")
+
+    for h in hits:
+        title = h.get("thread_title") or f"thread {h['thread_id']}"
+        ct = h.get("content_type") or h["event_type"]
+        head = f"[{h['thread_id']}/{h['event_id']}] {title} · {ct}"
+        if n_terms:
+            k = _term_hit_count(_hit_text(h), terms)
+            head += f" · {k}/{n_terms}"
+            if k == 0:
+                head += " (semantic)"
+        lines.append(head)
+
+        context = h.get("context")
+        if context:  # context_lines: a numbered multi-line block replaces the snippet
+            lines.extend(f"    {ln}" for ln in context.split("\n"))
+        else:
+            snippet = " ".join((h.get("snippet") or "").split())
+            if snippet:
+                lines.append(f"    {snippet}")
+
+        ctx_events = h.get("context_events") or {}
+        for direction in ("before", "after"):
+            for ev in ctx_events.get(direction, []):
+                body = " ".join((ev.get("content") or "")[:200].split())
+                lines.append(f"      ({direction} · {ev.get('content_type')}) {body}")
+
+    return "\n".join(lines)