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}
|
||||
}
|
||||
Reference in new issue
Block a user