Context Engineering with a Streaming Database

-- Scenario: an airline disruption agent helping passengers whose flights change.
CREATE SOURCE flight_events (
  flight_id VARCHAR, flight_date DATE, status VARCHAR, gate VARCHAR,
  est_departure TIMESTAMPTZ, delay_minutes INT, reason VARCHAR, event_ts TIMESTAMPTZ,
  WATERMARK FOR event_ts AS event_ts - INTERVAL '30 seconds'
) WITH (connector = 'kafka', topic = 'flight_events',
        properties.bootstrap.server = 'broker:9092')
FORMAT PLAIN ENCODE JSON;

CREATE SOURCE res_pg WITH (
  connector = 'postgres-cdc', hostname = 'pg', port = '5432',
  username = 'rw', password = '...', database.name = 'reservations', schema.name = 'public'
);

CREATE TABLE flights (flight_id VARCHAR, flight_date DATE, origin VARCHAR, destination VARCHAR,
                      sched_departure TIMESTAMPTZ, seats_available INT,
                      PRIMARY KEY (flight_id, flight_date))
  FROM res_pg TABLE 'public.flights';

CREATE TABLE bookings (pnr VARCHAR PRIMARY KEY, passenger_id BIGINT, flight_id VARCHAR,
                       flight_date DATE, cabin VARCHAR, status VARCHAR, booked_at TIMESTAMPTZ)
  FROM res_pg TABLE 'public.bookings';

CREATE TABLE passengers (passenger_id BIGINT PRIMARY KEY, name VARCHAR, tier VARCHAR,
                         language VARCHAR, needs_assistance BOOLEAN)
  FROM res_pg TABLE 'public.passengers';

-- Context object 1: current status per flight. Budget: one row. Latest event per key.
CREATE MATERIALIZED VIEW flight_now AS
SELECT flight_id, flight_date, status, gate, est_departure, delay_minutes, reason,
       event_ts AS status_as_of
FROM (SELECT *, ROW_NUMBER() OVER (PARTITION BY flight_id, flight_date ORDER BY event_ts DESC) AS rn
      FROM flight_events) t
WHERE rn = 1;

-- Context object 2: the passenger's trip. Budget: one row per booking, joined and flagged.
CREATE MATERIALIZED VIEW passenger_trip AS
SELECT b.pnr, b.passenger_id, p.name, p.tier, p.language, p.needs_assistance,
       b.flight_id, b.flight_date, b.cabin, b.status AS booking_status,
       f.origin, f.destination, f.sched_departure,
       n.status AS flight_status, n.gate, n.est_departure, n.delay_minutes, n.reason,
       n.status_as_of,
       (n.status IN ('cancelled', 'diverted') OR n.delay_minutes >= 120) AS disrupted
FROM bookings b
JOIN passengers p ON p.passenger_id = b.passenger_id
JOIN flights f ON f.flight_id = b.flight_id AND f.flight_date = b.flight_date
LEFT JOIN flight_now n ON n.flight_id = b.flight_id AND n.flight_date = b.flight_date
WHERE b.status = 'confirmed';

-- Context object 3: rebooking options. Budget: top 3 per disrupted flight, ranked in SQL.
CREATE MATERIALIZED VIEW rebooking_options AS
SELECT * FROM (
  SELECT d.flight_id AS from_flight, d.flight_date AS from_date,
         a.flight_id AS alt_flight, a.flight_date AS alt_date,
         COALESCE(an.est_departure, a.sched_departure) AS alt_departure,
         a.seats_available,
         ROW_NUMBER() OVER (PARTITION BY d.flight_id, d.flight_date
                            ORDER BY COALESCE(an.est_departure, a.sched_departure)) AS rn
  FROM flights d
  JOIN flight_now dn ON dn.flight_id = d.flight_id AND dn.flight_date = d.flight_date
  JOIN flights a ON a.origin = d.origin AND a.destination = d.destination
                AND a.sched_departure > d.sched_departure
                AND a.sched_departure < d.sched_departure + INTERVAL '12 hours'
  LEFT JOIN flight_now an ON an.flight_id = a.flight_id AND an.flight_date = a.flight_date
  WHERE dn.status IN ('cancelled', 'diverted')
    AND a.seats_available > 0
    AND COALESCE(an.status, 'scheduled') NOT IN ('cancelled', 'diverted')
) t WHERE rn <= 3;

