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")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
// Package diag holds the strict diagnosis schema of experiment design 6.3 and the enum sources.
|
||||
package diag
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Diagnosis is the submit_diagnosis payload (design 6.3). Every field is required under strict mode.
|
||||
type Diagnosis struct {
|
||||
RootCauseLabel string `json:"root_cause_label"`
|
||||
Confidence float64 `json:"confidence"`
|
||||
Evidence []string `json:"evidence"`
|
||||
ProposedAction string `json:"proposed_action"`
|
||||
Target string `json:"target"`
|
||||
Rationale string `json:"rationale"`
|
||||
}
|
||||
|
||||
// DefaultLabels mirrors the ground_truth_label values of the 13 scenario files (design 7); the agent refreshes the
|
||||
// list from the harness GET /scenarios at start-up and falls back to this copy.
|
||||
var DefaultLabels = []string{
|
||||
"api_pod_failure", "api_pod_killed", "db_primary_down", "gitops_drift", "legacy_instance_lost", "legacy_process_killed",
|
||||
"memory_pressure_api", "network_latency_api_to_db", "network_latency_legacy", "replication_paused", "schema_mismatch",
|
||||
"transient_loss_legacy", "transient_partition",
|
||||
}
|
||||
|
||||
// DefaultActions mirrors the layer-2 action_ids of design 2.3; refreshed from the harness GET /actions/catalog (layer contains "2").
|
||||
var DefaultActions = []string{
|
||||
"k8s.rollout_restart", "k8s.rollout_undo", "k8s.scale", "k8s.node_drain", "k8s.node_uncordon",
|
||||
"gce.instance_reset", "mig.recreate_instance", "mig.resize", "db.replay_resume", "net.route_change",
|
||||
}
|
||||
|
||||
// Enums are the closed sets the schema is built from.
|
||||
type Enums struct {
|
||||
Labels []string // + "unknown"
|
||||
Actions []string // + "none", "escalate"
|
||||
}
|
||||
|
||||
func (e Enums) LabelEnum() []string { return append(append([]string{}, e.Labels...), "unknown") }
|
||||
func (e Enums) ActionEnum() []string { return append(append([]string{}, e.Actions...), "none", "escalate") }
|
||||
|
||||
// Schema is the JSON schema of the submit_diagnosis tool parameters (strict: additionalProperties false, all required).
|
||||
func (e Enums) Schema() map[string]any {
|
||||
return map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"root_cause_label": map[string]any{"type": "string", "enum": e.LabelEnum(), "description": "Root cause label; unknown when the evidence does not support any label."},
|
||||
"confidence": map[string]any{"type": "number", "description": "0.0 to 1.0. Below 0.5 the harness routes the case to layer 3 (human)."},
|
||||
"evidence": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "Quotations of tool results that support the label. At least one."},
|
||||
"proposed_action": map[string]any{"type": "string", "enum": e.ActionEnum(), "description": "Catalogue action_id (layer 2), none when no action is needed, escalate when a human must decide."},
|
||||
"target": map[string]any{"type": "string", "description": "Action target: deployment name, instance name, legacy:N, or empty."},
|
||||
"rationale": map[string]any{"type": "string", "description": "Why this label and action; at most 200 characters."},
|
||||
},
|
||||
"required": []string{"root_cause_label", "confidence", "evidence", "proposed_action", "target", "rationale"},
|
||||
"additionalProperties": false,
|
||||
}
|
||||
}
|
||||
|
||||
// Validate enforces the schema constraints on the model output (strict mode guarantees shape; this checks the values).
|
||||
func (e Enums) Validate(d Diagnosis) error {
|
||||
if !contains(e.LabelEnum(), d.RootCauseLabel) {
|
||||
return fmt.Errorf("root_cause_label %q not in enum", d.RootCauseLabel)
|
||||
}
|
||||
if d.Confidence < 0 || d.Confidence > 1 {
|
||||
return fmt.Errorf("confidence %v out of [0,1]", d.Confidence)
|
||||
}
|
||||
if len(d.Evidence) == 0 {
|
||||
return fmt.Errorf("evidence must have at least one item")
|
||||
}
|
||||
if !contains(e.ActionEnum(), d.ProposedAction) {
|
||||
return fmt.Errorf("proposed_action %q not in enum", d.ProposedAction)
|
||||
}
|
||||
if n := len([]rune(d.Rationale)); n > 200 {
|
||||
return fmt.Errorf("rationale is %d characters, limit 200", n)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func Parse(args string, e Enums) (Diagnosis, error) {
|
||||
var d Diagnosis
|
||||
dec := json.NewDecoder(strings.NewReader(args))
|
||||
dec.DisallowUnknownFields()
|
||||
if err := dec.Decode(&d); err != nil {
|
||||
return d, fmt.Errorf("submit_diagnosis arguments: %w", err)
|
||||
}
|
||||
return d, e.Validate(d)
|
||||
}
|
||||
|
||||
func contains(xs []string, s string) bool {
|
||||
for _, x := range xs {
|
||||
if x == s {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
package diag
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSchemaStrict(t *testing.T) {
|
||||
e := Enums{Labels: DefaultLabels, Actions: DefaultActions}
|
||||
s := e.Schema()
|
||||
if s["additionalProperties"] != false {
|
||||
t.Fatal("additionalProperties must be false")
|
||||
}
|
||||
props := s["properties"].(map[string]any)
|
||||
req := s["required"].([]string)
|
||||
if len(props) != len(req) {
|
||||
t.Fatalf("strict mode: every property must be required (%d props, %d required)", len(props), len(req))
|
||||
}
|
||||
if _, err := json.Marshal(s); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseValidate(t *testing.T) {
|
||||
e := Enums{Labels: DefaultLabels, Actions: DefaultActions}
|
||||
good := `{"root_cause_label":"api_pod_killed","confidence":0.9,"evidence":["pods: api-xyz restarts=3"],"proposed_action":"k8s.rollout_restart","target":"demo/api","rationale":"restart loop"}`
|
||||
d, err := Parse(good, e)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if d.ProposedAction != "k8s.rollout_restart" {
|
||||
t.Fatal("wrong action")
|
||||
}
|
||||
bad := []string{
|
||||
`{"root_cause_label":"nope","confidence":0.9,"evidence":["x"],"proposed_action":"none","target":"","rationale":"r"}`,
|
||||
`{"root_cause_label":"unknown","confidence":1.5,"evidence":["x"],"proposed_action":"none","target":"","rationale":"r"}`,
|
||||
`{"root_cause_label":"unknown","confidence":0.5,"evidence":[],"proposed_action":"none","target":"","rationale":"r"}`,
|
||||
`{"root_cause_label":"unknown","confidence":0.5,"evidence":["x"],"proposed_action":"db.promote","target":"","rationale":"r"}`,
|
||||
`{"root_cause_label":"unknown","confidence":0.5,"evidence":["x"],"proposed_action":"none","target":"","rationale":"` + strings.Repeat("a", 201) + `"}`,
|
||||
`{"root_cause_label":"unknown","confidence":0.5,"evidence":["x"],"proposed_action":"none","target":"","rationale":"r","extra":1}`,
|
||||
}
|
||||
for i, b := range bad {
|
||||
if _, err := Parse(b, e); err == nil {
|
||||
t.Fatalf("case %d should fail", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,122 @@
|
||||
// Package harness talks to the harness runner (design 5.2): proposals, approval, execution, catalogue, scenarios.
|
||||
package harness
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Client struct {
|
||||
BaseURL string // http://10.20.0.11:8080
|
||||
HTTP *http.Client
|
||||
}
|
||||
|
||||
func New(base string) *Client {
|
||||
return &Client{BaseURL: strings.TrimRight(base, "/"), HTTP: &http.Client{Timeout: 11 * time.Minute}}
|
||||
}
|
||||
|
||||
func (c *Client) call(ctx context.Context, method, path string, body any, out any) (int, error) {
|
||||
var rd io.Reader
|
||||
if body != nil {
|
||||
b, _ := json.Marshal(body)
|
||||
rd = bytes.NewReader(b)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, method, c.BaseURL+path, rd)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if body != nil {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
resp, err := c.HTTP.Do(req)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
rb, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
|
||||
if out != nil && len(rb) > 0 {
|
||||
if err := json.Unmarshal(rb, out); err != nil {
|
||||
return resp.StatusCode, fmt.Errorf("harness %s %s: status %d, body not JSON: %s", method, path, resp.StatusCode, truncate(string(rb), 300))
|
||||
}
|
||||
}
|
||||
if resp.StatusCode/100 != 2 {
|
||||
return resp.StatusCode, fmt.Errorf("harness %s %s: status %d: %s", method, path, resp.StatusCode, truncate(string(rb), 300))
|
||||
}
|
||||
return resp.StatusCode, nil
|
||||
}
|
||||
|
||||
type Proposal struct {
|
||||
Token string `json:"token"`
|
||||
ExpiresIn float64 `json:"expires_in_s"`
|
||||
Approved bool `json:"approved"`
|
||||
}
|
||||
|
||||
// Propose registers a layer-2 proposal (POST /proposals); the harness records the proposal event and returns the token.
|
||||
func (c *Client) Propose(ctx context.Context, runID, scenarioID *string, actionID string, target, diagnosis map[string]any, ttl time.Duration) (*Proposal, error) {
|
||||
body := map[string]any{"run_id": runID, "scenario_id": scenarioID, "action_id": actionID, "target": target, "diagnosis": diagnosis}
|
||||
if ttl > 0 {
|
||||
body["ttl_seconds"] = int(ttl.Seconds())
|
||||
}
|
||||
var p Proposal
|
||||
if _, err := c.call(ctx, "POST", "/proposals", body, &p); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &p, nil
|
||||
}
|
||||
|
||||
// Approve marks the proposal approved (POST /proposals/{token}/approve); the harness records the approval event.
|
||||
func (c *Client) Approve(ctx context.Context, token, by string) error {
|
||||
_, err := c.call(ctx, "POST", "/proposals/"+token+"/approve", map[string]any{"approved_by": by}, nil)
|
||||
return err
|
||||
}
|
||||
|
||||
// Reject records the rejection (POST /proposals/{token}/reject).
|
||||
func (c *Client) Reject(ctx context.Context, token, by, reason string) error {
|
||||
_, err := c.call(ctx, "POST", "/proposals/"+token+"/reject", map[string]any{"rejected_by": by, "reason": reason}, nil)
|
||||
return err
|
||||
}
|
||||
|
||||
// Execute runs the catalogue action with the approved token (POST /actions/execute). The body is returned as-is for the
|
||||
// Discord message; a non-2xx status is an error with the body text.
|
||||
func (c *Client) Execute(ctx context.Context, token, actionID string, target map[string]any, runID, scenarioID *string, note string) (map[string]any, error) {
|
||||
body := map[string]any{"action_id": actionID, "target": target, "token": token, "run_id": runID, "scenario_id": scenarioID, "note": note}
|
||||
var out map[string]any
|
||||
_, err := c.call(ctx, "POST", "/actions/execute", body, &out)
|
||||
return out, err
|
||||
}
|
||||
|
||||
type CatalogEntry struct {
|
||||
ActionID string `json:"action_id"`
|
||||
Layer string `json:"layer"`
|
||||
Exec bool `json:"harness_executes"`
|
||||
}
|
||||
|
||||
func (c *Client) Catalog(ctx context.Context) ([]CatalogEntry, error) {
|
||||
var out []CatalogEntry
|
||||
_, err := c.call(ctx, "GET", "/actions/catalog", nil, &out)
|
||||
return out, err
|
||||
}
|
||||
|
||||
type Scenario struct {
|
||||
ID string `json:"scenario_id"`
|
||||
Label string `json:"ground_truth_label"`
|
||||
}
|
||||
|
||||
func (c *Client) Scenarios(ctx context.Context) ([]Scenario, error) {
|
||||
var out []Scenario
|
||||
_, err := c.call(ctx, "GET", "/scenarios", nil, &out)
|
||||
return out, err
|
||||
}
|
||||
|
||||
func truncate(s string, n int) string {
|
||||
if len(s) <= n {
|
||||
return s
|
||||
}
|
||||
return s[:n] + "..."
|
||||
}
|
||||
@@ -0,0 +1,358 @@
|
||||
// Package hitl is the Discord human-in-the-loop bot of experiment design 6.5/6.6: it posts layer-2 proposals with
|
||||
// Approve/Reject buttons, turns a button press into harness approval + execution, expires silent proposals into
|
||||
// the layer-3 (escalation) channel, and posts escalations with the matching runbook in a thread.
|
||||
package hitl
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/bwmarrin/discordgo"
|
||||
|
||||
"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/harness"
|
||||
"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/rag"
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
GuildID string
|
||||
ApprovalsChannel string
|
||||
EscalationChannel string
|
||||
ApproverRoleID string // optional: only members with this role may decide
|
||||
}
|
||||
|
||||
// Discord abstracts the discordgo calls so the decision logic is testable without a gateway.
|
||||
type Discord interface {
|
||||
Send(channelID string, msg *discordgo.MessageSend) (*discordgo.Message, error)
|
||||
Edit(channelID, messageID string, embed *discordgo.MessageEmbed, components []discordgo.MessageComponent) error
|
||||
Thread(channelID, messageID, name, text string) error
|
||||
}
|
||||
|
||||
type Bot struct {
|
||||
cfg Config
|
||||
dc Discord
|
||||
harness *harness.Client
|
||||
store *rag.Store
|
||||
mu sync.Mutex
|
||||
timers map[string]*time.Timer
|
||||
}
|
||||
|
||||
func New(cfg Config, dc Discord, hc *harness.Client, store *rag.Store) *Bot {
|
||||
return &Bot{cfg: cfg, dc: dc, harness: hc, store: store, timers: map[string]*time.Timer{}}
|
||||
}
|
||||
|
||||
// PostRequest is what the agent sends to POST /post.
|
||||
type PostRequest struct {
|
||||
Token string `json:"token"`
|
||||
ActionID string `json:"action_id"`
|
||||
Target map[string]any `json:"target"`
|
||||
Diagnosis map[string]any `json:"diagnosis"`
|
||||
RunID *string `json:"run_id"`
|
||||
ScenarioID *string `json:"scenario_id"`
|
||||
SessionID string `json:"session_id"`
|
||||
ExpiresIn float64 `json:"expires_in_s"`
|
||||
}
|
||||
|
||||
func deref(s *string) string {
|
||||
if s == nil {
|
||||
return "-"
|
||||
}
|
||||
return *s
|
||||
}
|
||||
|
||||
func embedFor(d map[string]any, actionID string, target map[string]any, runID, scenarioID *string, status string, color int) *discordgo.MessageEmbed {
|
||||
ev, _ := d["evidence"].([]any)
|
||||
var evLines []string
|
||||
for i, e := range ev {
|
||||
if i >= 5 {
|
||||
break
|
||||
}
|
||||
evLines = append(evLines, "- "+fmt.Sprint(e))
|
||||
}
|
||||
tgt, _ := json.Marshal(target)
|
||||
return &discordgo.MessageEmbed{
|
||||
Title: fmt.Sprintf("Layer-2 proposal: %s", actionID),
|
||||
Color: color,
|
||||
Fields: []*discordgo.MessageEmbedField{
|
||||
{Name: "Root cause", Value: fmt.Sprintf("%v (confidence %.2f)", d["root_cause_label"], toF(d["confidence"])), Inline: false},
|
||||
{Name: "Target", Value: "`" + string(tgt) + "`", Inline: false},
|
||||
{Name: "Evidence", Value: nonEmpty(strings.Join(evLines, "\n")), Inline: false},
|
||||
{Name: "Rationale", Value: nonEmpty(fmt.Sprint(d["rationale"])), Inline: false},
|
||||
{Name: "Run / scenario", Value: deref(runID) + " / " + deref(scenarioID), Inline: true},
|
||||
{Name: "Status", Value: status, Inline: true},
|
||||
},
|
||||
Footer: &discordgo.MessageEmbedFooter{Text: fmt.Sprintf("model %v, prompt %v, session %v", d["model"], d["prompt_version"], d["session_id"])},
|
||||
}
|
||||
}
|
||||
|
||||
func nonEmpty(s string) string {
|
||||
if strings.TrimSpace(s) == "" {
|
||||
return "-"
|
||||
}
|
||||
if len(s) > 1000 {
|
||||
return s[:1000] + "..."
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func toF(v any) float64 {
|
||||
f, _ := v.(float64)
|
||||
return f
|
||||
}
|
||||
|
||||
const (
|
||||
colorPending = 0x5865f2
|
||||
colorApproved = 0x57f287
|
||||
colorRejected = 0xed4245
|
||||
colorExpired = 0x99aab5
|
||||
colorEscalate = 0xfee75c
|
||||
)
|
||||
|
||||
func buttons(token string, disabled bool) []discordgo.MessageComponent {
|
||||
return []discordgo.MessageComponent{discordgo.ActionsRow{Components: []discordgo.MessageComponent{
|
||||
discordgo.Button{Label: "Approve", Style: discordgo.SuccessButton, CustomID: "approve:" + token, Disabled: disabled},
|
||||
discordgo.Button{Label: "Reject", Style: discordgo.DangerButton, CustomID: "reject:" + token, Disabled: disabled},
|
||||
}}}
|
||||
}
|
||||
|
||||
// Post publishes the proposal to the approvals channel and arms the expiry timer.
|
||||
func (b *Bot) Post(ctx context.Context, req PostRequest) error {
|
||||
ttl := time.Duration(req.ExpiresIn) * time.Second
|
||||
if ttl <= 0 {
|
||||
ttl = 15 * time.Minute
|
||||
}
|
||||
expires := time.Now().Add(ttl)
|
||||
em := embedFor(req.Diagnosis, req.ActionID, req.Target, req.RunID, req.ScenarioID, fmt.Sprintf("waiting for approval until <t:%d:T>", expires.Unix()), colorPending)
|
||||
msg, err := b.dc.Send(b.cfg.ApprovalsChannel, &discordgo.MessageSend{Embeds: []*discordgo.MessageEmbed{em}, Components: buttons(req.Token, false)})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := b.store.InsertHitl(ctx, rag.HitlMessage{Token: req.Token, ChannelID: msg.ChannelID, MessageID: msg.ID, ActionID: req.ActionID, Target: req.Target, RunID: req.RunID, ScenarioID: req.ScenarioID, SessionID: req.SessionID, ExpiresAt: expires}); err != nil {
|
||||
return err
|
||||
}
|
||||
b.arm(req.Token, time.Until(expires))
|
||||
return nil
|
||||
}
|
||||
|
||||
func (b *Bot) arm(token string, d time.Duration) {
|
||||
if d < 0 {
|
||||
d = 0
|
||||
}
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
if t, ok := b.timers[token]; ok {
|
||||
t.Stop()
|
||||
}
|
||||
b.timers[token] = time.AfterFunc(d, func() { b.expire(token) })
|
||||
}
|
||||
|
||||
func (b *Bot) disarm(token string) {
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
if t, ok := b.timers[token]; ok {
|
||||
t.Stop()
|
||||
delete(b.timers, token)
|
||||
}
|
||||
}
|
||||
|
||||
// RearmPending re-creates expiry timers after a restart.
|
||||
func (b *Bot) RearmPending(ctx context.Context) {
|
||||
list, err := b.store.Pending(ctx)
|
||||
if err != nil {
|
||||
log.Printf("pending: %v", err)
|
||||
return
|
||||
}
|
||||
for _, m := range list {
|
||||
b.arm(m.Token, time.Until(m.ExpiresAt))
|
||||
}
|
||||
log.Printf("re-armed %d pending proposals", len(list))
|
||||
}
|
||||
|
||||
// Decide handles a button press. by is "username (id)". Returns the status text for logging.
|
||||
func (b *Bot) Decide(ctx context.Context, token, decision, by string) (string, error) {
|
||||
m, err := b.store.GetHitl(ctx, token)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("unknown proposal %s", token)
|
||||
}
|
||||
first, err := b.store.Decide(ctx, token, decision, by, nil)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if !first {
|
||||
return "already decided", nil
|
||||
}
|
||||
b.disarm(token)
|
||||
switch decision {
|
||||
case "reject":
|
||||
if err := b.harness.Reject(ctx, token, by, "rejected in Discord"); err != nil {
|
||||
log.Printf("harness reject: %v", err)
|
||||
}
|
||||
_ = b.edit(m, "rejected by "+by, colorRejected, nil)
|
||||
return "rejected", nil
|
||||
case "approve":
|
||||
if time.Now().After(m.ExpiresAt) {
|
||||
_ = b.edit(m, "expired before approval", colorExpired, nil)
|
||||
return "expired", nil
|
||||
}
|
||||
if err := b.harness.Approve(ctx, token, by); err != nil {
|
||||
_ = b.edit(m, "approval failed: "+err.Error(), colorRejected, nil)
|
||||
_ = b.store.SetHitlResult(ctx, token, map[string]any{"error": err.Error()})
|
||||
return "approve failed", err
|
||||
}
|
||||
_ = b.edit(m, "approved by "+by+", executing...", colorApproved, nil)
|
||||
res, err := b.harness.Execute(ctx, token, m.ActionID, m.Target, m.RunID, m.ScenarioID, "approved in Discord by "+by)
|
||||
if err != nil {
|
||||
_ = b.store.SetHitlResult(ctx, token, map[string]any{"error": err.Error(), "result": res})
|
||||
_ = b.edit(m, "approved by "+by+"; execution failed: "+truncate(err.Error(), 300), colorRejected, res)
|
||||
return "execute failed", err
|
||||
}
|
||||
_ = b.store.SetHitlResult(ctx, token, res)
|
||||
_ = b.edit(m, "approved by "+by+"; executed", colorApproved, res)
|
||||
return "executed", nil
|
||||
}
|
||||
return "", fmt.Errorf("unknown decision %s", decision)
|
||||
}
|
||||
|
||||
func (b *Bot) edit(m *rag.HitlMessage, status string, color int, result map[string]any) error {
|
||||
ctx := context.Background()
|
||||
d := map[string]any{}
|
||||
if se, err := b.store.GetHitl(ctx, m.Token); err == nil {
|
||||
m = se
|
||||
}
|
||||
em := embedFor(d, m.ActionID, m.Target, m.RunID, m.ScenarioID, status, color)
|
||||
em.Fields = em.Fields[1:] // the diagnosis fields are already in the original message; keep the edit compact
|
||||
if result != nil {
|
||||
rb, _ := json.Marshal(result)
|
||||
em.Fields = append(em.Fields, &discordgo.MessageEmbedField{Name: "Result", Value: "```json\n" + truncate(string(rb), 900) + "\n```"})
|
||||
}
|
||||
return b.dc.Edit(m.ChannelID, m.MessageID, em, buttons(m.Token, true))
|
||||
}
|
||||
|
||||
func (b *Bot) expire(token string) {
|
||||
ctx := context.Background()
|
||||
m, err := b.store.GetHitl(ctx, token)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
first, err := b.store.Decide(ctx, token, "expired", "timer", nil)
|
||||
if err != nil || !first {
|
||||
return
|
||||
}
|
||||
_ = b.edit(m, "no decision within the window; moved to layer 3", colorExpired, nil)
|
||||
text := fmt.Sprintf("Proposal `%s` for `%s` (run %s, scenario %s) received no decision within the approval window. Layer 3: follow the runbook or call the harness `/actions/execute` with the operator token.",
|
||||
m.ActionID, mustJSON(m.Target), deref(m.RunID), deref(m.ScenarioID))
|
||||
if _, err := b.dc.Send(b.cfg.EscalationChannel, &discordgo.MessageSend{Content: text}); err != nil {
|
||||
log.Printf("escalation post: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// EscalateRequest is what the agent sends to POST /escalate (design 6.6).
|
||||
type EscalateRequest struct {
|
||||
Diagnosis map[string]any `json:"diagnosis"`
|
||||
Runbook string `json:"runbook"`
|
||||
RunID *string `json:"run_id"`
|
||||
ScenarioID *string `json:"scenario_id"`
|
||||
Alert map[string]any `json:"alert"`
|
||||
}
|
||||
|
||||
func (b *Bot) Escalate(ctx context.Context, req EscalateRequest) error {
|
||||
em := embedFor(req.Diagnosis, "escalate", map[string]any{"target": req.Diagnosis["target"]}, req.RunID, req.ScenarioID, "layer 3: human decision required", colorEscalate)
|
||||
em.Title = fmt.Sprintf("Escalation: %v", req.Diagnosis["root_cause_label"])
|
||||
if an, ok := req.Alert["alertname"]; ok {
|
||||
em.Fields = append(em.Fields, &discordgo.MessageEmbedField{Name: "Alert", Value: fmt.Sprint(an), Inline: true})
|
||||
}
|
||||
msg, err := b.dc.Send(b.cfg.EscalationChannel, &discordgo.MessageSend{Embeds: []*discordgo.MessageEmbed{em}})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
name := fmt.Sprintf("%s %s", deref(req.ScenarioID), time.Now().UTC().Format("01-02 15:04"))
|
||||
text := "Runbook:\n" + nonEmpty(truncate(req.Runbook, 1800))
|
||||
if err := b.dc.Thread(msg.ChannelID, msg.ID, name, text); err != nil {
|
||||
log.Printf("thread: %v", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Allowed checks the optional approver role.
|
||||
func (b *Bot) Allowed(roles []string) bool {
|
||||
if b.cfg.ApproverRoleID == "" {
|
||||
return true
|
||||
}
|
||||
for _, r := range roles {
|
||||
if r == b.cfg.ApproverRoleID {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func mustJSON(v any) string {
|
||||
b, _ := json.Marshal(v)
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func truncate(s string, n int) string {
|
||||
if len(s) <= n {
|
||||
return s
|
||||
}
|
||||
return s[:n] + "..."
|
||||
}
|
||||
|
||||
// ---- discordgo adapter ----
|
||||
|
||||
type Session struct{ S *discordgo.Session }
|
||||
|
||||
func (d Session) Send(channelID string, msg *discordgo.MessageSend) (*discordgo.Message, error) {
|
||||
return d.S.ChannelMessageSendComplex(channelID, msg)
|
||||
}
|
||||
|
||||
func (d Session) Edit(channelID, messageID string, embed *discordgo.MessageEmbed, components []discordgo.MessageComponent) error {
|
||||
embeds := []*discordgo.MessageEmbed{embed}
|
||||
_, err := d.S.ChannelMessageEditComplex(&discordgo.MessageEdit{Channel: channelID, ID: messageID, Embeds: &embeds, Components: &components})
|
||||
return err
|
||||
}
|
||||
|
||||
func (d Session) Thread(channelID, messageID, name, text string) error {
|
||||
th, err := d.S.MessageThreadStartComplex(channelID, messageID, &discordgo.ThreadStart{Name: truncate(name, 90), AutoArchiveDuration: 1440})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = d.S.ChannelMessageSend(th.ID, text)
|
||||
return err
|
||||
}
|
||||
|
||||
// HandleInteraction is the discordgo InteractionCreate handler: DEFERRED_UPDATE_MESSAGE within 3 s, then the decision.
|
||||
func (b *Bot) HandleInteraction(s *discordgo.Session, i *discordgo.InteractionCreate) {
|
||||
if i.Type != discordgo.InteractionMessageComponent {
|
||||
return
|
||||
}
|
||||
cid := i.MessageComponentData().CustomID
|
||||
parts := strings.SplitN(cid, ":", 2)
|
||||
if len(parts) != 2 || (parts[0] != "approve" && parts[0] != "reject") {
|
||||
return
|
||||
}
|
||||
var roles []string
|
||||
by := "unknown"
|
||||
if i.Member != nil && i.Member.User != nil {
|
||||
roles = i.Member.Roles
|
||||
by = fmt.Sprintf("%s (%s)", i.Member.User.Username, i.Member.User.ID)
|
||||
}
|
||||
if !b.Allowed(roles) {
|
||||
_ = s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{Type: discordgo.InteractionResponseChannelMessageWithSource,
|
||||
Data: &discordgo.InteractionResponseData{Content: "You are not in the approver role.", Flags: discordgo.MessageFlagsEphemeral}})
|
||||
return
|
||||
}
|
||||
if err := s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{Type: discordgo.InteractionResponseDeferredMessageUpdate}); err != nil {
|
||||
log.Printf("interaction respond: %v", err)
|
||||
}
|
||||
go func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 11*time.Minute)
|
||||
defer cancel()
|
||||
st, err := b.Decide(ctx, parts[1], parts[0], by)
|
||||
log.Printf("decision %s on %s by %s: %s err=%v", parts[0], parts[1][len(parts[1])-6:], by, st, err)
|
||||
}()
|
||||
}
|
||||
@@ -0,0 +1,184 @@
|
||||
// Package openai is a minimal client for the OpenAI Responses API (function calling) and the Embeddings API.
|
||||
// Shapes follow https://developers.openai.com/api/docs/guides/function-calling and /guides/embeddings (checked 2026-10-08).
|
||||
package openai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Client struct {
|
||||
BaseURL string // default https://api.openai.com/v1
|
||||
KeyFile string // read per call so a rotated key file is picked up without restart
|
||||
HTTP *http.Client
|
||||
keyFn func() (string, error)
|
||||
}
|
||||
|
||||
func New(baseURL string, keyFn func() (string, error)) *Client {
|
||||
if baseURL == "" {
|
||||
baseURL = "https://api.openai.com/v1"
|
||||
}
|
||||
return &Client{BaseURL: baseURL, keyFn: keyFn, HTTP: &http.Client{Timeout: 120 * time.Second}}
|
||||
}
|
||||
|
||||
// Tool is a function tool definition (strict mode: additionalProperties false and every property required).
|
||||
type Tool struct {
|
||||
Type string `json:"type"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Parameters map[string]any `json:"parameters"`
|
||||
Strict bool `json:"strict"`
|
||||
}
|
||||
|
||||
// Item is one input/output item of the Responses API. Only the fields this agent uses are modelled.
|
||||
type Item struct {
|
||||
Type string `json:"type"` // message | function_call | function_call_output
|
||||
Role string `json:"role,omitempty"` // message: system | user | assistant
|
||||
Content any `json:"content,omitempty"` // message: string (input) or []ContentPart (output)
|
||||
CallID string `json:"call_id,omitempty"` // function_call / function_call_output
|
||||
Name string `json:"name,omitempty"` // function_call
|
||||
Arguments string `json:"arguments,omitempty"` // function_call (JSON string)
|
||||
Output string `json:"output,omitempty"` // function_call_output
|
||||
ID string `json:"id,omitempty"`
|
||||
Status string `json:"status,omitempty"`
|
||||
}
|
||||
|
||||
type ContentPart struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
}
|
||||
|
||||
type Request struct {
|
||||
Model string `json:"model"`
|
||||
Input []Item `json:"input"`
|
||||
Tools []Tool `json:"tools,omitempty"`
|
||||
ToolChoice any `json:"tool_choice,omitempty"` // "auto" | "required" | {"type":"function","name":...}
|
||||
PreviousResponseID string `json:"previous_response_id,omitempty"`
|
||||
MaxOutputTokens int `json:"max_output_tokens,omitempty"`
|
||||
ParallelToolCalls *bool `json:"parallel_tool_calls,omitempty"`
|
||||
Store *bool `json:"store,omitempty"`
|
||||
}
|
||||
|
||||
type Response struct {
|
||||
ID string `json:"id"`
|
||||
Status string `json:"status"`
|
||||
Model string `json:"model"`
|
||||
Output []Item `json:"output"`
|
||||
Usage struct {
|
||||
InputTokens int `json:"input_tokens"`
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
} `json:"usage"`
|
||||
Error *struct {
|
||||
Message string `json:"message"`
|
||||
Type string `json:"type"`
|
||||
} `json:"error"`
|
||||
}
|
||||
|
||||
// FunctionCalls returns the function_call items of the output.
|
||||
func (r *Response) FunctionCalls() []Item {
|
||||
var out []Item
|
||||
for _, it := range r.Output {
|
||||
if it.Type == "function_call" {
|
||||
out = append(out, it)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Text concatenates assistant message text parts.
|
||||
func (r *Response) Text() string {
|
||||
var b bytes.Buffer
|
||||
for _, it := range r.Output {
|
||||
if it.Type != "message" {
|
||||
continue
|
||||
}
|
||||
if parts, ok := it.Content.([]any); ok {
|
||||
for _, p := range parts {
|
||||
if m, ok := p.(map[string]any); ok {
|
||||
if t, ok := m["text"].(string); ok {
|
||||
b.WriteString(t)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func (c *Client) do(ctx context.Context, path string, body any, out any) error {
|
||||
key, err := c.keyFn()
|
||||
if err != nil {
|
||||
return fmt.Errorf("openai key: %w", err)
|
||||
}
|
||||
if key == "" {
|
||||
return fmt.Errorf("openai key is empty (secret not registered yet)")
|
||||
}
|
||||
b, _ := json.Marshal(body)
|
||||
req, err := http.NewRequestWithContext(ctx, "POST", c.BaseURL+path, bytes.NewReader(b))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := c.HTTP.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
rb, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
|
||||
if resp.StatusCode/100 != 2 {
|
||||
return fmt.Errorf("openai %s: status %d: %s", path, resp.StatusCode, truncate(string(rb), 600))
|
||||
}
|
||||
return json.Unmarshal(rb, out)
|
||||
}
|
||||
|
||||
func (c *Client) CreateResponse(ctx context.Context, req Request) (*Response, error) {
|
||||
var r Response
|
||||
if err := c.do(ctx, "/responses", req, &r); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if r.Error != nil {
|
||||
return nil, fmt.Errorf("openai response error: %s", r.Error.Message)
|
||||
}
|
||||
return &r, nil
|
||||
}
|
||||
|
||||
type embedRequest struct {
|
||||
Model string `json:"model"`
|
||||
Input []string `json:"input"`
|
||||
}
|
||||
|
||||
type embedResponse struct {
|
||||
Data []struct {
|
||||
Index int `json:"index"`
|
||||
Embedding []float32 `json:"embedding"`
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
// Embed returns one vector per input (text-embedding-3-small: 1536 dimensions by default).
|
||||
func (c *Client) Embed(ctx context.Context, model string, inputs []string) ([][]float32, error) {
|
||||
var r embedResponse
|
||||
if err := c.do(ctx, "/embeddings", embedRequest{Model: model, Input: inputs}, &r); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([][]float32, len(inputs))
|
||||
for _, d := range r.Data {
|
||||
if d.Index >= 0 && d.Index < len(out) {
|
||||
out[d.Index] = d.Embedding
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func truncate(s string, n int) string {
|
||||
if len(s) <= n {
|
||||
return s
|
||||
}
|
||||
return s[:n] + "..."
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
package rag
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestChunk(t *testing.T) {
|
||||
doc := strings.Repeat("para one line.\n", 20) + "\n\n" + strings.Repeat("para two line.\n", 20) + "\n\n" + strings.Repeat("x", 5000)
|
||||
cs := Chunk(doc, 2000)
|
||||
if len(cs) < 3 {
|
||||
t.Fatalf("expected several chunks, got %d", len(cs))
|
||||
}
|
||||
for _, c := range cs {
|
||||
if len(c) > 2000 {
|
||||
t.Fatalf("chunk exceeds limit: %d", len(c))
|
||||
}
|
||||
if strings.TrimSpace(c) == "" {
|
||||
t.Fatal("empty chunk")
|
||||
}
|
||||
}
|
||||
if PromptVersion("a") == PromptVersion("b") || len(PromptVersion("a")) != 12 {
|
||||
t.Fatal("prompt version")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,367 @@
|
||||
// Package rag is the AI VM state store (experiment design 6.4): runbook chunks with pgvector embeddings, prompt
|
||||
// versions, diagnosis sessions and the Discord HITL message ledger. Schema is created at start-up; idempotent.
|
||||
package rag
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
pgvector "github.com/pgvector/pgvector-go"
|
||||
)
|
||||
|
||||
type Store struct{ pool *pgxpool.Pool }
|
||||
|
||||
const schema = `
|
||||
CREATE EXTENSION IF NOT EXISTS vector;
|
||||
CREATE SCHEMA IF NOT EXISTS ai;
|
||||
CREATE TABLE IF NOT EXISTS ai.rag_docs (
|
||||
id bigserial PRIMARY KEY,
|
||||
source text NOT NULL,
|
||||
chunk text NOT NULL,
|
||||
chunk_hash text NOT NULL UNIQUE,
|
||||
embedding vector(1536) NOT NULL,
|
||||
created_at timestamptz NOT NULL DEFAULT now()
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS rag_docs_embedding_idx ON ai.rag_docs USING hnsw (embedding vector_cosine_ops);
|
||||
CREATE TABLE IF NOT EXISTS ai.prompt_versions (
|
||||
version text PRIMARY KEY,
|
||||
body text NOT NULL,
|
||||
created_at timestamptz NOT NULL DEFAULT now()
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS ai.sessions (
|
||||
session_id text PRIMARY KEY,
|
||||
run_id text,
|
||||
scenario_id text,
|
||||
model text NOT NULL,
|
||||
prompt_version text NOT NULL,
|
||||
input jsonb NOT NULL,
|
||||
output jsonb,
|
||||
tool_calls jsonb NOT NULL DEFAULT '[]'::jsonb,
|
||||
usage jsonb,
|
||||
error text,
|
||||
duration_ms bigint,
|
||||
created_at timestamptz NOT NULL DEFAULT now()
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS ai.hitl_messages (
|
||||
token text PRIMARY KEY,
|
||||
channel_id text NOT NULL,
|
||||
message_id text NOT NULL,
|
||||
action_id text NOT NULL,
|
||||
target jsonb NOT NULL DEFAULT '{}'::jsonb,
|
||||
run_id text,
|
||||
scenario_id text,
|
||||
session_id text,
|
||||
expires_at timestamptz NOT NULL,
|
||||
decided_at timestamptz,
|
||||
decision text,
|
||||
decided_by text,
|
||||
result jsonb,
|
||||
created_at timestamptz NOT NULL DEFAULT now()
|
||||
);`
|
||||
|
||||
func New(ctx context.Context, dsn string) (*Store, error) {
|
||||
pool, err := pgxpool.New(ctx, dsn)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var lastErr error
|
||||
for i := 0; i < 30; i++ { // rag-db container may still be starting
|
||||
if lastErr = pool.Ping(ctx); lastErr == nil {
|
||||
break
|
||||
}
|
||||
time.Sleep(2 * time.Second)
|
||||
}
|
||||
if lastErr != nil {
|
||||
return nil, fmt.Errorf("rag db: %w", lastErr)
|
||||
}
|
||||
if _, err := pool.Exec(ctx, schema); err != nil {
|
||||
return nil, fmt.Errorf("rag schema: %w", err)
|
||||
}
|
||||
return &Store{pool: pool}, nil
|
||||
}
|
||||
|
||||
func (s *Store) Ping(ctx context.Context) error { return s.pool.Ping(ctx) }
|
||||
|
||||
// ---- prompt versions ----
|
||||
|
||||
// PromptVersion is the first 12 hex characters of sha256(body); the same body always maps to the same version.
|
||||
func PromptVersion(body string) string {
|
||||
h := sha256.Sum256([]byte(body))
|
||||
return hex.EncodeToString(h[:])[:12]
|
||||
}
|
||||
|
||||
func (s *Store) UpsertPrompt(ctx context.Context, body string) (string, error) {
|
||||
v := PromptVersion(body)
|
||||
_, err := s.pool.Exec(ctx, `INSERT INTO ai.prompt_versions (version, body) VALUES ($1,$2) ON CONFLICT (version) DO NOTHING`, v, body)
|
||||
return v, err
|
||||
}
|
||||
|
||||
// ---- runbook chunks ----
|
||||
|
||||
// Chunk splits a document into pieces of about maxChars characters on paragraph boundaries (design 6.4: ~500 tokens).
|
||||
func Chunk(text string, maxChars int) []string {
|
||||
paras := strings.Split(strings.ReplaceAll(text, "\r\n", "\n"), "\n\n")
|
||||
var out []string
|
||||
var cur strings.Builder
|
||||
flush := func() {
|
||||
if t := strings.TrimSpace(cur.String()); t != "" {
|
||||
out = append(out, t)
|
||||
}
|
||||
cur.Reset()
|
||||
}
|
||||
for _, p := range paras {
|
||||
p = strings.TrimSpace(p)
|
||||
if p == "" {
|
||||
continue
|
||||
}
|
||||
if cur.Len() > 0 && cur.Len()+len(p)+2 > maxChars {
|
||||
flush()
|
||||
}
|
||||
for len(p) > maxChars { // a single oversized paragraph is cut hard
|
||||
cut := strings.LastIndex(p[:maxChars], "\n")
|
||||
if cut < maxChars/2 {
|
||||
cut = maxChars
|
||||
}
|
||||
if cur.Len() > 0 {
|
||||
flush()
|
||||
}
|
||||
out = append(out, strings.TrimSpace(p[:cut]))
|
||||
p = strings.TrimSpace(p[cut:])
|
||||
}
|
||||
if cur.Len() > 0 {
|
||||
cur.WriteString("\n\n")
|
||||
}
|
||||
cur.WriteString(p)
|
||||
}
|
||||
flush()
|
||||
return out
|
||||
}
|
||||
|
||||
type Embedder interface {
|
||||
Embed(ctx context.Context, model string, inputs []string) ([][]float32, error)
|
||||
}
|
||||
|
||||
// IngestDir embeds every *.md under dir whose chunk is not stored yet and removes chunks of files that disappeared.
|
||||
// Returns (new chunks, removed chunks). Unchanged files cost nothing (hash lookup only).
|
||||
func (s *Store) IngestDir(ctx context.Context, dir string, emb Embedder, model string) (int, int, error) {
|
||||
files, err := filepath.Glob(filepath.Join(dir, "*.md"))
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
sort.Strings(files)
|
||||
type pending struct{ source, chunk, hash string }
|
||||
var todo []pending
|
||||
keep := map[string]bool{}
|
||||
sources := map[string]bool{}
|
||||
for _, f := range files {
|
||||
b, err := os.ReadFile(f)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
src := filepath.Base(f)
|
||||
sources[src] = true
|
||||
for _, c := range Chunk(string(b), 2000) {
|
||||
h := sha256.Sum256([]byte(src + "\x00" + c))
|
||||
hash := hex.EncodeToString(h[:])
|
||||
keep[hash] = true
|
||||
var exists bool
|
||||
if err := s.pool.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM ai.rag_docs WHERE chunk_hash=$1)`, hash).Scan(&exists); err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
if !exists {
|
||||
todo = append(todo, pending{src, c, hash})
|
||||
}
|
||||
}
|
||||
}
|
||||
added := 0
|
||||
for i := 0; i < len(todo); i += 32 {
|
||||
j := i + 32
|
||||
if j > len(todo) {
|
||||
j = len(todo)
|
||||
}
|
||||
batch := todo[i:j]
|
||||
texts := make([]string, len(batch))
|
||||
for k, p := range batch {
|
||||
texts[k] = p.chunk
|
||||
}
|
||||
vecs, err := emb.Embed(ctx, model, texts)
|
||||
if err != nil {
|
||||
return added, 0, err
|
||||
}
|
||||
for k, p := range batch {
|
||||
if len(vecs[k]) == 0 {
|
||||
return added, 0, fmt.Errorf("empty embedding for chunk of %s", p.source)
|
||||
}
|
||||
if _, err := s.pool.Exec(ctx, `INSERT INTO ai.rag_docs (source, chunk, chunk_hash, embedding) VALUES ($1,$2,$3,$4) ON CONFLICT (chunk_hash) DO NOTHING`,
|
||||
p.source, p.chunk, p.hash, pgvector.NewVector(vecs[k])); err != nil {
|
||||
return added, 0, err
|
||||
}
|
||||
added++
|
||||
}
|
||||
}
|
||||
// drop stale chunks (file removed or content changed)
|
||||
rows, err := s.pool.Query(ctx, `SELECT chunk_hash, source FROM ai.rag_docs`)
|
||||
if err != nil {
|
||||
return added, 0, err
|
||||
}
|
||||
var stale []string
|
||||
for rows.Next() {
|
||||
var h, src string
|
||||
if err := rows.Scan(&h, &src); err != nil {
|
||||
rows.Close()
|
||||
return added, 0, err
|
||||
}
|
||||
if !keep[h] {
|
||||
stale = append(stale, h)
|
||||
}
|
||||
}
|
||||
rows.Close()
|
||||
for _, h := range stale {
|
||||
if _, err := s.pool.Exec(ctx, `DELETE FROM ai.rag_docs WHERE chunk_hash=$1`, h); err != nil {
|
||||
return added, len(stale), err
|
||||
}
|
||||
}
|
||||
return added, len(stale), nil
|
||||
}
|
||||
|
||||
type Hit struct {
|
||||
Source string `json:"source"`
|
||||
Chunk string `json:"chunk"`
|
||||
Distance float64 `json:"distance"`
|
||||
}
|
||||
|
||||
// Search returns the k nearest chunks by cosine distance (design 6.4: ORDER BY embedding <=> $1 LIMIT k).
|
||||
func (s *Store) Search(ctx context.Context, vec []float32, k int) ([]Hit, error) {
|
||||
if k <= 0 || k > 20 {
|
||||
k = 5
|
||||
}
|
||||
rows, err := s.pool.Query(ctx, `SELECT source, chunk, embedding <=> $1 AS d FROM ai.rag_docs ORDER BY d LIMIT $2`, pgvector.NewVector(vec), k)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []Hit
|
||||
for rows.Next() {
|
||||
var h Hit
|
||||
if err := rows.Scan(&h.Source, &h.Chunk, &h.Distance); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, h)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) DocCount(ctx context.Context) (int, error) {
|
||||
var n int
|
||||
err := s.pool.QueryRow(ctx, `SELECT count(*) FROM ai.rag_docs`).Scan(&n)
|
||||
return n, err
|
||||
}
|
||||
|
||||
// ---- sessions ----
|
||||
|
||||
type Session struct {
|
||||
SessionID string
|
||||
RunID *string
|
||||
ScenarioID *string
|
||||
Model string
|
||||
PromptVersion string
|
||||
Input any
|
||||
Output any
|
||||
ToolCalls any
|
||||
Usage any
|
||||
Error string
|
||||
Duration time.Duration
|
||||
}
|
||||
|
||||
func (s *Store) SaveSession(ctx context.Context, se Session) error {
|
||||
in, _ := json.Marshal(se.Input)
|
||||
out, _ := json.Marshal(se.Output)
|
||||
tc, _ := json.Marshal(se.ToolCalls)
|
||||
us, _ := json.Marshal(se.Usage)
|
||||
_, err := s.pool.Exec(ctx, `INSERT INTO ai.sessions (session_id,run_id,scenario_id,model,prompt_version,input,output,tool_calls,usage,error,duration_ms)
|
||||
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,NULLIF($10,''),$11)
|
||||
ON CONFLICT (session_id) DO UPDATE SET output=EXCLUDED.output, tool_calls=EXCLUDED.tool_calls, usage=EXCLUDED.usage, error=EXCLUDED.error, duration_ms=EXCLUDED.duration_ms`,
|
||||
se.SessionID, se.RunID, se.ScenarioID, se.Model, se.PromptVersion, in, out, tc, us, se.Error, se.Duration.Milliseconds())
|
||||
return err
|
||||
}
|
||||
|
||||
// ---- HITL ledger ----
|
||||
|
||||
type HitlMessage struct {
|
||||
Token string
|
||||
ChannelID string
|
||||
MessageID string
|
||||
ActionID string
|
||||
Target map[string]any
|
||||
RunID *string
|
||||
ScenarioID *string
|
||||
SessionID string
|
||||
ExpiresAt time.Time
|
||||
DecidedAt *time.Time
|
||||
Decision *string
|
||||
DecidedBy *string
|
||||
}
|
||||
|
||||
func (s *Store) InsertHitl(ctx context.Context, m HitlMessage) error {
|
||||
t, _ := json.Marshal(m.Target)
|
||||
_, err := s.pool.Exec(ctx, `INSERT INTO ai.hitl_messages (token,channel_id,message_id,action_id,target,run_id,scenario_id,session_id,expires_at)
|
||||
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9)`, m.Token, m.ChannelID, m.MessageID, m.ActionID, t, m.RunID, m.ScenarioID, m.SessionID, m.ExpiresAt)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Store) GetHitl(ctx context.Context, token string) (*HitlMessage, error) {
|
||||
var m HitlMessage
|
||||
var t []byte
|
||||
err := s.pool.QueryRow(ctx, `SELECT token,channel_id,message_id,action_id,target,run_id,scenario_id,COALESCE(session_id,''),expires_at,decided_at,decision,decided_by FROM ai.hitl_messages WHERE token=$1`, token).
|
||||
Scan(&m.Token, &m.ChannelID, &m.MessageID, &m.ActionID, &t, &m.RunID, &m.ScenarioID, &m.SessionID, &m.ExpiresAt, &m.DecidedAt, &m.Decision, &m.DecidedBy)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = json.Unmarshal(t, &m.Target)
|
||||
return &m, nil
|
||||
}
|
||||
|
||||
// Decide records the first decision only; returns false when one was already recorded (double click / race).
|
||||
func (s *Store) Decide(ctx context.Context, token, decision, by string, result any) (bool, error) {
|
||||
r, _ := json.Marshal(result)
|
||||
tag, err := s.pool.Exec(ctx, `UPDATE ai.hitl_messages SET decided_at=now(), decision=$2, decided_by=$3, result=$4 WHERE token=$1 AND decided_at IS NULL`, token, decision, by, r)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return tag.RowsAffected() == 1, nil
|
||||
}
|
||||
|
||||
func (s *Store) SetHitlResult(ctx context.Context, token string, result any) error {
|
||||
r, _ := json.Marshal(result)
|
||||
_, err := s.pool.Exec(ctx, `UPDATE ai.hitl_messages SET result=$2 WHERE token=$1`, token, r)
|
||||
return err
|
||||
}
|
||||
|
||||
// Pending returns undecided messages (restart safety: expiry timers are re-armed from this list).
|
||||
func (s *Store) Pending(ctx context.Context) ([]HitlMessage, error) {
|
||||
rows, err := s.pool.Query(ctx, `SELECT token,channel_id,message_id,action_id,target,run_id,scenario_id,COALESCE(session_id,''),expires_at FROM ai.hitl_messages WHERE decided_at IS NULL`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []HitlMessage
|
||||
for rows.Next() {
|
||||
var m HitlMessage
|
||||
var t []byte
|
||||
if err := rows.Scan(&m.Token, &m.ChannelID, &m.MessageID, &m.ActionID, &t, &m.RunID, &m.ScenarioID, &m.SessionID, &m.ExpiresAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = json.Unmarshal(t, &m.Target)
|
||||
out = append(out, m)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
@@ -0,0 +1,560 @@
|
||||
// Package tools implements the read-only allow-list of experiment design 6.2. Every tool returns a compact JSON
|
||||
// string for the model; nothing here writes to any system.
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/aws"
|
||||
"github.com/aws/aws-sdk-go-v2/credentials"
|
||||
"github.com/aws/aws-sdk-go-v2/service/ec2"
|
||||
ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
|
||||
elbv2 "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2"
|
||||
"github.com/aws/aws-sdk-go-v2/service/sts"
|
||||
"google.golang.org/api/compute/v1"
|
||||
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
|
||||
"k8s.io/client-go/kubernetes"
|
||||
"k8s.io/client-go/rest"
|
||||
|
||||
"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/openai"
|
||||
"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/rag"
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
MimirURL string // http://10.20.0.10:9009/prometheus
|
||||
LokiURL string // http://10.20.0.10:3100
|
||||
K8sAPI string // https://10.10.0.10:6443
|
||||
K8sTokenFile string
|
||||
K8sCAFile string
|
||||
AWSRoleARN string
|
||||
AWSRegion string
|
||||
GCPProject string
|
||||
EmbedModel string
|
||||
MaxOutput int // bytes per tool result handed to the model
|
||||
}
|
||||
|
||||
type Runner struct {
|
||||
cfg Config
|
||||
http *http.Client
|
||||
rag *rag.Store
|
||||
oai *openai.Client
|
||||
}
|
||||
|
||||
func New(cfg Config, store *rag.Store, oai *openai.Client) *Runner {
|
||||
if cfg.MaxOutput == 0 {
|
||||
cfg.MaxOutput = 6000
|
||||
}
|
||||
return &Runner{cfg: cfg, http: &http.Client{Timeout: 30 * time.Second}, rag: store, oai: oai}
|
||||
}
|
||||
|
||||
// Definitions are the function tools exposed to the model (strict schemas, every property required).
|
||||
func Definitions() []openai.Tool {
|
||||
obj := func(props map[string]any, req ...string) map[string]any {
|
||||
if req == nil {
|
||||
req = []string{}
|
||||
}
|
||||
return map[string]any{"type": "object", "properties": props, "required": req, "additionalProperties": false}
|
||||
}
|
||||
return []openai.Tool{
|
||||
{Type: "function", Strict: true, Name: "query_metrics", Description: "Run a PromQL range query against Mimir (Prometheus API) over the last range_minutes. Metrics: probe_success{target=alb|db-primary}, demo_http_requests_total, demo_http_request_duration_seconds_bucket, up{job=legacy-api}, kube_pod_container_status_restarts_total, pg_replication_lag_seconds, pg_is_wal_replay_paused, argocd_app_info, harness_*.",
|
||||
Parameters: obj(map[string]any{"promql": map[string]any{"type": "string"}, "range_minutes": map[string]any{"type": "integer", "description": "1 to 120"}}, "promql", "range_minutes")},
|
||||
{Type: "function", Strict: true, Name: "query_logs", Description: "Run a LogQL query against Loki over the last range_minutes, newest first. Streams: {job=\"demo-api\"}, {job=\"legacy-api\"}, {namespace=\"demo\"}.",
|
||||
Parameters: obj(map[string]any{"logql": map[string]any{"type": "string"}, "range_minutes": map[string]any{"type": "integer"}, "limit": map[string]any{"type": "integer", "description": "1 to 100"}}, "logql", "range_minutes", "limit")},
|
||||
{Type: "function", Strict: true, Name: "get_probe_status", Description: "Latest synthetic probe results (probe_success per target: alb = demo-app through the ALB, db-primary = PostgreSQL TCP).",
|
||||
Parameters: obj(map[string]any{})},
|
||||
{Type: "function", Strict: true, Name: "get_k8s", Description: "List Kubernetes objects in a namespace (read-only). kind: pods | events | deployments | replicasets | nodes. namespace is ignored for nodes.",
|
||||
Parameters: obj(map[string]any{"kind": map[string]any{"type": "string", "enum": []string{"pods", "events", "deployments", "replicasets", "nodes"}}, "namespace": map[string]any{"type": "string"}}, "kind", "namespace")},
|
||||
{Type: "function", Strict: true, Name: "get_gcp_instances", Description: "GCE instances (selfheal-* pet VMs and legacy-* MIG instances) with status, zone, internal IP, plus the legacy MIG target size and instance health.",
|
||||
Parameters: obj(map[string]any{})},
|
||||
{Type: "function", Strict: true, Name: "get_aws_state", Description: "AWS state: EC2 instances tagged selfheal-* (state, private IP) and ALB target group health per worker.",
|
||||
Parameters: obj(map[string]any{})},
|
||||
{Type: "function", Strict: true, Name: "search_runbooks", Description: "Semantic search over the runbooks (one per catalogue action), the topology document and past verdict summaries. Returns the k closest chunks.",
|
||||
Parameters: obj(map[string]any{"query": map[string]any{"type": "string"}, "k": map[string]any{"type": "integer", "description": "1 to 8"}}, "query", "k")},
|
||||
}
|
||||
}
|
||||
|
||||
// Call executes one tool by name with JSON arguments. Errors are returned as JSON so the model can react.
|
||||
func (r *Runner) Call(ctx context.Context, name, args string) string {
|
||||
var a map[string]any
|
||||
if err := json.Unmarshal([]byte(args), &a); err != nil {
|
||||
return errJSON(fmt.Errorf("arguments: %w", err))
|
||||
}
|
||||
var (
|
||||
out any
|
||||
err error
|
||||
)
|
||||
switch name {
|
||||
case "query_metrics":
|
||||
out, err = r.queryMetrics(ctx, str(a, "promql"), num(a, "range_minutes"))
|
||||
case "query_logs":
|
||||
out, err = r.queryLogs(ctx, str(a, "logql"), num(a, "range_minutes"), num(a, "limit"))
|
||||
case "get_probe_status":
|
||||
out, err = r.probeStatus(ctx)
|
||||
case "get_k8s":
|
||||
out, err = r.getK8s(ctx, str(a, "kind"), str(a, "namespace"))
|
||||
case "get_gcp_instances":
|
||||
out, err = r.gcpInstances(ctx)
|
||||
case "get_aws_state":
|
||||
out, err = r.awsState(ctx)
|
||||
case "search_runbooks":
|
||||
out, err = r.SearchRunbooks(ctx, str(a, "query"), int(num(a, "k")))
|
||||
default:
|
||||
err = fmt.Errorf("unknown tool %s", name)
|
||||
}
|
||||
if err != nil {
|
||||
return errJSON(err)
|
||||
}
|
||||
b, _ := json.Marshal(out)
|
||||
if len(b) > r.cfg.MaxOutput {
|
||||
return string(b[:r.cfg.MaxOutput]) + `..."truncated":true}`
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func errJSON(err error) string {
|
||||
b, _ := json.Marshal(map[string]string{"error": err.Error()})
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// ---- Mimir / Loki ----
|
||||
|
||||
type promResp struct {
|
||||
Status string `json:"status"`
|
||||
Data struct {
|
||||
ResultType string `json:"resultType"`
|
||||
Result []struct {
|
||||
Metric map[string]string `json:"metric"`
|
||||
Values [][]any `json:"values"`
|
||||
Value []any `json:"value"`
|
||||
} `json:"result"`
|
||||
} `json:"data"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
|
||||
func (r *Runner) getJSON(ctx context.Context, u string, out any) error {
|
||||
req, _ := http.NewRequestWithContext(ctx, "GET", u, nil)
|
||||
resp, err := r.http.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
|
||||
if resp.StatusCode != 200 {
|
||||
return fmt.Errorf("status %d: %s", resp.StatusCode, truncate(string(b), 300))
|
||||
}
|
||||
return json.Unmarshal(b, out)
|
||||
}
|
||||
|
||||
func clampMinutes(m int64) int64 {
|
||||
if m <= 0 {
|
||||
return 10
|
||||
}
|
||||
if m > 120 {
|
||||
return 120
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func (r *Runner) queryMetrics(ctx context.Context, promql string, minutes int64) (any, error) {
|
||||
minutes = clampMinutes(minutes)
|
||||
end := time.Now()
|
||||
start := end.Add(-time.Duration(minutes) * time.Minute)
|
||||
step := minutes * 60 / 20
|
||||
if step < 15 {
|
||||
step = 15
|
||||
}
|
||||
q := url.Values{"query": {promql}, "start": {strconv.FormatInt(start.Unix(), 10)}, "end": {strconv.FormatInt(end.Unix(), 10)}, "step": {strconv.FormatInt(step, 10)}}
|
||||
var pr promResp
|
||||
if err := r.getJSON(ctx, strings.TrimRight(r.cfg.MimirURL, "/")+"/api/v1/query_range?"+q.Encode(), &pr); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if pr.Status != "success" {
|
||||
return nil, fmt.Errorf("mimir: %s", pr.Error)
|
||||
}
|
||||
type series struct {
|
||||
Metric map[string]string `json:"metric"`
|
||||
Points []string `json:"points"` // "t=HH:MM:SS v=..."
|
||||
}
|
||||
var out []series
|
||||
for i, s := range pr.Data.Result {
|
||||
if i >= 20 {
|
||||
break
|
||||
}
|
||||
var pts []string
|
||||
vals := s.Values
|
||||
if len(vals) > 12 { // keep the last 12 points
|
||||
vals = vals[len(vals)-12:]
|
||||
}
|
||||
for _, v := range vals {
|
||||
if len(v) == 2 {
|
||||
ts, _ := v[0].(float64)
|
||||
pts = append(pts, fmt.Sprintf("t=%s v=%v", time.Unix(int64(ts), 0).UTC().Format("15:04:05"), v[1]))
|
||||
}
|
||||
}
|
||||
out = append(out, series{Metric: s.Metric, Points: pts})
|
||||
}
|
||||
return map[string]any{"series_count": len(pr.Data.Result), "series": out, "range_minutes": minutes}, nil
|
||||
}
|
||||
|
||||
func (r *Runner) probeStatus(ctx context.Context) (any, error) {
|
||||
q := url.Values{"query": {"probe_success"}}
|
||||
var pr promResp
|
||||
if err := r.getJSON(ctx, strings.TrimRight(r.cfg.MimirURL, "/")+"/api/v1/query?"+q.Encode(), &pr); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := map[string]any{}
|
||||
for _, s := range pr.Data.Result {
|
||||
if len(s.Value) == 2 {
|
||||
out[s.Metric["target"]] = s.Value[1]
|
||||
}
|
||||
}
|
||||
return map[string]any{"probe_success": out, "checked_at": time.Now().UTC().Format(time.RFC3339)}, nil
|
||||
}
|
||||
|
||||
type lokiResp struct {
|
||||
Status string `json:"status"`
|
||||
Data struct {
|
||||
Result []struct {
|
||||
Stream map[string]string `json:"stream"`
|
||||
Values [][]string `json:"values"`
|
||||
} `json:"result"`
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
func (r *Runner) queryLogs(ctx context.Context, logql string, minutes, limit int64) (any, error) {
|
||||
minutes = clampMinutes(minutes)
|
||||
if limit <= 0 || limit > 100 {
|
||||
limit = 30
|
||||
}
|
||||
end := time.Now()
|
||||
start := end.Add(-time.Duration(minutes) * time.Minute)
|
||||
q := url.Values{"query": {logql}, "start": {strconv.FormatInt(start.UnixNano(), 10)}, "end": {strconv.FormatInt(end.UnixNano(), 10)}, "limit": {strconv.FormatInt(limit, 10)}, "direction": {"backward"}}
|
||||
var lr lokiResp
|
||||
if err := r.getJSON(ctx, strings.TrimRight(r.cfg.LokiURL, "/")+"/loki/api/v1/query_range?"+q.Encode(), &lr); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var lines []string
|
||||
for _, s := range lr.Data.Result {
|
||||
for _, v := range s.Values {
|
||||
if len(v) == 2 {
|
||||
ns, _ := strconv.ParseInt(v[0], 10, 64)
|
||||
lines = append(lines, fmt.Sprintf("%s %s %s", time.Unix(0, ns).UTC().Format("15:04:05"), labelsOf(s.Stream), truncate(v[1], 300)))
|
||||
}
|
||||
}
|
||||
}
|
||||
sort.Sort(sort.Reverse(sort.StringSlice(lines)))
|
||||
if int64(len(lines)) > limit {
|
||||
lines = lines[:limit]
|
||||
}
|
||||
return map[string]any{"lines": lines, "count": len(lines)}, nil
|
||||
}
|
||||
|
||||
func labelsOf(m map[string]string) string {
|
||||
keys := make([]string, 0, len(m))
|
||||
for k := range m {
|
||||
if k == "job" || k == "pod" || k == "instance" || k == "container" || k == "level" {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
}
|
||||
sort.Strings(keys)
|
||||
parts := make([]string, 0, len(keys))
|
||||
for _, k := range keys {
|
||||
parts = append(parts, k+"="+m[k])
|
||||
}
|
||||
return "{" + strings.Join(parts, ",") + "}"
|
||||
}
|
||||
|
||||
// ---- Kubernetes (SA ai-reader, ClusterRole view) ----
|
||||
|
||||
func (r *Runner) k8s() (kubernetes.Interface, error) {
|
||||
tok, err := os.ReadFile(r.cfg.K8sTokenFile)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("k8s token: %w", err)
|
||||
}
|
||||
ca, err := os.ReadFile(r.cfg.K8sCAFile)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("k8s ca: %w", err)
|
||||
}
|
||||
rc := &rest.Config{Host: r.cfg.K8sAPI, BearerToken: strings.TrimSpace(string(tok)), TLSClientConfig: rest.TLSClientConfig{CAData: ca}, Timeout: 20 * time.Second}
|
||||
return kubernetes.NewForConfig(rc)
|
||||
}
|
||||
|
||||
func (r *Runner) getK8s(ctx context.Context, kind, ns string) (any, error) {
|
||||
cs, err := r.k8s()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if ns == "" {
|
||||
ns = "demo"
|
||||
}
|
||||
switch kind {
|
||||
case "pods":
|
||||
l, err := cs.CoreV1().Pods(ns).List(ctx, metav1.ListOptions{})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var out []map[string]any
|
||||
for _, p := range l.Items {
|
||||
restarts := int32(0)
|
||||
ready := 0
|
||||
var states []string
|
||||
for _, c := range p.Status.ContainerStatuses {
|
||||
restarts += c.RestartCount
|
||||
if c.Ready {
|
||||
ready++
|
||||
}
|
||||
switch {
|
||||
case c.State.Waiting != nil:
|
||||
states = append(states, c.Name+":"+c.State.Waiting.Reason)
|
||||
case c.State.Terminated != nil:
|
||||
states = append(states, c.Name+":"+c.State.Terminated.Reason)
|
||||
}
|
||||
}
|
||||
out = append(out, map[string]any{"name": p.Name, "node": p.Spec.NodeName, "phase": string(p.Status.Phase), "ready": fmt.Sprintf("%d/%d", ready, len(p.Spec.Containers)), "restarts": restarts, "states": states, "age_s": int(time.Since(p.CreationTimestamp.Time).Seconds())})
|
||||
}
|
||||
return out, nil
|
||||
case "events":
|
||||
l, err := cs.CoreV1().Events(ns).List(ctx, metav1.ListOptions{})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items := l.Items
|
||||
sort.Slice(items, func(i, j int) bool { return items[i].LastTimestamp.After(items[j].LastTimestamp.Time) })
|
||||
var out []string
|
||||
for i, e := range items {
|
||||
if i >= 30 {
|
||||
break
|
||||
}
|
||||
out = append(out, fmt.Sprintf("%s %s %s/%s %s: %s (x%d)", e.LastTimestamp.UTC().Format("15:04:05"), e.Type, e.InvolvedObject.Kind, e.InvolvedObject.Name, e.Reason, truncate(e.Message, 160), e.Count))
|
||||
}
|
||||
return out, nil
|
||||
case "deployments":
|
||||
l, err := cs.AppsV1().Deployments(ns).List(ctx, metav1.ListOptions{})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var out []map[string]any
|
||||
for _, d := range l.Items {
|
||||
img := ""
|
||||
if len(d.Spec.Template.Spec.Containers) > 0 {
|
||||
img = d.Spec.Template.Spec.Containers[0].Image
|
||||
}
|
||||
out = append(out, map[string]any{"name": d.Name, "desired": *d.Spec.Replicas, "ready": d.Status.ReadyReplicas, "updated": d.Status.UpdatedReplicas, "available": d.Status.AvailableReplicas, "image": img, "generation": d.Generation, "observed": d.Status.ObservedGeneration})
|
||||
}
|
||||
return out, nil
|
||||
case "replicasets":
|
||||
l, err := cs.AppsV1().ReplicaSets(ns).List(ctx, metav1.ListOptions{})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var out []map[string]any
|
||||
for _, rs := range l.Items {
|
||||
out = append(out, map[string]any{"name": rs.Name, "desired": *rs.Spec.Replicas, "ready": rs.Status.ReadyReplicas, "revision": rs.Annotations["deployment.kubernetes.io/revision"]})
|
||||
}
|
||||
return out, nil
|
||||
case "nodes":
|
||||
l, err := cs.CoreV1().Nodes().List(ctx, metav1.ListOptions{})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var out []map[string]any
|
||||
for _, n := range l.Items {
|
||||
ready := "Unknown"
|
||||
for _, c := range n.Status.Conditions {
|
||||
if c.Type == "Ready" {
|
||||
ready = string(c.Status)
|
||||
}
|
||||
}
|
||||
out = append(out, map[string]any{"name": n.Name, "ready": ready, "unschedulable": n.Spec.Unschedulable, "taints": len(n.Spec.Taints)})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
return nil, fmt.Errorf("kind %q not allowed", kind)
|
||||
}
|
||||
|
||||
// ---- GCP (VM SA, roles/compute.viewer) ----
|
||||
|
||||
func (r *Runner) gcpInstances(ctx context.Context) (any, error) {
|
||||
svc, err := compute.NewService(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var vms []map[string]any
|
||||
err = svc.Instances.AggregatedList(r.cfg.GCPProject).Filter("name eq ^(selfheal-|legacy).*").Pages(ctx, func(l *compute.InstanceAggregatedList) error {
|
||||
for _, s := range l.Items {
|
||||
for _, in := range s.Instances {
|
||||
ip := ""
|
||||
if len(in.NetworkInterfaces) > 0 {
|
||||
ip = in.NetworkInterfaces[0].NetworkIP
|
||||
}
|
||||
vms = append(vms, map[string]any{"name": in.Name, "zone": in.Zone[strings.LastIndex(in.Zone, "/")+1:], "status": in.Status, "ip": ip})
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
sort.Slice(vms, func(i, j int) bool { return vms[i]["name"].(string) < vms[j]["name"].(string) })
|
||||
var migs []map[string]any
|
||||
_ = svc.RegionInstanceGroupManagers.List(r.cfg.GCPProject, "asia-northeast1").Pages(ctx, func(l *compute.RegionInstanceGroupManagerList) error {
|
||||
for _, m := range l.Items {
|
||||
entry := map[string]any{"name": m.Name, "target_size": m.TargetSize}
|
||||
if h, err := svc.RegionInstanceGroupManagers.ListManagedInstances(r.cfg.GCPProject, "asia-northeast1", m.Name).Context(ctx).Do(); err == nil {
|
||||
var inst []string
|
||||
for _, mi := range h.ManagedInstances {
|
||||
hs := "UNKNOWN"
|
||||
if len(mi.InstanceHealth) > 0 {
|
||||
hs = mi.InstanceHealth[0].DetailedHealthState
|
||||
}
|
||||
inst = append(inst, fmt.Sprintf("%s status=%s health=%s action=%s", mi.Instance[strings.LastIndex(mi.Instance, "/")+1:], mi.InstanceStatus, hs, mi.CurrentAction))
|
||||
}
|
||||
entry["instances"] = inst
|
||||
}
|
||||
migs = append(migs, entry)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
return map[string]any{"instances": vms, "migs": migs}, nil
|
||||
}
|
||||
|
||||
// ---- AWS (GCE ID token -> AssumeRoleWithWebIdentity, role selfheal-ai-readonly) ----
|
||||
|
||||
func gceIDToken(ctx context.Context, audience string) (string, error) {
|
||||
u := "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity?format=full&audience=" + url.QueryEscape(audience)
|
||||
req, _ := http.NewRequestWithContext(ctx, "GET", u, nil)
|
||||
req.Header.Set("Metadata-Flavor", "Google")
|
||||
resp, err := (&http.Client{Timeout: 10 * time.Second}).Do(req)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != 200 {
|
||||
return "", fmt.Errorf("metadata identity: %s", resp.Status)
|
||||
}
|
||||
return string(b), nil
|
||||
}
|
||||
|
||||
func (r *Runner) awsCfg(ctx context.Context) (aws.Config, error) {
|
||||
if r.cfg.AWSRoleARN == "" {
|
||||
return aws.Config{}, fmt.Errorf("AWS_ROLE_ARN not set")
|
||||
}
|
||||
tok, err := gceIDToken(ctx, r.cfg.AWSRoleARN)
|
||||
if err != nil {
|
||||
return aws.Config{}, err
|
||||
}
|
||||
anon := aws.Config{Region: r.cfg.AWSRegion, Credentials: aws.AnonymousCredentials{}}
|
||||
out, err := sts.NewFromConfig(anon).AssumeRoleWithWebIdentity(ctx, &sts.AssumeRoleWithWebIdentityInput{
|
||||
RoleArn: aws.String(r.cfg.AWSRoleARN), RoleSessionName: aws.String("ai-agent"), WebIdentityToken: aws.String(tok), DurationSeconds: aws.Int32(900),
|
||||
})
|
||||
if err != nil {
|
||||
return aws.Config{}, fmt.Errorf("assume-role-with-web-identity: %w", err)
|
||||
}
|
||||
c := out.Credentials
|
||||
return aws.Config{Region: r.cfg.AWSRegion, Credentials: credentials.NewStaticCredentialsProvider(*c.AccessKeyId, *c.SecretAccessKey, *c.SessionToken)}, nil
|
||||
}
|
||||
|
||||
// AWSCallerIdentity is used by the self-check.
|
||||
func (r *Runner) AWSCallerIdentity(ctx context.Context) (string, error) {
|
||||
cfg, err := r.awsCfg(ctx)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
id, err := sts.NewFromConfig(cfg).GetCallerIdentity(ctx, &sts.GetCallerIdentityInput{})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return aws.ToString(id.Arn), nil
|
||||
}
|
||||
|
||||
func (r *Runner) awsState(ctx context.Context) (any, error) {
|
||||
cfg, err := r.awsCfg(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ec := ec2.NewFromConfig(cfg)
|
||||
inst, err := ec.DescribeInstances(ctx, &ec2.DescribeInstancesInput{Filters: []ec2types.Filter{{Name: aws.String("tag:Name"), Values: []string{"selfheal-*"}}}})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var nodes []map[string]any
|
||||
for _, res := range inst.Reservations {
|
||||
for _, i := range res.Instances {
|
||||
name := ""
|
||||
for _, t := range i.Tags {
|
||||
if aws.ToString(t.Key) == "Name" {
|
||||
name = aws.ToString(t.Value)
|
||||
}
|
||||
}
|
||||
nodes = append(nodes, map[string]any{"name": name, "id": aws.ToString(i.InstanceId), "state": string(i.State.Name), "private_ip": aws.ToString(i.PrivateIpAddress)})
|
||||
}
|
||||
}
|
||||
sort.Slice(nodes, func(a, b int) bool { return nodes[a]["name"].(string) < nodes[b]["name"].(string) })
|
||||
lb := elbv2.NewFromConfig(cfg)
|
||||
tgs, err := lb.DescribeTargetGroups(ctx, &elbv2.DescribeTargetGroupsInput{Names: []string{"selfheal-tg"}})
|
||||
var targets []map[string]any
|
||||
if err == nil && len(tgs.TargetGroups) > 0 {
|
||||
th, err := lb.DescribeTargetHealth(ctx, &elbv2.DescribeTargetHealthInput{TargetGroupArn: tgs.TargetGroups[0].TargetGroupArn})
|
||||
if err == nil {
|
||||
for _, t := range th.TargetHealthDescriptions {
|
||||
targets = append(targets, map[string]any{"target": aws.ToString(t.Target.Id), "port": aws.ToInt32(t.Target.Port), "state": string(t.TargetHealth.State), "reason": string(t.TargetHealth.Reason)})
|
||||
}
|
||||
}
|
||||
}
|
||||
return map[string]any{"ec2": nodes, "alb_targets": targets}, nil
|
||||
}
|
||||
|
||||
// ---- RAG ----
|
||||
|
||||
func (r *Runner) SearchRunbooks(ctx context.Context, query string, k int) (any, error) {
|
||||
if r.rag == nil || r.oai == nil {
|
||||
return nil, fmt.Errorf("rag store not available")
|
||||
}
|
||||
if k <= 0 || k > 8 {
|
||||
k = 4
|
||||
}
|
||||
vecs, err := r.oai.Embed(ctx, r.cfg.EmbedModel, []string{query})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
hits, err := r.rag.Search(ctx, vecs[0], k)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return hits, nil
|
||||
}
|
||||
|
||||
func str(m map[string]any, k string) string {
|
||||
v, _ := m[k].(string)
|
||||
return v
|
||||
}
|
||||
|
||||
func num(m map[string]any, k string) int64 {
|
||||
switch v := m[k].(type) {
|
||||
case float64:
|
||||
return int64(v)
|
||||
case int64:
|
||||
return v
|
||||
case int:
|
||||
return int64(v)
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func truncate(s string, n int) string {
|
||||
if len(s) <= n {
|
||||
return s
|
||||
}
|
||||
return s[:n] + "..."
|
||||
}
|
||||
Reference in new issue
Block a user