feat: add provider-neutral LLM layer with secrets, budgets, recording and prompts
Phase 7 (ADR-0004). Cognitive code asks the router for a role; the router maps roles to configured models with ordered fallbacks, applies global and per-mind daily token budgets and a concurrency limit, records every call (request, response, usage, latency, cost, seed, request hash) and can serve recordings back for exact replay. Providers are raw HTTP clients for OpenAI-compatible, Anthropic and Ollama endpoints, with embedding support for the embed role wired into memory retrieval. Secrets are references only: env:, file:, or enc: values encrypted with AES-256-GCM under NOEMA_MASTER_KEY; keys are resolved at call time and never rendered, logged or exported. Outbound HTTP refuses private and loopback addresses at dial time unless a provider explicitly allows it, bounds redirects and caps bodies. Provider errors are scrubbed. Prompts are named, versioned rows with system and user templates and an optional JSON schema; untrusted text enters only through delimited data blocks. Six defaults are seeded and never overwritten. The UI adds the LLM page (providers, models, role matrix, usage, errors, budgets), a provider page with connectivity test, and a prompt manager with editor, test runner, history, restore, clone and diff. CLI: provider, prompt.
This commit is contained in:
Andy Dixon
@@ -0,0 +1,389 @@
|
||||
// Package router maps cognitive roles to configured providers and models,
|
||||
// applies budgets and concurrency limits, records every call and serves
|
||||
// recordings during replay (ADR-0004).
|
||||
package router
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/andydixon/noema/internal/embeddings"
|
||||
"github.com/andydixon/noema/internal/llm"
|
||||
"github.com/andydixon/noema/internal/llm/providers"
|
||||
"github.com/andydixon/noema/internal/observability"
|
||||
"github.com/andydixon/noema/internal/scheduler"
|
||||
"github.com/andydixon/noema/internal/security/secrets"
|
||||
"github.com/andydixon/noema/pkg/ids"
|
||||
)
|
||||
|
||||
// Router resolves roles to configured models, applies budgets and
|
||||
// concurrency limits, records every call and emits metrics. Cognitive code
|
||||
// depends on this type only.
|
||||
type Router struct {
|
||||
store *Store
|
||||
secrets *secrets.Resolver
|
||||
log *slog.Logger
|
||||
metrics *observability.Metrics
|
||||
pool *scheduler.Pool
|
||||
|
||||
mu sync.RWMutex
|
||||
cache routing
|
||||
cachedAt time.Time
|
||||
|
||||
DailyTokenBudget int64 // global hard stop (0 = unlimited)
|
||||
MindBudget func(mind ids.ID) int64 // per-mind daily budget lookup (nil = unlimited)
|
||||
Replay bool // serve recorded responses when a matching call exists
|
||||
Override map[llm.Role]llm.LanguageModel // tests inject models per role
|
||||
now func() time.Time
|
||||
}
|
||||
|
||||
type routing struct {
|
||||
providers map[ids.ID]Provider
|
||||
models map[ids.ID]Model
|
||||
roles map[llm.Role]RoleAssignment
|
||||
}
|
||||
|
||||
// NewRouter creates a router.
|
||||
func NewRouter(store *Store, sec *secrets.Resolver, log *slog.Logger, metrics *observability.Metrics, maxConcurrent int) *Router {
|
||||
return &Router{store: store, secrets: sec, log: log, metrics: metrics, pool: scheduler.NewPool(maxConcurrent), now: time.Now}
|
||||
}
|
||||
|
||||
// Invalidate forces the configuration cache to reload on the next call.
|
||||
func (r *Router) Invalidate() {
|
||||
r.mu.Lock()
|
||||
r.cachedAt = time.Time{}
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
func (r *Router) load(ctx context.Context) (routing, error) {
|
||||
r.mu.RLock()
|
||||
if time.Since(r.cachedAt) < 30*time.Second {
|
||||
c := r.cache
|
||||
r.mu.RUnlock()
|
||||
return c, nil
|
||||
}
|
||||
r.mu.RUnlock()
|
||||
provs, err := r.store.ListProviders(ctx)
|
||||
if err != nil {
|
||||
return routing{}, err
|
||||
}
|
||||
models, err := r.store.ListModels(ctx)
|
||||
if err != nil {
|
||||
return routing{}, err
|
||||
}
|
||||
roles, err := r.store.ListRoles(ctx)
|
||||
if err != nil {
|
||||
return routing{}, err
|
||||
}
|
||||
c := routing{providers: map[ids.ID]Provider{}, models: map[ids.ID]Model{}, roles: map[llm.Role]RoleAssignment{}}
|
||||
for _, p := range provs {
|
||||
c.providers[p.ID] = p
|
||||
}
|
||||
for _, m := range models {
|
||||
c.models[m.ID] = m
|
||||
}
|
||||
for _, a := range roles {
|
||||
c.roles[a.Role] = a
|
||||
}
|
||||
r.mu.Lock()
|
||||
r.cache, r.cachedAt = c, time.Now()
|
||||
r.mu.Unlock()
|
||||
return c, nil
|
||||
}
|
||||
|
||||
// Candidates returns the ordered models (primary then fallbacks) for a role.
|
||||
func (r *Router) Candidates(ctx context.Context, role llm.Role) ([]Model, error) {
|
||||
c, err := r.load(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
a, ok := c.roles[role]
|
||||
if !ok || a.ModelID == nil {
|
||||
return nil, llm.ErrNoModel
|
||||
}
|
||||
var out []Model
|
||||
for _, id := range append([]ids.ID{*a.ModelID}, a.FallbackIDs...) {
|
||||
m, ok := c.models[id]
|
||||
if !ok || !m.Enabled {
|
||||
continue
|
||||
}
|
||||
p, ok := c.providers[m.ProviderID]
|
||||
if !ok || !p.Enabled {
|
||||
continue
|
||||
}
|
||||
out = append(out, m)
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return nil, llm.ErrNoModel
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Available reports whether a role has at least one usable model.
|
||||
func (r *Router) Available(ctx context.Context, role llm.Role) bool {
|
||||
if r.Override != nil {
|
||||
if _, ok := r.Override[role]; ok {
|
||||
return true
|
||||
}
|
||||
}
|
||||
_, err := r.Candidates(ctx, role)
|
||||
return err == nil
|
||||
}
|
||||
|
||||
// build constructs a provider client for a model, resolving the key at call time.
|
||||
func (r *Router) build(c routing, m Model) (llm.LanguageModel, error) {
|
||||
p := c.providers[m.ProviderID]
|
||||
key, err := r.secrets.Resolve(p.APIKeyRef)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("provider %s: %w", p.Name, err)
|
||||
}
|
||||
cfg := providers.Config{Name: p.Name, Endpoint: p.Endpoint, APIKey: key, Model: m.Name, Timeout: p.Timeout, AllowPrivate: p.AllowPrivate, Headers: p.Headers}
|
||||
switch p.Kind {
|
||||
case llm.KindOpenAI:
|
||||
return providers.NewOpenAI(cfg), nil
|
||||
case llm.KindAnthropic:
|
||||
return providers.NewAnthropic(cfg), nil
|
||||
case llm.KindOllama:
|
||||
return providers.NewOllama(cfg), nil
|
||||
}
|
||||
return nil, fmt.Errorf("provider %s: unsupported kind %s", p.Name, p.Kind)
|
||||
}
|
||||
|
||||
// Embedder returns the embedder configured for llm.RoleEmbed, or embeddings.None.
|
||||
func (r *Router) Embedder(ctx context.Context) embeddings.Embedder {
|
||||
c, err := r.load(ctx)
|
||||
if err != nil {
|
||||
return embeddings.None{}
|
||||
}
|
||||
a, ok := c.roles[llm.RoleEmbed]
|
||||
if !ok || a.ModelID == nil {
|
||||
return embeddings.None{}
|
||||
}
|
||||
m, ok := c.models[*a.ModelID]
|
||||
if !ok || !m.Enabled {
|
||||
return embeddings.None{}
|
||||
}
|
||||
p, ok := c.providers[m.ProviderID]
|
||||
if !ok || !p.Enabled {
|
||||
return embeddings.None{}
|
||||
}
|
||||
key, err := r.secrets.Resolve(p.APIKeyRef)
|
||||
if err != nil {
|
||||
return embeddings.None{}
|
||||
}
|
||||
cfg := providers.Config{Name: p.Name, Endpoint: p.Endpoint, APIKey: key, Model: m.Name, Timeout: p.Timeout, AllowPrivate: p.AllowPrivate, Headers: p.Headers}
|
||||
switch p.Kind {
|
||||
case llm.KindOllama:
|
||||
return providers.NewOllamaEmbedder(cfg, m.Dims)
|
||||
case llm.KindOpenAI:
|
||||
return providers.NewOpenAIEmbedder(cfg, m.Dims)
|
||||
}
|
||||
return embeddings.None{}
|
||||
}
|
||||
|
||||
// Hash returns the deterministic request hash used for replay matching. It
|
||||
// covers everything that influences the model's answer.
|
||||
func Hash(model string, req llm.Request) string {
|
||||
h := sha256.New()
|
||||
enc := json.NewEncoder(h)
|
||||
enc.Encode(struct { //nolint:errcheck
|
||||
Model string
|
||||
Req llm.Request
|
||||
}{model, req})
|
||||
return hex.EncodeToString(h.Sum(nil))[:32]
|
||||
}
|
||||
|
||||
// Complete routes a request for a role: budget check, primary then fallbacks,
|
||||
// recording and metrics. The returned response carries the call id so events
|
||||
// can cite it.
|
||||
func (r *Router) Complete(ctx context.Context, req llm.Request) (llm.Response, error) {
|
||||
if req.Role == "" {
|
||||
return llm.Response{}, fmt.Errorf("%w: role required", ErrInvalid)
|
||||
}
|
||||
if r.Override != nil {
|
||||
if lm, ok := r.Override[req.Role]; ok {
|
||||
return r.call(ctx, lm, "override", "override", 0, 0, req)
|
||||
}
|
||||
}
|
||||
c, err := r.load(ctx)
|
||||
if err != nil {
|
||||
return llm.Response{}, err
|
||||
}
|
||||
cands, err := r.Candidates(ctx, req.Role)
|
||||
if err != nil {
|
||||
return llm.Response{}, err
|
||||
}
|
||||
if err := r.checkBudget(ctx, req.MindID); err != nil {
|
||||
return llm.Response{}, err
|
||||
}
|
||||
var lastErr error
|
||||
for i, m := range cands {
|
||||
lm, err := r.build(c, m)
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
applyDefaults(&req, c.roles[req.Role], m)
|
||||
resp, err := r.call(ctx, lm, m.Provider, m.Name, m.InputCostPerM, m.OutputCostPerM, req)
|
||||
if err == nil {
|
||||
return resp, nil
|
||||
}
|
||||
lastErr = err
|
||||
if errors.Is(err, context.Canceled) || errors.Is(err, llm.ErrBudgetExceeded) {
|
||||
break
|
||||
}
|
||||
if i < len(cands)-1 {
|
||||
r.log.Warn("model failed; trying fallback", "role", req.Role, "model", m.Name, "err", err)
|
||||
}
|
||||
}
|
||||
return llm.Response{}, lastErr
|
||||
}
|
||||
|
||||
func applyDefaults(req *llm.Request, a RoleAssignment, m Model) {
|
||||
if req.Temperature == nil {
|
||||
if a.Temperature != nil {
|
||||
req.Temperature = a.Temperature
|
||||
} else if t, ok := m.Params["temperature"].(float64); ok {
|
||||
req.Temperature = &t
|
||||
}
|
||||
}
|
||||
if req.MaxTokens == 0 {
|
||||
if a.MaxTokens != nil {
|
||||
req.MaxTokens = *a.MaxTokens
|
||||
} else if t, ok := m.Params["max_tokens"].(float64); ok {
|
||||
req.MaxTokens = int(t)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Router) checkBudget(ctx context.Context, mind ids.ID) error {
|
||||
dayStart := r.now().Truncate(24 * time.Hour)
|
||||
if r.DailyTokenBudget > 0 {
|
||||
used, err := r.store.TokensSince(ctx, dayStart, ids.Nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if used >= r.DailyTokenBudget {
|
||||
return fmt.Errorf("%w: global daily budget %d reached", llm.ErrBudgetExceeded, r.DailyTokenBudget)
|
||||
}
|
||||
}
|
||||
if r.MindBudget != nil && !mind.IsZero() {
|
||||
if b := r.MindBudget(mind); b > 0 {
|
||||
used, err := r.store.TokensSince(ctx, dayStart, mind)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if used >= b {
|
||||
return fmt.Errorf("%w: mind daily budget %d reached", llm.ErrBudgetExceeded, b)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *Router) call(ctx context.Context, lm llm.LanguageModel, provider, model string, inCost, outCost float64, req llm.Request) (llm.Response, error) {
|
||||
hash := Hash(model, req)
|
||||
reqJSON, _ := json.Marshal(req)
|
||||
if r.Replay {
|
||||
if rec, err := r.store.FindRecorded(ctx, hash); err == nil && len(rec.Response) > 0 {
|
||||
var resp llm.Response
|
||||
if json.Unmarshal(rec.Response, &resp) == nil {
|
||||
resp.Replayed, resp.CallID = true, ids.New()
|
||||
r.store.RecordCall(ctx, Call{ID: resp.CallID, MindID: req.MindID, CorrelationID: req.CorrelationID, Role: req.Role, Purpose: req.Purpose, Provider: provider, Model: model, //nolint:errcheck
|
||||
PromptName: req.PromptName, PromptVersion: req.PromptVersion, RequestHash: hash, Request: reqJSON, Response: rec.Response, Replayed: true, Seed: req.Seed})
|
||||
return resp, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
var resp llm.Response
|
||||
var callErr error
|
||||
start := time.Now()
|
||||
err := r.pool.Run(ctx, func(ctx context.Context) error {
|
||||
cctx := ctx
|
||||
if req.Timeout > 0 {
|
||||
var cancel context.CancelFunc
|
||||
cctx, cancel = context.WithTimeout(ctx, req.Timeout)
|
||||
defer cancel()
|
||||
}
|
||||
resp, callErr = lm.Complete(cctx, req)
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return llm.Response{}, err
|
||||
}
|
||||
latency := time.Since(start)
|
||||
call := Call{ID: ids.New(), MindID: req.MindID, CorrelationID: req.CorrelationID, Role: req.Role, Purpose: req.Purpose, Provider: provider, Model: model,
|
||||
PromptName: req.PromptName, PromptVersion: req.PromptVersion, RequestHash: hash, Request: reqJSON, LatencyMS: int(latency.Milliseconds()), Seed: req.Seed}
|
||||
if r.metrics != nil {
|
||||
r.metrics.LLMRequests.WithLabelValues(provider, string(req.Role)).Inc()
|
||||
}
|
||||
if callErr != nil {
|
||||
call.Error = callErr.Error()
|
||||
if r.metrics != nil {
|
||||
r.metrics.LLMFailures.WithLabelValues(provider).Inc()
|
||||
}
|
||||
r.store.RecordCall(ctx, call) //nolint:errcheck
|
||||
return llm.Response{}, callErr
|
||||
}
|
||||
resp.CallID, resp.Latency = call.ID, latency
|
||||
resp.Cost = float64(resp.Usage.PromptTokens)*inCost/1e6 + float64(resp.Usage.CompletionTokens)*outCost/1e6
|
||||
call.PromptTokens, call.OutputTokens, call.Cost = resp.Usage.PromptTokens, resp.Usage.CompletionTokens, resp.Cost
|
||||
call.Response, _ = json.Marshal(resp)
|
||||
if r.metrics != nil {
|
||||
r.metrics.LLMTokens.WithLabelValues(provider, "prompt").Add(float64(resp.Usage.PromptTokens))
|
||||
r.metrics.LLMTokens.WithLabelValues(provider, "completion").Add(float64(resp.Usage.CompletionTokens))
|
||||
}
|
||||
if err := r.store.RecordCall(ctx, call); err != nil {
|
||||
r.log.Error("record llm call failed", "err", err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// TestProvider performs a tiny completion against the first enabled chat model
|
||||
// of a provider (or the given model) and records the outcome. It never returns
|
||||
// secret material.
|
||||
func (r *Router) TestProvider(ctx context.Context, providerID ids.ID, modelName string) (llm.Response, error) {
|
||||
c, err := r.load(ctx)
|
||||
if err != nil {
|
||||
return llm.Response{}, err
|
||||
}
|
||||
p, ok := c.providers[providerID]
|
||||
if !ok {
|
||||
return llm.Response{}, fmt.Errorf("%w: provider not found", ErrInvalid)
|
||||
}
|
||||
var m Model
|
||||
for _, cand := range c.models {
|
||||
if cand.ProviderID == providerID && cand.Kind == "chat" && (modelName == "" || cand.Name == modelName) {
|
||||
m = cand
|
||||
break
|
||||
}
|
||||
}
|
||||
if m.ID.IsZero() {
|
||||
m = Model{ProviderID: providerID, Provider: p.Name, Name: modelName}
|
||||
}
|
||||
if m.Name == "" {
|
||||
return llm.Response{}, fmt.Errorf("%w: add a chat model to test this provider", ErrInvalid)
|
||||
}
|
||||
lm, err := r.build(c, m)
|
||||
if err != nil {
|
||||
r.store.RecordProviderTest(ctx, providerID, false, err.Error()) //nolint:errcheck
|
||||
return llm.Response{}, err
|
||||
}
|
||||
tctx, cancel := context.WithTimeout(ctx, 30*time.Second)
|
||||
defer cancel()
|
||||
resp, err := r.call(tctx, lm, p.Name, m.Name, m.InputCostPerM, m.OutputCostPerM, llm.Request{Role: llm.RoleJudge, Purpose: "connectivity test", MaxTokens: 8, Messages: []llm.Message{{Role: "user", Content: "Reply with the single word OK."}}})
|
||||
if err != nil {
|
||||
r.store.RecordProviderTest(ctx, providerID, false, err.Error()) //nolint:errcheck
|
||||
return llm.Response{}, err
|
||||
}
|
||||
r.store.RecordProviderTest(ctx, providerID, true, fmt.Sprintf("%s answered in %dms", resp.Model, resp.Latency.Milliseconds())) //nolint:errcheck
|
||||
r.Invalidate()
|
||||
return resp, nil
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/andydixon/noema/internal/llm"
|
||||
"github.com/andydixon/noema/internal/minds"
|
||||
"github.com/andydixon/noema/internal/security/secrets"
|
||||
"github.com/andydixon/noema/internal/storage"
|
||||
"github.com/andydixon/noema/migrations"
|
||||
"github.com/andydixon/noema/pkg/ids"
|
||||
)
|
||||
|
||||
func TestRouterOverrideBudgetRecordingAndReplay(t *testing.T) {
|
||||
pool := storage.TestPool(t, migrations.FS)
|
||||
ctx := context.Background()
|
||||
st := NewStore(pool)
|
||||
sec, _ := secrets.New("")
|
||||
r := NewRouter(st, sec, slog.New(slog.DiscardHandler), nil, 2)
|
||||
m, _ := minds.NewStore(pool).Create(ctx, minds.CreateParams{Slug: "r"})
|
||||
|
||||
if _, err := r.Complete(ctx, llm.Request{Role: llm.RoleJudge, Messages: []llm.Message{{Role: "user", Content: "x"}}}); !errors.Is(err, llm.ErrNoModel) {
|
||||
t.Fatalf("expected no model, got %v", err)
|
||||
}
|
||||
fake := &llm.Fake{Script: []string{"first", "second"}}
|
||||
r.Override = map[llm.Role]llm.LanguageModel{llm.RoleJudge: fake}
|
||||
resp, err := r.Complete(ctx, llm.Request{Role: llm.RoleJudge, MindID: m.ID, Messages: []llm.Message{{Role: "user", Content: "x"}}})
|
||||
if err != nil || resp.Text != "first" || resp.CallID.IsZero() {
|
||||
t.Fatalf("override call: %v %+v", err, resp)
|
||||
}
|
||||
calls, _ := st.ListCalls(ctx, CallFilter{MindID: m.ID})
|
||||
if len(calls) != 1 || calls[0].Role != llm.RoleJudge || calls[0].PromptTokens == 0 {
|
||||
t.Fatalf("call not recorded: %+v", calls)
|
||||
}
|
||||
// Replay serves the recording for an identical request without calling the model.
|
||||
r.Replay = true
|
||||
resp2, err := r.Complete(ctx, llm.Request{Role: llm.RoleJudge, MindID: m.ID, Messages: []llm.Message{{Role: "user", Content: "x"}}})
|
||||
if err != nil || !resp2.Replayed || resp2.Text != "first" || len(fake.Calls) != 1 {
|
||||
t.Fatalf("replay wrong: %v %+v calls=%d", err, resp2, len(fake.Calls))
|
||||
}
|
||||
r.Replay = false
|
||||
// Budget: exhaust and expect refusal.
|
||||
r.DailyTokenBudget = 1
|
||||
// Override path bypasses configured routing; budget applies to configured models, so configure one.
|
||||
r.Override = nil
|
||||
p, err := st.CreateProvider(ctx, ProviderInput{Name: "local", Kind: llm.KindOllama, Endpoint: "http://127.0.0.1:1", AllowPrivate: true, Enabled: true})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
mdl, _ := st.CreateModel(ctx, ModelInput{ProviderID: p.ID, Name: "m", Kind: "chat", Enabled: true})
|
||||
st.SetRole(ctx, RoleAssignment{Role: llm.RoleJudge, ModelID: &mdl.ID}) //nolint:errcheck
|
||||
r.Invalidate()
|
||||
if _, err := r.Complete(ctx, llm.Request{Role: llm.RoleJudge, Messages: []llm.Message{{Role: "user", Content: "x"}}}); !errors.Is(err, llm.ErrBudgetExceeded) {
|
||||
t.Fatalf("expected budget error, got %v", err)
|
||||
}
|
||||
r.DailyTokenBudget = 0
|
||||
// Provider unreachable → provider error recorded with error text.
|
||||
_, err = r.Complete(ctx, llm.Request{Role: llm.RoleJudge, Messages: []llm.Message{{Role: "user", Content: "x"}}, Timeout: time.Second})
|
||||
if err == nil {
|
||||
t.Fatal("unreachable provider should fail")
|
||||
}
|
||||
errs, _ := st.ListCalls(ctx, CallFilter{ErrorsOnly: true})
|
||||
if len(errs) != 1 {
|
||||
t.Fatalf("error not recorded: %d", len(errs))
|
||||
}
|
||||
// Provider input validation: credentials in headers are rejected; key refs kept on update.
|
||||
if _, err := st.CreateProvider(ctx, ProviderInput{Name: "bad", Kind: llm.KindOpenAI, Headers: map[string]string{"Authorization": "x"}}); err == nil {
|
||||
t.Fatal("credential header accepted")
|
||||
}
|
||||
p2, _ := st.UpdateProvider(ctx, p.ID, ProviderInput{Name: "local", Kind: llm.KindOllama, Endpoint: p.Endpoint, APIKeyRef: "env:X", Enabled: true, AllowPrivate: true})
|
||||
p3, _ := st.UpdateProvider(ctx, p.ID, ProviderInput{Name: "local", Kind: llm.KindOllama, Endpoint: p.Endpoint, APIKeyRef: "", Enabled: true, AllowPrivate: true})
|
||||
if p2.APIKeyRef != "env:X" || p3.APIKeyRef != "env:X" {
|
||||
t.Fatalf("key reference handling wrong: %q %q", p2.APIKeyRef, p3.APIKeyRef)
|
||||
}
|
||||
if h := Hash("m", llm.Request{Role: "a"}); h != Hash("m", llm.Request{Role: "a"}) || h == Hash("m2", llm.Request{Role: "a"}) {
|
||||
t.Fatal("hash not stable/discriminating")
|
||||
}
|
||||
_ = ids.Nil
|
||||
}
|
||||
@@ -0,0 +1,537 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
|
||||
"github.com/andydixon/noema/internal/llm"
|
||||
"github.com/andydixon/noema/internal/storage"
|
||||
"github.com/andydixon/noema/pkg/ids"
|
||||
)
|
||||
|
||||
// Provider is a configured endpoint. APIKeyRef is a secret reference; the
|
||||
// struct never carries plaintext.
|
||||
type Provider struct {
|
||||
ID ids.ID `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Kind llm.ProviderKind `json:"kind"`
|
||||
Endpoint string `json:"endpoint"`
|
||||
APIKeyRef string `json:"-"`
|
||||
KeyDescribed string `json:"api_key"` // human description only
|
||||
AllowPrivate bool `json:"allow_private"`
|
||||
Headers map[string]string `json:"headers"`
|
||||
Timeout time.Duration `json:"timeout"`
|
||||
Enabled bool `json:"enabled"`
|
||||
LastTestAt *time.Time `json:"last_test_at,omitempty"`
|
||||
LastTestOK *bool `json:"last_test_ok,omitempty"`
|
||||
LastTestMsg string `json:"last_test_msg"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// Model is a named model on a provider.
|
||||
type Model struct {
|
||||
ID ids.ID `json:"id"`
|
||||
ProviderID ids.ID `json:"provider_id"`
|
||||
Provider string `json:"provider"`
|
||||
Name string `json:"name"`
|
||||
Kind string `json:"kind"` // chat | embedding
|
||||
ContextWindow int `json:"context_window"`
|
||||
Dims int `json:"dims"`
|
||||
InputCostPerM float64 `json:"input_cost_per_m"`
|
||||
OutputCostPerM float64 `json:"output_cost_per_m"`
|
||||
Params map[string]any `json:"params"`
|
||||
Enabled bool `json:"enabled"`
|
||||
}
|
||||
|
||||
// RoleAssignment maps a role to a primary model and fallbacks.
|
||||
type RoleAssignment struct {
|
||||
Role llm.Role `json:"role"`
|
||||
ModelID *ids.ID `json:"model_id,omitempty"`
|
||||
FallbackIDs []ids.ID `json:"fallback_ids"`
|
||||
Temperature *float64 `json:"temperature,omitempty"`
|
||||
MaxTokens *int `json:"max_tokens,omitempty"`
|
||||
UpdatedBy string `json:"updated_by"`
|
||||
}
|
||||
|
||||
// Call is a recorded model invocation.
|
||||
type Call struct {
|
||||
ID ids.ID `json:"id"`
|
||||
MindID ids.ID `json:"mind_id,omitempty"`
|
||||
CorrelationID ids.ID `json:"correlation_id,omitempty"`
|
||||
Role llm.Role `json:"role"`
|
||||
Purpose string `json:"purpose"`
|
||||
Provider string `json:"provider"`
|
||||
Model string `json:"model"`
|
||||
PromptName string `json:"prompt_name"`
|
||||
PromptVersion int `json:"prompt_version"`
|
||||
RequestHash string `json:"request_hash"`
|
||||
Request json.RawMessage `json:"request"`
|
||||
Response json.RawMessage `json:"response,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
Cost float64 `json:"cost"`
|
||||
LatencyMS int `json:"latency_ms"`
|
||||
Seed *int64 `json:"seed,omitempty"`
|
||||
Replayed bool `json:"replayed"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
// ErrInvalid marks bad configuration input.
|
||||
var ErrInvalid = errors.New("llm: invalid")
|
||||
|
||||
// Store persists LLM configuration and calls.
|
||||
type Store struct{ db *pgxpool.Pool }
|
||||
|
||||
// NewStore creates a store.
|
||||
func NewStore(db *pgxpool.Pool) *Store { return &Store{db: db} }
|
||||
|
||||
// ---- providers ----
|
||||
|
||||
const provCols = `id, name, kind, endpoint, api_key_ref, allow_private, headers, timeout_ms, enabled, last_test_at, last_test_ok, last_test_msg, created_at, updated_at`
|
||||
|
||||
func scanProvider(row pgx.Row) (Provider, error) {
|
||||
var p Provider
|
||||
var kind string
|
||||
var headers []byte
|
||||
var timeoutMS int
|
||||
if err := row.Scan(&p.ID, &p.Name, &kind, &p.Endpoint, &p.APIKeyRef, &p.AllowPrivate, &headers, &timeoutMS, &p.Enabled, &p.LastTestAt, &p.LastTestOK, &p.LastTestMsg, &p.CreatedAt, &p.UpdatedAt); err != nil {
|
||||
return p, storage.Translate(err)
|
||||
}
|
||||
p.Kind = llm.ProviderKind(kind)
|
||||
p.Timeout = time.Duration(timeoutMS) * time.Millisecond
|
||||
json.Unmarshal(headers, &p.Headers) //nolint:errcheck
|
||||
if p.Headers == nil {
|
||||
p.Headers = map[string]string{}
|
||||
}
|
||||
return p, nil
|
||||
}
|
||||
|
||||
// ProviderInput is the editable part of a provider.
|
||||
type ProviderInput struct {
|
||||
Name string
|
||||
Kind llm.ProviderKind
|
||||
Endpoint string
|
||||
APIKeyRef string // "" keeps the existing reference on update
|
||||
AllowPrivate bool
|
||||
Headers map[string]string
|
||||
Timeout time.Duration
|
||||
Enabled bool
|
||||
}
|
||||
|
||||
func (in *ProviderInput) validate() error {
|
||||
in.Name = strings.TrimSpace(in.Name)
|
||||
if in.Name == "" || len(in.Name) > 64 {
|
||||
return fmt.Errorf("%w: name required (max 64)", ErrInvalid)
|
||||
}
|
||||
if !in.Kind.Valid() {
|
||||
return fmt.Errorf("%w: unknown provider kind %q", ErrInvalid, in.Kind)
|
||||
}
|
||||
in.Endpoint = strings.TrimSpace(in.Endpoint)
|
||||
if in.Endpoint == "" {
|
||||
in.Endpoint = in.Kind.DefaultEndpoint()
|
||||
}
|
||||
if !strings.HasPrefix(in.Endpoint, "http://") && !strings.HasPrefix(in.Endpoint, "https://") {
|
||||
return fmt.Errorf("%w: endpoint must be an http(s) URL", ErrInvalid)
|
||||
}
|
||||
if in.Timeout <= 0 {
|
||||
in.Timeout = 60 * time.Second
|
||||
}
|
||||
if in.Headers == nil {
|
||||
in.Headers = map[string]string{}
|
||||
}
|
||||
for k := range in.Headers {
|
||||
if strings.EqualFold(k, "authorization") || strings.EqualFold(k, "x-api-key") {
|
||||
return fmt.Errorf("%w: put credentials in the API key field, not headers", ErrInvalid)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// CreateProvider adds a provider.
|
||||
func (s *Store) CreateProvider(ctx context.Context, in ProviderInput) (Provider, error) {
|
||||
if err := in.validate(); err != nil {
|
||||
return Provider{}, err
|
||||
}
|
||||
id := ids.New()
|
||||
headers, _ := json.Marshal(in.Headers)
|
||||
_, err := s.db.Exec(ctx, `INSERT INTO llm_providers (id, name, kind, endpoint, api_key_ref, allow_private, headers, timeout_ms, enabled) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9)`,
|
||||
id, in.Name, string(in.Kind), in.Endpoint, in.APIKeyRef, in.AllowPrivate, headers, in.Timeout.Milliseconds(), in.Enabled)
|
||||
if err != nil {
|
||||
return Provider{}, storage.Translate(err)
|
||||
}
|
||||
return s.GetProvider(ctx, id)
|
||||
}
|
||||
|
||||
// UpdateProvider edits a provider; an empty APIKeyRef keeps the stored one.
|
||||
func (s *Store) UpdateProvider(ctx context.Context, id ids.ID, in ProviderInput) (Provider, error) {
|
||||
if err := in.validate(); err != nil {
|
||||
return Provider{}, err
|
||||
}
|
||||
headers, _ := json.Marshal(in.Headers)
|
||||
tag, err := s.db.Exec(ctx, `UPDATE llm_providers SET name=$2, kind=$3, endpoint=$4, api_key_ref=CASE WHEN $5='' THEN api_key_ref ELSE $5 END, allow_private=$6, headers=$7, timeout_ms=$8, enabled=$9, updated_at=now() WHERE id=$1`,
|
||||
id, in.Name, string(in.Kind), in.Endpoint, in.APIKeyRef, in.AllowPrivate, headers, in.Timeout.Milliseconds(), in.Enabled)
|
||||
if err != nil {
|
||||
return Provider{}, storage.Translate(err)
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return Provider{}, storage.ErrNotFound
|
||||
}
|
||||
return s.GetProvider(ctx, id)
|
||||
}
|
||||
|
||||
// ClearProviderKey removes the stored key reference.
|
||||
func (s *Store) ClearProviderKey(ctx context.Context, id ids.ID) error {
|
||||
_, err := s.db.Exec(ctx, `UPDATE llm_providers SET api_key_ref='', updated_at=now() WHERE id=$1`, id)
|
||||
return err
|
||||
}
|
||||
|
||||
// DeleteProvider removes a provider and its models (calls keep provider name).
|
||||
func (s *Store) DeleteProvider(ctx context.Context, id ids.ID) error {
|
||||
tag, err := s.db.Exec(ctx, `DELETE FROM llm_providers WHERE id=$1`, id)
|
||||
if err == nil && tag.RowsAffected() == 0 {
|
||||
return storage.ErrNotFound
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// RecordProviderTest stores the outcome of a connectivity test.
|
||||
func (s *Store) RecordProviderTest(ctx context.Context, id ids.ID, ok bool, msg string) error {
|
||||
_, err := s.db.Exec(ctx, `UPDATE llm_providers SET last_test_at=now(), last_test_ok=$2, last_test_msg=$3 WHERE id=$1`, id, ok, msg)
|
||||
return err
|
||||
}
|
||||
|
||||
// GetProvider fetches one provider.
|
||||
func (s *Store) GetProvider(ctx context.Context, id ids.ID) (Provider, error) {
|
||||
return scanProvider(s.db.QueryRow(ctx, `SELECT `+provCols+` FROM llm_providers WHERE id=$1`, id))
|
||||
}
|
||||
|
||||
// ListProviders lists providers by name.
|
||||
func (s *Store) ListProviders(ctx context.Context) ([]Provider, error) {
|
||||
rows, err := s.db.Query(ctx, `SELECT `+provCols+` FROM llm_providers ORDER BY lower(name)`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []Provider
|
||||
for rows.Next() {
|
||||
p, err := scanProvider(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, p)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// ---- models ----
|
||||
|
||||
const modelCols = `m.id, m.provider_id, p.name, m.name, m.kind, m.context_window, m.dims, m.input_cost_per_m, m.output_cost_per_m, m.params, m.enabled`
|
||||
|
||||
func scanModel(row pgx.Row) (Model, error) {
|
||||
var m Model
|
||||
var params []byte
|
||||
if err := row.Scan(&m.ID, &m.ProviderID, &m.Provider, &m.Name, &m.Kind, &m.ContextWindow, &m.Dims, &m.InputCostPerM, &m.OutputCostPerM, ¶ms, &m.Enabled); err != nil {
|
||||
return m, storage.Translate(err)
|
||||
}
|
||||
json.Unmarshal(params, &m.Params) //nolint:errcheck
|
||||
if m.Params == nil {
|
||||
m.Params = map[string]any{}
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// ModelInput is the editable part of a model.
|
||||
type ModelInput struct {
|
||||
ProviderID ids.ID
|
||||
Name string
|
||||
Kind string
|
||||
ContextWindow int
|
||||
Dims int
|
||||
InputCostPerM float64
|
||||
OutputCostPerM float64
|
||||
Params map[string]any
|
||||
Enabled bool
|
||||
}
|
||||
|
||||
// CreateModel adds a model to a provider.
|
||||
func (s *Store) CreateModel(ctx context.Context, in ModelInput) (Model, error) {
|
||||
in.Name = strings.TrimSpace(in.Name)
|
||||
if in.Name == "" || in.ProviderID.IsZero() {
|
||||
return Model{}, fmt.Errorf("%w: provider and model name required", ErrInvalid)
|
||||
}
|
||||
if in.Kind != "chat" && in.Kind != "embedding" {
|
||||
in.Kind = "chat"
|
||||
}
|
||||
if in.Params == nil {
|
||||
in.Params = map[string]any{}
|
||||
}
|
||||
params, _ := json.Marshal(in.Params)
|
||||
id := ids.New()
|
||||
_, err := s.db.Exec(ctx, `INSERT INTO llm_models (id, provider_id, name, kind, context_window, dims, input_cost_per_m, output_cost_per_m, params, enabled) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10)`,
|
||||
id, in.ProviderID, in.Name, in.Kind, in.ContextWindow, in.Dims, in.InputCostPerM, in.OutputCostPerM, params, in.Enabled)
|
||||
if err != nil {
|
||||
return Model{}, storage.Translate(err)
|
||||
}
|
||||
return s.GetModel(ctx, id)
|
||||
}
|
||||
|
||||
// UpdateModel edits a model.
|
||||
func (s *Store) UpdateModel(ctx context.Context, id ids.ID, in ModelInput) (Model, error) {
|
||||
params, _ := json.Marshal(in.Params)
|
||||
tag, err := s.db.Exec(ctx, `UPDATE llm_models SET name=$2, kind=$3, context_window=$4, dims=$5, input_cost_per_m=$6, output_cost_per_m=$7, params=$8, enabled=$9 WHERE id=$1`,
|
||||
id, strings.TrimSpace(in.Name), in.Kind, in.ContextWindow, in.Dims, in.InputCostPerM, in.OutputCostPerM, params, in.Enabled)
|
||||
if err != nil {
|
||||
return Model{}, storage.Translate(err)
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return Model{}, storage.ErrNotFound
|
||||
}
|
||||
return s.GetModel(ctx, id)
|
||||
}
|
||||
|
||||
// DeleteModel removes a model.
|
||||
func (s *Store) DeleteModel(ctx context.Context, id ids.ID) error {
|
||||
tag, err := s.db.Exec(ctx, `DELETE FROM llm_models WHERE id=$1`, id)
|
||||
if err == nil && tag.RowsAffected() == 0 {
|
||||
return storage.ErrNotFound
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// GetModel fetches a model with its provider name.
|
||||
func (s *Store) GetModel(ctx context.Context, id ids.ID) (Model, error) {
|
||||
return scanModel(s.db.QueryRow(ctx, `SELECT `+modelCols+` FROM llm_models m JOIN llm_providers p ON p.id=m.provider_id WHERE m.id=$1`, id))
|
||||
}
|
||||
|
||||
// ListModels lists all models.
|
||||
func (s *Store) ListModels(ctx context.Context) ([]Model, error) {
|
||||
rows, err := s.db.Query(ctx, `SELECT `+modelCols+` FROM llm_models m JOIN llm_providers p ON p.id=m.provider_id ORDER BY lower(p.name), lower(m.name)`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []Model
|
||||
for rows.Next() {
|
||||
m, err := scanModel(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, m)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// ---- roles ----
|
||||
|
||||
// ListRoles returns assignments for every known role (unassigned roles have nil ModelID).
|
||||
func (s *Store) ListRoles(ctx context.Context) ([]RoleAssignment, error) {
|
||||
rows, err := s.db.Query(ctx, `SELECT role, model_id, fallback_ids, temperature, max_tokens, updated_by FROM llm_roles`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
byRole := map[llm.Role]RoleAssignment{}
|
||||
for rows.Next() {
|
||||
var a RoleAssignment
|
||||
var role string
|
||||
var fallbacks []string
|
||||
if err := rows.Scan(&role, &a.ModelID, &fallbacks, &a.Temperature, &a.MaxTokens, &a.UpdatedBy); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
a.Role = llm.Role(role)
|
||||
for _, f := range fallbacks {
|
||||
a.FallbackIDs = append(a.FallbackIDs, ids.ID(f))
|
||||
}
|
||||
byRole[a.Role] = a
|
||||
}
|
||||
out := make([]RoleAssignment, 0, len(llm.Roles))
|
||||
for _, r := range llm.Roles {
|
||||
a, ok := byRole[r.Role]
|
||||
if !ok {
|
||||
a = RoleAssignment{Role: r.Role}
|
||||
}
|
||||
if a.FallbackIDs == nil {
|
||||
a.FallbackIDs = []ids.ID{}
|
||||
}
|
||||
out = append(out, a)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// SetRole assigns a role.
|
||||
func (s *Store) SetRole(ctx context.Context, a RoleAssignment) error {
|
||||
known := false
|
||||
for _, r := range llm.Roles {
|
||||
if r.Role == a.Role {
|
||||
known = true
|
||||
}
|
||||
}
|
||||
if !known {
|
||||
return fmt.Errorf("%w: unknown role %q", ErrInvalid, a.Role)
|
||||
}
|
||||
fallbacks := make([]string, 0, len(a.FallbackIDs))
|
||||
for _, f := range a.FallbackIDs {
|
||||
fallbacks = append(fallbacks, string(f))
|
||||
}
|
||||
_, err := s.db.Exec(ctx, `INSERT INTO llm_roles (role, model_id, fallback_ids, temperature, max_tokens, updated_by) VALUES ($1,$2,$3,$4,$5,$6)
|
||||
ON CONFLICT (role) DO UPDATE SET model_id=excluded.model_id, fallback_ids=excluded.fallback_ids, temperature=excluded.temperature, max_tokens=excluded.max_tokens, updated_by=excluded.updated_by, updated_at=now()`,
|
||||
string(a.Role), a.ModelID, fallbacks, a.Temperature, a.MaxTokens, a.UpdatedBy)
|
||||
return err
|
||||
}
|
||||
|
||||
// ---- calls ----
|
||||
|
||||
// RecordCall stores a call.
|
||||
func (s *Store) RecordCall(ctx context.Context, c Call) error {
|
||||
if c.ID.IsZero() {
|
||||
c.ID = ids.New()
|
||||
}
|
||||
_, err := s.db.Exec(ctx, `INSERT INTO llm_calls (id, mind_id, correlation_id, role, purpose, provider, model, prompt_name, prompt_version, request_hash, request, response, error, prompt_tokens, output_tokens, cost, latency_ms, seed, replayed)
|
||||
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19)`,
|
||||
c.ID, nullID(c.MindID), nullID(c.CorrelationID), string(c.Role), c.Purpose, c.Provider, c.Model, c.PromptName, c.PromptVersion, c.RequestHash, c.Request, nullRaw(c.Response), c.Error, c.PromptTokens, c.OutputTokens, c.Cost, c.LatencyMS, c.Seed, c.Replayed)
|
||||
return err
|
||||
}
|
||||
|
||||
func nullID(id ids.ID) *string {
|
||||
if id.IsZero() {
|
||||
return nil
|
||||
}
|
||||
s := string(id)
|
||||
return &s
|
||||
}
|
||||
|
||||
func nullRaw(r json.RawMessage) *json.RawMessage {
|
||||
if len(r) == 0 {
|
||||
return nil
|
||||
}
|
||||
return &r
|
||||
}
|
||||
|
||||
// CallFilter selects calls.
|
||||
type CallFilter struct {
|
||||
MindID ids.ID
|
||||
CorrelationID ids.ID
|
||||
Role llm.Role
|
||||
ErrorsOnly bool
|
||||
Limit int
|
||||
}
|
||||
|
||||
// ListCalls returns calls newest first.
|
||||
func (s *Store) ListCalls(ctx context.Context, f CallFilter) ([]Call, error) {
|
||||
where := []string{"TRUE"}
|
||||
args := []any{}
|
||||
arg := func(v any) string { args = append(args, v); return fmt.Sprintf("$%d", len(args)) }
|
||||
if !f.MindID.IsZero() {
|
||||
where = append(where, "mind_id = "+arg(f.MindID))
|
||||
}
|
||||
if !f.CorrelationID.IsZero() {
|
||||
where = append(where, "correlation_id = "+arg(f.CorrelationID))
|
||||
}
|
||||
if f.Role != "" {
|
||||
where = append(where, "role = "+arg(string(f.Role)))
|
||||
}
|
||||
if f.ErrorsOnly {
|
||||
where = append(where, "error <> ''")
|
||||
}
|
||||
limit := f.Limit
|
||||
if limit <= 0 || limit > 500 {
|
||||
limit = 50
|
||||
}
|
||||
rows, err := s.db.Query(ctx, `SELECT id, mind_id, correlation_id, role, purpose, provider, model, prompt_name, prompt_version, request_hash, request, response, error, prompt_tokens, output_tokens, cost, latency_ms, seed, replayed, created_at
|
||||
FROM llm_calls WHERE `+strings.Join(where, " AND ")+` ORDER BY created_at DESC LIMIT `+arg(limit), args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []Call
|
||||
for rows.Next() {
|
||||
c, err := scanCall(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, c)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func scanCall(row pgx.Row) (Call, error) {
|
||||
var c Call
|
||||
var mind, corr *string
|
||||
var role string
|
||||
var resp *json.RawMessage
|
||||
if err := row.Scan(&c.ID, &mind, &corr, &role, &c.Purpose, &c.Provider, &c.Model, &c.PromptName, &c.PromptVersion, &c.RequestHash, &c.Request, &resp, &c.Error, &c.PromptTokens, &c.OutputTokens, &c.Cost, &c.LatencyMS, &c.Seed, &c.Replayed, &c.CreatedAt); err != nil {
|
||||
return c, storage.Translate(err)
|
||||
}
|
||||
c.Role = llm.Role(role)
|
||||
if mind != nil {
|
||||
c.MindID = ids.ID(*mind)
|
||||
}
|
||||
if corr != nil {
|
||||
c.CorrelationID = ids.ID(*corr)
|
||||
}
|
||||
if resp != nil {
|
||||
c.Response = *resp
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
// GetCall fetches a call.
|
||||
func (s *Store) GetCall(ctx context.Context, id ids.ID) (Call, error) {
|
||||
return scanCall(s.db.QueryRow(ctx, `SELECT id, mind_id, correlation_id, role, purpose, provider, model, prompt_name, prompt_version, request_hash, request, response, error, prompt_tokens, output_tokens, cost, latency_ms, seed, replayed, created_at FROM llm_calls WHERE id=$1`, id))
|
||||
}
|
||||
|
||||
// FindRecorded returns the most recent successful call with the same request
|
||||
// hash, for exact replay.
|
||||
func (s *Store) FindRecorded(ctx context.Context, hash string) (Call, error) {
|
||||
return scanCall(s.db.QueryRow(ctx, `SELECT id, mind_id, correlation_id, role, purpose, provider, model, prompt_name, prompt_version, request_hash, request, response, error, prompt_tokens, output_tokens, cost, latency_ms, seed, replayed, created_at
|
||||
FROM llm_calls WHERE request_hash=$1 AND error='' AND NOT replayed ORDER BY created_at DESC LIMIT 1`, hash))
|
||||
}
|
||||
|
||||
// UsageRow is aggregated usage.
|
||||
type UsageRow struct {
|
||||
Provider string `json:"provider"`
|
||||
Model string `json:"model"`
|
||||
Role string `json:"role"`
|
||||
Calls int64 `json:"calls"`
|
||||
Errors int64 `json:"errors"`
|
||||
PromptTokens int64 `json:"prompt_tokens"`
|
||||
OutputTokens int64 `json:"output_tokens"`
|
||||
Cost float64 `json:"cost"`
|
||||
AvgLatencyMS float64 `json:"avg_latency_ms"`
|
||||
}
|
||||
|
||||
// Usage aggregates calls since a time, grouped by provider, model and role.
|
||||
func (s *Store) Usage(ctx context.Context, since time.Time, mind ids.ID) ([]UsageRow, error) {
|
||||
rows, err := s.db.Query(ctx, `SELECT provider, model, role, count(*), count(*) FILTER (WHERE error <> ''), coalesce(sum(prompt_tokens),0), coalesce(sum(output_tokens),0), coalesce(sum(cost),0), coalesce(avg(latency_ms),0)
|
||||
FROM llm_calls WHERE created_at >= $1 AND ($2::uuid IS NULL OR mind_id = $2) GROUP BY provider, model, role ORDER BY sum(prompt_tokens)+sum(output_tokens) DESC`, since, nullID(mind))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []UsageRow
|
||||
for rows.Next() {
|
||||
var u UsageRow
|
||||
if err := rows.Scan(&u.Provider, &u.Model, &u.Role, &u.Calls, &u.Errors, &u.PromptTokens, &u.OutputTokens, &u.Cost, &u.AvgLatencyMS); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, u)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// TokensSince sums tokens used since a time, optionally for one mind. Used
|
||||
// for budget enforcement.
|
||||
func (s *Store) TokensSince(ctx context.Context, since time.Time, mind ids.ID) (int64, error) {
|
||||
var n int64
|
||||
err := s.db.QueryRow(ctx, `SELECT coalesce(sum(prompt_tokens + output_tokens),0) FROM llm_calls WHERE created_at >= $1 AND NOT replayed AND ($2::uuid IS NULL OR mind_id = $2)`, since, nullID(mind)).Scan(&n)
|
||||
return n, err
|
||||
}
|
||||
Reference in New Issue
Block a user