-- Context object 4: how bad is it right now at the airport? Budget: one row, sliding 30 minutes.
CREATE MATERIALIZED VIEW airport_disruption_30m AS
SELECT f.origin AS airport, window_start, window_end,
       COUNT(*) FILTER (WHERE e.status = 'cancelled')        AS cancellations,
       COUNT(*) FILTER (WHERE e.delay_minutes >= 60)         AS long_delays,
       MAX(e.delay_minutes)                                   AS worst_delay
FROM HOP(flight_events, event_ts, INTERVAL '5 minutes', INTERVAL '30 minutes') e
JOIN flights f ON f.flight_id = e.flight_id AND f.flight_date = e.flight_date
GROUP BY f.origin, window_start, window_end;

-- Context object 5: memory. The agent writes; a view keeps the last five per passenger.
CREATE TABLE agent_notes (
  note_id BIGINT PRIMARY KEY, passenger_id BIGINT, kind VARCHAR, note VARCHAR, created_at TIMESTAMPTZ
);
CREATE MATERIALIZED VIEW passenger_memory AS
SELECT passenger_id, ARRAY_AGG(kind || ': ' || note ORDER BY created_at DESC) AS notes
FROM (SELECT *, ROW_NUMBER() OVER (PARTITION BY passenger_id ORDER BY created_at DESC) AS rn
      FROM agent_notes) t
WHERE rn <= 5
GROUP BY passenger_id;

CREATE INDEX passenger_trip_by_pnr ON passenger_trip (pnr);
CREATE INDEX rebooking_by_flight ON rebooking_options (from_flight, from_date);
CREATE INDEX disruption_by_airport ON airport_disruption_30m (airport);

-- context.sql: the live layer as one JSON document, from one snapshot. $1 = pnr.
SELECT jsonb_build_object(
  'trip',        to_jsonb(t) - 'passenger_id',
  'rebooking',   CASE WHEN t.disrupted THEN (
                   SELECT jsonb_agg(jsonb_build_object(
                            'flight', r.alt_flight, 'departs', r.alt_departure,
                            'seats', r.seats_available))
                   FROM rebooking_options r
                   WHERE r.from_flight = t.flight_id AND r.from_date = t.flight_date)
                 ELSE NULL END,
  'airport',     (SELECT to_jsonb(a) - 'airport' FROM airport_disruption_30m a
                   WHERE a.airport = t.origin
                   ORDER BY a.window_end DESC LIMIT 1),
  'memory',      COALESCE(m.notes, ARRAY[]::VARCHAR[]),
  'freshness',   jsonb_build_object(
                   'flight_status_as_of', t.status_as_of,
                   'generated_at', now(),
                   'status_lag_seconds', EXTRACT(EPOCH FROM now() - t.status_as_of))
) AS context
FROM passenger_trip t
LEFT JOIN passenger_memory m ON m.passenger_id = t.passenger_id
WHERE t.pnr = $1;

import json, re, time, psycopg, anthropic

rw  = psycopg.connect("postgresql://disruption_agent@rw:4566/dev", autocommit=True)
llm = anthropic.Anthropic()
CONTEXT_SQL = open("context.sql").read()

def bind(sql: str) -> str:
    """The SQL files use Postgres $n placeholders; psycopg wants %(name)s."""
    return re.sub(r"\$(\d+)", r"%(p\1)s", sql)
CONTEXT_VERSION = "v3"           # bump when a view or the query changes
LIVE_BUDGET_TOKENS = 1500

SYSTEM = """You are an airline disruption agent. <context> holds the passenger's trip, rebooking
options ranked by departure time, airport conditions and notes from earlier conversations.
Every fact has a timestamp under freshness; if status_lag_seconds exceeds 120, say the status
may be out of date. Use the lookup tool only for details not in the context."""

LOOKUP_TOOL = {
    "name": "lookup",
    "description": "Read more detail from live views. kind is one of: all_options (every rebooking "
                   "option, not just the top 3), flight_history (last 20 status events for a flight), "
                   "other_bookings (the passenger's other confirmed trips).",
    "input_schema": {"type": "object",
                     "properties": {"kind": {"type": "string"}, "key": {"type": "string"}},
                     "required": ["kind", "key"]},
}

def fetch_context(pnr: str) -> dict:
    with rw.cursor() as cur:
        cur.execute(bind(CONTEXT_SQL), {"p1": pnr})
        row = cur.fetchone()
    return row[0] if row else {}

