Database Specialist

Overview

database_specialist is the second agent in the SequentialAgent pipeline (query_resolution). It acts as the primary data gatherer and SQL executor for the Kavion AI Assistant. It receives the structured, fully-resolved query output from resolver_classifier and translates it into precise SQL queries executed against Supabase, or delegates to rgb_surface_analyzer for surface analysis turns (it acts as a coordinator: it validates the target image via SQL, then delegates vision analysis to the rgb_surface_analyzer AgentTool.)


Role and Position

  • Role: Data retrieval and tool delegation sub-agent within the SequentialAgent pipeline (query_resolution).
  • Position: Second agent in the pipeline. Runs after resolver_classifier and before report_generator.
  • Caller: Called by the SequentialAgent pipeline (query_resolution) when the pipeline is executed.
  • Callee: rgb_surface_analyzer (via AgentTool delegation on surface_analysis turns).
  • Input: Reads resolver_classifier output from ctx.state["resolver_classifier"].
  • Output: Its output is written to ctx.state["database_specialist"] and merged into ctx.state["global"] by the database_specialist_after_agent callback before the downstream report_generator agent runs.

Tools

Tool Type Role in the System When Invoked
execute_sql Plain Python callable (closure-wrapped by tool_adapter.py) Executes raw SQL queries against Supabase behind an AST-based SQL firewall. All non-surface-analysis query types, and for image validation/resolution on surface-analysis turns.
rgb_surface_analyzer AgentTool (wraps a full LlmAgent) Standalone LLM agent that analyzes RGB images for surface conditions (stains, corrosion). query_type == "surface_analysis" only.

LLM Context Design

This section details every item that ends up in the actual request sent to the LLM, categorized by its lifecycle and source.

Context Lifecycle Categorization

  • Static (fixed at code-authoring time, identical every call):
    • Base Instruction Template: The core instruction text defined in adk_agents_config.yaml (excluding placeholders).
  • Dynamic · session (varies by session/user state but stable within a session):
    • {database_context}: Formatted markdown of the database schema and metadata, loaded at server start.
    • {multi_image_cap}: Maximum images per surface analysis turn (configured via multi_image_cap in adk_agents_config.yaml).
  • Dynamic · per-turn (recomputed every turn):
    • {resolved_query}: Sourced from resolver_classifier.resolved_query in session state.
    • {active_filters}: Sourced from global.active_filters in session state.
    • {user_db_access}: Sourced from global.user_db_access in session state (substituted by factory.py but not present in the database_specialist instruction template — no-op substitution).
    • {sql_tool_result}: Sourced from global.sql_tool_result in session state (substituted by factory.py but not present in the database_specialist instruction template — no-op substitution).
  • ADK-injected / default (added automatically by the ADK framework itself):
    • Conversation History: Pulled in automatically by ADK because include_contents: default is set. This includes user messages, assistant responses, tool calls, and tool responses from previous turns in the session.
    • Output Schema Formatting Instructions: Automatically appended by ADK to the end of the prompt because output_schema: DatabaseSpecialistResponse is set. This instructs the model to output JSON matching the Pydantic schema.
  • Callback-mutated (present, changed, or removed because a callback touched it):
    • temp:jwt_token: Injected into tool_context.state by database_specialist_before_tool before calling rgb_surface_analyzer AgentTool.

History Inclusion (include_contents)

  • Setting: Not explicitly set in adk_agents_config.yaml; factory default resolves to "default".
  • Behavior: Instructs the ADK framework to automatically include the conversation history (user messages, assistant responses, tool calls, and tool responses) in the LLM request. The history is formatted as a sequence of turns.

Instruction Template and Placeholders

The instruction template is defined in src/kavai/config/adk_agents_config.yaml. The factory (factory.py) substitutes placeholders into the instruction text at agent build time:

Placeholder Source What it injects Insertion Count
{database_context} Database schema config Formatted markdown of the database schema and metadata 1
{multi_image_cap} adk_agents_config.yaml Maximum images per surface analysis turn (configured via multi_image_cap) 6

