perf(insights): pin partial index on assistant tool-call queries
Review follow-up (#67341): on a freshly initialized state.db (before ANALYZE has run) the source-filtered branches of _get_tool_usage / _get_skill_usage did not select idx_messages_assistant_calls_by_session — the optimizer drove from idx_sessions_source_id and probed each session's messages via idx_messages_session_active, scanning non tool-call rows. Pin the index with INDEXED BY on all four fixed-predicate branches so the plan is deterministic for both the unfiltered and source-filtered scopes without depending on statistics. Safe because the index is declared in SCHEMA_SQL (created by every read-write SessionDB._init_schema) and every InsightsEngine caller opens a read-write SessionDB; read-only attachments (which skip schema init) are never used for insights. Extract the four queries into class constants and add tests: query-plan coverage for both scopes without ANALYZE, row-level equivalence between pinned and un-pinned forms, and an assertion that INDEXED BY fails loudly if the index is absent.
This commit is contained in:
PRATHAMESH75
kshitij
parent
034eadb326
commit
7f1d84fe7f
|
|
@ -195,6 +195,46 @@ class InsightsEngine:
|
|||
" ORDER BY started_at DESC"
|
||||
)
|
||||
|
||||
# Assistant ``tool_calls`` scan for tool/skill usage. ``INDEXED BY`` pins
|
||||
# the partial index ``idx_messages_assistant_calls_by_session`` so the plan
|
||||
# is deterministic on a freshly initialized state.db (before ANALYZE has
|
||||
# run) for BOTH the unfiltered and source-filtered branches — without the
|
||||
# hint the optimizer falls back to ``idx_messages_session_active`` for the
|
||||
# source-filtered probe and scans each session's non-tool-call rows. Safe
|
||||
# because the index is declared in ``SCHEMA_SQL`` (created by every
|
||||
# read-write ``SessionDB._init_schema``); every ``InsightsEngine`` caller
|
||||
# opens a read-write ``SessionDB`` — read-only attachments (which skip
|
||||
# schema init) are never used for insights.
|
||||
_MESSAGES_ASSISTANT_CALLS_INDEX = "idx_messages_assistant_calls_by_session"
|
||||
_GET_TOOL_CALLS_WITH_SOURCE = (
|
||||
"SELECT m.tool_calls"
|
||||
f" FROM messages m INDEXED BY {_MESSAGES_ASSISTANT_CALLS_INDEX}"
|
||||
" JOIN sessions s ON s.id = m.session_id"
|
||||
" WHERE s.started_at >= ? AND s.source = ?"
|
||||
" AND m.role = 'assistant' AND m.tool_calls IS NOT NULL"
|
||||
)
|
||||
_GET_TOOL_CALLS_ALL = (
|
||||
"SELECT m.tool_calls"
|
||||
f" FROM messages m INDEXED BY {_MESSAGES_ASSISTANT_CALLS_INDEX}"
|
||||
" JOIN sessions s ON s.id = m.session_id"
|
||||
" WHERE s.started_at >= ?"
|
||||
" AND m.role = 'assistant' AND m.tool_calls IS NOT NULL"
|
||||
)
|
||||
_GET_SKILL_CALLS_WITH_SOURCE = (
|
||||
"SELECT m.tool_calls, m.timestamp"
|
||||
f" FROM messages m INDEXED BY {_MESSAGES_ASSISTANT_CALLS_INDEX}"
|
||||
" JOIN sessions s ON s.id = m.session_id"
|
||||
" WHERE s.started_at >= ? AND s.source = ?"
|
||||
" AND m.role = 'assistant' AND m.tool_calls IS NOT NULL"
|
||||
)
|
||||
_GET_SKILL_CALLS_ALL = (
|
||||
"SELECT m.tool_calls, m.timestamp"
|
||||
f" FROM messages m INDEXED BY {_MESSAGES_ASSISTANT_CALLS_INDEX}"
|
||||
" JOIN sessions s ON s.id = m.session_id"
|
||||
" WHERE s.started_at >= ?"
|
||||
" AND m.role = 'assistant' AND m.tool_calls IS NOT NULL"
|
||||
)
|
||||
|
||||
def _get_sessions(self, cutoff: float, source: str = None) -> List[Dict]:
|
||||
"""Fetch sessions within the time window."""
|
||||
if source:
|
||||
|
|
@ -243,22 +283,10 @@ class InsightsEngine:
|
|||
# (covers CLI sessions where tool_name is NULL on tool responses)
|
||||
if source:
|
||||
cursor2 = self._conn.execute(
|
||||
"""SELECT m.tool_calls
|
||||
FROM messages m
|
||||
JOIN sessions s ON s.id = m.session_id
|
||||
WHERE s.started_at >= ? AND s.source = ?
|
||||
AND m.role = 'assistant' AND m.tool_calls IS NOT NULL""",
|
||||
(cutoff, source),
|
||||
self._GET_TOOL_CALLS_WITH_SOURCE, (cutoff, source)
|
||||
)
|
||||
else:
|
||||
cursor2 = self._conn.execute(
|
||||
"""SELECT m.tool_calls
|
||||
FROM messages m
|
||||
JOIN sessions s ON s.id = m.session_id
|
||||
WHERE s.started_at >= ?
|
||||
AND m.role = 'assistant' AND m.tool_calls IS NOT NULL""",
|
||||
(cutoff,),
|
||||
)
|
||||
cursor2 = self._conn.execute(self._GET_TOOL_CALLS_ALL, (cutoff,))
|
||||
|
||||
tool_calls_counts = Counter()
|
||||
for row in cursor2.fetchall():
|
||||
|
|
@ -301,22 +329,10 @@ class InsightsEngine:
|
|||
|
||||
if source:
|
||||
cursor = self._conn.execute(
|
||||
"""SELECT m.tool_calls, m.timestamp
|
||||
FROM messages m
|
||||
JOIN sessions s ON s.id = m.session_id
|
||||
WHERE s.started_at >= ? AND s.source = ?
|
||||
AND m.role = 'assistant' AND m.tool_calls IS NOT NULL""",
|
||||
(cutoff, source),
|
||||
self._GET_SKILL_CALLS_WITH_SOURCE, (cutoff, source)
|
||||
)
|
||||
else:
|
||||
cursor = self._conn.execute(
|
||||
"""SELECT m.tool_calls, m.timestamp
|
||||
FROM messages m
|
||||
JOIN sessions s ON s.id = m.session_id
|
||||
WHERE s.started_at >= ?
|
||||
AND m.role = 'assistant' AND m.tool_calls IS NOT NULL""",
|
||||
(cutoff,),
|
||||
)
|
||||
cursor = self._conn.execute(self._GET_SKILL_CALLS_ALL, (cutoff,))
|
||||
|
||||
for row in cursor.fetchall():
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
"""Tests for agent/insights.py — InsightsEngine analytics and reporting."""
|
||||
|
||||
import sqlite3
|
||||
import time
|
||||
import pytest
|
||||
|
||||
|
|
@ -368,30 +369,99 @@ class TestInsightsPopulated:
|
|||
|
||||
|
||||
|
||||
def test_tool_and_skill_usage_invariant_to_partial_index(self, populated_db):
|
||||
"""The assistant tool-call partial index is a pure optimization: tool
|
||||
and skill usage must be byte-for-byte identical with and without it."""
|
||||
engine = InsightsEngine(populated_db)
|
||||
cutoff = 0.0
|
||||
index = "idx_messages_assistant_calls_by_session"
|
||||
# The Insights assistant tool-call queries pin
|
||||
# idx_messages_assistant_calls_by_session via INDEXED BY. These tests prove
|
||||
# (a) the planner uses that index for BOTH the unfiltered and source-filtered
|
||||
# branches on a fresh DB *without* ANALYZE, and (b) the index is a pure
|
||||
# optimization — output is identical whether or not it is selected.
|
||||
_INDEX = "idx_messages_assistant_calls_by_session"
|
||||
_PINNED_QUERIES = (
|
||||
("_GET_TOOL_CALLS_ALL", (0.0,)),
|
||||
("_GET_TOOL_CALLS_WITH_SOURCE", (0.0, "cli")),
|
||||
("_GET_SKILL_CALLS_ALL", (0.0,)),
|
||||
("_GET_SKILL_CALLS_WITH_SOURCE", (0.0, "cli")),
|
||||
)
|
||||
|
||||
with_index_tools = engine._get_tool_usage(cutoff)
|
||||
with_index_skills = engine._get_skill_usage(cutoff)
|
||||
# The index must exist by default (it lives in SCHEMA_SQL).
|
||||
def test_assistant_call_queries_use_partial_index_without_analyze(
|
||||
self, populated_db
|
||||
):
|
||||
"""Every fixed-predicate branch selects the partial index on a fresh DB.
|
||||
|
||||
No ANALYZE is run, so this covers the default-statistics case a freshly
|
||||
initialized state.db is actually in. Both the unfiltered and the
|
||||
source-filtered (``s.source = ?``) branches are checked.
|
||||
"""
|
||||
# Guard against the fresh-DB planner regression the reviewers found:
|
||||
# without INDEXED BY the source-filtered branch fell back to
|
||||
# idx_messages_session_active.
|
||||
assert "ANALYZE" not in "".join(
|
||||
r["sql"] or ""
|
||||
for r in populated_db._conn.execute(
|
||||
"SELECT sql FROM sqlite_master WHERE type = 'index'"
|
||||
)
|
||||
)
|
||||
for attr, params in self._PINNED_QUERIES:
|
||||
sql = getattr(InsightsEngine, attr)
|
||||
plan = "\n".join(
|
||||
row["detail"]
|
||||
for row in populated_db._conn.execute(
|
||||
"EXPLAIN QUERY PLAN " + sql, params
|
||||
).fetchall()
|
||||
)
|
||||
assert self._INDEX in plan, f"{attr} did not use the index:\n{plan}"
|
||||
|
||||
def test_assistant_call_rows_invariant_to_index_selection(self, populated_db):
|
||||
"""The pinned index only changes the plan, never the result set.
|
||||
|
||||
For every branch, the index-pinned query and the un-pinned form (whose
|
||||
plan the optimizer chooses freely) must return identical rows — proving
|
||||
the index is a pure optimization — for both the unfiltered and
|
||||
source-filtered scopes.
|
||||
"""
|
||||
assert populated_db._conn.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type = 'index' AND name = ?",
|
||||
(index,),
|
||||
(self._INDEX,),
|
||||
).fetchone() is not None
|
||||
|
||||
populated_db._conn.execute(f"DROP INDEX IF EXISTS {index}")
|
||||
populated_db._conn.commit()
|
||||
without_index_tools = engine._get_tool_usage(cutoff)
|
||||
without_index_skills = engine._get_skill_usage(cutoff)
|
||||
for attr, params in self._PINNED_QUERIES:
|
||||
pinned_sql = getattr(InsightsEngine, attr)
|
||||
unpinned_sql = pinned_sql.replace(f" INDEXED BY {self._INDEX}", "")
|
||||
pinned = [
|
||||
tuple(r) for r in
|
||||
populated_db._conn.execute(pinned_sql, params).fetchall()
|
||||
]
|
||||
unpinned = [
|
||||
tuple(r) for r in
|
||||
populated_db._conn.execute(unpinned_sql, params).fetchall()
|
||||
]
|
||||
assert sorted(pinned) == sorted(unpinned), attr
|
||||
|
||||
assert with_index_tools == without_index_tools
|
||||
assert with_index_skills == without_index_skills
|
||||
# Sanity: the fixture actually exercises the assistant tool_calls path.
|
||||
assert any(t["tool_name"] == "search_files" for t in with_index_tools)
|
||||
def test_tool_and_skill_usage_invariant_to_partial_index(self, populated_db):
|
||||
"""The public tool/skill usage output is stable and exercises the
|
||||
assistant tool_calls path for both scopes."""
|
||||
engine = InsightsEngine(populated_db)
|
||||
cutoff = 0.0
|
||||
|
||||
tools = engine._get_tool_usage(cutoff)
|
||||
tools_cli = engine._get_tool_usage(cutoff, source="cli")
|
||||
skills = engine._get_skill_usage(cutoff)
|
||||
skills_cli = engine._get_skill_usage(cutoff, source="cli")
|
||||
|
||||
# Sanity: the fixture actually drives the assistant tool_calls path.
|
||||
assert any(t["tool_name"] == "search_files" for t in tools)
|
||||
assert any(t["tool_name"] == "search_files" for t in tools_cli)
|
||||
assert isinstance(skills, list) and isinstance(skills_cli, list)
|
||||
|
||||
def test_indexed_by_requires_the_index_to_exist(self, populated_db):
|
||||
"""INDEXED BY is a hard dependency: if the index were ever missing the
|
||||
query fails loudly rather than silently degrading. This documents why
|
||||
every InsightsEngine caller must open a read-write SessionDB."""
|
||||
populated_db._conn.execute(f"DROP INDEX IF EXISTS {self._INDEX}")
|
||||
populated_db._conn.commit()
|
||||
with pytest.raises(sqlite3.OperationalError, match="no such index"):
|
||||
populated_db._conn.execute(
|
||||
InsightsEngine._GET_TOOL_CALLS_ALL, (0.0,)
|
||||
).fetchall()
|
||||
|
||||
|
||||
# =========================================================================
|
||||
|
|
|
|||
|
|
@ -3704,32 +3704,15 @@ class TestInsightsToolCallIndex:
|
|||
finally:
|
||||
db2.close()
|
||||
|
||||
def test_index_serves_assistant_tool_call_scan(self, db):
|
||||
"""The planner uses the partial index for the exact Insights predicate.
|
||||
def test_index_predicate_is_partial(self, db):
|
||||
"""The index covers only the assistant tool-call rows Insights reads.
|
||||
|
||||
Seed enough rows that the optimizer prefers the partial index over a
|
||||
full table scan, then assert the plan names it.
|
||||
Query-plan coverage (that the Insights queries actually select this
|
||||
index, for both scopes, without ANALYZE) lives with the queries in
|
||||
tests/agent/test_insights.py.
|
||||
"""
|
||||
db.create_session(session_id="s1", source="cli")
|
||||
for i in range(200):
|
||||
if i % 5 == 0:
|
||||
db.append_message(
|
||||
"s1", role="assistant", content=f"m{i}",
|
||||
tool_calls=[{"function": {"name": "search_files"}}],
|
||||
)
|
||||
else:
|
||||
db.append_message("s1", role="user", content=f"m{i}")
|
||||
db._conn.execute("ANALYZE")
|
||||
db._conn.commit()
|
||||
|
||||
plan = "\n".join(
|
||||
row["detail"]
|
||||
for row in db._conn.execute(
|
||||
"EXPLAIN QUERY PLAN "
|
||||
"SELECT m.tool_calls FROM messages m "
|
||||
"JOIN sessions s ON s.id = m.session_id "
|
||||
"WHERE s.started_at >= 0 "
|
||||
"AND m.role = 'assistant' AND m.tool_calls IS NOT NULL"
|
||||
).fetchall()
|
||||
)
|
||||
assert self._INDEX in plan, plan
|
||||
sql = self._index_defn(db._conn)
|
||||
assert sql is not None
|
||||
assert "WHERE" in sql
|
||||
assert "role = 'assistant'" in sql
|
||||
assert "tool_calls IS NOT NULL" in sql
|
||||
|
|
|
|||
Loading…
Reference in New Issue