def fit_budget(ctx: dict) -> dict:
    """Drop the least important layer until the live layer fits. Order matters."""
    for layer in ("airport", "memory", "rebooking"):
        tokens = llm.messages.count_tokens(
            model="claude-sonnet-5-5", system=SYSTEM,
            messages=[{"role": "user", "content": json.dumps(ctx, default=str)}]).input_tokens
        if tokens <= LIVE_BUDGET_TOKENS:
            break
        ctx.pop(layer, None)
    return ctx

def answer(pnr: str, question: str) -> str:
    t0 = time.time()
    ctx = fit_budget(fetch_context(pnr))
    resp = llm.messages.create(
        model="claude-sonnet-5-5", max_tokens=700, system=SYSTEM, tools=[LOOKUP_TOOL],
        messages=[{"role": "user",
                   "content": f"<context>{json.dumps(ctx, default=str)}</context>\n\n{question}"}])
    log_turn(pnr, resp, time.time() - t0, ctx)
    return resp.content[0].text if resp.stop_reason == "end_turn" else handle_tool_call(resp, pnr, question)

LOOKUPS = {
    "all_options":    "SELECT alt_flight, alt_departure, seats_available FROM rebooking_options "
                      "WHERE from_flight = split_part($1, '/', 1) AND from_date = split_part($1, '/', 2)::date "
                      "ORDER BY alt_departure",
    "flight_history": "SELECT status, gate, delay_minutes, reason, event_ts FROM flight_events "
                      "WHERE flight_id = split_part($1, '/', 1) AND flight_date = split_part($1, '/', 2)::date "
                      "ORDER BY event_ts DESC LIMIT 20",
    "other_bookings": "SELECT pnr, flight_id, flight_date, origin, destination, flight_status "
                      "FROM passenger_trip WHERE passenger_id = $1::bigint",
}

def lookup(kind: str, key: str) -> str:
    sql = LOOKUPS.get(kind)
    if sql is None:
        return json.dumps({"error": f"unknown kind {kind}"})
    with rw.cursor() as cur:
        cur.execute(bind(sql), {"p1": key})
        cols = [d.name for d in cur.description]
        return json.dumps([dict(zip(cols, r)) for r in cur.fetchall()], default=str)

# Memory is a write to a table; the passenger_memory view picks it up on the next barrier.
def remember(passenger_id: int, kind: str, note: str):
    with rw.cursor() as cur:
        cur.execute("INSERT INTO agent_notes VALUES (%s, %s, %s, %s, now())",
                    (int(time.time() * 1000), passenger_id, kind, note))

-- The application writes one row per turn.
CREATE TABLE turn_log (
  turn_id VARCHAR PRIMARY KEY, agent VARCHAR, context_version VARCHAR,
  pnr VARCHAR, input_tokens INT, output_tokens INT, tool_calls INT,
  latency_ms INT, outcome VARCHAR, logged_at TIMESTAMPTZ
);

-- Cost and behaviour per context version, hourly. Compare v2 with v3 the hour after you ship it.
CREATE MATERIALIZED VIEW context_cost_1h AS
SELECT agent, context_version, window_start, window_end,
       COUNT(*)                                            AS turns,
       AVG(input_tokens)                                   AS avg_input_tokens,
       MAX(input_tokens)                                   AS max_input_tokens,
       AVG(latency_ms)                                     AS avg_latency_ms,
       SUM(tool_calls)::float / COUNT(*)                   AS tool_calls_per_turn,
       COUNT(*) FILTER (WHERE outcome = 'resolved')::float / COUNT(*) AS resolved_rate,
       COUNT(*) FILTER (WHERE outcome = 'corrected')       AS corrections
FROM TUMBLE(turn_log, logged_at, INTERVAL '1 hour')
GROUP BY agent, context_version, window_start, window_end;

-- Freshness SLA: how far behind is the live layer? Read query, run by the monitor every minute.
SELECT MAX(event_ts) AS latest_event, now() - MAX(event_ts) AS lag FROM flight_events;
SELECT MAX(status_as_of) AS latest_status, now() - MAX(status_as_of) AS lag FROM flight_now;

-- Which context objects are big? State per view; a growing one needs a tighter definition.
SELECT * FROM rw_catalog.rw_table_stats ORDER BY total_key_count DESC;