merge: T106 search.py N+1 batching + k constant
This commit is contained in:
Joseph Doherty
+133
-9
@@ -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"],
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user