perf: search.py N+1 batching + k constant extraction (T106)

chat/web/search.py

@@ -14,6 +14,12 @@ For each match we hydrate just enough metadata to render a row:
 * the originating scene title when one exists,
 * and the ``pov_summary`` itself.
 
+T106 (Phase 4.5): hydration is batched. Pre-T106 the route called
+``get_bot``/``get_chat``/``get_scene`` once per result row — N+1 with
+``DEFAULT_SEARCH_K=50`` meaning up to 150 individual SELECTs per page
+load. We now collect distinct ids first and fan-in via three
+``WHERE id IN (...)`` queries, then map back per row.
+
 We deliberately keep this module synchronous and template-only — no
 HTMX swaps, no JSON API — because the search box is a "leave the
 current chat to look something up" surface, not an inline drawer.
@@ -21,7 +27,9 @@ current chat to look something up" surface, not an inline drawer.
 
 from __future__ import annotations
 
+import json
 from pathlib import Path
+from sqlite3 import Connection
 
 from fastapi import APIRouter, Depends, Request
 from fastapi.responses import HTMLResponse
@@ -36,29 +44,145 @@ TEMPLATES = Jinja2Templates(
     directory=str(Path(__file__).resolve().parent.parent / "templates")
 )
 
+#: Maximum cross-chat FTS matches surfaced per ``/search`` page load.
+#: Extracted as a module-level constant (T106) so the cap is tunable
+#: without touching the route body. ``search_all_memories`` itself
+#: defaults to a smaller ``k=20``; we override here because the
+#: top-bar search is a "scan everything I've seen" surface, not an
+#: inline drawer.
+DEFAULT_SEARCH_K = 50
+
 router = APIRouter()
 
 
+def _fetch_bots_by_ids(conn: Connection, ids: set[str]) -> dict[str, dict]:
+    """Batched sibling of :func:`chat.state.entities.get_bot`.
+
+    Inlined here (not exported from ``state.entities``) to keep T106's
+    scope confined to ``search.py`` per the Phase 4.5 plan. Returns
+    ``{bot_id: bot_dict}`` for every id present in ``ids``; ids with
+    no matching row are simply absent from the map (the caller falls
+    back to the raw id string the same way it did pre-T106).
+
+    Empty ``ids`` short-circuits to ``{}`` because SQLite rejects
+    ``WHERE id IN ()`` as a syntax error.
+    """
+    if not ids:
+        return {}
+    placeholders = ",".join("?" * len(ids))
+    cols = [c[1] for c in conn.execute("PRAGMA table_info(bots)").fetchall()]
+    rows = conn.execute(
+        f"SELECT * FROM bots WHERE id IN ({placeholders})",
+        tuple(ids),
+    ).fetchall()
+    out: dict[str, dict] = {}
+    for row in rows:
+        d = dict(zip(cols, row))
+        d["voice_samples"] = json.loads(d.pop("voice_samples_json"))
+        d["traits"] = json.loads(d.pop("traits_json"))
+        out[d["id"]] = d
+    return out
+
+
+def _fetch_chats_by_ids(conn: Connection, ids: set[str]) -> dict[str, dict]:
+    """Batched sibling of :func:`chat.state.world.get_chat`.
+
+    Mirrors that helper's ``chats``/``chat_state`` JOIN so the returned
+    dicts have the same shape (``narrative_anchor``, ``time``,
+    ``weather``, ``active_scene_id``, etc.). Empty ``ids`` returns
+    ``{}`` to dodge the ``IN ()`` syntax error.
+    """
+    if not ids:
+        return {}
+    placeholders = ",".join("?" * len(ids))
+    rows = conn.execute(
+        "SELECT c.id, c.host_bot_id, c.guest_bot_id, c.created_at, "
+        "       s.time, s.weather, s.active_scene_id, s.narrative_anchor "
+        f"FROM chats c JOIN chat_state s ON s.chat_id = c.id "
+        f"WHERE c.id IN ({placeholders})",
+        tuple(ids),
+    ).fetchall()
+    return {
+        row[0]: {
+            "id": row[0],
+            "host_bot_id": row[1],
+            "guest_bot_id": row[2],
+            "created_at": row[3],
+            "time": row[4],
+            "weather": row[5],
+            "active_scene_id": row[6],
+            "narrative_anchor": row[7],
+        }
+        for row in rows
+    }
+
+
+def _fetch_scenes_by_ids(conn: Connection, ids: set[int]) -> dict[int, dict]:
+    """Batched sibling of :func:`chat.state.world.get_scene`.
+
+    Returns ``{scene_id: scene_dict}`` with ``participants`` already
+    JSON-decoded so callers see the same shape as the per-row helper.
+    Empty ``ids`` returns ``{}``.
+    """
+    if not ids:
+        return {}
+    placeholders = ",".join("?" * len(ids))
+    cols = [c[1] for c in conn.execute("PRAGMA table_info(scenes)").fetchall()]
+    rows = conn.execute(
+        f"SELECT * FROM scenes WHERE id IN ({placeholders})",
+        tuple(ids),
+    ).fetchall()
+    out: dict[int, dict] = {}
+    for row in rows:
+        d = dict(zip(cols, row))
+        d["participants"] = json.loads(d.pop("participants_json"))
+        out[d["id"]] = d
+    return out
+
+
 @router.get("/search", response_class=HTMLResponse)
 async def search(request: Request, q: str = "", conn=Depends(get_conn)):