Prompt Scope & Retrieval/Delegation Rules

The agent executes different protocols based on the resolver_classifier.query_type and pre-flight checks:

Pre-Flight Check (Mandatory)

Before executing any SQL or tool invocation, the agent checks:

  1. resolver_classifier.is_in_scope == false → Return immediately with either:
    • Security Threat Indicators: If the query contains schema introspection attempts (information_schema, pg_catalog, pg_*, auth.*, profiles, users), system metadata queries, privilege escalation, data exfiltration patterns, injection attempts, or environment probing:
      • agent_result: "Security violation detected. This request attempts to access restricted system resources."
    • Generic Out-of-Scope: If the query is causal reasoning, predictions, suggestions, or professional advice:
      • agent_result: "Out of scope. Cannot assist with non-database queries."
    • executed_queries: [], is_complete: false, retrieved_rows: [], errors: [].
  2. resolver_classifier.is_ambiguous == true → Return immediately with:
    • agent_result: "Query is ambiguous. Please clarify your request."
    • executed_queries: [], is_complete: false, retrieved_rows: [], errors: [].
  3. Otherwise → Proceed to intent routing.

Retrieval Intents (metadata_lookup, record_filtering, anomaly_retrieval, aggregation, deterministic_linking, general)

  • Two-Query Protocol (Mandatory):
    1. COUNT First (No Limit): Run SELECT COUNT(*) first to get the true total count.
    2. Fetch with Limit (1-50): Fetch sample rows for display with LIMIT 50 or less.
    3. Report Both Counts: agent_result must report both the true total count and the displayed subset (e.g., “Found 500 images with X anomaly. Retrieved first 50 for display.”).
  • Mandatory Fields for Images: Always include id, filename, storage_path, thumbnail_url, sensor_modality, and metadata. NEVER use SELECT * on the images table. NEVER omit thumbnail_url.
  • Modality Filtering (Strict - Zero Tolerance):
    • Use ONLY filename patterns for modality filtering (e.g., filename ILIKE '%_T%.jpg' for Thermal, filename ILIKE '%_V%.jpg' for RGB).
    • NEVER use sensor_modality column or metadata->>'ImageType'.
    • Include all case variations (.jpg, .JPG, .jpeg, .JPEG).
  • Location Filtering (Strict - Mandatory):
    • ONLY use metadata->>'PhysicalLocation' for location filtering.
    • If location filter returns 0 results, return 0 results (do NOT return unfiltered data).
  • ILIKE Only (Mandatory): ALWAYS use ILIKE for string matches on name, slug, filename, category, and metadata->>'PhysicalLocation'. NEVER use = or LIKE for strings.
  • Anomaly Joins (Critical - Use DISTINCT): Always use SELECT DISTINCT when joining with the annotations table to prevent duplicate images.
  • Golden Gas Readings Rule: Use v_dataset_gas_readings view or fallback to spatial join with ST_DWithin().

Surface Analysis Delegation (query_type == 'surface_analysis')

  1. Read target_image_ids: Read image identifiers EXCLUSIVELY from resolver_classifier.target_image_ids in session state. NEVER derive image IDs from the user query text.
  2. Validate: If target_image_ids is null, missing, or empty → return error immediately.
  3. Cap Enforcement: If target_image_ids contains more than {multi_image_cap} identifiers, use only the first {multi_image_cap}.
  4. ⛔ SQL Gate — Mandatory Decision Point: Classify each identifier into:
    • Placeholder (starts with "__", e.g. "__first_1__") → Fetch first N RGB image IDs from the active dataset.
    • Filename (contains a . extension) → Resolve filename to its UUID.
    • UUID (standard UUID format) → Validate UUID exists in the dataset and is accessible under the user’s RLS policy.
    • Note: SQL is NEVER skipped for surface_analysis. Every identifier type requires a database round-trip to enforce RLS and catch non-existent images before calling the model.
  5. Invoke the Analyzer Once: After all IDs are resolved and validated, invoke rgb_surface_analyzer EXACTLY ONCE with ALL resolved UUIDs in a single call: rgb_surface_analyzer(image_ids=["<uuid1>", "<uuid2>", ...], dataset_slug="<slug>")
  6. Forward Results: Capture the structured JSON output of the rgb_surface_analyzer tool and include it in the output JSON under the rgb_surface_analyzer key. Keep retrieved_rows EMPTY.
  7. Agent Result Format: agent_result must use the exact delegation format: "Surface analysis delegated to rgb_surface_analyzer for <N> image(s) in <dataset_slug>." and must NOT contain image URLs, thumbnail_url values, or storage_path values.

