feat(ai): add AI VM services (diagnosis agent, pgvector RAG, Discord HITL bot) and harness approval endpoints
실험 설계서 6장(WP-12)의 AI VM 2층 구현.
- cmd/agent: 하네스 /diagnose 수신 → OpenAI Responses API 함수 호출 루프(읽기 전용 도구 7종:
query_metrics, query_logs, get_probe_status, get_k8s, get_gcp_instances, get_aws_state,
search_runbooks) → strict 스키마 submit_diagnosis(6.3) → 하네스 /proposals 등록 + HITL 게시,
또는 신뢰도 0.5 미만·escalate 시 3층 큐 에스컬레이션. 세션·프롬프트 버전·러너북 벡터는
pgvector(ai 스키마)에 저장. 라벨·조치 enum은 하네스 /scenarios·/actions/catalog에서 갱신.
- cmd/hitl: Discord Gateway 봇(intents GUILDS만). 승인/거부 버튼(custom_id approve:/reject:<token>),
3초 내 DEFERRED_UPDATE_MESSAGE, 승인 시 하네스 승인 → /actions/execute → 원 메시지 편집,
15분 무응답 시 에스컬레이션 채널로 이동. 재기동 시 미결 제안 타이머 복원.
- harness: POST /proposals/{token}/approve, /reject 추가. 제안은 진단 시점에 등록하고 승인은
버튼 시점에 기록해 MTTA(승인-제안)가 실제 대기 시간이 되도록 함.
OpenAI·Discord 토큰은 Secret Manager 파일로만 읽고 코드·저장소에 두지 않는다.
This commit is contained in:
seon committed
1 parent
de076a79e4
commit
f30af2a4c3
21 files changed
+2875
-1
No files matched your search
@@ -0,0 +1,303 @@
|
||||
// Package agent runs one diagnosis (experiment design 6.3/6.4): a Responses API tool loop over the read-only tools,
|
||||
// ending in a strict submit_diagnosis call. After the diagnosis it registers a layer-2 proposal with the harness and
|
||||
// hands it to the HITL bot, or sends the case to the escalation channel (design 6.5/6.6).
|
||||
package agent
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"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"
|
||||
"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/rag"
|
||||
"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/tools"
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
Model string
|
||||
MaxTurns int
|
||||
MaxOutputTokens int
|
||||
SystemPrompt string // body of prompts/system.md
|
||||
Topology string // body of prompts/topology.md (design 1 IP table)
|
||||
SLO string // body of prompts/slo.md (design 4.2)
|
||||
HitlURL string // http://hitl:8081
|
||||
ProposalTTL time.Duration
|
||||
ConfidenceMin float64 // 0.5 (design 6.3)
|
||||
}
|
||||
|
||||
type Agent struct {
|
||||
cfg Config
|
||||
oai *openai.Client
|
||||
tools *tools.Runner
|
||||
harness *harness.Client
|
||||
store *rag.Store
|
||||
enums diag.Enums
|
||||
promptV string
|
||||
http *http.Client
|
||||
}
|
||||
|
||||
func New(cfg Config, oai *openai.Client, tr *tools.Runner, hc *harness.Client, store *rag.Store, enums diag.Enums, promptVersion string) *Agent {
|
||||
if cfg.MaxTurns == 0 {
|
||||
cfg.MaxTurns = 12
|
||||
}
|
||||
if cfg.MaxOutputTokens == 0 {
|
||||
cfg.MaxOutputTokens = 4000
|
||||
}
|
||||
if cfg.ProposalTTL == 0 {
|
||||
cfg.ProposalTTL = 15 * time.Minute
|
||||
}
|
||||
if cfg.ConfidenceMin == 0 {
|
||||
cfg.ConfidenceMin = 0.5
|
||||
}
|
||||
return &Agent{cfg: cfg, oai: oai, tools: tr, harness: hc, store: store, enums: enums, promptV: promptVersion, http: &http.Client{Timeout: 30 * time.Second}}
|
||||
}
|
||||
|
||||
// Request is the harness -> AI VM contract (harness alerts.go diagnose()).
|
||||
type Request struct {
|
||||
Alert map[string]any `json:"alert"`
|
||||
RunID *string `json:"run_id"`
|
||||
ScenarioID *string `json:"scenario_id"`
|
||||
}
|
||||
|
||||
// Result is returned to the harness: the diagnosis fields at top level (the harness stores the whole object as the
|
||||
// diagnosis event payload and reads root_cause_label / proposed_action from it) plus bookkeeping.
|
||||
type Result struct {
|
||||
diag.Diagnosis
|
||||
SessionID string `json:"session_id"`
|
||||
Model string `json:"model"`
|
||||
PromptVersion string `json:"prompt_version"`
|
||||
Turns int `json:"turns"`
|
||||
ToolCalls int `json:"tool_calls"`
|
||||
DurationMs int64 `json:"duration_ms"`
|
||||
Disposition string `json:"disposition"` // proposal | escalated | none
|
||||
ProposalToken string `json:"proposal_token,omitempty"`
|
||||
}
|
||||
|
||||
type toolCallLog struct {
|
||||
Name string `json:"name"`
|
||||
Args string `json:"args"`
|
||||
Output string `json:"output"`
|
||||
Ms int64 `json:"ms"`
|
||||
}
|
||||
|
||||
func (a *Agent) toolDefs() []openai.Tool {
|
||||
defs := tools.Definitions()
|
||||
defs = append(defs, openai.Tool{Type: "function", Strict: true, Name: "submit_diagnosis",
|
||||
Description: "Submit the final diagnosis. Call exactly once, after gathering evidence. This records the result; it does not execute anything.",
|
||||
Parameters: a.enums.Schema()})
|
||||
return defs
|
||||
}
|
||||
|
||||
func (a *Agent) userMessage(ctx context.Context, req Request) string {
|
||||
alert, _ := json.MarshalIndent(req.Alert, "", " ")
|
||||
var b strings.Builder
|
||||
b.WriteString("## Topology\n")
|
||||
b.WriteString(a.cfg.Topology)
|
||||
b.WriteString("\n\n## SLO\n")
|
||||
b.WriteString(a.cfg.SLO)
|
||||
b.WriteString("\n\n## Alert payload\n```json\n")
|
||||
b.Write(alert)
|
||||
b.WriteString("\n```\n\n## Last 10 minutes (pre-fetched)\n")
|
||||
for _, q := range []struct{ name, promql string }{
|
||||
{"probe_success", `probe_success`},
|
||||
{"api_p95_latency_s", `histogram_quantile(0.95, sum by (le) (rate(demo_http_request_duration_seconds_bucket[5m])))`},
|
||||
{"api_error_ratio", `sum(rate(demo_http_requests_total{status=~"5.."}[5m])) / sum(rate(demo_http_requests_total[5m]))`},
|
||||
{"legacy_up", `up{job="legacy-api"}`},
|
||||
} {
|
||||
out := a.tools.Call(ctx, "query_metrics", fmt.Sprintf(`{"promql":%q,"range_minutes":10}`, q.promql))
|
||||
if len(out) > 1200 {
|
||||
out = out[:1200] + "..."
|
||||
}
|
||||
fmt.Fprintf(&b, "- %s: %s\n", q.name, out)
|
||||
}
|
||||
if req.RunID != nil {
|
||||
fmt.Fprintf(&b, "\nrun_id=%s", *req.RunID)
|
||||
}
|
||||
if req.ScenarioID != nil {
|
||||
fmt.Fprintf(&b, " scenario_id=%s", *req.ScenarioID)
|
||||
}
|
||||
b.WriteString("\n\nInvestigate with the tools, then call submit_diagnosis exactly once.")
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// Diagnose runs the tool loop and returns the validated diagnosis. The session is stored whatever the outcome.
|
||||
func (a *Agent) Diagnose(ctx context.Context, req Request) (*Result, error) {
|
||||
start := time.Now()
|
||||
sid := "ses-" + uuid.NewString()[:8]
|
||||
res := &Result{SessionID: sid, Model: a.cfg.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)},
|
||||
}
|
||||
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}}
|
||||
finish := func(err error) {
|
||||
session.Output = res
|
||||
session.ToolCalls = calls
|
||||
session.Duration = time.Since(start)
|
||||
if err != nil {
|
||||
session.Error = err.Error()
|
||||
}
|
||||
if a.store != nil {
|
||||
if e := a.store.SaveSession(context.Background(), session); e != nil {
|
||||
log.Printf("session save: %v", e)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
defs := a.toolDefs()
|
||||
var prevID string
|
||||
var diagnosis *diag.Diagnosis
|
||||
pf := false
|
||||
var usage any
|
||||
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}
|
||||
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)
|
||||
if err != nil {
|
||||
finish(err)
|
||||
return nil, err
|
||||
}
|
||||
usage = resp.Usage
|
||||
prevID = resp.ID
|
||||
fcs := resp.FunctionCalls()
|
||||
if len(fcs) == 0 {
|
||||
// text only: nudge once towards the tool
|
||||
input = []openai.Item{{Type: "message", Role: "user", Content: "Use the tools to gather evidence and finish with submit_diagnosis."}}
|
||||
continue
|
||||
}
|
||||
input = input[:0]
|
||||
for _, fc := range fcs {
|
||||
if fc.Name == "submit_diagnosis" {
|
||||
d, err := diag.Parse(fc.Arguments, a.enums)
|
||||
if err != nil {
|
||||
input = append(input, openai.Item{Type: "function_call_output", CallID: fc.CallID, Output: fmt.Sprintf(`{"error":%q}`, err.Error())})
|
||||
calls = append(calls, toolCallLog{Name: fc.Name, Args: fc.Arguments, Output: "rejected: " + err.Error()})
|
||||
continue
|
||||
}
|
||||
diagnosis = &d
|
||||
calls = append(calls, toolCallLog{Name: fc.Name, Args: fc.Arguments, Output: "accepted"})
|
||||
break
|
||||
}
|
||||
t0 := time.Now()
|
||||
out := a.tools.Call(ctx, fc.Name, fc.Arguments)
|
||||
calls = append(calls, toolCallLog{Name: fc.Name, Args: fc.Arguments, Output: out, Ms: time.Since(t0).Milliseconds()})
|
||||
input = append(input, openai.Item{Type: "function_call_output", CallID: fc.CallID, Output: out})
|
||||
}
|
||||
}
|
||||
session.Usage = usage
|
||||
res.ToolCalls = len(calls)
|
||||
res.DurationMs = time.Since(start).Milliseconds()
|
||||
if diagnosis == nil {
|
||||
err := fmt.Errorf("no valid submit_diagnosis after %d turns", res.Turns)
|
||||
finish(err)
|
||||
return nil, err
|
||||
}
|
||||
res.Diagnosis = *diagnosis
|
||||
finish(nil)
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// Dispose applies design 6.3/6.5/6.6 after the diagnosis: proposal + HITL message, or escalation, or nothing.
|
||||
func (a *Agent) Dispose(ctx context.Context, req Request, res *Result) {
|
||||
d := res.Diagnosis
|
||||
diagMap := map[string]any{"root_cause_label": d.RootCauseLabel, "confidence": d.Confidence, "evidence": d.Evidence, "proposed_action": d.ProposedAction, "target": d.Target, "rationale": d.Rationale, "session_id": res.SessionID, "model": res.Model, "prompt_version": res.PromptVersion}
|
||||
switch {
|
||||
case d.ProposedAction == "none":
|
||||
res.Disposition = "none"
|
||||
return
|
||||
case d.ProposedAction == "escalate" || d.Confidence < a.cfg.ConfidenceMin:
|
||||
res.Disposition = "escalated"
|
||||
a.escalate(ctx, req, diagMap)
|
||||
return
|
||||
}
|
||||
target := targetFor(d.ProposedAction, d.Target)
|
||||
p, err := a.harness.Propose(ctx, req.RunID, req.ScenarioID, d.ProposedAction, target, diagMap, a.cfg.ProposalTTL)
|
||||
if err != nil {
|
||||
log.Printf("propose: %v", err)
|
||||
res.Disposition = "escalated"
|
||||
diagMap["note"] = "proposal registration failed: " + err.Error()
|
||||
a.escalate(ctx, req, diagMap)
|
||||
return
|
||||
}
|
||||
res.Disposition = "proposal"
|
||||
res.ProposalToken = p.Token
|
||||
body := map[string]any{"token": p.Token, "action_id": d.ProposedAction, "target": target, "diagnosis": diagMap, "run_id": req.RunID, "scenario_id": req.ScenarioID, "session_id": res.SessionID, "expires_in_s": p.ExpiresIn}
|
||||
if err := a.postJSON(ctx, a.cfg.HitlURL+"/post", body); err != nil {
|
||||
log.Printf("hitl post: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Agent) escalate(ctx context.Context, req Request, diagMap map[string]any) {
|
||||
runbook := ""
|
||||
if hits, err := a.tools.SearchRunbooks(ctx, fmt.Sprint(diagMap["root_cause_label"])+" "+fmt.Sprint(diagMap["rationale"]), 2); err == nil {
|
||||
if hs, ok := hits.([]rag.Hit); ok {
|
||||
var parts []string
|
||||
for _, h := range hs {
|
||||
parts = append(parts, "["+h.Source+"]\n"+h.Chunk)
|
||||
}
|
||||
runbook = strings.Join(parts, "\n\n")
|
||||
}
|
||||
}
|
||||
body := map[string]any{"diagnosis": diagMap, "runbook": runbook, "run_id": req.RunID, "scenario_id": req.ScenarioID, "alert": req.Alert}
|
||||
if err := a.postJSON(ctx, a.cfg.HitlURL+"/escalate", body); err != nil {
|
||||
log.Printf("hitl escalate: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Agent) postJSON(ctx context.Context, u string, body any) error {
|
||||
b, _ := json.Marshal(body)
|
||||
req, _ := http.NewRequestWithContext(ctx, "POST", u, bytes.NewReader(b))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := a.http.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode/100 != 2 {
|
||||
return fmt.Errorf("%s: status %d", u, resp.StatusCode)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// targetFor maps the free-text target of the diagnosis onto the harness action target keys (harness actions.go).
|
||||
func targetFor(actionID, target string) map[string]any {
|
||||
t := strings.TrimSpace(target)
|
||||
switch {
|
||||
case strings.HasPrefix(actionID, "k8s.rollout"), actionID == "k8s.scale":
|
||||
m := map[string]any{"namespace": "demo", "deployment": t}
|
||||
if actionID == "k8s.scale" {
|
||||
m["replicas"] = 3
|
||||
}
|
||||
if i := strings.Index(t, "/"); i > 0 {
|
||||
m["namespace"], m["deployment"] = t[:i], t[i+1:]
|
||||
}
|
||||
return m
|
||||
case actionID == "k8s.node_drain", actionID == "k8s.node_uncordon":
|
||||
return map[string]any{"node": t}
|
||||
case actionID == "gce.instance_reset":
|
||||
return map[string]any{"instance": t}
|
||||
case actionID == "mig.recreate_instance":
|
||||
return map[string]any{"instance": t}
|
||||
case actionID == "mig.resize":
|
||||
return map[string]any{"mig": "selfheal-legacy-mig", "size": 2}
|
||||
case actionID == "db.replay_resume":
|
||||
return map[string]any{"host": "db-replica"}
|
||||
}
|
||||
return map[string]any{"target": t}
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"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"
|
||||
"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/tools"
|
||||
)
|
||||
|
||||
// A mock Responses API: turn 1 asks for get_probe_status, turn 2 submits the diagnosis. Verifies the loop wiring
|
||||
// (function_call_output with the call_id, previous_response_id chaining, strict submit parsing).
|
||||
func TestDiagnoseLoop(t *testing.T) {
|
||||
var turns int
|
||||
var sawOutput, sawPrev bool
|
||||
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
|
||||
}
|
||||
var req openai.Request
|
||||
_ = json.NewDecoder(r.Body).Decode(&req)
|
||||
turns++
|
||||
for _, it := range req.Input {
|
||||
if it.Type == "function_call_output" && it.CallID == "call_1" && strings.Contains(it.Output, "probe_success") {
|
||||
sawOutput = true
|
||||
}
|
||||
}
|
||||
if req.PreviousResponseID == "resp_1" {
|
||||
sawPrev = true
|
||||
}
|
||||
if r.Header.Get("Authorization") != "Bearer test-key" {
|
||||
t.Errorf("missing bearer")
|
||||
}
|
||||
switch turns {
|
||||
case 1:
|
||||
w.Write([]byte(`{"id":"resp_1","status":"completed","output":[{"type":"function_call","call_id":"call_1","name":"get_probe_status","arguments":"{}"}],"usage":{"input_tokens":10,"output_tokens":5}}`))
|
||||
default:
|
||||
args := `{\"root_cause_label\":\"db_primary_down\",\"confidence\":0.8,\"evidence\":[\"probe_success db-primary=0\"],\"proposed_action\":\"escalate\",\"target\":\"db-primary\",\"rationale\":\"primary unreachable\"}`
|
||||
w.Write([]byte(`{"id":"resp_2","status":"completed","output":[{"type":"function_call","call_id":"call_2","name":"submit_diagnosis","arguments":"` + args + `"}],"usage":{"input_tokens":20,"output_tokens":9}}`))
|
||||
}
|
||||
}))
|
||||
defer oaiSrv.Close()
|
||||
mimir := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if strings.HasSuffix(r.URL.Path, "/api/v1/query") {
|
||||
w.Write([]byte(`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"target":"db-primary"},"value":[1,"0"]},{"metric":{"target":"alb"},"value":[1,"1"]}]}}`))
|
||||
return
|
||||
}
|
||||
w.Write([]byte(`{"status":"success","data":{"resultType":"matrix","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: "m", SystemPrompt: "sys", Topology: "topo", SLO: "slo", HitlURL: "http://127.0.0.1:1", MaxTurns: 5}, oai, tr, harness.New("http://127.0.0.1:1"), nil, enums, "pv1")
|
||||
res, err := ag.Diagnose(context.Background(), Request{Alert: map[string]any{"alertname": "DbPrimaryDown"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if res.RootCauseLabel != "db_primary_down" || res.Turns != 2 || !sawOutput || !sawPrev {
|
||||
t.Fatalf("unexpected result: %+v sawOutput=%v sawPrev=%v", res, sawOutput, sawPrev)
|
||||
}
|
||||
ag.Dispose(context.Background(), Request{Alert: map[string]any{}}, res)
|
||||
if res.Disposition != "escalated" {
|
||||
t.Fatalf("escalate action must escalate, got %s", res.Disposition)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTargetFor(t *testing.T) {
|
||||
m := targetFor("k8s.rollout_restart", "demo/api")
|
||||
if m["namespace"] != "demo" || m["deployment"] != "api" {
|
||||
t.Fatalf("%v", m)
|
||||
}
|
||||
if targetFor("mig.resize", "")["mig"] != "selfheal-legacy-mig" {
|
||||
t.Fatal("mig")
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user