feat(llm): add profile failover and pool health
This commit is contained in:
pphua committed
1 parent
d4b21f0bdb
commit
18926c129f
15 files changed
+2507
-761
No files matched your search
+136
-39
@@ -11,25 +11,55 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
"github.com/Autumn-27/norma/transcript"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
type taskIDContextKey struct{}
|
||||
|
||||
// WithTaskID attaches the owning task registry id to an LLM call. Session ids
|
||||
// are based on exploration ids, which are not interchangeable with task ids.
|
||||
func WithTaskID(ctx context.Context, taskID string) context.Context {
|
||||
taskID = strings.TrimSpace(taskID)
|
||||
if taskID == "" {
|
||||
return ctx
|
||||
}
|
||||
return context.WithValue(ctx, taskIDContextKey{}, taskID)
|
||||
}
|
||||
|
||||
// TaskIDFrom returns the explicit task registry id attached by the task runtime.
|
||||
func TaskIDFrom(ctx context.Context) string {
|
||||
if ctx == nil {
|
||||
return ""
|
||||
}
|
||||
taskID, _ := ctx.Value(taskIDContextKey{}).(string)
|
||||
return strings.TrimSpace(taskID)
|
||||
}
|
||||
|
||||
// Recorder wraps an llm.Provider and records every completion call.
|
||||
type Recorder struct {
|
||||
inner llm.Provider
|
||||
pg *db.DB
|
||||
model string // model name (from config, not in CompletionRequest)
|
||||
prof string // LLM profile name (from llm_profiles)
|
||||
enabled func() bool // reports whether recording is currently on; nil = always record
|
||||
model string // model name (from config, not in CompletionRequest)
|
||||
prof string // LLM profile name (from llm_profiles)
|
||||
// thinkingType / reasoningEffort 是配置级思考参数(思考开关 / 思考强度)。它们在
|
||||
// norma 的 buildBody() 里从 provider 配置注入真正的 HTTP body,不出现在
|
||||
// CompletionRequest 上,故 Recorder 需在此单独带一份,序列化时写进录制。
|
||||
thinkingType string
|
||||
reasoningEffort string
|
||||
enabled func() bool // reports whether recording is currently on; nil = always record
|
||||
}
|
||||
|
||||
// Wrap returns a Provider that records calls to pg when enabled() reports true.
|
||||
// profName is the LLM profile name (e.g. "default"); may be empty. A nil enabled
|
||||
// predicate records unconditionally (backward-compatible default).
|
||||
func Wrap(inner llm.Provider, pg *db.DB, model, profName string, enabled func() bool) *Recorder {
|
||||
return &Recorder{inner: inner, pg: pg, model: model, prof: profName, enabled: enabled}
|
||||
// profName is the LLM profile name (e.g. "default"); may be empty. thinkingType /
|
||||
// reasoningEffort are the config-level thinking params actually sent to the API
|
||||
// (empty = not sent). A nil enabled predicate records unconditionally.
|
||||
func Wrap(inner llm.Provider, pg *db.DB, model, profName, thinkingType, reasoningEffort string, enabled func() bool) *Recorder {
|
||||
return &Recorder{
|
||||
inner: inner, pg: pg, model: model, prof: profName,
|
||||
thinkingType: thinkingType, reasoningEffort: reasoningEffort, enabled: enabled,
|
||||
}
|
||||
}
|
||||
|
||||
// parseSession extracts task id and worker role from session strings like
|
||||
@@ -63,64 +93,123 @@ func parseSession(s string) (taskID, worker string) {
|
||||
// harness via transcript.WithSessionID; reading it per-call is race-free even when
|
||||
// planner and multiple workers share one Recorder instance.
|
||||
func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
|
||||
// Recording off → passthrough with zero overhead: no request serialization,
|
||||
// no response accumulation, no DB insert.
|
||||
if r.enabled != nil && !r.enabled() {
|
||||
return r.inner.Stream(ctx, req)
|
||||
}
|
||||
// Body recording (the heavy debug trace: full request/response) is gated by the
|
||||
// llm_record setting. Lightweight usage metering always runs — it powers token
|
||||
// stats and must be complete even for interrupted/failed runs, so it is NOT
|
||||
// gated. Only body serialization + response accumulation are skipped when off.
|
||||
recordBodies := r.enabled == nil || r.enabled()
|
||||
start := time.Now()
|
||||
session := transcript.SessionIDFrom(ctx)
|
||||
taskID, worker := parseSession(session)
|
||||
parsedID, worker := parseSession(session)
|
||||
expID := db.ParseExpID(parsedID)
|
||||
taskID := TaskIDFrom(ctx)
|
||||
if taskID == "" {
|
||||
// Backward compatibility for non-task callers. For task calls, the task
|
||||
// runtime always supplies the registry id explicitly.
|
||||
taskID = parsedID
|
||||
}
|
||||
|
||||
// Serialize the request for storage.
|
||||
reqBody := serializeRequest(req)
|
||||
// Serialize the request only when storing bodies (this is the expensive part).
|
||||
reqBody := ""
|
||||
if recordBodies {
|
||||
reqBody = r.serializeRequest(req)
|
||||
}
|
||||
|
||||
return func(yield func(llm.StreamEvent, error) bool) {
|
||||
var (
|
||||
textBuf strings.Builder
|
||||
textBuf strings.Builder
|
||||
thinkingBuf strings.Builder
|
||||
usage llm.Usage
|
||||
stopReason string
|
||||
streamErr error
|
||||
usage llm.Usage
|
||||
stopReason string
|
||||
streamErr error
|
||||
)
|
||||
|
||||
finish := func(err error) {
|
||||
status := "ok"
|
||||
if err != nil {
|
||||
status = "error"
|
||||
}
|
||||
// Lightweight metering row — always written.
|
||||
r.recordUsage(taskID, expID, worker, usage, int(time.Since(start).Milliseconds()), status)
|
||||
// Heavy trace row — only when body recording is on.
|
||||
if recordBodies {
|
||||
r.record(req, session, taskID, worker, reqBody, start, textBuf.String(), thinkingBuf.String(), usage, stopReason, err)
|
||||
}
|
||||
}
|
||||
|
||||
for ev, err := range r.inner.Stream(ctx, req) {
|
||||
if err != nil {
|
||||
streamErr = err
|
||||
// Record the failed call, then propagate.
|
||||
r.record(req, session, taskID, worker, reqBody, start, textBuf.String(), thinkingBuf.String(), usage, "", err)
|
||||
finish(streamErr)
|
||||
if !yield(ev, err) {
|
||||
return
|
||||
}
|
||||
return
|
||||
}
|
||||
// Accumulate response content.
|
||||
// Always track usage (cheap); accumulate text/thinking only for bodies.
|
||||
// Anthropic (and the other providers) split token usage across events:
|
||||
// message_start carries ONLY the input side (input + cache), message_delta
|
||||
// ONLY the output. They must be FOLDED with Add — overwriting on delta
|
||||
// would zero out the input already counted at start (mirrors the SDK's own
|
||||
// llm.Accumulator; see norma/llm/accumulate.go).
|
||||
switch ev.Type {
|
||||
case llm.SETextDelta:
|
||||
textBuf.WriteString(ev.Text)
|
||||
if recordBodies {
|
||||
textBuf.WriteString(ev.Text)
|
||||
}
|
||||
case llm.SEThinkingDelta:
|
||||
thinkingBuf.WriteString(ev.Text)
|
||||
if recordBodies {
|
||||
thinkingBuf.WriteString(ev.Text)
|
||||
}
|
||||
case llm.SEMessageStart:
|
||||
usage.Add(ev.Usage)
|
||||
case llm.SEMessageDelta:
|
||||
if ev.StopReason != "" {
|
||||
stopReason = ev.StopReason
|
||||
}
|
||||
usage = ev.Usage
|
||||
case llm.SEMessageStart:
|
||||
if ev.Usage.InputTokens > 0 {
|
||||
usage = ev.Usage
|
||||
}
|
||||
usage.Add(ev.Usage)
|
||||
}
|
||||
if !yield(ev, err) {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Stream completed normally — record it.
|
||||
r.record(req, session, taskID, worker, reqBody, start, textBuf.String(), thinkingBuf.String(), usage, stopReason, streamErr)
|
||||
// Stream completed normally.
|
||||
finish(streamErr)
|
||||
}
|
||||
}
|
||||
|
||||
// record persists one LLM call to PostgreSQL (async, best-effort).
|
||||
// recordUsage appends one lightweight metering row to llm_usage (no bodies). Skips
|
||||
// zero-token calls with no model, which carry nothing worth metering.
|
||||
func (r *Recorder) recordUsage(taskID string, expID int64, worker string, usage llm.Usage, latencyMs int, status string) {
|
||||
if r.pg == nil {
|
||||
return
|
||||
}
|
||||
if r.model == "" && usage.InputTokens == 0 && usage.OutputTokens == 0 &&
|
||||
usage.CacheReadTokens == 0 && usage.CacheWriteTokens == 0 {
|
||||
return
|
||||
}
|
||||
err := r.pg.InsertLLMUsage(&db.LLMUsage{
|
||||
TaskID: taskID,
|
||||
ExplorationID: expID,
|
||||
Worker: worker,
|
||||
Model: r.model,
|
||||
ProfileName: r.prof,
|
||||
LatencyMs: latencyMs,
|
||||
InputTokens: usage.InputTokens,
|
||||
OutputTokens: usage.OutputTokens,
|
||||
CacheRead: usage.CacheReadTokens,
|
||||
CacheWrite: usage.CacheWriteTokens,
|
||||
Status: status,
|
||||
})
|
||||
if err != nil {
|
||||
log.Printf("[llmusage] insert: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// record persists one LLM call to PostgreSQL before the provider stream returns.
|
||||
// Keeping the write inside the owning task operation means task deletion can
|
||||
// drain calls and then remove records without a late async insert recreating one.
|
||||
func (r *Recorder) record(req llm.CompletionRequest, session, taskID, worker, reqBody string, start time.Time, text, thinking string, usage llm.Usage, stopReason string, streamErr error) {
|
||||
latency := int(time.Since(start).Milliseconds())
|
||||
status := "ok"
|
||||
@@ -169,20 +258,28 @@ func (r *Recorder) record(req llm.CompletionRequest, session, taskID, worker, re
|
||||
ResponseBody: string(respBody),
|
||||
}
|
||||
|
||||
go func() {
|
||||
if err := r.pg.InsertLLMRecord(rec); err != nil {
|
||||
log.Printf("[llmrec] insert: %v", err)
|
||||
}
|
||||
}()
|
||||
if err := r.pg.InsertLLMRecord(rec); err != nil {
|
||||
log.Printf("[llmrec] insert: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// serializeRequest builds a JSON representation of the completion request.
|
||||
func serializeRequest(req llm.CompletionRequest) string {
|
||||
func (r *Recorder) serializeRequest(req llm.CompletionRequest) string {
|
||||
m := map[string]any{
|
||||
"system": req.System,
|
||||
"messages": req.Messages,
|
||||
"max_tokens": req.MaxTokens,
|
||||
}
|
||||
// 记录本次调用实际发出的思考参数。type 采用「有效值」：每请求覆盖 req.Thinking
|
||||
// 优先于配置级 thinkingType(与 norma buildBody 的判定一致，如 compaction 摘要会
|
||||
// 强制 disabled)；effort 无每请求覆盖，直接取配置值。两者皆空则不写 thinking 字段。
|
||||
effType := r.thinkingType
|
||||
if req.Thinking != "" {
|
||||
effType = req.Thinking
|
||||
}
|
||||
if effType != "" || r.reasoningEffort != "" {
|
||||
m["thinking"] = map[string]string{"type": effType, "effort": r.reasoningEffort}
|
||||
}
|
||||
if len(req.Tools) > 0 {
|
||||
// Store tool names only (full schemas are huge).
|
||||
names := make([]string, len(req.Tools))
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
package llmrec
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestTaskIDContextUsesExplicitRegistryID(t *testing.T) {
|
||||
ctx := WithTaskID(context.Background(), " 42 ")
|
||||
if got := TaskIDFrom(ctx); got != "42" {
|
||||
t.Fatalf("TaskIDFrom()=%q want 42", got)
|
||||
}
|
||||
if got := TaskIDFrom(WithTaskID(ctx, " ")); got != "42" {
|
||||
t.Fatalf("blank task id should preserve parent context, got %q", got)
|
||||
}
|
||||
if got := TaskIDFrom(nil); got != "" {
|
||||
t.Fatalf("nil context returned %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseSessionFallbackIsExplorationScoped(t *testing.T) {
|
||||
taskID, worker := parseSession("exp12-worker-i99")
|
||||
if taskID != "12" || worker != "worker" {
|
||||
t.Fatalf("parseSession()=(%q,%q)", taskID, worker)
|
||||
}
|
||||
if taskID, worker := parseSession("not-a-task"); taskID != "" || worker != "" {
|
||||
t.Fatalf("unexpected non-session parse: (%q,%q)", taskID, worker)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user