Skip to content

Commit fb781a9

Browse files
moshloopmoshloop
committed
feat(ai): add native OpenAI provider, SQLite response cache, and OpenRouter pricing
- Add BackendOpenAI with native openai-go SDK provider (pkg/ai/provider/openai.go) - Route GPT/o-series models to native API instead of Codex CLI - Add SQLite-based response cache (pkg/ai/cache/) with WAL mode, TTL expiration, hourly cleanup, and daily stats aggregation - Replace in-memory cache interface with SQLite-backed middleware - Keep grok-* models on BackendCodexCLI, OpenRouter pricing was already present
1 parent 5024df7 commit fb781a9

10 files changed

Lines changed: 604 additions & 232 deletions

File tree

‎go.mod‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ require (
1212
github.com/flanksource/clicky v1.19.0
1313
github.com/flanksource/commons v1.48.1
1414
github.com/flanksource/sandbox-runtime v1.0.2
15+
github.com/mattn/go-sqlite3 v1.14.38
16+
github.com/openai/openai-go v1.12.0
1517
github.com/samber/lo v1.53.0
1618
github.com/sergi/go-diff v1.4.0
1719
github.com/spf13/cobra v1.10.2

‎go.sum‎

Lines changed: 4 additions & 177 deletions
Large diffs are not rendered by default.

‎pkg/ai/cache/cache.go‎

Lines changed: 307 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,307 @@
1+
package cache
2+
3+
import (
4+
"crypto/sha256"
5+
"database/sql"
6+
"errors"
7+
"fmt"
8+
"os"
9+
"path/filepath"
10+
"time"
11+
12+
"github.com/flanksource/commons/logger"
13+
_ "github.com/mattn/go-sqlite3"
14+
"github.com/samber/lo"
15+
)
16+
17+
var (
18+
ErrCacheDisabled = errors.New("caching is disabled")
19+
ErrNotFound = errors.New("cache entry not found")
20+
)
21+
22+
type Config struct {
23+
DBPath string
24+
TTL time.Duration
25+
NoCache bool
26+
}
27+
28+
type Entry struct {
29+
ID int64
30+
CacheKey string
31+
PromptHash string
32+
Model string
33+
Prompt string
34+
Response string
35+
Error string
36+
TokensInput int
37+
TokensOutput int
38+
TokensReasoning int
39+
TokensCacheRead int
40+
TokensCacheWrite int
41+
TokensTotal int
42+
CostUSD float64
43+
DurationMS int64
44+
Provider string
45+
Temperature float64
46+
MaxTokens int
47+
CreatedAt time.Time
48+
AccessedAt time.Time
49+
ExpiresAt *time.Time
50+
}
51+
52+
type StatsEntry struct {
53+
Model string
54+
Provider string
55+
TotalRequests int64
56+
CacheHits int64
57+
CacheMisses int64
58+
ErrorCount int64
59+
TotalInputTokens int64
60+
TotalOutputTokens int64
61+
TotalReasoningTokens int64
62+
TotalCacheReadTokens int64
63+
TotalCacheWriteTokens int64
64+
TotalCost float64
65+
AvgDurationMS int64
66+
FirstRequest time.Time
67+
LastRequest time.Time
68+
}
69+
70+
type Cache struct {
71+
db *sql.DB
72+
config Config
73+
}
74+
75+
func (c Cache) GetTTL() time.Duration { return c.config.TTL }
76+
77+
func New(config Config) (*Cache, error) {
78+
if config.DBPath == "" {
79+
homeDir, err := os.UserHomeDir()
80+
if err != nil {
81+
return nil, fmt.Errorf("failed to get home directory: %w", err)
82+
}
83+
config.DBPath = filepath.Join(homeDir, ".cache", "captain-llm.db")
84+
}
85+
86+
if err := os.MkdirAll(filepath.Dir(config.DBPath), 0o750); err != nil {
87+
return nil, fmt.Errorf("failed to create cache directory: %w", err)
88+
}
89+
90+
db, err := sql.Open("sqlite3", config.DBPath)
91+
if err != nil {
92+
return nil, fmt.Errorf("failed to open database: %w", err)
93+
}
94+
95+
for _, pragma := range []string{
96+
"PRAGMA journal_mode = WAL",
97+
"PRAGMA synchronous = NORMAL",
98+
"PRAGMA cache_size = -64000",
99+
"PRAGMA busy_timeout = 5000",
100+
} {
101+
if _, err := db.Exec(pragma); err != nil {
102+
db.Close()
103+
return nil, fmt.Errorf("failed to set pragma %s: %w", pragma, err)
104+
}
105+
}
106+
107+
cache := &Cache{db: db, config: config}
108+
if _, err := db.Exec(embeddedSchema); err != nil {
109+
db.Close()
110+
return nil, fmt.Errorf("failed to initialize schema: %w", err)
111+
}
112+
113+
go cache.cleanupExpired()
114+
return cache, nil
115+
}
116+
117+
func (c *Cache) Close() error { return c.db.Close() }
118+
119+
func generateCacheKey(prompt, model string) string {
120+
h := sha256.Sum256([]byte(fmt.Sprintf("%s|%s", prompt, model)))
121+
return fmt.Sprintf("%x", h)
122+
}
123+
124+
func (c *Cache) Get(prompt, model string) (*Entry, error) {
125+
if c.config.NoCache {
126+
return nil, ErrCacheDisabled
127+
}
128+
129+
cacheKey := generateCacheKey(prompt, model)
130+
131+
var entry Entry
132+
var expiresAt sql.NullTime
133+
err := c.db.QueryRow(`
134+
SELECT id, cache_key, prompt_hash, model, prompt, response, error,
135+
tokens_input, tokens_output, tokens_reasoning,
136+
tokens_cache_read, tokens_cache_write, tokens_total,
137+
cost_usd, duration_ms, provider, temperature, max_tokens,
138+
created_at, accessed_at, expires_at
139+
FROM llm_cache
140+
WHERE cache_key = ?
141+
AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
142+
ORDER BY created_at DESC
143+
LIMIT 1
144+
`, cacheKey).Scan(
145+
&entry.ID, &entry.CacheKey, &entry.PromptHash, &entry.Model,
146+
&entry.Prompt, &entry.Response, &entry.Error,
147+
&entry.TokensInput, &entry.TokensOutput, &entry.TokensReasoning,
148+
&entry.TokensCacheRead, &entry.TokensCacheWrite, &entry.TokensTotal,
149+
&entry.CostUSD, &entry.DurationMS, &entry.Provider,
150+
&entry.Temperature, &entry.MaxTokens,
151+
&entry.CreatedAt, &entry.AccessedAt, &expiresAt,
152+
)
153+
154+
if err == sql.ErrNoRows {
155+
return nil, ErrNotFound
156+
}
157+
if err != nil {
158+
return nil, fmt.Errorf("failed to get cache entry: %w", err)
159+
}
160+
161+
if expiresAt.Valid {
162+
entry.ExpiresAt = &expiresAt.Time
163+
if time.Now().After(expiresAt.Time) {
164+
return nil, ErrNotFound
165+
}
166+
}
167+
168+
_, _ = c.db.Exec("UPDATE llm_cache SET accessed_at = CURRENT_TIMESTAMP WHERE id = ?", entry.ID)
169+
return &entry, nil
170+
}
171+
172+
func (c *Cache) Set(entry *Entry) error {
173+
if c.config.NoCache {
174+
return nil
175+
}
176+
177+
entry.CreatedAt = time.Now()
178+
entry.CacheKey = generateCacheKey(entry.Prompt, entry.Model)
179+
entry.PromptHash = entry.CacheKey
180+
181+
logger.Tracef("[%s] caching response for %s (hash:%s)", entry.Model, lo.Ellipsis(entry.Prompt, 20), entry.PromptHash)
182+
183+
var expiresAt *time.Time
184+
if c.config.TTL > 0 {
185+
exp := time.Now().Add(c.config.TTL)
186+
expiresAt = &exp
187+
}
188+
189+
_, err := c.db.Exec(`
190+
INSERT OR REPLACE INTO llm_cache (
191+
cache_key, prompt_hash, model, prompt, response, error,
192+
tokens_input, tokens_output, tokens_reasoning,
193+
tokens_cache_read, tokens_cache_write, tokens_total,
194+
cost_usd, duration_ms, provider, temperature, max_tokens, expires_at
195+
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
196+
entry.CacheKey, entry.PromptHash, entry.Model, entry.Prompt, entry.Response, entry.Error,
197+
entry.TokensInput, entry.TokensOutput, entry.TokensReasoning,
198+
entry.TokensCacheRead, entry.TokensCacheWrite, entry.TokensTotal,
199+
entry.CostUSD, entry.DurationMS, entry.Provider, entry.Temperature, entry.MaxTokens,
200+
expiresAt,
201+
)
202+
if err != nil {
203+
return fmt.Errorf("failed to set cache entry: %w", err)
204+
}
205+
206+
c.updateStats(entry)
207+
return nil
208+
}
209+
210+
func (c *Cache) Clear() error {
211+
_, err := c.db.Exec("DELETE FROM llm_cache")
212+
return err
213+
}
214+
215+
func (c *Cache) GetStats() ([]StatsEntry, error) {
216+
rows, err := c.db.Query(`
217+
SELECT model, provider,
218+
COUNT(*) as total_requests,
219+
SUM(CASE WHEN error IS NULL OR error = '' THEN 1 ELSE 0 END),
220+
SUM(CASE WHEN error IS NOT NULL AND error != '' THEN 1 ELSE 0 END),
221+
SUM(tokens_input), SUM(tokens_output), SUM(tokens_reasoning),
222+
SUM(tokens_cache_read), SUM(tokens_cache_write),
223+
SUM(cost_usd), AVG(duration_ms),
224+
MIN(created_at), MAX(created_at)
225+
FROM llm_cache
226+
GROUP BY model, provider
227+
ORDER BY total_requests DESC
228+
`)
229+
if err != nil {
230+
return nil, fmt.Errorf("failed to get stats: %w", err)
231+
}
232+
defer rows.Close()
233+
234+
var stats []StatsEntry
235+
for rows.Next() {
236+
var s StatsEntry
237+
var provider sql.NullString
238+
var inputTok, outputTok, reasonTok, cacheReadTok, cacheWriteTok sql.NullInt64
239+
240+
if err := rows.Scan(
241+
&s.Model, &provider,
242+
&s.TotalRequests, &s.CacheHits, &s.ErrorCount,
243+
&inputTok, &outputTok, &reasonTok,
244+
&cacheReadTok, &cacheWriteTok,
245+
&s.TotalCost, &s.AvgDurationMS,
246+
&s.FirstRequest, &s.LastRequest,
247+
); err != nil {
248+
return nil, fmt.Errorf("failed to scan stats: %w", err)
249+
}
250+
251+
if provider.Valid {
252+
s.Provider = provider.String
253+
}
254+
if inputTok.Valid {
255+
s.TotalInputTokens = inputTok.Int64
256+
}
257+
if outputTok.Valid {
258+
s.TotalOutputTokens = outputTok.Int64
259+
}
260+
if reasonTok.Valid {
261+
s.TotalReasoningTokens = reasonTok.Int64
262+
}
263+
if cacheReadTok.Valid {
264+
s.TotalCacheReadTokens = cacheReadTok.Int64
265+
}
266+
if cacheWriteTok.Valid {
267+
s.TotalCacheWriteTokens = cacheWriteTok.Int64
268+
}
269+
270+
stats = append(stats, s)
271+
}
272+
return stats, nil
273+
}
274+
275+
func (c *Cache) cleanupExpired() {
276+
ticker := time.NewTicker(1 * time.Hour)
277+
defer ticker.Stop()
278+
279+
for range ticker.C {
280+
_, _ = c.db.Exec("DELETE FROM llm_cache WHERE expires_at IS NOT NULL AND expires_at < CURRENT_TIMESTAMP")
281+
}
282+
}
283+
284+
func (c *Cache) updateStats(entry *Entry) {
285+
date := entry.CreatedAt.Format("2006-01-02")
286+
_, _ = c.db.Exec(`
287+
INSERT INTO llm_stats (
288+
date, model, provider, request_count,
289+
total_input_tokens, total_output_tokens, total_reasoning_tokens,
290+
total_cache_read_tokens, total_cache_write_tokens,
291+
total_cost_usd
292+
) VALUES (?, ?, ?, 1, ?, ?, ?, ?, ?, ?)
293+
ON CONFLICT(date, model, provider) DO UPDATE SET
294+
request_count = request_count + 1,
295+
total_input_tokens = total_input_tokens + excluded.total_input_tokens,
296+
total_output_tokens = total_output_tokens + excluded.total_output_tokens,
297+
total_reasoning_tokens = total_reasoning_tokens + excluded.total_reasoning_tokens,
298+
total_cache_read_tokens = total_cache_read_tokens + excluded.total_cache_read_tokens,
299+
total_cache_write_tokens = total_cache_write_tokens + excluded.total_cache_write_tokens,
300+
total_cost_usd = total_cost_usd + excluded.total_cost_usd,
301+
updated_at = CURRENT_TIMESTAMP`,
302+
date, entry.Model, entry.Provider,
303+
entry.TokensInput, entry.TokensOutput, entry.TokensReasoning,
304+
entry.TokensCacheRead, entry.TokensCacheWrite,
305+
entry.CostUSD,
306+
)
307+
}

‎pkg/ai/cache/schema.go‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
package cache
2+
3+
const embeddedSchema = `
4+
CREATE TABLE IF NOT EXISTS llm_cache (
5+
id INTEGER PRIMARY KEY AUTOINCREMENT,
6+
cache_key TEXT NOT NULL,
7+
prompt_hash TEXT NOT NULL,
8+
model TEXT NOT NULL,
9+
prompt TEXT NOT NULL,
10+
response TEXT NOT NULL,
11+
error TEXT,
12+
tokens_input INTEGER DEFAULT 0,
13+
tokens_output INTEGER DEFAULT 0,
14+
tokens_reasoning INTEGER DEFAULT 0,
15+
tokens_cache_read INTEGER DEFAULT 0,
16+
tokens_cache_write INTEGER DEFAULT 0,
17+
tokens_total INTEGER DEFAULT 0,
18+
cost_usd REAL DEFAULT 0.0,
19+
duration_ms INTEGER DEFAULT 0,
20+
provider TEXT,
21+
temperature REAL DEFAULT 0.2,
22+
max_tokens INTEGER,
23+
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
24+
accessed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
25+
expires_at TIMESTAMP,
26+
UNIQUE(cache_key, model)
27+
);
28+
29+
CREATE INDEX IF NOT EXISTS idx_cache_lookup ON llm_cache(cache_key, model, expires_at);
30+
CREATE INDEX IF NOT EXISTS idx_prompt_hash ON llm_cache(prompt_hash);
31+
CREATE INDEX IF NOT EXISTS idx_created_at ON llm_cache(created_at DESC);
32+
CREATE INDEX IF NOT EXISTS idx_provider ON llm_cache(provider);
33+
CREATE INDEX IF NOT EXISTS idx_model ON llm_cache(model);
34+
35+
CREATE TABLE IF NOT EXISTS llm_stats (
36+
id INTEGER PRIMARY KEY AUTOINCREMENT,
37+
date DATE NOT NULL,
38+
model TEXT NOT NULL,
39+
provider TEXT,
40+
request_count INTEGER DEFAULT 0,
41+
cache_hit_count INTEGER DEFAULT 0,
42+
cache_miss_count INTEGER DEFAULT 0,
43+
error_count INTEGER DEFAULT 0,
44+
total_input_tokens INTEGER DEFAULT 0,
45+
total_output_tokens INTEGER DEFAULT 0,
46+
total_reasoning_tokens INTEGER DEFAULT 0,
47+
total_cache_read_tokens INTEGER DEFAULT 0,
48+
total_cache_write_tokens INTEGER DEFAULT 0,
49+
total_cost_usd REAL DEFAULT 0.0,
50+
avg_duration_ms INTEGER DEFAULT 0,
51+
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
52+
UNIQUE(date, model, provider)
53+
);
54+
`

‎pkg/ai/client.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ func NewProvider(cfg Config) (Provider, error) {
2424
}
2525

2626
if cfg.Model == "" {
27-
return nil, fmt.Errorf("model cannot be empty; specify a model or backend (available: anthropic, gemini, codex-cli, claude-cli, gemini-cli)")
27+
return nil, fmt.Errorf("model cannot be empty; specify a model or backend (available: anthropic, gemini, openai, codex-cli, claude-cli, gemini-cli)")
2828
}
2929

3030
if backend == "" {
@@ -52,6 +52,7 @@ func GetAPIKeyFromEnv(backend Backend) string {
5252
envVars := map[Backend][]string{
5353
BackendAnthropic: {"ANTHROPIC_API_KEY"},
5454
BackendGemini: {"GEMINI_API_KEY", "GOOGLE_API_KEY"},
55+
BackendOpenAI: {"OPENAI_API_KEY"},
5556
BackendGeminiCLI: {},
5657
BackendClaudeCLI: {},
5758
BackendCodexCLI: {},

0 commit comments

Comments
 (0)