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