diff --git a/pyproject.toml b/pyproject.toml index 10c0d05..de2e437 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "zero-shot-agent" version = "0.1.0" -description = "Baseline agent skeleton — FastAPI + LangGraph + SQLite, provider-agnostic LLM. The capability slot is transform_text." +description = "Baseline agent skeleton — FastAPI + LangGraph + MSSQL, provider-agnostic LLM. The capability slot is for police data analysis." readme = "README.md" requires-python = ">=3.11" dependencies = [ @@ -14,11 +14,14 @@ dependencies = [ "langgraph>=0.2.28", "httpx>=0.27", "structlog>=24.1", + "pyodbc>=5.0", + "aioodbc>=0.16.0", ] [dependency-groups] dev = [ "pytest>=8.2", + "docker>=7.0.0", ] [tool.pytest.ini_options] @@ -27,4 +30,4 @@ pythonpath = ["."] addopts = "-q" [tool.uv] -package = false +package = false \ No newline at end of file diff --git a/spec/agent.md b/spec/agent.md index 49f33fa..e38bff0 100644 --- a/spec/agent.md +++ b/spec/agent.md @@ -1,218 +1,214 @@ # Agent -> Required when the project uses an agent framework. Delete this file if your project has no agent framework. -> -> If your project has no agent framework (e.g., a simple script or single-LLM API call), delete this file. -> - ---- - ## Agent Architecture Pattern - - - -| Pattern | Use when | -|---------|----------| -| **Single-agent loop** | One LLM drives a deterministic tool-call loop. No branches, no handoffs. | -| **Graph (LangGraph)** | Multi-step pipeline with conditional edges, checkpointing, or parallel nodes. | -| **Multi-agent** | Specialised sub-agents with distinct roles; orchestrator routes between them. | -| **Supervisor** | One supervisor LLM dispatches to worker agents based on task type. | -| **Human-in-the-loop** | Execution pauses at defined checkpoints for user review or approval. | - -**Chosen:** - ---- +**Chosen:** Graph (LangGraph) — The agent requires a multi-step pipeline with conditional edges for validation, retries, and error handling. A simple linear chain would not allow for re-planning when the validator rejects SQL or when the executor returns an error. The graph pattern gives us the needed control flow while keeping the implementation clear. ## LLM Provider & Model - - - | Agent / Node | Provider | Model ID | Rationale | -|-------------|----------|----------|-----------| -| | Anthropic | | | - -**Fallback behaviour:** +|--------------|----------|----------|-----------| +| Planner | OpenRouter | anthropic/claude-sonnet-3.5 | Strong reasoning for question refinement and ambiguity detection. | +| SQL Writer | OpenRouter | anthropic/claude-sonnet-3.5 | High-quality SQL generation; Sonnet 3.5 is strong at structured output. | +| Validator | OpenRouter | anthropic/claude-3-haiku | Fast, cheap check for SQL safety; does not need heavy reasoning. | +| Executor | OpenRouter | anthropic/claude-3-haiku | Not an LLM node; this node executes SQL. (Kept for completeness; actually a tool call.) | +| Answer Writer | OpenRouter | anthropic/claude-sonnet-3.5 | Good at natural-language summarization and optional chart suggestion. | -**Prompt strategy:** +**Fallback behaviour:** On 429/5xx, retry with exponential backoff (max 2 retries). On 401/403, surface an actionable error to check the API key. The agent does not degrade to a stub; it returns an error to the user. ---- +**Prompt strategy:** System/user split with structured output (JSON mode) for Planner, SQL Writer, and Validator. Answer Writer uses plain text with optional markdown for chart specifications. ## Tools & Tool Calling - - - | Tool name | Description | Inputs | Output | Side-effects | |-----------|-------------|--------|--------|--------------| -| | | | | | +| `introspect_schema` | Returns the database schema (tables and columns only) for the agent to use in SQL generation. | None | JSON: `{ tables: { name: string, columns: [{name: string, type: string}][] } }` | Read-only query to the database's information schema. | +| `execute_sql` | Runs a parameterized SQL query against the MSSQL database and returns the result set. | SQL string, parameters (optional) | `{ rows: any[], row_count: int, column_names: string[] }` | Read-only database query. Commits nothing; uses a read-only transaction. | +| `log_audit` | Writes an immutable audit record for the agent run. | `run_id`, `officer_id`, `question`, `refined_question`, `sql`, `row_count`, `latency_ms`, `result_hash` | None | Inserts into the audit table. | -**Tool selection strategy:** +**Tool selection strategy:** The LLM (Planner, SQL Writer, Validator) decides which tool to call based on the current state and the node's purpose. The graph wiring forces the sequence: Planner may call `introspect_schema`; SQL Writer calls `introspect_schema` to ground the SQL; Validator does not call tools (it validates the SQL string); Executor calls `execute_sql`; Answer Writer does not call tools. -**Tool failure handling:** - ---- +**Tool failure handling:** +- `introspect_schema`: On failure, set `state.error` and route to `handle_error`. +- `execute_sql`: On SQL error (e.g., syntax, permission), capture the error message in `state.error` and route to `handle_error`. On transient errors (timeout, deadlock), retry up to 2 times with backoff before setting error. +- `log_audit`: On failure, log the error but do not fail the run (audit is important but not critical to the answer). ## Agent State - - - ```python class AgentState(TypedDict): # Identity run_id: int # set at initialisation - + officer_id: str | None # from header or login (Phase 3) # Input - # ... # fields populated from the trigger - + question: str # raw user question + refined_question: str | None # after Planner (if clarification not needed) + clarification_needed: bool # True if Planner could not confidently refine + clarification_question: str | None # what to ask the user # Pipeline data (populated progressively by nodes) - # ... - + sql: str | None # generated SQL + sql_safe: bool # after Validator + execution_result: dict | None # from Executor: rows, count, etc. # Output - # ... # final result fields - + answer: str | None # natural-language answer + chart_spec: dict | None # optional Vega-Lite or similar spec # Control error: str | None # set by any node on fatal failure - checkpoint: str | None # last completed node (for resume) ``` ---- - ## Nodes / Steps - - - -### `node_[name]` - -**Reads from state:** - -**Writes to state:** - -**LLM call:** - -**External calls:** - -| System | Operation | On Failure | -|--------|-----------|------------| -| | | | - -**Behaviour:** - ---- +### `planner_node` +**Reads from state:** `question` +**Writes to state:** `refined_question`, `clarification_needed`, `clarification_question` +**LLM call:** Yes; prompt: "Given the user's question about police data, determine if it is clear enough to generate SQL. If ambiguous or missing key details (like time frame, location, or specific data type), produce a clarification question. Otherwise, produce a refined question that removes ambiguity and is ready for SQL generation. Output JSON with keys: `refined_question` (string), `clarification_needed` (boolean), `clarification_question` (string or null)." +**External calls:** May call `introspect_schema` to ground the refinement in available tables/columns. +**Behaviour:** The Planner decides whether the user's question is answerable as-is or needs clarification. It outputs a refined question if confident, or a clarification question to ask the user. This prevents the SQL Writer from guessing and producing unsafe or incorrect SQL. + +### `sql_writer_node` +**Reads from state:** `refined_question` (if `clarification_needed` is False) +**Writes to state:** `sql` +**LLM call:** Yes; prompt: "You are a SQL expert for a Microsoft SQL Server database with tables for FIR/CCTNS, HR/personnel, and logistics/property. Given the following schema: {schema} and the user's question: {refined_question}, generate a single, safe, read-only SQL query that answers the question. The query must use parameterized placeholders for any user-specific values (like officer_id) and must include a LIMIT or TOP clause to prevent large result sets. Output ONLY the SQL query in a code block marked as sql." +**External calls:** Calls `introspect_schema` to get the schema (tables and columns only). +**Behaviour:** The SQL Writer generates a parameterized SQL query using the provided schema. It is instructed to avoid `SELECT *`, to use explicit column lists, and to add a `LIMIT` or `TOP` clause (configurable, default 1000). The query is designed to be run by a read-only database user. + +### `validator_node` +**Reads from state:** `sql` +**Writes to state:** `sql_safe` +**LLM call:** Yes; prompt: "You are a SQL safety expert. Review the following SQL query for a read-only reporting tool. Determine if it is safe to execute. It must: 1) be a SELECT statement (no INSERT/UPDATE/DELETE/DROP/etc.), 2) not use `SELECT *`, 3) include a LIMIT or TOP clause, 4) only reference tables and columns from the provided schema, and 5) not contain any user-submitted strings directly in the SQL (must use parameters). Output JSON with a single key `sql_safe` (boolean)." +**External calls:** None (validates the SQL string directly). +**Behaviour:** The Validator checks the generated SQL for safety compliance. If it passes, `sql_safe` is set to True; if not, it is False and the error is captured in `state.error` (via an edge) to route to the handler. + +### `executor_node` +**Reads from state:** `sql` (only if `sql_safe` is True) +**Writes to state:** `execution_result` +**LLM call:** No; this is a tool call node. +**External calls:** Calls `execute_sql` with the SQL query. +**Behaviour:** Executes the SQL against the MSSQL database using SQLAlchemy with a read-only connection. Returns the result set, row count, and column names. On success, populates `execution_result`. On failure (SQL error, timeout, etc.), sets `state.error` with the error message. + +### `answer_writer_node` +**Reads from state:** `question`, `refined_question`, `execution_result` +**Writes to state:** `answer`, `chart_spec` +**LLM call:** Yes; prompt: "You are a helpful police data analyst. Given the original question: {original_question}, the refined question: {refined_question}, and the query results: {result_summary} (including row count and a sample of the data), produce a natural-language answer that directly addresses the user's question. If the result is suitable for a visualization (e.g., time series, categorical counts, geographic data), suggest a chart type and provide a Vega-Lite specification. Output JSON with keys: `answer` (string) and `chart_spec` (object or null)." +**External calls:** None (uses the execution result from state). +**Behaviour:** The Answer Writer crafts a final answer in plain language, optionally suggesting a visualization. It does not see the raw rows (only a summary and sample to avoid leaking data), but can suggest charts based on the shape of the data (e.g., if there's a date column and a numeric column, suggest a line chart). + +### `handle_error_node` +**Reads from state:** `error` +**Writes to state:** None (terminal node) +**LLM call:** No +**External calls:** May call `log_audit` to record the failed run. +**Behaviour:** Sets the final answer to an user-friendly error message (e.g., "I encountered an error while processing your request: {error}. Please try rephrasing your question.") and logs the audit entry with the error. ## Graph / Flow Topology - - - -``` -START - │ - ▼ -node_a ──(error)──► node_handle_error ──► END - │ - ▼ -node_b ──(condition)──► node_c - │ │ - │ ▼ - └──────────────────► node_finalize - │ - ▼ - END +```mermaid +graph TD + A[Start] --> B[planner_node] + B -->|clarification_needed=true| C[Ask User for Clarification] + C -->|User responds| B + B -->|clarification_needed=false| D[sql_writer_node] + D --> E[validator_node] + E -->|sql_safe=true| F[executor_node] + E -->|sql_safe=false| G[handle_error_node] + F --> H[answer_writer_node] + H --> I[End] + G --> I + style C fill:#f9f,stroke:#333,stroke-dasharray: 5 5 ``` **Conditional edges:** - -| Source node | Condition | Target | -|-------------|-----------|--------| -| | | | - ---- +- From `planner_node`: if `state.clarification_needed` is True, go to clarification step (which is outside the graph; the API will return a clarification request to the user and wait for a new input that resets the graph with the new question). +- From `validator_node`: if `state.sql_safe` is True, proceed to `executor_node`; else, go to `handle_error_node`. +- All other edges are sequential. ## Memory & Context - - - | Scope | Mechanism | What is stored | |-------|-----------|----------------| -| **Within a run** | LangGraph state | All in-progress data | -| **Across runs** | | | -| **Conversation** | | | +| **Within a run** | LangGraph state | All in-progress data (question, refined question, SQL, execution result, etc.) | +| **Across runs** | Database table (`officer_reports`) | Pinned reports (saved queries) and their results (optional cache) for the officer. | +| **Conversation** | None (stateless by design; each question is independent) | Not applicable; the agent does not maintain chat history across turns in Phase 1. (Phase 2 may add sidebar recency.) | -**Context window management:** - ---- +**Context window management:** The agent uses schema-only context (tables and columns) for the LLM, which is typically under 2k tokens. No conversation history is included in the prompt, keeping the context window lean. ## Human-in-the-Loop Checkpoints - - - | Checkpoint | What is shown to the user | Expected user action | Timeout / default | -|------------|--------------------------|----------------------|-------------------| -| | | | | - ---- +|------------|---------------------------|----------------------|-------------------| +| **Clarification** | A plain-language question asking for missing details (e.g., "To answer that, I need to know: which district are you interested in?") | Provide the missing information or rephrase the question. | 2 minutes; after timeout, the agent proceeds with the original question and may produce a best-effort answer or error. | ## Error Handling & Recovery - - - -**Node-level:** +**Node-level:** Each node catches its own exceptions; fatal errors set `state["error"]` and route to `handle_error_node` via the graph's error edges (implicit in the conditional edges above; any node can set error and the next conditional edge will route to handler based on error presence). **Graph-level (handle_error node):** - Reads: `state.error`, `state.run_id` -- Updates DB: run status → "failed", `error_message`, `completed_at` +- Updates DB: audit table record status → "failed", `error_message`, `completed_at` - Logs error with `run_id` context -- Terminates graph - -**Resume / retry strategy:** +- Terminates graph (returns an error answer to the user) -**Partial failure:** +**Resume / retry strategy:** The agent does not support resuming a failed run from a checkpoint; each question is a new run. Retries are handled at the tool level (e.g., `execute_sql` retries on transient DB errors). ---- +**Partial failure:** If a non-critical step fails (e.g., suggesting a chart spec fails), the agent logs the error but continues to provide the answer. The chart spec is optional. ## Observability - - - | Signal | What | Where | |--------|------|-------| -| **Trace** | One trace per run, one span per node | | -| **LLM calls** | Prompt tokens, completion tokens, latency, model | | -| **Tool calls** | Tool name, inputs, success/error, latency | Structured log | -| **Run outcome** | Status, total duration, error if any | DB + structured log | - ---- +| **Trace** | One trace per run, one span per node | stdout via structlog (JSON format) | +| **LLM calls** | Prompt tokens, completion tokens, latency, model | Structured log (includes `llm_call` event) | +| **Tool calls** | Tool name, inputs, success/error, latency | Structured log (includes `tool_call` event) | +| **Run outcome** | Status, total duration, error if any | Database audit table + structured log | ## Concurrency Model +- **Run isolation:** Each API call (`/api/ask`) creates a new `run_id` and runs the agent graph in isolation. Concurrent runs are allowed; they are distinguished by `run_id`. +- **Parallel nodes within a run:** The current graph is linear; no nodes run in parallel. (Future phases could parallelize schema introspection and LLM calls if needed.) +- **Checkpointing:** None required for Phase 1 (stateless runs). If we add conversation history in Phase 2, we may add a checkpointer. - - -- **Run isolation:** -- **Parallel nodes within a run:** -- **Checkpointing:** - ---- - -## Graph Assembly (`agent/graph.py`) - - - +## Graph Assembly (`src/graph/agent.py`) ```python -graph = StateGraph(AgentState) - -graph.add_node("node_a", node_a) -graph.add_node("node_b", node_b) -graph.add_node("finalize", node_finalize) -graph.add_node("handle_error", node_handle_error) - -graph.set_entry_point("node_a") - -graph.add_conditional_edges( - "node_a", - lambda s: "handle_error" if s.get("error") else "node_b", +from langgraph.graph import StateGraph, END +from .state import AgentState +from .nodes import ( + planner_node, + sql_writer_node, + validator_node, + executor_node, + answer_writer_node, + handle_error_node, ) -graph.add_edge("node_b", "finalize") -graph.add_edge("finalize", END) -graph.add_edge("handle_error", END) - -compiled_graph = graph.compile() +def build_graph(): + graph = StateGraph(AgentState) + + # Add nodes + graph.add_node("planner", planner_node) + graph.add_node("sql_writer", sql_writer_node) + graph.add_node("validator", validator_node) + graph.add_node("executor", executor_node) + graph.add_node("answer_writer", answer_writer_node) + graph.add_node("handle_error", handle_error_node) + + # Set entry point + graph.set_entry_point("planner") + + # Add edges + graph.add_conditional_edges( + "planner", + lambda state: "ask_user" if state.get("clarification_needed") else "sql_writer", + { + "ask_user": "ask_user", # This is handled outside the graph by the API + "sql_writer": "sql_writer", + }, + ) + graph.add_edge("sql_writer", "validator") + graph.add_conditional_edges( + "validator", + lambda state: "executor" if state.get("sql_safe") else "handle_error", + { + "executor": "executor", + "handle_error": "handle_error", + }, + ) + graph.add_edge("executor", "answer_writer") + graph.add_edge("answer_writer", END) + graph.add_edge("handle_error", END) + + return graph.compile() + +# The actual graph instance used by the runner +agent_graph = build_graph() ``` + +Note: The `ask_user` step is not a node in the graph; the API layer handles it by returning a clarification request to the user and waiting for a new POST that restarts the graph with the updated question. \ No newline at end of file diff --git a/spec/api.md b/spec/api.md index 442a58b..ca289a9 100644 --- a/spec/api.md +++ b/spec/api.md @@ -1,41 +1,117 @@ # API -> Fill in this section — see comments below. - ---- - ## API Style - - +REST ## Endpoints / Commands +### `POST /api/ask` +**Purpose:** Accept a natural-language question from an officer, run it through the agent pipeline, and return the answer with optional SQL and chart. - +**Request:** +```json +{ + "question": "string", + "officer_id": "string (optional, from header in Phase 3)" +} +``` -### `` +**Response:** +```json +{ + "answer": "string (natural-language response)", + "sql": "string (the executed SQL, for transparency)", + "execution_time_ms": "integer", + "row_count": "integer", + "chart_spec": "object (optional Vega-Lite specification)", + "status": "string: "success" | "clarification_needed" | "error" +} +``` -**Purpose:** +**Error cases:** +| Status | Condition | +|--------|-----------| +| 400 | Missing or invalid `question` in request body. | +| 401 | Missing or invalid `officer_id` (when required by phase). | +| 429 | Agent rate-limited (upstream LLM or DB); retry after delay. | +| 500 | Internal agent error (see `answer` field for message). | +### `POST /api/pin` +**Purpose:** Save a question and its result as a pinned report for quick reuse. **Request:** ```json +```json { - "": "" + "question": "string", + "officer_id": "string" + 1": "{ + "answer": "string": "integer", + "string: "string thetime_ms": integer", + "row_count": "integer", + "created_at": "ISO timestamp" +}' + } } ``` **Response:** ```json { - "": "" + "id": "string (uuid)", + "message": "Report pinned successfully" } ``` **Error cases:** | Status | Condition | |--------|-----------| -| 400 | | -| 500 | | +| 400 | Missing required fields. | +| 401 | Invalid or missing officer_id. | +| 500 | Database error while saving. | +### `GET /api/reports` +**Purpose:** List recent and pinned reports for an officer. -## Authentication +**Request:** None (officer_id from header or query) - +**Response:** +```json +{ + "recent": [ + { + "id": "string", + "question": "string", + "answer": "string", + "created_at": "ISO timestamp" + } + ], + "pinned": [ + { + "id": "string", + "question": "string", + "answer": "string", + "pinned_at": "ISO timestamp" + } + ] +} +``` + +**Error cases:** +| Status | Condition | +|--------|-----------| +| 401 | Invalid or missing officer_id. | +| 500 | Database error while fetching. | +### `GET /api/health` +**Purpose:** Simple health check; returns 200 if the service is up and can reach the DB and LLM API (without validating keys). + +**Response:** +```json +{ + "status": "ok", + "timestamp": "ISO timestamp", + "services": { + "database": "reachable" | "unreachable", + "llm_api": "reachable" | "unreachable" + } +} +``` +## Authentication +In Phase 1, the officer_id is accepted as an optional header (`X-Officer-Id`) or query parameter for audit logging; no real authentication is performed. In Phase 3, this will be replaced with proper JWT validation against an identity provider. \ No newline at end of file diff --git a/spec/architecture.md b/spec/architecture.md index e5e3bd9..c161e40 100644 --- a/spec/architecture.md +++ b/spec/architecture.md @@ -1,68 +1,73 @@ # Architecture -> Fill in this section — see comments below. - ---- - ## System Overview - - +The UP Police Data Analyst agent is a read-only interface over a Microsoft SQL Server database containing FIR/CCTNS, HR/personnel, and logistics/property data. Authorized officers interact via a chat-based web UI to ask natural-language questions about crime, personnel, or logistics data. The agent translates questions into safe SQL queries, executes them against the database, and returns answers with the executed SQL, timing, and optional visualizations. All interactions are logged immutably for audit and compliance. ## Component Map - - - -``` -[Component A] - ↓ -[Component B] ←→ [External Service] - ↓ -[Component C] +```mermaid +graph TD + A[Officer's Browser] -->|HTTPS/WSS| B(FastAPI Backend) + B -->|SQLAlchemy + pyodbc| C[Microsoft SQL Server] + B -->|HTTP (JSON)| D[OpenRouter LLM API] + B -->|Append-only| E[Audit Table (SQL Server)] + B -->|Serve at /app| F[Static Frontend (HTML/JS)] ``` ## Layers - - - | Layer | Responsibility | |-------|----------------| -| | | +| Presentation | Static HTML/JS frontend served at `/app`; handles user input, displays answers with collapsible SQL/chart/table panels, and shows recent/pinned reports sidebar. | +| API | FastAPI endpoints (`/api/ask`, `/api/pin`, `/api/recent`) that validate input, invoke the agent loop, and return structured responses. | +| Agent Loop | LangGraph-based reasoning loop: Planner → SQL Writer → Validator → Executor → Answer Writer. Each step is a node with conditional edges for error handling and retries. | +| Tools & Storage | SQLAlchemy ORM with MSSQL driver (`pyodbc`/`aioodbc`) for schema introspection and query execution; audit table for immutable logging. | +| External | OpenRouter LLM API (primary: Anthropic Claude Sonnet 3.5, fallback: Claude Haiku) for reasoning and SQL generation. | ## Data Flow - - - -1. Trigger: -2. -3. -4. Output: +1. Trigger: Officer types a question in the chat input and presses Enter. +2. Frontend sends POST `/api/ask` with `{ question, officer_id }` (officer_id from headers or login context). +3. API validates input, creates an agent run record, and invokes the LangGraph agent loop. +4. Planner node refines the question (if ambiguous) and outputs a refined question or clarification request. +5. SQL Writer node generates a parameterized SQL query using the cached schema (tables/columns only, no row data). +6. Validator node checks the SQL for safety: ensures read-only, mandates `LIMIT/TOP`, blocks `SELECT *`, and validates against a read-only DB user policy. +7. Executor node runs the SQL against the MSSQL database via SQLAlchemy, capturing row count and latency. +8. Answer Writer node formats the result into a natural-language answer, optionally generating a chart/table if the data is suitable. +9. The run record is updated with the answer, SQL, timing, and an immutable audit log entry is written. +10. API returns the answer, SQL, timing, and optional chart spec to the frontend. +11. Frontend displays the answer with collapsible panels for SQL and chart/table, and updates the recent/pinned sidebar. ## External Dependencies - - - | Dependency | Purpose | Failure Mode | |------------|---------|--------------| -| | | | +| Microsoft SQL Server | Stores FIR/CCTNS, HR/personnel, and logistics/property data; source of truth for all queries. | Database downtime causes 503 errors; agent returns "Database unavailable" error. | +| OpenRouter LLM API | Provides reasoning and SQL generation via Anthropic Claude models. | API downtime or rate limiting causes fallback to cheaper model or explicit error; agent may ask user to rephrase or try later. | +| Officer Auth System (future) | Provides JWT tokens for officer identification (Phase 3). | Missing/invalid token results in 401 error; Phase 1 uses stubbed officer_id from header. | ## Stack - -> This project's concrete technology choices (captured at intake, filled by the spec-writer). The generic, every-project rules — model-naming, DB driver, dev port, test environment — live in `harness/patterns/tech-stack.md`; this section is only what **this** project picked. - -- **Language:** -- **Agent framework:** -- **LLM provider + model:** -- **Backend:** -- **Database + ORM:** -- **Frontend:** -- **Dependency management:** +- **Language:** Python 3.11 +- **Agent framework:** LangGraph +- **LLM provider + model:** OpenRouter (Anthropic Claude Sonnet 3.5 for planning, Claude Haiku for cheap fallback) +- **Backend:** FastAPI +- **Database + ORM:** Microsoft SQL Server + SQLAlchemy 2.0 (with `pyodbc`/`aioodbc` driver) +- **Frontend:** Zero-build static HTML/CSS/JS (served at `/app`) +- **Dependency management:** uv + pyproject.toml | Key library | Version | Purpose | |-------------|---------|---------| -| | | | - -**Avoid:** +| langgraph | 0.1.x | Agent reasoning loop | +| fastapi | 0.109.x | Web API and server | +| sqlalchemy | 2.0.x | ORM and DB abstraction | +| pyodbc | 5.x | MSSQL driver (sync) | +| aioodbc | 1.x | Async MSSQL driver (if needed) | +| pydantic | 2.x | Data validation and settings | +| python-dotenv | 1.x | Load environment variables from `.env` | +| structlog | 23.x | Structured logging | +| alembic | 1.13.x | Database migrations | + +**Avoid:** +- Raw SQL string concatenation (use ORM or parameterized queries only). +- Storing secrets in code or logs (use `.env` and secret hygiene rules). +- Using SQLite as a substitute for production PostgreSQL/MSSQL in tests (tests must use the real driver). +- Hardcoding model names or DB connection strings (load from config/env). ## Deployment Model - - +The agent runs as a long-running service via `uv run python -m src` (which starts Uvicorn on port 8001). It is designed for deployment behind a reverse proxy (NGINX/Apache) with TLS termination in production. In development, it runs locally on `http://localhost:8001`. \ No newline at end of file diff --git a/spec/capabilities/ask_question.md b/spec/capabilities/ask_question.md new file mode 100644 index 0000000..beb161c --- /dev/null +++ b/spec/capabilities/ask_question.md @@ -0,0 +1,45 @@ +# Capability: ask_question +## What It Does +Converts a natural-language question about police data into a safe SQL query, executes it against the read-only MSSQL database, and returns an answer with the executed SQL, timing, and optional visualizations. + +## Inputs +| Input | Type | Source | Required | +|-------|------|--------|----------| +| question | string | user (chat input) | yes | +| officer_id | string | header or auth (Phase 3) | no (Phase 1: optional for audit) | + +## Outputs +| Output | Type | Destination | +|--------|------|-------------| +| answer | string | agent message bubble | +| sql | string | collapsible "View SQL" panel | +| row_count | integer | part of answer text | +| latency_ms | integer | part of answer text | +| chart_spec | dict (optional) | collapsible "View Chart" panel | +| status | string ("success", "clarification_needed", "error") | determines message bubble type | + +## External Calls +| System | Operation | On Failure | +|--------|-----------|------------| +| MSSQL (read-only) | introspect_schema (tables/columns only) | set error, route to handle_error | +| MSSQL (read-only) | execute_sql (parameterized SELECT) | capture error, set state.error | +| OpenRouter LLM API | planner (question refinement) | retry 2x with backoff, then error | +| OpenRouter LLM API | sql_writer (SQL generation) | retry 2x with backoff, then error | +| OpenRouter LLM API | validator (SQL safety check) | retry 2x with backoff, then error | +| OpenRouter LLM API | answer_writer (natural-language summary) | retry 2x with backoff, then error | +| AuditLog table | insert immutable audit log | log error but do not fail run | + +## Business Rules +- The agent must never include raw row data in any LLM prompt (schema-only by default). +- All SQL queries must be read-only (SELECT only) and include a LIMIT or TOP clause. +- The agent must use parameterized queries; never concatenate user input into SQL. +- If the planner determines the question is ambiguous, it must ask for clarification before proceeding. +- Every query must be logged immutably in the audit table for compliance. +- The officer_id (when available) must be included in the audit log for accountability. + +## Success Criteria +- [ ] Given a clear natural-language question, the agent returns a correct answer with SQL and timing within 15 seconds p50. +- [ ] Given an ambiguous question, the agent asks for clarification instead of guessing. +- [ ] Given a potentially unsafe question (e.g., attempting INSERT), the agent rejects it with an error. +- [ ] The audit log contains an entry for every query with officer_id, question, SQL, row count, latency, and result hash. +- [ ] Raw data never appears in LLM prompts or logs (verified via observability). \ No newline at end of file diff --git a/spec/capabilities/index.md b/spec/capabilities/index.md index 7455bda..64f878a 100644 --- a/spec/capabilities/index.md +++ b/spec/capabilities/index.md @@ -11,13 +11,12 @@ A capability is a single, discrete action or behavior the agent performs. Exampl - "Draft a personalized email given a lead profile" - "Send a Slack notification when a threshold is crossed" -## Capabilities in This Project - - +## Capabilities In This Project | Capability | File | |-----------|------| -| | [name.md](name.md) | +| ask_question | [ask_question.md](ask_question.md) | +| pin_report | [pin_report.md](pin_report.md) | ## How to Add a New Capability @@ -35,4 +34,4 @@ Each capability file should answer: - **Outputs** (what it produces) - **External calls** (APIs, LLMs, databases it touches) - **Error cases** (what can go wrong and how it's handled) -- **Success criteria** (how we test it) +- **Success criteria** (how we test it) \ No newline at end of file diff --git a/spec/capabilities/pin_report.md b/spec/capabilities/pin_report.md new file mode 100644 index 0000000..f957b5c --- /dev/null +++ b/spec/capabilities/pin_report.md @@ -0,0 +1,38 @@ +# Capability: pin_report +## What It Does +Allows an officer to save a question and its result (or a link to re-run it) for quick access later, pinned to their sidebar. + +## Inputs +| Input | Type | Source | Required | +|-------|------|--------|----------| +| question | string | from a prior ask_question turn | yes | +| officer_id | string | header or auth | yes | +| nickname | string | user-defined label for the pin | no | + +## Outputs +| Output | Type | Destination | +|--------|------|-------------| +| pin_id | integer | returned to frontend | +| nickname | string | displayed in sidebar | +| pinned_at | timestamp | stored in database | +| status | string | "success" or "error" | + +## External Calls +| System | Operation | On Failure | +|--------|-----------|------------| +| OfficerReport table | INSERT or UPDATE | set error, return error response | +| AuditLog table | INSERT (pin event) | log error but do not fail pin | + +## Business Rules +- A pinned report stores the original question (not the result set) so it can be re-run against fresh data. +- Officers can pin the same question multiple times with different nicknames. +- Pinned reports are private to the officer unless sharing is enabled (Phase 2+). +- The pin does not store the result set to avoid stale data; re-running gets current data. +- Each pin action is logged in the audit trail for accountability. + +## Success Criteria +- [ ] After a successful ask_question, the officer can click a pin icon to save the query. +- [ ] The pinned query appears in the sidebar under "Pinned" with the given nickname (or a default). +- [ ] Clicking the pinned query re-runs the question and displays the current result. +- [ ] The audit log records the pin action with officer_id and question. +- [ ] Pins persist across page refreshes and sessions. \ No newline at end of file diff --git a/spec/data.md b/spec/data.md index e331007..2fa6db0 100644 --- a/spec/data.md +++ b/spec/data.md @@ -1,34 +1,54 @@ # Data Model -> Fill in this section — see comments below. - ---- - ## Storage Technology - - +Microsoft SQL Server (MSSQL) is used as the source of truth for all police data (FIR/CCTNS, HR/personnel, logistics/property). The agent connects to a read-only database user with permissions to SELECT from the relevant schemas. An audit table is added to log all queries for compliance. ## Entities - - - -### Entity: - - - -| Field | Type | Required | Description | -|-------|------|----------|-------------| -| id | | yes | Primary key | -| | | | | - -### Relationships - - +### Entity: AuditLog +Purpose: Immutable record of every analyst query for audit and compliance (PVP/IT Act). +|| Field | Type | Required | Description || +|-------|------|----------|-------------|| +| id | BIGINT IDENTITY(1,1) | yes | Primary key || +| officer_id | VARCHAR(50) | yes | Identifier of the officer who made the request (from header or auth) || +| question_text | TEXT | yes | The natural-language question posed by the officer || +| refined_question | TEXT | no | The question after any clarification steps (if applicable) || +| executed_sql | TEXT | yes | The SQL query that was run against the database || +| row_count | INT | yes | Number of rows returned by the query || +| execution_time_ms | INT | yes | Time taken to execute the SQL query (in milliseconds) || +| result_hash | CHAR(64) | yes | SHA-256 hash of the query result (for tamper evidence) || +| llm_model_used | VARCHAR(50) | yes | The LLM model that generated the SQL (e.g., "anthropic/claude-sonnet-3.5") || +| llm_prompt_tokens | INT | yes | Number of prompt tokens sent to the LLM || +| llm_completion_tokens | INT | yes | Number of completion tokens returned by the LLM || +| status | VARCHAR(20) | yes | Outcome: "success", "clarification_needed", "error" || +| error_message | TEXT | no | If status is "error", the error message from the agent || +| created_at | DATETIME2(3) | yes | Timestamp when the query was received (UTC) || +| completed_at | DATETIME2(3) | no | Timestamp when the query finished processing || + +### Entity: OfficerReport +Purpose: Stores pinned (saved) reports for officers to reuse quickly. +|| Field | Type | Required | Description || +|-------|------|----------|-------------|| +| id | UNIQUEIDENTIFIER | yes | Primary key (NEWID()) || +| officer_id | VARCHAR(50) | yes | Identifier of the officer who pinned the report || +| question_text | TEXT | yes | The natural-language question that was pinned || +| answer_text | TEXT | yes | The answer generated at the time of pinning (for quick view) || +| pinned_at | DATETIME2(3) | yes | Timestamp when the report was pinned || +| expires_at | DATETIME2(3) | no | Optional expiry for temporary pins (null = permanent) || + +## Relationships +- AuditLog.officer_id references OfficerReport.officer_id (logical; no FK as officer info may come from external auth). +- One officer can have many audit log entries and many pinned reports. ## Data Lifecycle - - +- **AuditLog**: Insert-only; rows are never updated or deleted. Partitioning by month/year may be applied for performance in production. +- **OfficerReport**: Insert when a user pins a report; update if the user refreshes the pin (e.g., with new data); delete when the user unpins or when expires_at is passed (cleanup job). ## Sensitive Data +The agent is designed to **never** store or transmit raw row data from the protected tables (FIR, HR, logistics). The AuditLog stores only: +- Aggregates (row count, execution time) +- A hash of the result set (for tamper evidence, not reconstruction) +- The SQL query (which is read-only and parameterized) + +Fields that contain PII (e.g., victim names, addresses, officer personal details) are **never** included in LLM prompts, logs, or API responses beyond what is strictly necessary for the aggregate answer. The database user used by the agent has SELECT access only to the necessary views/tables, and row-level security (to be added in Phase 3) will further restrict data to the officer's jurisdiction. - +All logs (application and audit) avoid storing raw PII; any error messages that might contain data are sanitized before logging. \ No newline at end of file diff --git a/spec/roadmap.md b/spec/roadmap.md index 03ae242..9a9dd2b 100644 --- a/spec/roadmap.md +++ b/spec/roadmap.md @@ -1,62 +1,92 @@ # Roadmap -> Fill in each section. Run `/zero-shot-build [your idea]` to have it filled automatically. - ---- - ## What This Agent Does - - +The UP Police Data Analyst agent enables authorized officers to query FIR/CCTNS, HR/personnel, and logistics/property data in a read-only Microsoft SQL Server database using natural language. It returns answers with the executed SQL, timing, and optional visualizations (charts/tables), while ensuring raw data never leaves the DB and all interactions are audited per PVP/IT Act requirements. ## Who Uses It - - +Primary users are IPS officers (senior leadership), SHOs (station house officers), and beat constables who need ad-hoc insights during shifts or for reports. The analytics team uses it for scheduled queries and dashboard building. ## Core Problem Being Solved - - +Manual SQL querying is slow, error-prone, and limited to technical users. Officers currently rely on the analytics team for data requests, causing delays. This agent puts safe, governed data access directly in officers' hands via a chat interface. ## Success Criteria - - - -- [ ] -- [ ] -- [ ] +- [ ] Officers can ask natural-language questions about crime, personnel, or logistics data and receive accurate answers with SQL and timing within 15 seconds p50. +- [ ] Raw database rows are never included in LLM prompts; only schema and aggregates are used unless explicitly opted-in for disambiguation. +- [ ] Every query is logged immutably (user ID, question, SQL, row count, latency, result hash) for audit and compliance. +- [ ] The agent handles ambiguous queries by requesting clarification or showing a refined question before execution. +- [ ] Saved/pinned reports persist across sessions and can be shared within teams. ## What This Agent Does NOT Do (Out of Scope) - - +- Write or modify data in the database (strictly read-only). +- Provide real-time alerting or predictive analytics (Phase 2+). +- Integrate with external systems like crime mapping or facial recognition (future phases). +- Replace the need for approved analytics team reports; it supplements ad-hoc inquiry. ## Key Constraints - - +- Latency: p50 ≤ 15 seconds end-to-end (user question to answer display). +- DB load: Minimize via schema-only LLM context by default, mandatory LIMIT/TOP clauses, and connection pooling (max 5). +- Privacy: Sensitive PII (victims, addresses, chargesheets) never leaves the SQL Server; audit logs are tamper-evident. +- Access: Role-based row-level filtering (district/unit) enforced at the DB layer via service policies; officer identity via header. +- Compliance: Full audit trail meets PVP Act and IT Act requirements for data access logging. ## Phases of Development - - -> **Phase 1 is the smallest first-time-right user-testable win.** It must work perfectly the first time the user tests it — zero rough edges on the tested path. Its backend is minimal but REAL on the one core path (no fake data on the tested path). Its frontend is visually complete: real UI for the one working path PLUS clearly-labelled NON-FUNCTIONAL stubs for everything coming later, so the user sees the vision (a stub must never be mistaken for a bug). Each later phase wires those stubs into real functionality, one increment at a time. - -### Phase 1 — - -- **Goal:** -- **Independent slices (parallel build units):** - - `slice-a` (backend) — - - `slice-b` (frontend) — -- **Key surfaces / files:** -- **Gate command:** -- **How the user tests it (handoff seed):** - -### Phase 2 — - -- **Goal:** +### Phase 1 — Core Inquiry & Pinning +- **Goal:** Enable a single natural-language question → SQL → executed → answered flow with collapsible SQL/chart/table panels and a sidebar of recent/pinned reports. - **Independent slices (parallel build units):** - - `slice-a` (backend) — - - `slice-b` (frontend) — -- **Key surfaces / files:** -- **Gate command:** -- **How the user tests it (handoff seed):** - - - + - `slice-schema` (backend) — introspect MSSQL schema and cache it for LLM context; deps: none + - `slice-agent-loop` (backend) — implement the planner→sql_writer→validator→executor→answer_writer graph; deps: slice-schema + - `slice-audit` (backend) — add immutable audit table and logging middleware; deps: none + - `slice-frontend` (frontend) — chat input, collapsible answer panels (SQL/chart/table), sidebar with recent queries and pinned reports; deps: none +- **Key surfaces / files:** + - Backend: `src/db/mssql/` (schema introspection, connection), `src/graph/` (agent nodes/edges), `src/prompts/data_analyst.md`, `src/api/analyst.py` (endpoint), `src/db/models.py` (audit table), `alembic/` (migration for audit table) + - Frontend: `frontend/public/` (index.html, styles.css, app.js for chat UI) +- **Gate command:** `uv run alembic upgrade head && uv run pytest tests/phase1 -v` +- **How the user tests it (handoff seed):** + 1. Start the server: `.venv/bin/python -m src` (or `uv run python -m src`) from repo root. + 2. Open `http://localhost:8001/app/` in a browser. + 3. Type a natural-language question like "Show total FIRs registered last week in Lucknow district". + 4. Verify the agent returns an answer with: + - The executed SQL (collapsible panel) with a `LIMIT` or `TOP` clause. + - Row count and timing. + - Optional chart/table if the result is suitable. + - The question and answer appear in the sidebar under "Recent". + 5. Click the pin icon to save the query; verify it appears under "Pinned" and persists after a page refresh. + 6. Confirm no raw data appears in the agent's thinking (visible via observability logs) — only schema or aggregates. + +### Phase 2 — Sharing & Visualization +- **Goal:** Wire the pinned reports into real functionality: export, scheduled queries, and richer visualizations (time-series, heatmaps). +- **Independent slices (parallel build units):** + - `slice-export` (backend) — add CSV/XLSX export for pinned reports; deps: slice-agent-loop + - `slice-scheduler` (backend) — add cron-like scheduling for pinned reports; deps: slice-agent-loop + - `slice-viz` (backend) — add chart-type selection (bar, line, pie, heatmap) and rendering; deps: slice-agent-loop + - `slice-frontend-viz` (frontend) — enhance answer panels with chart controls and export buttons; deps: slice-frontend +- **Key surfaces / files:** + - Backend: new API endpoints for export/scheduling, chart generation logic. + - Frontend: updated collapsible panels with chart type dropdown and export. +- **Gate command:** `uv run pytest tests/phase2 -v` +- **How the user tests it (handoff seed):** + 1. After Phase 1 is working, create a pinned report. + 2. Click "Export" to download CSV/XLSX; verify file contains expected data. + 3. Set a schedule (e.g., daily at 9 AM) and verify the system records it. + 4. Change chart type in the panel and verify the visualization updates accordingly. + 5. Confirm audit logs still capture all interactions. + +### Phase 3 — Role-Based Access & Advanced Audit +- **Goal:** Implement district/unit row-level security, JWT-based officer authentication, and real-time alerting on anomalous patterns (e.g., sudden spike in complaints). +- **Independent slices (parallel build units):** + - `slice-auth` (backend) — validate JWT tokens and map to officer ID/unit; deps: slice-agent-loop + - `slice-rls` (backend) — add row-level security via MSSQL session context or views; deps: slice-auth + - `slice-alert` (backend) — add anomaly detection rules and notification hooks; deps: slice-agent-loop + - `slice-frontend-auth` (frontend) — login page and unit/district selector; deps: slice-frontend +- **Key surfaces / files:** + - Backend: auth middleware, RLS implementation, alert rule engine. + - Frontend: login UI and context-aware query routing. +- **Gate command:** `uv run pytest tests/phase3 -v` +- **How the user tests it (handoff seed):** + 1. Log in as an officer from a specific district. + 2. Ask a question that should return only data from that district (e.g., "Show pending cases"). + 3. Verify results are filtered to the officer's unit. + 4. Attempt a cross-district query and confirm it returns only permitted data or an error. + 5. Trigger an anomaly (e.g., upload a test CSV with a spike) and verify an alert is generated. + 6. Confirm audit logs record the officer ID for every query. \ No newline at end of file diff --git a/spec/ui.md b/spec/ui.md index 15219c3..c8ab0c7 100644 --- a/spec/ui.md +++ b/spec/ui.md @@ -1,32 +1,33 @@ # UI -> **Boilerplate status:** Delete this file if the agent has no UI. Otherwise, filled in by the spec-writer sub-agent. - ---- - ## UI Type - - +Web dashboard (chat interface) served at `/app`. ## Views / Screens - - - -### Screen: - -**Purpose:** +### Screen: Analyst Chat +**Purpose:** The main interface for officers to ask questions about police data, view answers with supporting details (SQL, charts), and manage recent/pinned reports. **Key elements:** -- -- - -**Actions available:** -- +- Chat input box at the bottom: single-line text field with placeholder "Ask about FIRs, personnel, or logistics..." +- Send button (paper plane icon) to the right of the input. +- Chat history area above the input, displaying messages in bubbles: + - User messages: left-aligned, with a subtle background. + - Agent messages: right-aligned, containing: + - The answer in plain text. + - A collapsible section labeled "View SQL" that, when expanded, shows the executed SQL query in a monospace block with a copy button. + - A collapsible section labeled "View Chart" (if a chart is applicable) that shows a visualization (e.g., bar, line, pie) of the result data. + - A collapsible section labeled "View Table" (if the result is tabular) that shows the data in a searchable, sortable table with pagination. + - A pin icon (bookmark) in the top-right of the agent message bubble to save the query and result as a pinned report. +- Sidebar on the left (collapsible via a hamburger menu at top-left): + - Section "Recent": list of the last 10 questions asked by the officer (most recent at top), each showing a truncated question and timestamp; clicking re-populates the input with that question. + - Section "Pinned": list of pinned reports (saved by the officer), each showing a truncated question and a pin icon; clicking re-populates the input with that question. + - Section "Dashboard" (Phase 2+): placeholder for pre-built reports. ## Error States - - +- **Empty state:** When no messages have been exchanged, the chat area shows a friendly illustration and text: "Ask me anything about FIRs, personnel, or logistics data. For example: 'Show total FIRs registered last week in Lucknow district'." +- **Loading state:** When the agent is processing a question, the input is disabled, and a spinner appears inside the send button with the text "Thinking...". +- **Error state:** If the agent returns an error, the message bubble shows a red warning icon and the error message in plain text, with a suggestion to rephrase the question. +- **Clarification state:** If the agent needs more information, the message bubble shows a yellow info icon and a question asking for the missing details (e.g., "To answer that, I need to know: which district are you interested in?"), followed by suggested options or a text field. ## Tech Stack - - +HTML5, CSS3, vanilla JavaScript (no build step). Served as static files by the FastAPI backend at `/app`. \ No newline at end of file diff --git a/src/db/mssql/schema.py b/src/db/mssql/schema.py new file mode 100644 index 0000000..86cb368 --- /dev/null +++ b/src/db/mssql/schema.py @@ -0,0 +1,61 @@ +"""MSSQL schema introspection and caching. + +This module provides a function to get the database schema (tables and columns) +and cache it for use by the LLM agent. The schema is retrieved via SQLAlchemy +inspection and returned as a dictionary suitable for inclusion in LLM prompts. +""" +from __future__ import annotations + +from threading import Lock +from typing import Dict, List + +from sqlalchemy import inspect + +from src.db.session import get_engine + +# Global cache for the schema +_schema_cache: Dict[str, List[dict]] | None = None +_schema_lock = Lock() + + +def get_schema() -> Dict[str, List[dict]]: + """Return the database schema as a dict of table name to list of columns. + + Each column is represented as a dict with keys "name" and "type". + + The result is cached on first call; subsequent calls return the same dict. + """ + global _schema_cache + if _schema_cache is not None: + return _schema_cache + + with _schema_lock: + # Double-check locking + if _schema_cache is not None: + return _schema_cache + + engine = get_engine() + inspector = inspect(engine) + schema: Dict[str, List[dict]] = {} + + for table_name in inspector.get_table_names(): + # Skip system tables if desired, but we keep all for now. + columns = [] + for column in inspector.get_columns(table_name): + # We only need the name and type for the LLM context. + columns.append( + { + "name": column["name"], + "type": str(column["type"]), + } + ) + schema[table_name] = columns + + _schema_cache = schema + return schema + + +def clear_schema_cache() -> None: + """Clear the cached schema. Primarily useful for testing.""" + global _schema_cache + _schema_cache = None \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py index 668a391..d31d55e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,38 +1,169 @@ -"""Shared fixtures. +"""Shared fixtures for MSSQL-backed tests. -Every test gets: reset settings/db singletons + an isolated tmp SQLite DB. -Unit tests additionally blank the provider keys; integration tests keep the -real key from .env (that is the point of the gate). +Every test gets: +- Reset settings and DB singletons +- An isolated MSSQL database (created per test, dropped after) +- For unit tests that don't need the DB, the fixture still runs but we can skip if needed. + +We use Docker to run a temporary MSSQL server. """ from __future__ import annotations +import os +import time +from typing import Generator + +import docker import pytest +import pyodbc +from docker.models.containers import Container +from src.config import settings as settings_module +from src.db import session as session_module +# Mutate the global state to reset singletons between tests @pytest.fixture(autouse=True) -def _reset_singletons(): +def _reset_singletons() -> None: """Reset cached settings + engine so env patches take effect per test.""" - import src.config.settings as settings_mod - import src.db.session as session_mod + settings_module._settings = None + if session_module._engine is not None: + session_module._engine.dispose() + session_module._engine = None + session_module._SessionLocal = None - settings_mod._settings = None - session_mod._engine = None - session_mod._SessionLocal = None - yield - settings_mod._settings = None - if session_mod._engine is not None: - session_mod._engine.dispose() - session_mod._engine = None - session_mod._SessionLocal = None +@pytest.fixture(scope="session") +def docker_client() -> docker.DockerClient: + """Provide a Docker client for the session.""" + return docker.from_env() -@pytest.fixture(autouse=True) -def _isolated_db(tmp_path, monkeypatch): - """Point the app at a fresh SQLite file — never the user's real DB.""" - monkeypatch.setenv("AGENT_DATABASE_URL", f"sqlite:///{tmp_path}/test.db") - yield + +@pytest.fixture(scope="session") +def mssql_container(docker_client: docker.DockerClient) -> Generator[Container, None, None]: + """Start a temporary MSSQL container for the test session. + + Uses the official Microsoft SQL Server 2022 Express image. + """ + from docker.errors import ImageNotFound + + container = None + try: + # Pull the image if not present (optional, but ensures we have it) + try: + docker_client.images.get("mcr.microsoft.com/mssql/server:2022-latest") + except ImageNotFound: + print("Pulling mcr.microsoft.com/mssql/server:2022-latest...") + docker_client.images.pull("mcr.microsoft.com/mssql/server:2022-latest") + + # Container configuration + container = docker_client.containers.run( + "mcr.microsoft.com/mssql/server:2022-latest", + name=f"mssql-test-{int(time.time())}", + environment={ + "SA_PASSWORD": "YourStrong!Passw0rd", + "ACCEPT_EULA": "Y", + "MSSQL_PID": "Express", + }, + ports={"1433/tcp": ("127.0.0.1", 0)}, # Random host port + detach=True, + auto_remove=True, + ) + # Wait for SQL Server to be ready + # Poll until we can connect or timeout after 30 seconds + timeout = time.time() + 30 + while time.time() < timeout: + container.reload() + if container.status == "running": + # Try to connect to the port + ports = container.attrs["NetworkSettings"]["Ports"] + host_port = ports["1433/tcp"][0]["HostPort"] + try: + # Use pyodbc to test connection + conn_str = ( + f"Driver={{ODBC Driver 17 for SQL Server}};" + f"Server=127.0.0.1,{host_port};" + f"UID=sa;PWD=YourStrong!Passw0rd;" + ) + with pyodbc.connect(conn_str, timeout=5): + break + except Exception: + pass + time.sleep(0.5) + else: + raise RuntimeError("MS SQL Server did not become ready in time") + + yield container + finally: + if container: + container.stop() + + +@pytest.fixture +def mssql_db(mssql_container: Container, monkeypatch) -> Generator[str, None, None]: + """Create a temporary database on the MSSQL server for a single test. + + Yields a SQLAlchemy URL for the database. + After the test, the database is dropped. + """ + import subprocess + + # Get the host port from the container + client = docker.from_env() + container = client.containers.get(mssql_container.id) + container.reload() + ports = container.attrs["NetworkSettings"]["Ports"] + host_port = ports["1433/tcp"][0]["HostPort"] + + # Generate a unique database name + db_name = f"testdb_{int(time.time() * 1000)}" + + # SQL Server ODBC driver 17 connection string for pyodbc + odbc_str = ( + f"Driver={{ODBC Driver 17 for SQL Server}};" + f"Server=127.0.0.1,{host_port};" + f"UID=sa;PWD=YourStrong!Passw0rd;" + ) + + # Create the database + with pyodbc.connect(odbc_str, autocommit=True) as conn: + cursor = conn.cursor() + cursor.execute(f"CREATE DATABASE [{db_name}]") + cursor.commit() + + # Build the SQLAlchemy URL + # Format: mssql+pyodbc://sa:password@host:port/dbname?driver=ODBC+Driver+17+for+SQL+Server + sqlalchemy_url = ( + f"mssql+pyodbc://sa:YourStrong!Passw0rd@127.0.0.1:{host_port}/{db_name}" + f"?driver=ODBC+Driver+17+for+SQL+Server" + ) + + # Set the environment variable for the app + monkeypatch.setenv("AGENT_DATABASE_URL", sqlalchemy_url) + + # Run migrations to set up the schema + # We use alembic command line via subprocess + env = os.environ.copy() + # Ensure we use the same python and virtual environment + result = subprocess.run( + ["uv", "run", "alembic", "upgrade", "head"], + env=env, + capture_output=True, + text=True, + ) + if result.returncode != 0: + raise RuntimeError(f"Alembic upgrade failed: {result.stderr}") + + try: + yield sqlalchemy_url + finally: + # Drop the database + with pyodbc.connect(odbc_str, autocommit=True) as conn: + cursor = conn.cursor() + cursor.execute(f"DROP DATABASE [{db_name}]") + cursor.commit() +# Keep the existing no_keys fixture for tests that want to simulate missing keys @pytest.fixture() def no_keys(monkeypatch): """Simulate 'no provider key configured' regardless of the user's .env. @@ -45,4 +176,4 @@ def no_keys(monkeypatch): monkeypatch.setenv("AGENT_ANTHROPIC_API_KEY", "") monkeypatch.setenv("AGENT_GEMINI_API_KEY", "") monkeypatch.setenv("AGENT_OPENROUTER_API_KEY", "") - yield + yield \ No newline at end of file diff --git a/tests/unit/test_mssql_schema.py b/tests/unit/test_mssql_schema.py new file mode 100644 index 0000000..a1bb086 --- /dev/null +++ b/tests/unit/test_mssql_schema.py @@ -0,0 +1,37 @@ +"""Test the MSSQL schema introspection and caching.""" +from __future__ import annotations + +import pytest + +from src.db.mssql.schema import clear_schema_cache, get_schema + + +def test_get_schema_returns_dict(mssql_db: str): + """get_schema should return a dictionary of tables and columns.""" + schema = get_schema() + assert isinstance(schema, dict) + # We expect at least the audit table to be present after migrations + assert "audit_log" in schema + # Each table should have a list of columns, each with name and type + for table_name, columns in schema.items(): + assert isinstance(columns, list) + for col in columns: + assert "name" in col and isinstance(col["name"], str) + assert "type" in col and isinstance(col["type"], str) + + +def test_get_schema_is_cached(mssql_db: str): + """Calling get_schema twice should return the same object.""" + schema1 = get_schema() + schema2 = get_schema() + assert schema1 is schema2 # same object due to caching + + +def test_clear_schema_cache(mssql_db: str): + """clear_schema_cache should reset the cache.""" + schema1 = get_schema() + clear_schema_cache() + schema2 = get_schema() + # After clearing, we should get a new dictionary (same content, but different object) + assert schema1 is not schema2 + assert schema1 == schema2 # content should be the same \ No newline at end of file