mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
Patrick Buckley
215f7506ba
Reuse the shipped Cohere/Jina rerank client as an optional post-process on the BM25 surfaces (tool search, skill search, memory composition) via one seam: BM25Index gains an injected reranker + a two-stage search (BM25 recall top-50 -> rerank -> top-k). No new storage. Gated on a configured endpoint plus tools.rerank_bm25 (default on, matching rerank_web_search). tools.rerank_bm25_threshold (default 0.0 = off) is a relevance FLOOR for proactive memory surfacing: BM25 always returns something, so without a floor every-turn memory injection spends tokens on the top-k of whatever lexically matched; the reranker score is what makes a meaningful "inject nothing" gate possible. Two reranker modes (BM25Index rerank_filters): - REORDER (reactive tool/skill search): the reranker must never drop results -> fall back to BM25 order on empty, backfill omitted pool items, so a misbehaving endpoint can't silently lose tools. - FILTER (memory, rerank_filters = threshold > 0): a clean empty/short result is honoured (inject nothing) -- a deliberate divergence from web_search._rerank_results. Parse/endpoint failure is a discrete branch from the floor: an empty result for non-empty input means an unparseable response (a conforming reranker scores every doc), so the closure raises RerankError and BM25Index falls back to BM25 order in BOTH modes -- the floor only acts on valid scores. Also: cap the rerank client timeout at 15s (the per-turn memory path can't afford tools.timeout's 120s default); move the Reranker alias to rerank.py (shared, no import cycle); document the endpoint egress in the rerank_bm25 help, the admin Reranker-role description, and docs/tools.md; add scripts/bench_bm25_rerank.py (manual, needs a live endpoint) to measure precision@k/MRR lift and recommend a threshold default. Negative-tested: reorder fallback-on-empty and omitted-item backfill, filter-mode honor-empty, singleton-still-floored, the parse-fail RerankError raise, the >= floor boundary, and pool-position-to-doc-index mapping -- each guard reverted to confirm its test fails, then restored.
103 lines
3.3 KiB
Python
103 lines
3.3 KiB
Python
"""BM25-based memory relevance scoring and system message formatting."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from html import escape as _html_escape
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from turnstone.core.bm25 import BM25Index
|
|
|
|
if TYPE_CHECKING:
|
|
from turnstone.core.rerank import Reranker
|
|
|
|
|
|
@dataclass
|
|
class MemoryConfig:
|
|
"""Configuration for the structured memory system."""
|
|
|
|
relevance_k: int = 5
|
|
fetch_limit: int = 50
|
|
max_content: int = 32768
|
|
nudge_cooldown: int = 300
|
|
nudges: bool = True
|
|
|
|
|
|
def score_memories(
|
|
memories: list[dict[str, str]],
|
|
query: str,
|
|
k: int = 5,
|
|
reranker: Reranker | None = None,
|
|
rerank_filters: bool = False,
|
|
) -> list[dict[str, str]]:
|
|
"""Return the top-k memories most relevant to *query*.
|
|
|
|
Builds a BM25 index over ``name + description + content prefix``
|
|
for each memory and returns matches sorted by relevance. If *query*
|
|
is empty, returns the most recent *k* memories (they are already
|
|
ordered by ``updated DESC`` from storage).
|
|
"""
|
|
if not memories:
|
|
return []
|
|
if not query or not query.strip():
|
|
return memories[:k]
|
|
|
|
documents = [
|
|
f"{m.get('name', '')} {m.get('description', '')} {m.get('content', '')[:200]}"
|
|
for m in memories
|
|
]
|
|
index = BM25Index(documents, reranker=reranker, rerank_filters=rerank_filters)
|
|
top_indices = index.search(query, k)
|
|
return [memories[i] for i in top_indices]
|
|
|
|
|
|
def build_memory_context(memories: list[dict[str, str]]) -> str:
|
|
"""Format selected memories as an XML block for system message injection.
|
|
|
|
Produces a compact ``<memories>`` section matching the style used
|
|
for MCP resources (``<mcp-resources>``).
|
|
"""
|
|
if not memories:
|
|
return ""
|
|
lines = ["<memories>"]
|
|
for m in memories:
|
|
name = _html_escape(m.get("name", ""))
|
|
mem_type = _html_escape(m.get("type", "project"))
|
|
scope = _html_escape(m.get("scope", "global"))
|
|
desc = m.get("description", "")
|
|
content = m.get("content", "")
|
|
# Truncate content to avoid bloating system message
|
|
if len(content) > 500:
|
|
content = content[:500] + "..."
|
|
desc_attr = f' description="{_html_escape(desc)}"' if desc else ""
|
|
lines.append(
|
|
f' <memory name="{name}" type="{mem_type}" scope="{scope}"{desc_attr}>'
|
|
f"{_html_escape(content)}</memory>"
|
|
)
|
|
lines.append("</memories>")
|
|
return "\n".join(lines)
|
|
|
|
|
|
def extract_recent_context(messages: list[dict[str, Any]], max_messages: int = 3) -> str:
|
|
"""Extract text from the last N user messages for relevance scoring.
|
|
|
|
Handles both string and list content formats.
|
|
"""
|
|
user_texts: list[str] = []
|
|
for msg in reversed(messages):
|
|
if msg.get("role") != "user":
|
|
continue
|
|
content = msg.get("content", "")
|
|
if isinstance(content, str):
|
|
user_texts.append(content)
|
|
elif isinstance(content, list):
|
|
# Multi-part content (text + images)
|
|
for part in content:
|
|
if isinstance(part, dict) and part.get("type") == "text":
|
|
user_texts.append(part.get("text", ""))
|
|
elif isinstance(part, str):
|
|
user_texts.append(part)
|
|
if len(user_texts) >= max_messages:
|
|
break
|
|
return " ".join(user_texts)
|