-    """Render ``search.html`` with up to 50 cross-chat FTS matches.
+    """Render ``search.html`` with up to :data:`DEFAULT_SEARCH_K` matches.
 
     ``q`` is intentionally allowed to be empty — that path renders the
     page's "enter a query" placeholder rather than a 400, because the
     top-bar form submits to this URL even with an empty input. T93's
     service short-circuits whitespace-only queries to ``[]`` so there
     is no FTS5 ``MATCH ''`` syntax error to guard against here.
-    """
-    raw_results = search_all_memories(conn, query=q, k=50) if q else []
 
-    # Hydrate display fields per row. We do this in the route (not the
+    Hydration (T106) is batched: rather than calling ``get_bot`` /
-    # service) so the service stays a pure FTS shim that other UIs
+    ``get_chat`` / ``get_scene`` per row (worst case 3 * k individual
-    # can reuse.
+    SELECTs), we collect distinct ids and issue one ``IN (...)`` query
+    per entity kind, then map back during the row build. ``get_bot``
+    et al. remain imported for test-time monkeypatching but are no
+    longer invoked on the hot path.
+    """
+    raw_results = (
+        search_all_memories(conn, query=q, k=DEFAULT_SEARCH_K) if q else []
+    )
+
+    # Collect distinct ids up front so the IN-list queries dedupe (a
+    # popular bot or scene shows up many times across the result set).
+    bot_ids: set[str] = {r["owner_id"] for r in raw_results if r["owner_id"]}
+    chat_ids: set[str] = {r["chat_id"] for r in raw_results if r["chat_id"]}
+    scene_ids: set[int] = {
+        r["scene_id"] for r in raw_results if r["scene_id"]
+    }
+
+    bots_by_id = _fetch_bots_by_ids(conn, bot_ids)
+    chats_by_id = _fetch_chats_by_ids(conn, chat_ids)
+    scenes_by_id = _fetch_scenes_by_ids(conn, scene_ids)
+
+    # Hydrate display fields per row from the batched maps. We do this
+    # in the route (not the service) so the service stays a pure FTS
+    # shim that other UIs can reuse.
     results = []
     for row in raw_results:
-        bot = get_bot(conn, row["owner_id"])
+        bot = bots_by_id.get(row["owner_id"])
-        chat = get_chat(conn, row["chat_id"])
+        chat = chats_by_id.get(row["chat_id"])
-        scene = get_scene(conn, row["scene_id"]) if row["scene_id"] else None
+        scene = (
+            scenes_by_id.get(row["scene_id"]) if row["scene_id"] else None
+        )
         results.append(
             {
                 "memory_id": row["memory_id"],

tests/test_search_ux.py

@@ -16,6 +16,7 @@ Verifies the FastAPI ``/search`` route that wraps T93's
 from __future__ import annotations
 
 from pathlib import Path
+from unittest.mock import patch
 
 import pytest
 from fastapi.testclient import TestClient
@@ -133,3 +134,30 @@ def test_result_links_navigate_to_chat(client, tmp_path):
     # The link target is chat-level (memories don't carry an event_id
     # column today, so we don't deep-link to a specific turn).
     assert 'href="/chats/chat_a"' in resp.text
+
+
+def test_search_results_use_batched_lookups(client, tmp_path):
+    """T106: hydration must not fan out to per-row ``get_bot``/
+    ``get_chat``/``get_scene`` calls.
+
+    The previous implementation called each helper once per result row
+    (worst case 50 rows x 3 helpers = 150 individual queries). The
+    batched implementation collects distinct ids and issues at most one
+    query per entity kind via ``WHERE id IN (...)``, so the per-row
+    helpers should not be invoked at all when there are matches.
+
+    We seed two chats (so both ``get_bot`` and ``get_chat`` would have
+    been hit pre-T106) and assert each helper sees zero per-row calls.
+    """
+    _seed_two_chats_with_memories(tmp_path / "test.db")
+    with (
+        patch("chat.web.search.get_bot") as mock_get_bot,
+        patch("chat.web.search.get_chat") as mock_get_chat,
+        patch("chat.web.search.get_scene") as mock_get_scene,
+    ):
+        resp = client.get("/search?q=rabbit")
+    assert resp.status_code == 200
+    # Batched IN-list queries replace the per-row helpers entirely.
+    assert mock_get_bot.call_count == 0
+    assert mock_get_chat.call_count == 0
+    assert mock_get_scene.call_count == 0