State I/O

Reads

  • ctx.state["resolver_classifier"] (to read is_in_scope, is_ambiguous, query_type, target_image_ids, resolved_query)
  • ctx.state["global"]["active_filters"] (to read active_filters for multi-turn context retention and {active_filters} placeholder substitution)
  • ctx.state["global"]["query_type"] (to select the sliced database context schema injected via {database_context})
  • ctx.state["global"]["sql_tool_result"] (substituted into {sql_tool_result} by factory.py — no-op in current YAML template)
  • ctx.state["global"]["user_db_access"] (substituted into {user_db_access} by factory.py — no-op in current YAML template)
  • ctx.state["global"]["resolved_query"] (substituted into {resolved_query} placeholder)

Writes

  • ctx.state["database_specialist"]: The full validated DatabaseSpecialistResponse JSON dict.
  • ctx.state["rgb_surface_analyzer"]: Written by ADK’s own output_key mechanism for the rgb_surface_analyzer sub-agent (not by database_specialist_after_agent). The surface analysis data is accessible nested inside ctx.state["database_specialist"]["rgb_surface_analyzer"] after the callback runs.
  • ctx.state["global"]["generated_sql"]: Sourced from temp:real_execute_sql["query"] (captured by database_specialist_after_tool) or fallback to executed_queries[0].
  • ctx.state["global"]["sql_tool_result"]: Sourced from temp:real_execute_sql["rows"] (captured by database_specialist_after_tool) or fallback to retrieved_rows.
  • ctx.state["global"]["sql_execution_error"]: Sourced from errors[0] or None.
  • ctx.state["global"]["Executor_final_answer"]: Sourced from agent_result.

Scope and Persistence

  • Scope: Session-scoped. All state is written to ctx.state and manually persisted to session.db via persist_session_state in the after_agent callback.
  • Temp Prefixes: Ephemeral keys starting with temp: (like temp:real_execute_sql and temp:jwt_token) or __ (double underscore) are used in callbacks but are stripped before persistence to session.db to prevent leaking transient state or secrets.

Output Schema

  • Schema: DatabaseSpecialistResponse (defined in src/kavai/foundation_services/utility/models/database_models.py).
  • How it changes the request:
    • ADK automatically appends formatting instructions to the end of the prompt, forcing the model to output a JSON object matching the schema.
    • The schema is wrapped using make_gemini_compatible_model (which strips additionalProperties from the JSON schema to prevent Gemini constrained decoding crashes).

DatabaseSpecialistResponse Schema Fields

Field Type Default Description
executed_queries List[str] required The actual SQL queries executed to retrieve data. Empty for OOS/ambiguous.
agent_result str required A high-level summary of what was found in the database.
retrieved_rows List[Dict[str, Any]] [] The raw data rows retrieved from the database. Empty on surface-analysis turns.
is_complete bool required Whether the agent successfully executed all intended tasks/queries.
errors List[str] [] List of any errors or warnings encountered during database operations.
rgb_surface_analyzer Optional[RGBSurfaceAnalyzerResponse] None Optional surface analysis results from the rgb_surface_analyzer agent.

Callbacks

Callback Used? Behavior
before_agent ❌ N/A
before_model ❌ N/A
before_tool ✅ database_specialist_before_tool
after_tool ✅ database_specialist_after_tool
after_model ❌ N/A
after_agent ✅ database_specialist_after_agent
on_model_error ✅ on_model_error_callback

