File size: 6,207 Bytes
a845c8d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
CREATE EXTENSION IF NOT EXISTS "pg_trgm";

CREATE TABLE IF NOT EXISTS sessions (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    tenant_id VARCHAR(64) DEFAULT 'default',
    created_at TIMESTAMPTZ DEFAULT NOW(),
    last_active_at TIMESTAMPTZ DEFAULT NOW(),
    expires_at TIMESTAMPTZ DEFAULT NOW() + INTERVAL '7 days',
    budget_usd NUMERIC(10, 6) DEFAULT 10.0,
    spent_usd NUMERIC(10, 6) DEFAULT 0.0,
    status VARCHAR(32) DEFAULT 'active',
    metadata JSONB DEFAULT '{}',
    context_json JSONB DEFAULT '[]',
    model_name VARCHAR(128),
    reasoning_effort VARCHAR(32),
    max_iterations INTEGER DEFAULT 300,
    yolo_mode BOOLEAN DEFAULT FALSE,
    CONSTRAINT positive_budget CHECK (budget_usd >= 0),
    CONSTRAINT positive_spent CHECK (spent_usd >= 0)
);

CREATE TABLE IF NOT EXISTS requests (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    session_id UUID REFERENCES sessions(id) ON DELETE CASCADE,
    tenant_id VARCHAR(64) DEFAULT 'default',
    trace_id VARCHAR(128),
    model VARCHAR(128) NOT NULL,
    provider VARCHAR(64) NOT NULL,
    operation_type VARCHAR(32) DEFAULT 'chat',
    input_tokens INTEGER DEFAULT 0,
    output_tokens INTEGER DEFAULT 0,
    total_tokens INTEGER GENERATED ALWAYS AS (input_tokens + output_tokens) STORED,
    cost_usd NUMERIC(10, 6) DEFAULT 0.0,
    cached BOOLEAN DEFAULT FALSE,
    latency_ms INTEGER,
    queue_time_ms INTEGER DEFAULT 0,
    request_payload JSONB,
    response_payload JSONB,
    tool_calls JSONB DEFAULT '[]',
    finish_reason VARCHAR(32),
    rate_limited BOOLEAN DEFAULT FALSE,
    circuit_open BOOLEAN DEFAULT FALSE,
    retry_count INTEGER DEFAULT 0,
    created_at TIMESTAMPTZ DEFAULT NOW(),
    CONSTRAINT non_negative_tokens CHECK (input_tokens >= 0 AND output_tokens >= 0)
);

CREATE TABLE IF NOT EXISTS circuit_events (
    id SERIAL PRIMARY KEY,
    provider VARCHAR(64) NOT NULL,
    event_type VARCHAR(32) NOT NULL,
    failure_count INTEGER DEFAULT 0,
    details JSONB DEFAULT '{}',
    created_at TIMESTAMPTZ DEFAULT NOW()
);

CREATE TABLE IF NOT EXISTS budget_events (
    id SERIAL PRIMARY KEY,
    session_id UUID REFERENCES sessions(id) ON DELETE CASCADE,
    event_type VARCHAR(32) NOT NULL,
    threshold_percent INTEGER,
    spent_before NUMERIC(10, 6),
    spent_after NUMERIC(10, 6),
    budget NUMERIC(10, 6),
    created_at TIMESTAMPTZ DEFAULT NOW()
);

CREATE TABLE IF NOT EXISTS tool_executions (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    request_id UUID REFERENCES requests(id) ON DELETE CASCADE,
    tool_name VARCHAR(128) NOT NULL,
    tool_args JSONB,
    result_json JSONB,
    success BOOLEAN DEFAULT TRUE,
    error_message TEXT,
    latency_ms INTEGER,
    created_at TIMESTAMPTZ DEFAULT NOW()
);

CREATE TABLE IF NOT EXISTS cache_events (
    id SERIAL PRIMARY KEY,
    cache_key VARCHAR(256) NOT NULL,
    cache_type VARCHAR(32) NOT NULL,
    hit BOOLEAN NOT NULL,
    ttl_seconds INTEGER,
    size_bytes INTEGER,
    created_at TIMESTAMPTZ DEFAULT NOW()
);

