From 996a16cfb50f014580d912cc87f0dcc92912a849 Mon Sep 17 00:00:00 2001 From: Joseph Doherty Date: Mon, 27 Apr 2026 04:34:18 -0400 Subject: [PATCH] perf: search.py N+1 batching + k constant extraction (T106) --- chat/web/search.py | 142 +++++++++++++++++++++++++++++++++++++--- tests/test_search_ux.py | 28 ++++++++ 2 files changed, 161 insertions(+), 9 deletions(-) diff --git a/chat/web/search.py b/chat/web/search.py index 51d75ea..458c7c7 100644 --- a/chat/web/search.py +++ b/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 - # service) so the service stays a pure FTS shim that other UIs - # can reuse. + Hydration (T106) is batched: rather than calling ``get_bot`` / + ``get_chat`` / ``get_scene`` per row (worst case 3 * k individual + 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"]) - chat = get_chat(conn, row["chat_id"]) - scene = get_scene(conn, row["scene_id"]) if row["scene_id"] else None + bot = bots_by_id.get(row["owner_id"]) + chat = chats_by_id.get(row["chat_id"]) + scene = ( + scenes_by_id.get(row["scene_id"]) if row["scene_id"] else None + ) results.append( { "memory_id": row["memory_id"], diff --git a/tests/test_search_ux.py b/tests/test_search_ux.py index 7254549..013337b 100644 --- a/tests/test_search_ux.py +++ b/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