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 2026-10-14 23:21:46 +09:00
1 parent fa9c08c778
commit 53ff243183
8 files changed
+492 -12

No files matched your search

+6
View File
@@ -11,6 +11,7 @@ import (
"time" "time"
"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/agent" "github.com/selfheal-lab/selfheal-lab-apps/ai/internal/agent"
"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/diag"
"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/harness" "github.com/selfheal-lab/selfheal-lab-apps/ai/internal/harness"
"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/jev" "github.com/selfheal-lab/selfheal-lab-apps/ai/internal/jev"
@@ -95,6 +96,10 @@ func main() {
HitlURL: env("HITL_URL", "http://hitl:8081"), HitlURL: env("HITL_URL", "http://hitl:8081"),
} }
ag := agent.New(acfg, oai, tr, hc, store, enums, pv) ag := agent.New(acfg, oai, tr, hc, store, enums, pv)
// Claude Messages API for claude-* model ids chosen per campaign; the key may be registered later (empty = diagnoses with
// a claude-* model fail with "anthropic api key is empty", recorded in the harness diagnosis event).
anth := anthropic.New(env("ANTHROPIC_BASE_URL", ""), keyReader(env("ANTHROPIC_KEY_FILE", "/run/secrets/anthropic-api-key")))
ag.SetAnthropic(anth)
jevMode := env("JEV_MODE", "off") // off | shadow | enforce (shadow records Jev answers next to the diagnosis; enforce also applies the two review rules) jevMode := env("JEV_MODE", "off") // off | shadow | enforce (shadow records Jev answers next to the diagnosis; enforce also applies the two review rules)
var jc *jev.Client var jc *jev.Client
if jevMode == "shadow" || jevMode == "enforce" { if jevMode == "shadow" || jevMode == "enforce" {
@@ -130,6 +135,7 @@ func main() {
out := map[string]any{"model": acfg.Model, "prompt_version": pv, "labels": enums.Labels, "actions": enums.Actions} out := map[string]any{"model": acfg.Model, "prompt_version": pv, "labels": enums.Labels, "actions": enums.Actions}
key, kerr := keyReader(env("OPENAI_KEY_FILE", "/run/secrets/openai-api-key"))() key, kerr := keyReader(env("OPENAI_KEY_FILE", "/run/secrets/openai-api-key"))()
out["openai_key_present"] = kerr == nil && key != "" out["openai_key_present"] = kerr == nil && key != ""
out["anthropic_key_present"] = anth.KeyPresent()
out["jev_mode"] = jevMode out["jev_mode"] = jevMode
if jc != nil { if jc != nil {
out["jev_key_present"] = jc.KeyPresent() out["jev_key_present"] = jc.KeyPresent()
+38 -5
View File
@@ -35,9 +35,16 @@ type Config struct {
ConfidenceMin float64 // 0.5 (design 6.3) 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 { type Agent struct {
cfg Config cfg Config
oai *openai.Client oai *openai.Client
anth Responder // nil = claude-* models unavailable
tools *tools.Runner tools *tools.Runner
harness *harness.Client harness *harness.Client
store *rag.Store 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}} 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). // 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 } 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"` Alert map[string]any `json:"alert"`
RunID *string `json:"run_id"` RunID *string `json:"run_id"`
ScenarioID *string `json:"scenario_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 // 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"` Turns int `json:"turns"`
ToolCalls int `json:"tool_calls"` ToolCalls int `json:"tool_calls"`
DurationMs int64 `json:"duration_ms"` 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"` ProposalToken string `json:"proposal_token,omitempty"`
Jev map[string]any `json:"jev,omitempty"` // shadow gate answers (recorded, never acted on) 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) { func (a *Agent) DiagnoseWith(ctx context.Context, req Request, prefetched string) (*Result, error) {
start := time.Now() start := time.Now()
sid := "ses-" + uuid.NewString()[:8] 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 var calls []toolCallLog
input := []openai.Item{ input := []openai.Item{
{Type: "message", Role: "system", Content: a.cfg.SystemPrompt}, {Type: "message", Role: "system", Content: a.cfg.SystemPrompt},
{Type: "message", Role: "user", Content: a.userMessage(ctx, req, prefetched)}, {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) { finish := func(err error) {
session.Output = res session.Output = res
session.ToolCalls = calls 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() defs := a.toolDefs()
var prevID string var prevID string
var diagnosis *diag.Diagnosis 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} 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++ { for turn := 1; turn <= a.cfg.MaxTurns && diagnosis == nil; turn++ {
res.Turns = 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 != "" { if prevID != "" {
r.PreviousResponseID = prevID r.PreviousResponseID = prevID
} }
if turn == a.cfg.MaxTurns { // last chance: force the final call if turn == a.cfg.MaxTurns { // last chance: force the final call
r.ToolChoice = map[string]any{"type": "function", "name": "submit_diagnosis"} 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 { if err != nil {
finish(err) finish(err)
return nil, err return nil, err
@@ -228,6 +260,7 @@ func (a *Agent) DiagnoseWith(ctx context.Context, req Request, prefetched string
} }
} }
session.Usage = usage session.Usage = usage
res.Usage = usage
res.ToolCalls = len(calls) res.ToolCalls = len(calls)
res.DurationMs = time.Since(start).Milliseconds() res.DurationMs = time.Since(start).Milliseconds()
if diagnosis == nil { if diagnosis == nil {
+46
View File
@@ -8,6 +8,7 @@ import (
"strings" "strings"
"testing" "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/diag"
"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/harness" "github.com/selfheal-lab/selfheal-lab-apps/ai/internal/harness"
"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/openai" "github.com/selfheal-lab/selfheal-lab-apps/ai/internal/openai"
@@ -81,3 +82,48 @@ func TestTargetFor(t *testing.T) {
t.Fatal("mig") 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)
}
}
+262
View File
@@ -0,0 +1,262 @@
// Package anthropic adapts the Claude Messages API to the small Responses-API surface the agent loop uses
// (openai.Request / openai.Response with message, function_call and function_call_output items), so one loop
// serves both providers. The agent picks this client for model ids starting with "claude-".
//
// Wire format (Messages API reference, anthropic-version 2023-06-01): system prompt as the top-level `system`
// field, tools as {name, description, input_schema, strict}, tool calls as `tool_use` content blocks answered by
// `tool_result` blocks in the next user message. The Responses-API `previous_response_id` chain is emulated
// client-side: every response id maps to the full message list up to and including that assistant turn.
//
// Claude Opus 5.5 / Sonnet 5.5 reject forced tool use (tool_choice type "tool"/"any" -> 400), so a forced
// tool_choice from the loop becomes {type:"auto"} plus a one-line instruction in the user turn.
package anthropic
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"sync"
"time"
"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/openai"
)
const (
apiVersion = "2023-06-01"
maxConvCached = 64 // response ids kept for previous_response_id chaining (one diagnosis uses <= MaxTurns)
defaultTokens = 4000
)
type Client struct {
BaseURL string
HTTP *http.Client
keyFn func() (string, error)
mu sync.Mutex
conv map[string]conversation // response id -> messages (and system prompt) up to and including that assistant turn
order []string // insertion order for eviction
seq int
}
type conversation struct {
system string
msgs []message
}
func New(baseURL string, keyFn func() (string, error)) *Client {
if baseURL == "" {
baseURL = "https://api.anthropic.com"
}
return &Client{BaseURL: strings.TrimRight(baseURL, "/"), HTTP: &http.Client{Timeout: 150 * time.Second}, keyFn: keyFn, conv: map[string]conversation{}}
}
// KeyPresent reports whether the key file has a non-empty value (the secret may not be registered yet).
func (c *Client) KeyPresent() bool {
k, err := c.keyFn()
return err == nil && k != ""
}
type message struct {
Role string `json:"role"`
Content any `json:"content"` // string or []block
}
type block struct {
Type string `json:"type"`
Text string `json:"text,omitempty"`
ID string `json:"id,omitempty"` // tool_use
Name string `json:"name,omitempty"` // tool_use
Input json.RawMessage `json:"input,omitempty"` // tool_use
ToolUseID string `json:"tool_use_id,omitempty"` // tool_result
Content string `json:"content,omitempty"` // tool_result
}
type tool struct {
Name string `json:"name"`
Description string `json:"description,omitempty"`
InputSchema map[string]any `json:"input_schema"`
Strict bool `json:"strict,omitempty"`
}
type request struct {
Model string `json:"model"`
MaxTokens int `json:"max_tokens"`
System string `json:"system,omitempty"`
Messages []message `json:"messages"`
Tools []tool `json:"tools,omitempty"`
ToolChoice map[string]any `json:"tool_choice,omitempty"`
}
type response struct {
ID string `json:"id"`
Type string `json:"type"`
Model string `json:"model"`
StopReason string `json:"stop_reason"`
Content []block `json:"content"`
Usage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
CacheCreationInputTokens int `json:"cache_creation_input_tokens"`
CacheReadInputTokens int `json:"cache_read_input_tokens"`
} `json:"usage"`
Error *struct {
Type string `json:"type"`
Message string `json:"message"`
} `json:"error"`
}
// forcedNudge is appended to the user turn when the loop asked for a forced submit_diagnosis call.
const forcedNudge = "You must call the submit_diagnosis tool now with your final answer; do not call any other tool."
// CreateResponse translates one loop turn to a Messages API call and the answer back to Responses-API items.
func (c *Client) CreateResponse(ctx context.Context, req openai.Request) (*openai.Response, error) {
key, err := c.keyFn()
if err != nil {
return nil, fmt.Errorf("anthropic key: %w", err)
}
if key == "" {
return nil, fmt.Errorf("anthropic api key is empty (secret selfheal-anthropic-api-key not registered yet)")
}
var msgs []message
var system []string
if req.PreviousResponseID != "" {
c.mu.Lock()
prev, ok := c.conv[req.PreviousResponseID]
c.mu.Unlock()
if !ok {
return nil, fmt.Errorf("anthropic: unknown previous_response_id %q", req.PreviousResponseID)
}
msgs = append(msgs, prev.msgs...)
if prev.system != "" {
system = append(system, prev.system) // the loop sends the system message only on the first turn
}
}
var userBlocks []block
flush := func() {
if len(userBlocks) == 0 {
return
}
msgs = append(msgs, message{Role: "user", Content: userBlocks})
userBlocks = nil
}
for _, it := range req.Input {
switch it.Type {
case "message":
text, _ := it.Content.(string)
switch it.Role {
case "system":
system = append(system, text)
case "assistant":
flush()
msgs = append(msgs, message{Role: "assistant", Content: []block{{Type: "text", Text: text}}})
default:
userBlocks = append(userBlocks, block{Type: "text", Text: text})
}
case "function_call_output":
userBlocks = append(userBlocks, block{Type: "tool_result", ToolUseID: it.CallID, Content: it.Output})
}
}
forced := false
if m, ok := req.ToolChoice.(map[string]any); ok && m["type"] == "function" {
forced = true
}
if forced {
userBlocks = append(userBlocks, block{Type: "text", Text: forcedNudge})
}
flush()
if len(msgs) == 0 {
return nil, fmt.Errorf("anthropic: no messages to send")
}
body := request{Model: req.Model, MaxTokens: req.MaxOutputTokens, System: strings.Join(system, "\n\n"), Messages: msgs}
if body.MaxTokens <= 0 {
body.MaxTokens = defaultTokens
}
for _, t := range req.Tools {
body.Tools = append(body.Tools, tool{Name: t.Name, Description: t.Description, InputSchema: t.Parameters, Strict: t.Strict})
}
if len(body.Tools) > 0 {
tc := map[string]any{"type": "auto"}
if req.ParallelToolCalls != nil && !*req.ParallelToolCalls {
tc["disable_parallel_tool_use"] = true
}
body.ToolChoice = tc
}
b, _ := json.Marshal(body)
hreq, err := http.NewRequestWithContext(ctx, "POST", c.BaseURL+"/v1/messages", bytes.NewReader(b))
if err != nil {
return nil, err
}
hreq.Header.Set("x-api-key", key)
hreq.Header.Set("anthropic-version", apiVersion)
hreq.Header.Set("Content-Type", "application/json")
resp, err := c.HTTP.Do(hreq)
if err != nil {
return nil, err
}
defer resp.Body.Close()
rb, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
var r response
_ = json.Unmarshal(rb, &r)
if resp.StatusCode/100 != 2 {
msg := truncate(string(rb), 600)
if r.Error != nil {
msg = r.Error.Type + ": " + r.Error.Message
}
return nil, fmt.Errorf("anthropic /v1/messages: status %d: %s", resp.StatusCode, msg)
}
if r.Error != nil {
return nil, fmt.Errorf("anthropic response error: %s", r.Error.Message)
}
out := &openai.Response{ID: r.ID, Status: "completed", Model: r.Model}
if out.ID == "" {
c.mu.Lock()
c.seq++
out.ID = fmt.Sprintf("msg_local_%d", c.seq)
c.mu.Unlock()
}
var text strings.Builder
for _, bl := range r.Content {
switch bl.Type {
case "text":
text.WriteString(bl.Text)
case "tool_use":
args := string(bl.Input)
if args == "" {
args = "{}"
}
out.Output = append(out.Output, openai.Item{Type: "function_call", CallID: bl.ID, Name: bl.Name, Arguments: args, ID: bl.ID, Status: "completed"})
}
}
if text.Len() > 0 {
out.Output = append([]openai.Item{{Type: "message", Role: "assistant", Content: []any{map[string]any{"type": "output_text", "text": text.String()}}}}, out.Output...)
}
// input tokens include cache writes/reads so the sum matches what is billed
out.Usage.InputTokens = r.Usage.InputTokens + r.Usage.CacheCreationInputTokens + r.Usage.CacheReadInputTokens
out.Usage.OutputTokens = r.Usage.OutputTokens
out.Usage.TotalTokens = out.Usage.InputTokens + out.Usage.OutputTokens
// remember the conversation for the next turn (assistant content kept verbatim so tool_use ids line up)
full := append(append([]message{}, msgs...), message{Role: "assistant", Content: r.Content})
c.mu.Lock()
c.conv[out.ID] = conversation{system: body.System, msgs: full}
c.order = append(c.order, out.ID)
for len(c.order) > maxConvCached {
delete(c.conv, c.order[0])
c.order = c.order[1:]
}
c.mu.Unlock()
return out, nil
}
func truncate(s string, n int) string {
if len(s) <= n {
return s
}
return s[:n] + "..."
}
+118
View File
@@ -0,0 +1,118 @@
package anthropic
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/openai"
)
// Two turns through a fake Messages API: the first asks for a tool, the second (with previous_response_id and a
// forced submit) must carry the tool_result, the assistant tool_use block, the nudge text and an auto tool_choice.
func TestTranslateTurns(t *testing.T) {
var bodies []map[string]any
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/v1/messages" || r.Header.Get("x-api-key") != "k1" || r.Header.Get("anthropic-version") != "2023-06-01" {
t.Errorf("bad request line/headers: %s %v", r.URL.Path, r.Header)
}
var b map[string]any
_ = json.NewDecoder(r.Body).Decode(&b)
bodies = append(bodies, b)
switch len(bodies) {
case 1:
w.Write([]byte(`{"id":"msg_1","type":"message","model":"claude-sonnet-5-5","stop_reason":"tool_use","content":[{"type":"text","text":"checking"},{"type":"tool_use","id":"toolu_1","name":"get_probe_status","input":{}}],"usage":{"input_tokens":100,"output_tokens":20,"cache_creation_input_tokens":5,"cache_read_input_tokens":0}}`))
default:
w.Write([]byte(`{"id":"msg_2","type":"message","model":"claude-sonnet-5-5","stop_reason":"tool_use","content":[{"type":"tool_use","id":"toolu_2","name":"submit_diagnosis","input":{"root_cause_label":"x","confidence":0.9}}],"usage":{"input_tokens":200,"output_tokens":30}}`))
}
}))
defer srv.Close()
c := New(srv.URL, func() (string, error) { return "k1", nil })
pf := false
tools := []openai.Tool{{Type: "function", Strict: true, Name: "get_probe_status", Description: "probe", Parameters: map[string]any{"type": "object", "properties": map[string]any{}}}}
r1, err := c.CreateResponse(context.Background(), openai.Request{Model: "claude-sonnet-5-5", MaxOutputTokens: 1234, ParallelToolCalls: &pf, Tools: tools, ToolChoice: "auto",
Input: []openai.Item{{Type: "message", Role: "system", Content: "SYS"}, {Type: "message", Role: "user", Content: "alert"}}})
if err != nil {
t.Fatal(err)
}
fcs := r1.FunctionCalls()
if r1.ID != "msg_1" || len(fcs) != 1 || fcs[0].CallID != "toolu_1" || fcs[0].Name != "get_probe_status" || fcs[0].Arguments != "{}" {
t.Fatalf("turn 1 translation: %+v", r1)
}
if r1.Usage.InputTokens != 105 || r1.Usage.OutputTokens != 20 || r1.Usage.TotalTokens != 125 {
t.Fatalf("usage: %+v", r1.Usage)
}
b1 := bodies[0]
if b1["system"] != "SYS" || b1["max_tokens"].(float64) != 1234 || b1["model"] != "claude-sonnet-5-5" {
t.Fatalf("body 1: %v", b1)
}
tl := b1["tools"].([]any)[0].(map[string]any)
if tl["name"] != "get_probe_status" || tl["input_schema"] == nil || tl["strict"] != true {
t.Fatalf("tool translation: %v", tl)
}
if tc := b1["tool_choice"].(map[string]any); tc["type"] != "auto" || tc["disable_parallel_tool_use"] != true {
t.Fatalf("tool_choice 1: %v", tc)
}
r2, err := c.CreateResponse(context.Background(), openai.Request{Model: "claude-sonnet-5-5", Tools: tools, PreviousResponseID: "msg_1",
ToolChoice: map[string]any{"type": "function", "name": "submit_diagnosis"},
Input: []openai.Item{{Type: "function_call_output", CallID: "toolu_1", Output: `{"alb":1}`}}})
if err != nil {
t.Fatal(err)
}
if got := r2.FunctionCalls(); len(got) != 1 || got[0].Name != "submit_diagnosis" || !strings.Contains(got[0].Arguments, `"root_cause_label":"x"`) {
t.Fatalf("turn 2 translation: %+v", r2.Output)
}
b2 := bodies[1]
msgs := b2["messages"].([]any)
if len(msgs) != 3 { // user(alert) / assistant(text+tool_use) / user(tool_result+nudge)
t.Fatalf("messages: %v", msgs)
}
as := msgs[1].(map[string]any)
if as["role"] != "assistant" || !strings.Contains(mustJSON(as["content"]), `"tool_use"`) {
t.Fatalf("assistant turn not replayed: %v", as)
}
last := msgs[2].(map[string]any)["content"].([]any)
if tr := last[0].(map[string]any); tr["type"] != "tool_result" || tr["tool_use_id"] != "toolu_1" || tr["content"] != `{"alb":1}` {
t.Fatalf("tool_result: %v", tr)
}
if nd := last[1].(map[string]any); nd["type"] != "text" || !strings.Contains(nd["text"].(string), "submit_diagnosis") {
t.Fatalf("forced nudge missing: %v", nd)
}
if tc := b2["tool_choice"].(map[string]any); tc["type"] != "auto" {
t.Fatalf("forced tool_choice must become auto for Claude 5.5: %v", tc)
}
if b2["system"] != "SYS" {
t.Fatalf("system prompt must be carried on the chained turn, got %v", b2["system"])
}
}
func TestEmptyKeyAndErrors(t *testing.T) {
c := New("http://127.0.0.1:1", func() (string, error) { return "", nil })
if c.KeyPresent() {
t.Fatal("empty key reported present")
}
_, err := c.CreateResponse(context.Background(), openai.Request{Model: "claude-haiku-5-5", Input: []openai.Item{{Type: "message", Role: "user", Content: "x"}}})
if err == nil || !strings.Contains(err.Error(), "anthropic api key is empty") {
t.Fatalf("expected empty-key error, got %v", err)
}
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(400)
w.Write([]byte(`{"type":"error","error":{"type":"invalid_request_error","message":"tool_choice: type \"tool\" and \"any\" are not supported for this model."}}`))
}))
defer srv.Close()
c2 := New(srv.URL, func() (string, error) { return "k", nil })
_, err = c2.CreateResponse(context.Background(), openai.Request{Model: "claude-haiku-5-5", Input: []openai.Item{{Type: "message", Role: "user", Content: "x"}}})
if err == nil || !strings.Contains(err.Error(), "status 400") || !strings.Contains(err.Error(), "invalid_request_error") {
t.Fatalf("expected 400 error, got %v", err)
}
_, err = c2.CreateResponse(context.Background(), openai.Request{Model: "claude-haiku-5-5", PreviousResponseID: "nope", Input: []openai.Item{{Type: "function_call_output", CallID: "a", Output: "b"}}})
if err == nil || !strings.Contains(err.Error(), "unknown previous_response_id") {
t.Fatalf("expected unknown id error, got %v", err)
}
}
func mustJSON(v any) string { b, _ := json.Marshal(v); return string(b) }
+3
View File
@@ -231,6 +231,9 @@ func (r *AlertRouter) diagnose(alert map[string]any, runID, sid *string) {
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute) ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
defer cancel() defer cancel()
reqBody := map[string]any{"alert": alert, "run_id": runID, "scenario_id": sid} reqBody := map[string]any{"alert": alert, "run_id": runID, "scenario_id": sid}
if cur := r.app.runs.current(); cur != nil && runID != nil && cur.RunID == *runID && cur.Model != "" {
reqBody["model"] = cur.Model // campaign-level model selection (the AI VM falls back to its default when absent)
}
b, _ := json.Marshal(reqBody) b, _ := json.Marshal(reqBody)
req, _ := http.NewRequestWithContext(ctx, "POST", r.app.cfg.AIDiagnoseURL, bytes.NewReader(b)) req, _ := http.NewRequestWithContext(ctx, "POST", r.app.cfg.AIDiagnoseURL, bytes.NewReader(b))
req.Header.Set("Content-Type", "application/json") req.Header.Set("Content-Type", "application/json")
+13 -3
View File
@@ -9,6 +9,7 @@ import (
"fmt" "fmt"
"log" "log"
"net/http" "net/http"
"regexp"
"strings" "strings"
"sync" "sync"
"time" "time"
@@ -25,6 +26,7 @@ type RunManager struct {
type activeRun struct { type activeRun struct {
RunID string RunID string
Campaign string Campaign string
Model string // LLM for this run's diagnoses ("" = AI VM default); passed to the AI VM on every /diagnose
ScenarioID string // "" between scenarios ScenarioID string // "" between scenarios
Scenario *Scenario Scenario *Scenario
TInject time.Time TInject time.Time
@@ -33,6 +35,7 @@ type activeRun struct {
type RunRequest struct { type RunRequest struct {
Campaign string `json:"campaign"` Campaign string `json:"campaign"`
Model string `json:"model"` // optional: LLM id for the diagnoses (e.g. gpt-6-luna, claude-sonnet-5-5); "" = AI VM default
ScenarioIDs []string `json:"scenario_ids"` ScenarioIDs []string `json:"scenario_ids"`
BaselineS *int `json:"baseline_s"` // default 300 (design 8) BaselineS *int `json:"baseline_s"` // default 300 (design 8)
CooldownS *int `json:"cooldown_s"` // default 300 CooldownS *int `json:"cooldown_s"` // default 300
@@ -44,6 +47,9 @@ type RunRequest struct {
SkipRevert bool `json:"skip_revert"` SkipRevert bool `json:"skip_revert"`
} }
// modelRe bounds the model id (empty = AI VM default). The AI VM routes by prefix: claude-* -> Anthropic, else OpenAI.
var modelRe = regexp.MustCompile(`^[A-Za-z0-9._:-]{0,64}$`)
func NewRunManager(a *App) *RunManager { return &RunManager{app: a} } func NewRunManager(a *App) *RunManager { return &RunManager{app: a} }
func (m *RunManager) current() *activeRun { func (m *RunManager) current() *activeRun {
@@ -78,6 +84,10 @@ func (a *App) handleRunCreate(w http.ResponseWriter, r *http.Request) {
return return
} }
} }
if !modelRe.MatchString(req.Model) {
writeErr(w, 400, errors.New("model must match [A-Za-z0-9._:-]{1,64}"))
return
}
m := a.runs m := a.runs
m.mu.Lock() m.mu.Lock()
if m.active != nil { if m.active != nil {
@@ -88,9 +98,9 @@ func (a *App) handleRunCreate(w http.ResponseWriter, r *http.Request) {
} }
runID := newID("run-") runID := newID("run-")
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
m.active = &activeRun{RunID: runID, Campaign: req.Campaign, cancel: cancel} m.active = &activeRun{RunID: runID, Campaign: req.Campaign, Model: req.Model, cancel: cancel}
m.mu.Unlock() m.mu.Unlock()
if err := a.store.CreateRun(r.Context(), runID, req.Campaign, req.DriftBefore, req.GitSHA, req.Notes); err != nil { if err := a.store.CreateRun(r.Context(), runID, req.Campaign, req.DriftBefore, req.GitSHA, req.Notes, req.Model); err != nil {
m.mu.Lock() m.mu.Lock()
m.active = nil m.active = nil
m.mu.Unlock() m.mu.Unlock()
@@ -100,7 +110,7 @@ func (a *App) handleRunCreate(w http.ResponseWriter, r *http.Request) {
} }
mRuns.Set(1) mRuns.Set(1)
go m.execute(ctx, runID, req) go m.execute(ctx, runID, req)
writeJSON(w, 202, map[string]any{"run_id": runID, "scenario_ids": req.ScenarioIDs}) writeJSON(w, 202, map[string]any{"run_id": runID, "scenario_ids": req.ScenarioIDs, "model": req.Model})
} }
func (a *App) handleRunAbort(w http.ResponseWriter, r *http.Request) { func (a *App) handleRunAbort(w http.ResponseWriter, r *http.Request) {
+6 -4
View File
@@ -24,6 +24,7 @@ CREATE TABLE IF NOT EXISTS runs (
status text NOT NULL DEFAULT 'running', status text NOT NULL DEFAULT 'running',
notes text notes text
); );
ALTER TABLE runs ADD COLUMN IF NOT EXISTS model text; -- LLM requested for the run's diagnoses (campaign-level model selection, 2026-10-14)
CREATE TABLE IF NOT EXISTS scenarios ( CREATE TABLE IF NOT EXISTS scenarios (
scenario_id text PRIMARY KEY, scenario_id text PRIMARY KEY,
path text NOT NULL, path text NOT NULL,
@@ -119,8 +120,8 @@ func (s *Store) UpsertScenarios(ctx context.Context, m map[string]*Scenario) err
return nil return nil
} }
func (s *Store) CreateRun(ctx context.Context, id, campaign string, driftBefore *int, gitSHA, notes string) error { func (s *Store) CreateRun(ctx context.Context, id, campaign string, driftBefore *int, gitSHA, notes, model string) error {
_, err := s.pool.Exec(ctx, `INSERT INTO runs (run_id,campaign,drift_before,git_sha,notes) VALUES ($1,$2,$3,$4,$5)`, id, campaign, driftBefore, nz(gitSHA), nz(notes)) _, err := s.pool.Exec(ctx, `INSERT INTO runs (run_id,campaign,drift_before,git_sha,notes,model) VALUES ($1,$2,$3,$4,$5,$6)`, id, campaign, driftBefore, nz(gitSHA), nz(notes), nz(model))
return err return err
} }
@@ -222,10 +223,11 @@ type RunRow struct {
GitSHA *string `json:"git_sha"` GitSHA *string `json:"git_sha"`
Status string `json:"status"` Status string `json:"status"`
Notes *string `json:"notes"` Notes *string `json:"notes"`
Model *string `json:"model"` // nil for runs before campaign-level model selection
} }
func (s *Store) Runs(ctx context.Context, runID string) ([]RunRow, error) { func (s *Store) Runs(ctx context.Context, runID string) ([]RunRow, error) {
rows, err := s.pool.Query(ctx, `SELECT run_id,campaign,started_at,finished_at,drift_before,drift_after,git_sha,status,notes FROM runs WHERE ($1='' OR run_id=$1) ORDER BY started_at DESC LIMIT 100`, runID) rows, err := s.pool.Query(ctx, `SELECT run_id,campaign,started_at,finished_at,drift_before,drift_after,git_sha,status,notes,model FROM runs WHERE ($1='' OR run_id=$1) ORDER BY started_at DESC LIMIT 100`, runID)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@@ -233,7 +235,7 @@ func (s *Store) Runs(ctx context.Context, runID string) ([]RunRow, error) {
var out []RunRow var out []RunRow
for rows.Next() { for rows.Next() {
var r RunRow var r RunRow
if err := rows.Scan(&r.RunID, &r.Campaign, &r.StartedAt, &r.FinishedAt, &r.DriftBefore, &r.DriftAfter, &r.GitSHA, &r.Status, &r.Notes); err != nil { if err := rows.Scan(&r.RunID, &r.Campaign, &r.StartedAt, &r.FinishedAt, &r.DriftBefore, &r.DriftAfter, &r.GitSHA, &r.Status, &r.Notes, &r.Model); err != nil {
return nil, err return nil, err
} }
out = append(out, r) out = append(out, r)