oikos/internal/knowledge/seed.go

package knowledge

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
)

type SeedResult struct {
	Documents      int
	Investigations int
	Runbooks       int
}

func contentHash(s string) string {
	h := sha256.Sum256([]byte(s))
	return hex.EncodeToString(h[:])
}

func Ingest(ctx context.Context, tx pgx.Tx, data map[string]any) (*SeedResult, error) {
	r := &SeedResult{}

	docs, _ := data["documents"].([]any)
	for _, raw := range docs {
		d, _ := raw.(map[string]any)
		if err := ingestDocument(ctx, tx, d); err != nil {
			return nil, fmt.Errorf("document %v: %w", str(d, "slug"), err)
		}
		r.Documents++
	}

	invs, _ := data["investigations"].([]any)
	for _, raw := range invs {
		m, _ := raw.(map[string]any)
		if err := ingestInvestigation(ctx, tx, m); err != nil {
			return nil, fmt.Errorf("investigation %v: %w", str(m, "slug"), err)
		}
		r.Investigations++
	}

	rbs, _ := data["runbooks"].([]any)
	for _, raw := range rbs {
		m, _ := raw.(map[string]any)
		if err := ingestRunbook(ctx, tx, m); err != nil {
			return nil, fmt.Errorf("runbook %v: %w", str(m, "slug"), err)
		}
		r.Runbooks++
	}

	return r, nil
}

func str(m map[string]any, key string) string {
	s, _ := m[key].(string)
	return s
}

func strSlice(m map[string]any, key string) []string {
	raw, _ := m[key].([]any)
	var out []string
	for _, v := range raw {
		if s, ok := v.(string); ok {
			out = append(out, s)
		}
	}
	return out
}

func mapVal(m map[string]any, key string) map[string]any {
	v, _ := m[key].(map[string]any)
	return v
}

func ingestDocument(ctx context.Context, tx pgx.Tx, m map[string]any) error {
	slug := str(m, "slug")
	title := str(m, "title")
	content := str(m, "content")
	entitySlug := str(m, "entity_slug")
	tags := strSlice(m, "tags")
	atGlance, _ := m["at_glance"].(map[string]any)
	clRaw, _ := m["changelog"].([]any)

	entityDocSlug := "document:" + slug

	if err := upsertKnowledgeEntity(ctx, tx, entityDocSlug, "document", title, content, slug, tags); err != nil {
		return err
	}

	attrs := map[string]any{}
	if len(atGlance) > 0 {
		attrs["at_glance"] = atGlance
	}
	if len(clRaw) > 0 {
		attrs["changelog"] = clRaw
	}
	if len(attrs) > 0 {
		attrsBytes, _ := json.Marshal(attrs)
		_, err := tx.Exec(ctx,
			`UPDATE entities SET attributes = attributes || $1, updated_at = now()
			 WHERE slug = $2`, string(attrsBytes), entityDocSlug)
		if err != nil {
			return fmt.Errorf("update document attrs: %w", err)
		}
	}

	if entitySlug != "" {
		if err := createEdge(ctx, tx, entityDocSlug, entitySlug, "documents", nil); err != nil {
			return fmt.Errorf("link document: %w", err)
		}
	}

	if len(atGlance) > 0 && entitySlug != "" {
		backfillAttrs, _ := json.Marshal(atGlance)
		_, err := tx.Exec(ctx,
			`UPDATE entities SET attributes = $1 || attributes, updated_at = now()
			 WHERE slug = $2`, string(backfillAttrs), entitySlug)
		if err != nil {
			return fmt.Errorf("backfill entity attrs: %w", err)
		}
	}

	return nil
}

func ingestInvestigation(ctx context.Context, tx pgx.Tx, m map[string]any) error {
	slug := str(m, "slug")
	title := str(m, "title")
	content := str(m, "content")
	date := str(m, "date")
	status := str(m, "status")
	duration := str(m, "duration")
	aboutSlugs := strSlice(m, "about_slugs")
	tags := strSlice(m, "tags")

	entitySlug := "investigation:" + slug

	if err := upsertKnowledgeEntity(ctx, tx, entitySlug, "investigation", title, content, slug, tags); err != nil {
		return err
	}

	attrs := map[string]any{}
	if date != "" {
		attrs["date"] = date
	}
	if status != "" {
		attrs["status"] = status
	}
	if duration != "" {
		attrs["duration"] = duration
	}
	if len(attrs) > 0 {
		attrsBytes, _ := json.Marshal(attrs)
		_, err := tx.Exec(ctx,
			`UPDATE entities SET attributes = attributes || $1, updated_at = now()
			 WHERE slug = $2`, string(attrsBytes), entitySlug)
		if err != nil {
			return fmt.Errorf("update investigation attrs: %w", err)
		}
	}

	for _, aboutSlug := range aboutSlugs {
		if err := createEdge(ctx, tx, entitySlug, aboutSlug, "about", nil); err != nil {
			return fmt.Errorf("link investigation about %s: %w", aboutSlug, err)
		}
	}

	return nil
}