database_specialist_before_tool

Injects JWT into state before invoking rgb_surface_analyzer AgentTool.

  1. Isolated Runner Context: When database_specialist invokes rgb_surface_analyzer as an AgentTool, ADK creates an isolated Runner with a fresh async context where get_current_jwt() returns None.
  2. JWT Injection: We MUST pass the JWT through tool_context.state so the rgb_surface_analyzer_before_model callback can mint Azure SAS URLs.
  3. Ephemeral State: The JWT is written under the temp: prefix (temp:jwt_token), which is ADK’s canonical “ephemeral state” namespace. ADK’s session services (and our persist_session_state helper) skip temp: keys when committing to durable storage, so the token never lands in session.db.

database_specialist_after_tool

Captures the real execute_sql ground truth into ephemeral state.

  1. Ground Truth Capture: Writes temp:real_execute_sql = {"query": <actual SQL>, "rows": <actual rows>} so database_specialist_after_agent can populate global.generated_sql and global.sql_tool_result from the real tool response instead of the LLM’s self-reported executed_queries / retrieved_rows fields.
  2. Tool Filtering: Only fires for execute_sql; all other tools (like rgb_surface_analyzer) pass through unchanged.
  3. Robust Parsing: Parses the actual rows out of the untrusted-data wrapper using a two-tier strategy: Tier 1 tries json.loads() first; Tier 2 falls back to ast.literal_eval() to handle Python literals (since ExecuteSQLTool formats response data using repr(), producing Python literal syntax rather than valid JSON).

database_specialist_after_agent

Runs after the agent emits its response. Parses the JSON output, validates it, and writes to session state.

  1. Stale Analysis Bleed Guard:
    • Reads ctx.state["resolver_classifier"]["query_type"] (always fresh).
    • If query_type != "surface_analysis" but rgb_surface_analyzer is present in the output → strips it before persistence to prevent stale analysis results from bleeding into unrelated turns.
  2. State Updates:
    • Writes the validated dict to ctx.state["database_specialist"].
    • Extracts rgb_surface_analyzer to ctx.state["rgb_surface_analyzer"] if present.
    • Updates global state keys (generated_sql, sql_tool_result, sql_execution_error, and Executor_final_answer).
  3. Manual Persistence: Safely resolves the session ID and commits the updated state to session.db via persist_session_state.

on_model_error_callback

Handles raised Gemini exceptions (network, auth, quota, 5xx) after SDK-level retries are exhausted.

  1. Invocation Reset: Resets retry and model-swap counters at the start of each new pipeline invocation using a temporary invocation ID (temp:_llm_last_invocation_id).
  2. App-Level Retries: If the error is retryable, it increments the retry counter and directly retries the request with the same model (up to 2 retries).
  3. Model Swap: If retries are exhausted or the error is non-retryable but swap-eligible, it swaps the model to the next one in the fallback chain: gemini-3.1-flash-lite (default) → gemini-3.5-flash (first fallback) → gemini-2.5-pro (last resort).
  4. Graceful Degradation: If all models are exhausted or a terminal error occurs, it returns a graceful LlmResponse with a user-friendly error message, swallowing the exception.

Planner

  • Planner: None. database_specialist does not use a planner.

ADK Version Verified Against

  • Verified against: google-adk >= 1.26.0 (as specified in pyproject.toml).

Generation Config

Parameter Value Source / Notes
Model gemini-3.1-flash-lite Pinned in adk_agents_config.yaml
Temperature Default (1.0) Left to ADK/Gemini defaults
Max Output Tokens Default Left to ADK/Gemini defaults
Safety Settings Default Left to ADK/Gemini defaults
Stop Sequences None Left to ADK/Gemini defaults

Open Questions / Unverified Items

  • Unverified items:
    • The exact token limit or window size for include_contents: default in ADK 1.26.0+ is not explicitly documented in the codebase, but it is assumed to be managed dynamically by the ADK framework.
    • Whether the ADK framework performs any prompt-level optimization or compression on the conversation history before sending it to the model.