perf(session-search): adapt discovery result hydration

(cherry picked from commit 60a3530444f65c2cdcd4e5b983e4aa380ed651c7)
This commit is contained in:
blunkjamie-dev
2026-08-09 11:00:29 -05:00
committed by Teknium
parent ee9ec6164c
commit 6e1bdc0a18
2 changed files with 161 additions and 34 deletions
+84 -5
View File
@@ -1,12 +1,14 @@
"""Tests for the single-shape session_search tool.
Three calling shapes:
1. DISCOVERY — pass query → FTS5 + anchored window + bookends per hit
Four calling shapes:
1. DISCOVERY — pass query → FTS5 + adaptive/full hydration
2. SCROLL — pass session_id + around_message_id → just the window
3. BROWSE — no args → recent sessions chronologically
3. READ — pass session_id → whole or head/tail-truncated session
4. BROWSE — no args → recent sessions chronologically
All run zero LLM calls.
"""
import inspect
import json
import time
@@ -72,6 +74,8 @@ class TestSchema:
assert "query" in params
assert "limit" in params
assert params["sort"]["enum"] == ["newest", "oldest"]
assert params["detail"]["enum"] == ["adaptive", "full"]
assert params["detail"]["default"] == "adaptive"
# Scroll shape
assert "session_id" in params
assert "around_message_id" in params
@@ -81,6 +85,10 @@ class TestSchema:
# Mode is inferred from which args are set — no explicit mode param
assert "mode" not in params
def test_detail_parameter_is_appended_for_positional_compatibility(self):
parameters = list(inspect.signature(session_search).parameters)
assert parameters[-1] == "detail"
class TestFormatTimestamp:
def test_formats_unix_and_passes_through_the_rest(self):
@@ -176,17 +184,22 @@ class TestDiscoveryShape:
assert "context" not in requested_fields
assert len(result["results"]) == 1
hit = result["results"][0]
assert hit["detail"] == "full"
assert "bookend_start" in hit
assert hit["messages"]
assert "bookend_end" in hit
def test_discovery_result_has_bookends_and_window(self, db):
def test_full_detail_returns_bookends_and_window_for_every_hit(self, db):
_seed_modpack_sessions(db)
result = json.loads(session_search(query="modpack", limit=3, db=db))
result = json.loads(session_search(
query="modpack", limit=3, detail="full", db=db
))
assert result["success"] is True
assert result["mode"] == "discover"
assert result["detail"] == "full"
assert result["count"] >= 1
for hit in result["results"]:
assert hit["detail"] == "full"
assert "bookend_start" in hit
assert "messages" in hit
assert "bookend_end" in hit
@@ -195,6 +208,72 @@ class TestDiscoveryShape:
assert "messages_before" in hit
assert "messages_after" in hit
def test_default_discovery_keeps_top_full_and_compacts_lower_hits(self, db):
_seed_modpack_sessions(db)
result = json.loads(session_search(query="modpack", limit=3, db=db))
assert result["success"] is True
assert result["detail"] == "adaptive"
assert len(result["results"]) == 3
top, *lower = result["results"]
assert top["detail"] == "full"
assert "bookend_start" in top
assert len(top["messages"]) > 1
assert "bookend_end" in top
for hit in lower:
assert hit["detail"] == "compact"
assert hit["bookend_start"] == []
assert len(hit["messages"]) == 1
assert hit["messages"][0]["id"] == hit["match_message_id"]
assert hit["messages"][0]["anchor"] is True
assert hit["bookend_end"] == []
def test_adaptive_detail_preserves_ranking_and_reduces_payload(self, db):
now = int(time.time())
for session_index in range(3):
session_id = f"payload_{session_index}"
db.create_session(session_id, source="cli")
db._conn.execute(
"UPDATE sessions SET started_at = ? WHERE id = ?",
(now - session_index, session_id),
)
for message_index in range(8):
db.append_message(
session_id,
role="user" if message_index % 2 == 0 else "assistant",
content=f"opening {session_index}-{message_index} " + "o" * 2500,
)
db.append_message(
session_id,
role="user",
content=f"payloadneedle anchor {session_index} " + "a" * 3500,
)
for message_index in range(8):
db.append_message(
session_id,
role="assistant" if message_index % 2 == 0 else "user",
content=f"closing {session_index}-{message_index} " + "c" * 2500,
)
db._conn.commit()
adaptive_json = session_search(query="payloadneedle", limit=3, db=db)
full_json = session_search(
query="payloadneedle", limit=3, detail="full", db=db
)
adaptive = json.loads(adaptive_json)
full = json.loads(full_json)
assert [r["session_id"] for r in adaptive["results"]] == [
r["session_id"] for r in full["results"]
]
assert [r["match_message_id"] for r in adaptive["results"]] == [
r["match_message_id"] for r in full["results"]
]
assert len(adaptive_json.encode("utf-8")) < len(full_json.encode("utf-8")) * 0.6
def test_current_session_filtered_out(self, db):
_seed_modpack_sessions(db)
+77 -29
View File
@@ -2,23 +2,27 @@
"""
Session Search Tool - Long-Term Conversation Recall
Single-shape tool with three calling modes (inferred from args, no explicit
Single-shape tool with four calling modes (inferred from args, no explicit
mode parameter):
1. DISCOVERY — pass ``query``. Runs FTS5, dedupes hits by session lineage,
returns top N sessions each with: snippet, ±5 message window around the
match, plus bookend_start (first 3 user+assistant msgs of session) and
bookend_end (last 3). Zero LLM cost.
1. DISCOVERY — pass ``query``. Runs FTS5 and dedupes hits by session lineage.
Adaptive detail (the default) fully hydrates the top result with a ±5
message window and bookends, while lower-ranked results keep the exact
anchor message plus metadata. Pass ``detail="full"`` to fully hydrate
every result. Zero LLM cost.
2. SCROLL — pass ``session_id`` + ``around_message_id``. Returns a window
of ±window messages centered on the anchor, no FTS5, no bookends. To
scroll forward / backward, re-anchor on the last / first message id of
the returned window.
3. BROWSE — no args. Returns recent sessions chronologically (titles,
3. READ — pass ``session_id`` without an anchor. Returns the whole session,
or a bounded head/tail view for large sessions.
4. BROWSE — no args. Returns recent sessions chronologically (titles,
previews, timestamps).
All three modes operate on the SQLite session DB via the FTS5 index and
All four modes operate on the SQLite session DB via the FTS5 index and
the get_anchored_view / get_messages_around primitives in hermes_state.
No LLM calls anywhere — every shape returns actual messages from the DB.
@@ -740,6 +744,7 @@ def _title_match_result(
"bookend_end": [_shape_message(m) for m in (view.get("bookend_end") or messages[-3:])],
"messages_before": view.get("messages_before", 0),
"messages_after": view.get("messages_after", max(len(messages) - 5, 0)),
"detail": "full",
"_lineage_root": lineage_root,
}
if lineage_root and lineage_root != session_id:
@@ -753,10 +758,11 @@ def _discover(
role_filter: Optional[List[str]],
limit: int,
sort: Optional[str],
detail: str,
current_session_id: str = None,
link_profile: str = None,
) -> str:
"""Discovery shape: FTS5 + anchored window + bookends per hit. Single call."""
"""Discovery shape: FTS5 plus adaptive or full result hydration."""
role_list = role_filter if role_filter else ["user", "assistant"]
current_lineage_root = _resolve_lineage(db, current_session_id) if current_session_id else None
title_result = _title_match_result(db, query, current_lineage_root)
@@ -788,6 +794,7 @@ def _discover(
"success": True,
"mode": "discover",
"query": query,
"detail": detail,
"results": [],
"count": 0,
"message": "No matching sessions found.",
@@ -864,6 +871,11 @@ def _discover(
except Exception:
session_meta = {}
result_detail = "full" if detail == "full" or not results else "compact"
window_messages = view.get("window") or []
if result_detail == "compact":
window_messages = [m for m in window_messages if m.get("id") == msg_id]
entry = {
"session_id": hit_sid,
"when": _format_timestamp(
@@ -875,19 +887,31 @@ def _discover(
"matched_role": match_info.get("role"),
"match_message_id": msg_id,
"snippet": match_info.get("snippet") or "",
"bookend_start": [
_shape_message(m, max_content_len=1200)
for m in (view.get("bookend_start") or [])
if not _is_compaction_summary(m.get("content", ""))
],
"messages": [_shape_message(m, anchor_id=msg_id, max_content_len=4000) for m in (view.get("window") or [])],
"bookend_end": [
_shape_message(m, max_content_len=1200)
for m in (view.get("bookend_end") or [])
if not _is_compaction_summary(m.get("content", ""))
"bookend_start": (
[
_shape_message(m, max_content_len=1200)
for m in (view.get("bookend_start") or [])
if not _is_compaction_summary(m.get("content", ""))
]
if result_detail == "full"
else []
),
"messages": [
_shape_message(m, anchor_id=msg_id, max_content_len=4000)
for m in window_messages
],
"bookend_end": (
[
_shape_message(m, max_content_len=1200)
for m in (view.get("bookend_end") or [])
if not _is_compaction_summary(m.get("content", ""))
]
if result_detail == "full"
else []
),
"messages_before": view.get("messages_before", 0),
"messages_after": view.get("messages_after", 0),
"detail": result_detail,
}
if lineage_root and lineage_root != hit_sid:
entry["parent_session_id"] = lineage_root
@@ -900,6 +924,7 @@ def _discover(
"success": True,
"mode": "discover",
"query": query,
"detail": detail,
"results": results,
"count": len(results),
"sessions_searched": len(seen_sessions),
@@ -922,12 +947,14 @@ def _session_search_impl(
sort: str = None,
# Cross-profile (any shape)
profile: str = None,
# Discovery result shaping (appended to preserve positional compatibility)
detail: str = "adaptive",
*,
_owned_dbs: Optional[List[Any]] = None,
) -> str:
"""Single-shape tool. Mode inferred from which args are set.
Discovery: pass ``query``.
Discovery: pass ``query``; ``detail="full"`` hydrates every result.
Scroll: pass ``session_id`` + ``around_message_id``.
Read: pass ``session_id`` (no anchor) — dumps the whole session.
Browse: pass nothing.
@@ -1017,12 +1044,19 @@ def _session_search_impl(
if candidate in ("newest", "oldest"):
sort_norm = candidate
detail_norm = (
"full"
if isinstance(detail, str) and detail.strip().lower() == "full"
else "adaptive"
)
return _discover(
db=db,
query=query.strip(),
role_filter=role_list,
limit=limit,
sort=sort_norm,
detail=detail_norm,
current_session_id=current_session_id,
link_profile=profile,
)
@@ -1108,19 +1142,21 @@ SESSION_SEARCH_SCHEMA = {
"FOUR CALLING SHAPES\n\n"
" 1) DISCOVERY — pass `query`:\n"
" session_search(query=\"auth refactor\", limit=3)\n"
" Runs FTS5, dedupes hits by session lineage, returns the top N sessions. "
"Each result carries:\n"
" Runs FTS5, dedupes hits by session lineage, and returns the top N "
"sessions. Adaptive detail is the default: the top-ranked result carries "
"full context, while lower-ranked results stay compact. Pass `detail=\"full\"` "
"to fully hydrate every result. Every result carries:\n"
" - session_id, title, when, source\n"
" - snippet: FTS5-highlighted match excerpt\n"
" - bookend_start: first 3 user+assistant messages of the session "
"(the goal / kickoff)\n"
" - messages: ±5 messages around the FTS5 match, with the anchor message "
"flagged (the hit in context)\n"
" - bookend_end: last 3 user+assistant messages of the session "
"(the resolution / decisions)\n"
" - detail: `full` or `compact`\n"
" - bookend_start/bookend_end: the first/last 3 user+assistant messages "
"for full results; empty lists for compact results\n"
" - messages: ±5 messages around the FTS5 match for full results; only "
"the flagged anchor message for compact results\n"
" - match_message_id, messages_before, messages_after\n"
" Bookends + window together let you reconstruct goal → match → resolution "
"without paying for the whole transcript.\n\n"
" The top result's bookends + window let you reconstruct goal → match → "
"resolution immediately. Scroll a compact result when another session looks "
"more promising.\n\n"
" 2) SCROLL — pass `session_id` + `around_message_id`:\n"
" session_search(session_id=\"...\", around_message_id=12345, window=10)\n"
" Returns a window of ±`window` messages centered on the anchor. No FTS5, "
@@ -1197,6 +1233,17 @@ SESSION_SEARCH_SCHEMA = {
"and browse shapes."
),
},
"detail": {
"type": "string",
"enum": ["adaptive", "full"],
"description": (
"Discovery shape only. 'adaptive' (default) fully hydrates the "
"top-ranked result and returns only the exact anchor message for "
"lower-ranked results. 'full' returns bookends and the complete "
"anchored window for every result."
),
"default": "adaptive",
},
"session_id": {
"type": "string",
"description": (
@@ -1261,6 +1308,7 @@ registry.register(
around_message_id=args.get("around_message_id"),
window=args.get("window", 5),
sort=args.get("sort"),
detail=args.get("detail", "adaptive"),
profile=args.get("profile"),
db=kw.get("db"),
current_session_id=kw.get("current_session_id"),