func ingestRunbook(ctx context.Context, tx pgx.Tx, m map[string]any) error {
	slug := str(m, "slug")
	name := str(m, "name")
	riskClass := str(m, "risk_class")
	entityType := str(m, "entity_type")
	content := str(m, "content")
	tags := strSlice(m, "tags")
	procedure, _ := m["procedure"].(map[string]any)

	entitySlug := "runbook:" + slug

	if err := upsertKnowledgeEntity(ctx, tx, entitySlug, "runbook", name, content, slug, tags); err != nil {
		return err
	}

	attrs := map[string]any{}
	if riskClass != "" {
		attrs["risk_class"] = riskClass
	}
	if len(procedure) > 0 {
		attrs["procedure"] = procedure
	}
	if len(attrs) > 0 {
		attrsBytes, _ := json.Marshal(attrs)
		_, err := tx.Exec(ctx,
			`UPDATE entities SET attributes = attributes || $1, updated_at = now()
			 WHERE slug = $2`, string(attrsBytes), entitySlug)
		if err != nil {
			return fmt.Errorf("update runbook attrs: %w", err)
		}
	}

	if entityType != "" {
		if err := createEdge(ctx, tx, entitySlug, entityType, "procedure-for", nil); err != nil {
			return fmt.Errorf("link runbook: %w", err)
		}
	}

	return nil
}

func upsertKnowledgeEntity(ctx context.Context, tx pgx.Tx, slug, entityType, title, content, source string, tags []string) error {
	id, err := getOrCreateEntity(ctx, tx, slug, entityType, title)
	if err != nil {
		return err
	}

	hash := contentHash(content)
	tagsJSON, _ := json.Marshal(tags)

	_, err = tx.Exec(ctx,
		`INSERT INTO knowledge_entities (entity_id, title, content, source, tags, content_hash, created_at, updated_at)
		 VALUES ($1, $2, $3, $4, $5, $6, now(), now())
		 ON CONFLICT (entity_id) DO UPDATE SET
		   title = $2, content = $3, source = $4, tags = $5,
		   content_hash = $6, updated_at = now()`,
		id, title, content, source, string(tagsJSON), hash)
	return err
}

func getOrCreateEntity(ctx context.Context, tx pgx.Tx, slug, entityType, name string) (uuid.UUID, error) {
	var id uuid.UUID
	err := tx.QueryRow(ctx, "SELECT id FROM entities WHERE slug = $1", slug).Scan(&id)
	if err == nil {
		return id, nil
	}
	if err != pgx.ErrNoRows {
		return uuid.Nil, fmt.Errorf("lookup entity %s: %w", slug, err)
	}

	id, err = uuid.NewV7()
	if err != nil {
		return uuid.Nil, fmt.Errorf("generate uuid: %w", err)
	}
	_, err = tx.Exec(ctx,
		`INSERT INTO entities (id, slug, type, name, state, attributes, version, created_at, updated_at)
		 VALUES ($1, $2, $3, $4, NULL, '{}', 1, now(), now())`,
		id, slug, entityType, name)
	if err != nil {
		return uuid.Nil, fmt.Errorf("create entity %s: %w", slug, err)
	}
	return id, nil
}

func createEdge(ctx context.Context, tx pgx.Tx, sourceSlug, targetSlug, relType string, attrs map[string]any) error {
	var sourceID, targetID uuid.UUID
	if err := tx.QueryRow(ctx, "SELECT id FROM entities WHERE slug = $1", sourceSlug).Scan(&sourceID); err != nil {
		return fmt.Errorf("source %s: %w", sourceSlug, err)
	}
	if err := tx.QueryRow(ctx, "SELECT id FROM entities WHERE slug = $1", targetSlug).Scan(&targetID); err != nil {
		return fmt.Errorf("target %s: %w", targetSlug, err)
	}

	attrsBytes, _ := json.Marshal(attrs)
	_, err := tx.Exec(ctx,
		`INSERT INTO relationships (source_id, target_id, type, attributes, valid_from, valid_to)
		 VALUES ($1, $2, $3, $4, now(), NULL)
		 ON CONFLICT (source_id, target_id, type) WHERE valid_to IS NULL
		 DO UPDATE SET attributes = EXCLUDED.attributes`,
		sourceID, targetID, relType, string(attrsBytes))
	return err
}