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,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