trading-agent/state/db.py

"""SQLite state: audit log + context cache, dimensioned by account (persona).

Alpaca remains the source of truth for live positions; this DB records what each agent
proposed and did, and caches each account's latest snapshot so a nightly agent can see its
own holdings (tools/portfolio.py) and the dashboard can show all three books.
"""

from __future__ import annotations

import json
import sqlite3
from datetime import datetime, timezone
from pathlib import Path

SCHEMA = """
CREATE TABLE IF NOT EXISTS runs (
    run_id TEXT PRIMARY KEY, trade_date TEXT, started_at TEXT, ended_at TEXT, status TEXT
);
CREATE TABLE IF NOT EXISTS theses (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    trade_date TEXT, agent TEXT, ticker TEXT, direction TEXT, conviction REAL,
    thesis TEXT, requested_weight_pct REAL, stop_pct REAL, take_pct REAL, horizon TEXT,
    json_blob TEXT
);
CREATE TABLE IF NOT EXISTS orders (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    trade_date TEXT, agent TEXT, ticker TEXT, side TEXT, qty INTEGER, entry_type TEXT,
    est_price REAL, stop_price REAL, take_price REAL,
    alpaca_order_id TEXT, status TEXT, avg_fill_price REAL, submitted_at TEXT
);
CREATE TABLE IF NOT EXISTS snapshots (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    ts TEXT, account TEXT, equity REAL, cash REAL, buying_power REAL, positions_json TEXT
);
CREATE TABLE IF NOT EXISTS positions (
    symbol TEXT, account TEXT, qty REAL, avg_entry REAL, market_value REAL,
    unrealized_pl REAL, as_of TEXT, PRIMARY KEY (symbol, account)
);
"""


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def connect(path: str | Path) -> sqlite3.Connection:
    con = sqlite3.connect(str(path))
    con.row_factory = sqlite3.Row
    return con


def init_db(path: str | Path) -> None:
    con = connect(path)
    con.executescript(SCHEMA)
    con.commit()
    con.close()


def record_thesis(con, trade_date, agent, pick_dict: dict) -> None:
    con.execute(
        "INSERT INTO theses (trade_date, agent, ticker, direction, conviction, thesis, "
        "requested_weight_pct, stop_pct, take_pct, horizon, json_blob) "
        "VALUES (?,?,?,?,?,?,?,?,?,?,?)",
        (
            trade_date, agent, pick_dict.get("ticker"), pick_dict.get("direction"),
            pick_dict.get("conviction"), pick_dict.get("thesis"),
            pick_dict.get("requested_weight_pct"), pick_dict.get("stop_loss_pct"),
            pick_dict.get("take_profit_pct"), pick_dict.get("time_horizon"),
            json.dumps(pick_dict),
        ),
    )


def record_order(con, *, trade_date, agent, ticker, side, qty, entry_type, est_price,
                 stop_price, take_price, alpaca_order_id, status, avg_fill_price=None) -> None:
    con.execute(
        "INSERT INTO orders (trade_date, agent, ticker, side, qty, entry_type, est_price, "
        "stop_price, take_price, alpaca_order_id, status, avg_fill_price, submitted_at) "
        "VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)",
        (trade_date, agent, ticker, side, qty, entry_type, est_price, stop_price, take_price,
         alpaca_order_id, status, avg_fill_price, _now()),
    )


def record_snapshot(con, *, account, equity, cash, buying_power, positions: list[dict]) -> None:
    con.execute(
        "INSERT INTO snapshots (ts, account, equity, cash, buying_power, positions_json) "
        "VALUES (?,?,?,?,?,?)",
        (_now(), account, equity, cash, buying_power, json.dumps(positions)),
    )


def upsert_position(con, *, account, symbol, qty, avg_entry, market_value, unrealized_pl) -> None:
    con.execute(
        "INSERT INTO positions (symbol, account, qty, avg_entry, market_value, unrealized_pl, "
        "as_of) VALUES (?,?,?,?,?,?,?) ON CONFLICT(symbol, account) DO UPDATE SET "
        "qty=excluded.qty, avg_entry=excluded.avg_entry, market_value=excluded.market_value, "
        "unrealized_pl=excluded.unrealized_pl, as_of=excluded.as_of",
        (symbol, account, qty, avg_entry, market_value, unrealized_pl, _now()),
    )


def latest_snapshot(con, account: str | None = None) -> sqlite3.Row | None:
    if account:
        return con.execute(
            "SELECT * FROM snapshots WHERE account=? ORDER BY ts DESC LIMIT 1", (account,)
        ).fetchone()
    return con.execute("SELECT * FROM snapshots ORDER BY ts DESC LIMIT 1").fetchone()