feat: per-run model selection, usage in diagnosis results and a Claude Messages API provider
캠페인마다 진단 모델을 고를 수 있게 한다. 하네스는 runs.model 열을 추가해 POST /runs의 model을 저장하고, 활성 런의 model을 /diagnose 요청에 실어 보낸다. 에이전트는 요청의 model을 우선 사용하고(없으면 OPENAI_MODEL 기본값), 모델 id가 claude-로 시작하면 internal/anthropic(Claude Messages API 어댑터)로 라우팅한다. 어댑터는 Responses API 항목(message/function_call/function_call_output)을 system·tools(input_schema)·tool_use/ tool_result로 변환하고 previous_response_id 체인을 클라이언트 쪽에서 재현하며, Claude 5.5가 강제 tool_choice를 거부하므로 강제 호출은 auto + 지시문으로 바꾼다. 키는 ANTHROPIC_KEY_FILE(/run/secrets/anthropic-api-key)에서 읽고 비어 있으면 명확한 오류를 낸다. Result에 usage(입력·출력·합계 토큰)를 추가해 하네스 diagnosis 이벤트에 비용 근거를 남긴다.
This commit is contained in:
seon committed
1 parent
fa9c08c778
commit
53ff243183
8 files changed
+492
-12
No files matched your search
@@ -35,9 +35,16 @@ type Config struct {
|
||||
ConfidenceMin float64 // 0.5 (design 6.3)
|
||||
}
|
||||
|
||||
// Responder is the one call the diagnosis loop needs from an LLM provider (OpenAI Responses API natively,
|
||||
// the Claude Messages API through internal/anthropic).
|
||||
type Responder interface {
|
||||
CreateResponse(ctx context.Context, req openai.Request) (*openai.Response, error)
|
||||
}
|
||||
|
||||
type Agent struct {
|
||||
cfg Config
|
||||
oai *openai.Client
|
||||
anth Responder // nil = claude-* models unavailable
|
||||
tools *tools.Runner
|
||||
harness *harness.Client
|
||||
store *rag.Store
|
||||
@@ -64,6 +71,20 @@ func New(cfg Config, oai *openai.Client, tr *tools.Runner, hc *harness.Client, s
|
||||
return &Agent{cfg: cfg, oai: oai, tools: tr, harness: hc, store: store, enums: enums, promptV: promptVersion, http: &http.Client{Timeout: 30 * time.Second}}
|
||||
}
|
||||
|
||||
// SetAnthropic enables claude-* model ids (Claude Messages API adapter).
|
||||
func (a *Agent) SetAnthropic(c Responder) { a.anth = c }
|
||||
|
||||
// provider picks the LLM client for a model id: claude-* -> Anthropic adapter, everything else -> OpenAI.
|
||||
func (a *Agent) provider(model string) (Responder, error) {
|
||||
if strings.HasPrefix(model, "claude-") {
|
||||
if a.anth == nil {
|
||||
return nil, fmt.Errorf("model %s: anthropic provider not configured", model)
|
||||
}
|
||||
return a.anth, nil
|
||||
}
|
||||
return a.oai, nil
|
||||
}
|
||||
|
||||
// EnableShadow turns the Jev gates on (JEV_MODE=shadow records only; enforce also applies the review rules).
|
||||
func (a *Agent) EnableShadow(c *jev.Client, enforce bool) { a.jev, a.jevEnforce = c, enforce }
|
||||
|
||||
@@ -72,6 +93,7 @@ type Request struct {
|
||||
Alert map[string]any `json:"alert"`
|
||||
RunID *string `json:"run_id"`
|
||||
ScenarioID *string `json:"scenario_id"`
|
||||
Model string `json:"model,omitempty"` // campaign-level model selection (harness runs.model); "" = Config.Model
|
||||
}
|
||||
|
||||
// Result is returned to the harness: the diagnosis fields at top level (the harness stores the whole object as the
|
||||
@@ -84,7 +106,8 @@ type Result struct {
|
||||
Turns int `json:"turns"`
|
||||
ToolCalls int `json:"tool_calls"`
|
||||
DurationMs int64 `json:"duration_ms"`
|
||||
Disposition string `json:"disposition"` // proposal | escalated | none
|
||||
Usage map[string]int `json:"usage,omitempty"` // input_tokens / output_tokens / total_tokens summed over the turns
|
||||
Disposition string `json:"disposition"` // proposal | escalated | none
|
||||
ProposalToken string `json:"proposal_token,omitempty"`
|
||||
Jev map[string]any `json:"jev,omitempty"` // shadow gate answers (recorded, never acted on)
|
||||
}
|
||||
@@ -155,13 +178,17 @@ func (a *Agent) Diagnose(ctx context.Context, req Request) (*Result, error) {
|
||||
func (a *Agent) DiagnoseWith(ctx context.Context, req Request, prefetched string) (*Result, error) {
|
||||
start := time.Now()
|
||||
sid := "ses-" + uuid.NewString()[:8]
|
||||
res := &Result{SessionID: sid, Model: a.cfg.Model, PromptVersion: a.promptV}
|
||||
model := a.cfg.Model
|
||||
if req.Model != "" {
|
||||
model = req.Model
|
||||
}
|
||||
res := &Result{SessionID: sid, Model: model, PromptVersion: a.promptV}
|
||||
var calls []toolCallLog
|
||||
input := []openai.Item{
|
||||
{Type: "message", Role: "system", Content: a.cfg.SystemPrompt},
|
||||
{Type: "message", Role: "user", Content: a.userMessage(ctx, req, prefetched)},
|
||||
}
|
||||
session := rag.Session{SessionID: sid, RunID: req.RunID, ScenarioID: req.ScenarioID, Model: a.cfg.Model, PromptVersion: a.promptV, Input: map[string]any{"alert": req.Alert, "user_message": input[1].Content}}
|
||||
session := rag.Session{SessionID: sid, RunID: req.RunID, ScenarioID: req.ScenarioID, Model: model, PromptVersion: a.promptV, Input: map[string]any{"alert": req.Alert, "user_message": input[1].Content}}
|
||||
finish := func(err error) {
|
||||
session.Output = res
|
||||
session.ToolCalls = calls
|
||||
@@ -176,6 +203,11 @@ func (a *Agent) DiagnoseWith(ctx context.Context, req Request, prefetched string
|
||||
}
|
||||
}
|
||||
|
||||
llm, err := a.provider(model)
|
||||
if err != nil {
|
||||
finish(err)
|
||||
return nil, err
|
||||
}
|
||||
defs := a.toolDefs()
|
||||
var prevID string
|
||||
var diagnosis *diag.Diagnosis
|
||||
@@ -185,14 +217,14 @@ func (a *Agent) DiagnoseWith(ctx context.Context, req Request, prefetched string
|
||||
usage := map[string]int{"input_tokens": 0, "output_tokens": 0, "total_tokens": 0, "turns": 0}
|
||||
for turn := 1; turn <= a.cfg.MaxTurns && diagnosis == nil; turn++ {
|
||||
res.Turns = turn
|
||||
r := openai.Request{Model: a.cfg.Model, Input: input, Tools: defs, ToolChoice: "auto", MaxOutputTokens: a.cfg.MaxOutputTokens, ParallelToolCalls: &pf}
|
||||
r := openai.Request{Model: model, Input: input, Tools: defs, ToolChoice: "auto", MaxOutputTokens: a.cfg.MaxOutputTokens, ParallelToolCalls: &pf}
|
||||
if prevID != "" {
|
||||
r.PreviousResponseID = prevID
|
||||
}
|
||||
if turn == a.cfg.MaxTurns { // last chance: force the final call
|
||||
r.ToolChoice = map[string]any{"type": "function", "name": "submit_diagnosis"}
|
||||
}
|
||||
resp, err := a.oai.CreateResponse(ctx, r)
|
||||
resp, err := llm.CreateResponse(ctx, r)
|
||||
if err != nil {
|
||||
finish(err)
|
||||
return nil, err
|
||||
@@ -228,6 +260,7 @@ func (a *Agent) DiagnoseWith(ctx context.Context, req Request, prefetched string
|
||||
}
|
||||
}
|
||||
session.Usage = usage
|
||||
res.Usage = usage
|
||||
res.ToolCalls = len(calls)
|
||||
res.DurationMs = time.Since(start).Milliseconds()
|
||||
if diagnosis == nil {
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/anthropic"
|
||||
"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/diag"
|
||||
"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/harness"
|
||||
"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/openai"
|
||||
@@ -81,3 +82,48 @@ func TestTargetFor(t *testing.T) {
|
||||
t.Fatal("mig")
|
||||
}
|
||||
}
|
||||
|
||||
// A claude-* model id routes the same loop through the Anthropic adapter; the request model wins over Config.Model
|
||||
// and the summed usage is returned in the result.
|
||||
func TestDiagnoseRoutesClaudeModel(t *testing.T) {
|
||||
var models []string
|
||||
anth := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var b map[string]any
|
||||
_ = json.NewDecoder(r.Body).Decode(&b)
|
||||
models = append(models, b["model"].(string))
|
||||
args := `{"root_cause_label":"db_primary_down","confidence":0.8,"evidence":["e"],"proposed_action":"escalate","target":"db-primary","rationale":"r"}`
|
||||
w.Write([]byte(`{"id":"msg_a","type":"message","model":"claude-sonnet-5-5","stop_reason":"tool_use","content":[{"type":"tool_use","id":"toolu_a","name":"submit_diagnosis","input":` + args + `}],"usage":{"input_tokens":40,"output_tokens":8}}`))
|
||||
}))
|
||||
defer anth.Close()
|
||||
oaiSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/embeddings" {
|
||||
w.Write([]byte(`{"data":[{"index":0,"embedding":[0.1,0.2]}]}`))
|
||||
return
|
||||
}
|
||||
t.Errorf("openai must not be called for a claude-* model: %s", r.URL.Path)
|
||||
}))
|
||||
defer oaiSrv.Close()
|
||||
mimir := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Write([]byte(`{"status":"success","data":{"resultType":"vector","result":[]}}`))
|
||||
}))
|
||||
defer mimir.Close()
|
||||
oai := openai.New(oaiSrv.URL, func() (string, error) { return "test-key", nil })
|
||||
tr := tools.New(tools.Config{MimirURL: mimir.URL, LokiURL: mimir.URL, EmbedModel: "e"}, nil, oai)
|
||||
enums := diag.Enums{Labels: diag.DefaultLabels, Actions: diag.DefaultActions}
|
||||
ag := New(Config{Model: "gpt-default", SystemPrompt: "sys", Topology: "topo", SLO: "slo", HitlURL: "http://127.0.0.1:1", MaxTurns: 3}, oai, tr, harness.New("http://127.0.0.1:1"), nil, enums, "pv1")
|
||||
req := Request{Alert: map[string]any{"alertname": "DbPrimaryDown"}, Model: "claude-sonnet-5-5"}
|
||||
if _, err := ag.Diagnose(context.Background(), req); err == nil || !strings.Contains(err.Error(), "anthropic provider not configured") {
|
||||
t.Fatalf("without SetAnthropic a claude-* model must fail clearly, got %v", err)
|
||||
}
|
||||
ag.SetAnthropic(anthropic.New(anth.URL, func() (string, error) { return "k", nil }))
|
||||
res, err := ag.Diagnose(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if res.Model != "claude-sonnet-5-5" || len(models) != 1 || models[0] != "claude-sonnet-5-5" || res.RootCauseLabel != "db_primary_down" {
|
||||
t.Fatalf("routing: res=%+v models=%v", res, models)
|
||||
}
|
||||
if res.Usage["input_tokens"] != 40 || res.Usage["output_tokens"] != 8 || res.Usage["total_tokens"] != 48 || res.Usage["turns"] != 1 {
|
||||
t.Fatalf("usage: %v", res.Usage)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user