| 1 | /** |
| 2 | * Agents as stored in D1, and their spend. Shared by the RPC methods and |
| 3 | * the desk. |
| 4 | */ |
| 5 | import type { WorkspaceAgent } from "@g1t/contracts"; |
| 6 | |
| 7 | import { agentStatus, budgetBlock, dayKey, monthKey } from "./budget.ts"; |
| 8 | import { type Definition, DEFAULT_AUTONOMY, DEFAULT_BUDGET, DEFAULT_ROUTING, PRESETS, readJson } from "./definition.ts"; |
| 9 | |
| 10 | export type Row = { |
| 11 | id: string; |
| 12 | workspace_id: string; |
| 13 | handle: string; |
| 14 | display_name: string; |
| 15 | avatar: string | null; |
| 16 | role: string; |
| 17 | instructions: string; |
| 18 | personality_preset: string; |
| 19 | personality: string; |
| 20 | routing: string; |
| 21 | budget: string; |
| 22 | autonomy: string; |
| 23 | capacity: number; |
| 24 | template: string | null; |
| 25 | version: number; |
| 26 | busy_until: string | null; |
| 27 | created_by: string; |
| 28 | created_at: string; |
| 29 | updated_at: string; |
| 30 | archived_at: string | null; |
| 31 | /** Joined from agent_spend: this month and today. */ |
| 32 | spent_month?: number | null; |
| 33 | spent_day?: number | null; |
| 34 | }; |
| 35 | |
| 36 | /** |
| 37 | * Agents with this month's and today's spend: `where` filters `agents` |
| 38 | * (aliased `a`); its parameters follow the two periods. |
| 39 | */ |
| 40 | export function selectAgents(where: string): string { |
| 41 | return `SELECT a.*, m.micros AS spent_month, d.micros AS spent_day FROM agents a |
| 42 | LEFT JOIN agent_spend m ON m.agent_id = a.id AND m.period = ?1 |
| 43 | LEFT JOIN agent_spend d ON d.agent_id = a.id AND d.period = ?2 |
| 44 | WHERE ${where}`; |
| 45 | } |
| 46 | |
| 47 | export function periods(now: Date): [string, string] { |
| 48 | return [monthKey(now), dayKey(now)]; |
| 49 | } |
| 50 | |
| 51 | export function definitionOf(row: Row): Definition { |
| 52 | return { |
| 53 | handle: row.handle, |
| 54 | display_name: row.display_name, |
| 55 | role: row.role, |
| 56 | instructions: row.instructions, |
| 57 | personality_preset: PRESETS.includes(row.personality_preset as Definition["personality_preset"]) |
| 58 | ? (row.personality_preset as Definition["personality_preset"]) |
| 59 | : "crisp", |
| 60 | personality: row.personality, |
| 61 | routing: readJson(row.routing, DEFAULT_ROUTING), |
| 62 | budget: readJson(row.budget, DEFAULT_BUDGET), |
| 63 | autonomy: readJson(row.autonomy, DEFAULT_AUTONOMY), |
| 64 | capacity: row.capacity, |
| 65 | template: row.template, |
| 66 | }; |
| 67 | } |
| 68 | |
| 69 | export function toAgent(row: Row, now: Date): WorkspaceAgent { |
| 70 | const definition = definitionOf(row); |
| 71 | const spent = { month: row.spent_month ?? 0, day: row.spent_day ?? 0 }; |
| 72 | return { |
| 73 | id: row.id, |
| 74 | workspace_id: row.workspace_id, |
| 75 | avatar: row.avatar, |
| 76 | ...definition, |
| 77 | version: row.version, |
| 78 | status: row.archived_at ? "paused" : agentStatus({ busyUntil: row.busy_until, now, blocked: budgetBlock(definition.budget, spent, now) !== null }), |
| 79 | spent_month_micros: spent.month, |
| 80 | created_by: row.created_by, |
| 81 | created_at: row.created_at, |
| 82 | updated_at: row.updated_at, |
| 83 | archived_at: row.archived_at, |
| 84 | }; |
| 85 | } |
| 86 | |
| 87 | /** Adds a reply's charge to the agent's month and day. */ |
| 88 | export function spendStatements(db: D1Database, agentId: string, micros: number, now: Date): D1PreparedStatement[] { |
| 89 | return periods(now).map((period) => |
| 90 | db |
| 91 | .prepare( |
| 92 | `INSERT INTO agent_spend (agent_id, period, micros, replies) VALUES (?1, ?2, ?3, 1) |
| 93 | ON CONFLICT (agent_id, period) DO UPDATE SET micros = micros + ?3, replies = replies + 1`, |
| 94 | ) |
| 95 | .bind(agentId, period, Math.max(0, Math.ceil(micros))), |
| 96 | ); |
| 97 | } |