Agent-System/cmd/dataset/main.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"

	"chat-agent-system/internal/dataset"

	"github.com/jackc/pgx/v5/pgxpool"
)

func main() {
	if len(os.Args) < 2 {
		fmt.Fprintln(os.Stderr, "usage: dataset <tracks|registry-import|backfill|destroy-preflight> ...")
		os.Exit(2)
	}
	switch os.Args[1] {
	case "tracks":
		root, _ := os.Getwd()
		path := filepath.Join(root, "internal/dataset/testdata/registry.json")
		if len(os.Args) > 2 {
			path = os.Args[2]
		}
		rs := dataset.RunEightTracks(context.Background(), path)
		enc := json.NewEncoder(os.Stdout)
		enc.SetIndent("", "  ")
		_ = enc.Encode(rs)
		if err := dataset.AllTracksPass(rs); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
	case "destroy-preflight":
		raw, err := os.ReadFile(os.Args[2])
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		var req dataset.DestroyRequest
		if err := json.Unmarshal(raw, &req); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		if err := dataset.EvaluateDestroy(req); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		fmt.Println("destroy-preflight PASS")
	case "registry-import":
		if len(os.Args) < 4 {
			fmt.Fprintln(os.Stderr, "usage: dataset registry-import <dsn> <registry.json>")
			os.Exit(2)
		}
		dsn := os.Args[2]
		if err := dataset.RefuseOperationalDSN(dsn, dsn); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		reg, err := dataset.LoadRegistryFile(os.Args[3])
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		ctx := context.Background()
		pool, err := pgxpool.New(ctx, dsn)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		defer pool.Close()
		pg := &dataset.PG{Pool: pool}
		for _, d := range reg.ListApproved() {
			if err := pg.UpsertDataset(ctx, d); err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(1)
			}
		}
		fmt.Printf("imported %d approved datasets\n", len(reg.ListApproved()))
	default:
		fmt.Fprintln(os.Stderr, "unknown command")
		os.Exit(2)
	}
}