hobokenchicken
6b10d4249c
This commit replaces the Axum/Rust backend with a Gin/Go implementation. The original Rust code has been archived in the 'rust' branch.
114 lines
3.2 KiB
Go
114 lines
3.2 KiB
Go
package server
|
|
|
|
import (
|
|
"log"
|
|
"time"
|
|
|
|
"llm-proxy/internal/db"
|
|
)
|
|
|
|
type RequestLog struct {
|
|
Timestamp time.Time `json:"timestamp"`
|
|
ClientID string `json:"client_id"`
|
|
Provider string `json:"provider"`
|
|
Model string `json:"model"`
|
|
PromptTokens uint32 `json:"prompt_tokens"`
|
|
CompletionTokens uint32 `json:"completion_tokens"`
|
|
ReasoningTokens uint32 `json:"reasoning_tokens"`
|
|
TotalTokens uint32 `json:"total_tokens"`
|
|
CacheReadTokens uint32 `json:"cache_read_tokens"`
|
|
CacheWriteTokens uint32 `json:"cache_write_tokens"`
|
|
Cost float64 `json:"cost"`
|
|
HasImages bool `json:"has_images"`
|
|
Status string `json:"status"`
|
|
ErrorMessage string `json:"error_message,omitempty"`
|
|
DurationMS int64 `json:"duration_ms"`
|
|
}
|
|
|
|
type RequestLogger struct {
|
|
database *db.DB
|
|
hub *Hub
|
|
logChan chan RequestLog
|
|
}
|
|
|
|
func NewRequestLogger(database *db.DB, hub *Hub) *RequestLogger {
|
|
return &RequestLogger{
|
|
database: database,
|
|
hub: hub,
|
|
logChan: make(chan RequestLog, 100),
|
|
}
|
|
}
|
|
|
|
func (l *RequestLogger) Start() {
|
|
go func() {
|
|
for entry := range l.logChan {
|
|
l.processLog(entry)
|
|
}
|
|
}()
|
|
}
|
|
|
|
func (l *RequestLogger) LogRequest(entry RequestLog) {
|
|
select {
|
|
case l.logChan <- entry:
|
|
default:
|
|
log.Println("Request log channel full, dropping log entry")
|
|
}
|
|
}
|
|
|
|
func (l *RequestLogger) processLog(entry RequestLog) {
|
|
// Broadcast to dashboard
|
|
l.hub.broadcast <- map[string]interface{}{
|
|
"type": "request",
|
|
"channel": "requests",
|
|
"payload": entry,
|
|
}
|
|
|
|
// Insert into DB
|
|
tx, err := l.database.Begin()
|
|
if err != nil {
|
|
log.Printf("Failed to begin transaction for logging: %v", err)
|
|
return
|
|
}
|
|
defer tx.Rollback()
|
|
|
|
// Ensure client exists
|
|
_, _ = tx.Exec("INSERT OR IGNORE INTO clients (client_id, name, description) VALUES (?, ?, 'Auto-created from request')",
|
|
entry.ClientID, entry.ClientID)
|
|
|
|
// Insert log
|
|
_, err = tx.Exec(`
|
|
INSERT INTO llm_requests
|
|
(timestamp, client_id, provider, model, prompt_tokens, completion_tokens, reasoning_tokens, total_tokens, cache_read_tokens, cache_write_tokens, cost, has_images, status, error_message, duration_ms)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
`, entry.Timestamp, entry.ClientID, entry.Provider, entry.Model,
|
|
entry.PromptTokens, entry.CompletionTokens, entry.ReasoningTokens, entry.TotalTokens,
|
|
entry.CacheReadTokens, entry.CacheWriteTokens, entry.Cost, entry.HasImages,
|
|
entry.Status, entry.ErrorMessage, entry.DurationMS)
|
|
|
|
if err != nil {
|
|
log.Printf("Failed to insert request log: %v", err)
|
|
return
|
|
}
|
|
|
|
// Update client stats
|
|
_, _ = tx.Exec(`
|
|
UPDATE clients SET
|
|
total_requests = total_requests + 1,
|
|
total_tokens = total_tokens + ?,
|
|
total_cost = total_cost + ?,
|
|
updated_at = CURRENT_TIMESTAMP
|
|
WHERE client_id = ?
|
|
`, entry.TotalTokens, entry.Cost, entry.ClientID)
|
|
|
|
// Update provider balance
|
|
if entry.Cost > 0 {
|
|
_, _ = tx.Exec("UPDATE provider_configs SET credit_balance = credit_balance - ? WHERE id = ? AND (billing_mode IS NULL OR billing_mode != 'postpaid')",
|
|
entry.Cost, entry.Provider)
|
|
}
|
|
|
|
err = tx.Commit()
|
|
if err != nil {
|
|
log.Printf("Failed to commit logging transaction: %v", err)
|
|
}
|
|
}
|