From 7f1d84fe7f5b03ef9a3a37733957728b8e83db93 Mon Sep 17 00:00:00 2001 From: PRATHAMESH75 Date: Sun, 19 Jul 2026 12:33:45 +0530 Subject: [PATCH] perf(insights): pin partial index on assistant tool-call queries MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- agent/insights.py | 72 +++++++++++++++--------- tests/agent/test_insights.py | 106 +++++++++++++++++++++++++++++------ tests/test_hermes_state.py | 37 ++++-------- 3 files changed, 142 insertions(+), 73 deletions(-) diff --git a/agent/insights.py b/agent/insights.py index 9d148a15446bf..2e45b222d4152 100644 --- a/agent/insights.py +++ b/agent/insights.py @@ -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: diff --git a/tests/agent/test_insights.py b/tests/agent/test_insights.py index 23fd93fe9c6df..b638be70df4b5 100644 --- a/tests/agent/test_insights.py +++ b/tests/agent/test_insights.py @@ -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() # ========================================================================= diff --git a/tests/test_hermes_state.py b/tests/test_hermes_state.py index 928f7eb0ce3f8..eac6b2b8baa59 100644 --- a/tests/test_hermes_state.py +++ b/tests/test_hermes_state.py @@ -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