CREATE INDEX IF NOT EXISTS idx_sessions_tenant ON sessions(tenant_id);
CREATE INDEX IF NOT EXISTS idx_sessions_status ON sessions(status);
CREATE INDEX IF NOT EXISTS idx_sessions_expires ON sessions(expires_at);
CREATE INDEX IF NOT EXISTS idx_sessions_created ON sessions(created_at DESC);
CREATE INDEX IF NOT EXISTS idx_sessions_active ON sessions(tenant_id, status) WHERE status = 'active';
CREATE INDEX IF NOT EXISTS idx_sessions_metadata ON sessions USING GIN (metadata);

CREATE INDEX IF NOT EXISTS idx_requests_session ON requests(session_id);
CREATE INDEX IF NOT EXISTS idx_requests_tenant ON requests(tenant_id);
CREATE INDEX IF NOT EXISTS idx_requests_trace ON requests(trace_id);
CREATE INDEX IF NOT EXISTS idx_requests_provider ON requests(provider);
CREATE INDEX IF NOT EXISTS idx_requests_created ON requests(created_at DESC);
CREATE INDEX IF NOT EXISTS idx_requests_model ON requests(model);
CREATE INDEX IF NOT EXISTS idx_requests_payload ON requests USING GIN (request_payload);

CREATE INDEX IF NOT EXISTS idx_circuit_provider ON circuit_events(provider);
CREATE INDEX IF NOT EXISTS idx_circuit_created ON circuit_events(created_at DESC);
CREATE INDEX IF NOT EXISTS idx_budget_session ON budget_events(session_id);
CREATE INDEX IF NOT EXISTS idx_tool_request ON tool_executions(request_id);
CREATE INDEX IF NOT EXISTS idx_tool_name ON tool_executions(tool_name);
CREATE INDEX IF NOT EXISTS idx_cache_key ON cache_events(cache_key);
CREATE INDEX IF NOT EXISTS idx_cache_created ON cache_events(created_at DESC);

CREATE OR REPLACE FUNCTION update_last_active()
RETURNS TRIGGER AS $$
BEGIN
    NEW.last_active_at = NOW();
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

DROP TRIGGER IF EXISTS trigger_update_last_active ON sessions;
CREATE TRIGGER trigger_update_last_active
    BEFORE UPDATE ON sessions
    FOR EACH ROW
    EXECUTE FUNCTION update_last_active();

CREATE MATERIALIZED VIEW IF NOT EXISTS mv_daily_costs AS
SELECT 
    tenant_id,
    DATE(created_at) as date,
    provider,
    model,
    COUNT(*) as request_count,
    SUM(input_tokens) as total_input_tokens,
    SUM(output_tokens) as total_output_tokens,
    SUM(cost_usd) as total_cost_usd,
    AVG(latency_ms)::NUMERIC(10,2) as avg_latency_ms,
    SUM(CASE WHEN cached THEN 1 ELSE 0 END) as cache_hits,
    SUM(CASE WHEN NOT cached THEN 1 ELSE 0 END) as cache_misses
FROM requests
GROUP BY tenant_id, DATE(created_at), provider, model;

CREATE UNIQUE INDEX IF NOT EXISTS idx_mv_daily_costs ON mv_daily_costs(tenant_id, date, provider, model);

CREATE OR REPLACE FUNCTION refresh_daily_costs()
RETURNS void AS $$
BEGIN
    REFRESH MATERIALIZED VIEW CONCURRENTLY mv_daily_costs;
END;
$$ LANGUAGE plpgsql;

ALTER TABLE sessions ENABLE ROW LEVEL SECURITY;
ALTER TABLE requests ENABLE ROW LEVEL SECURITY;

CREATE POLICY tenant_isolation_sessions ON sessions
    USING (tenant_id = current_setting('app.current_tenant', true));
CREATE POLICY tenant_isolation_requests ON requests
    USING (tenant_id = current_setting('app.current_tenant', true));

GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA public TO ml_intern;
GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA public TO ml_intern;