trajectory.go

  1package cmd
  2
  3import (
  4	"encoding/json"
  5	"fmt"
  6	"os"
  7
  8	"github.com/charmbracelet/crush/internal/config"
  9	"github.com/charmbracelet/crush/internal/db"
 10	"github.com/charmbracelet/crush/internal/message"
 11	"github.com/charmbracelet/crush/internal/session"
 12	"github.com/charmbracelet/crush/internal/trajectory"
 13	"github.com/charmbracelet/crush/internal/version"
 14	"github.com/spf13/cobra"
 15)
 16
 17var trajectoryCmd = &cobra.Command{
 18	Use:   "trajectory",
 19	Short: "Trajectory export utilities",
 20	Long:  "Export session trajectories in Harbor ATIF format for analysis and sharing",
 21}
 22
 23var trajectoryExportCmd = &cobra.Command{
 24	Use:   "export",
 25	Short: "Export a session as ATIF trajectory",
 26	Long:  "Export a Crush session in Harbor ATIF (Agent Trajectory Interchange Format) v1.4",
 27	Example: `
 28# Export a session as JSON to stdout
 29crush trajectory export --session <session-id>
 30
 31# Export a session to a JSON file
 32crush trajectory export --session <session-id> --output trajectory.json
 33
 34# Export as HTML for visualization
 35crush trajectory export --session <session-id> --format html --output trajectory.html
 36
 37# Validate with Harbor validator
 38crush trajectory export --session <session-id> > out.json
 39python -m harbor.utils.trajectory_validator out.json
 40  `,
 41	RunE: func(cmd *cobra.Command, args []string) error {
 42		sessionID, _ := cmd.Flags().GetString("session")
 43		outputFile, _ := cmd.Flags().GetString("output")
 44		format, _ := cmd.Flags().GetString("format")
 45		dataDir, _ := cmd.Flags().GetString("data-dir")
 46
 47		if sessionID == "" {
 48			return fmt.Errorf("--session flag is required")
 49		}
 50
 51		ctx := cmd.Context()
 52
 53		cwd, err := ResolveCwd(cmd)
 54		if err != nil {
 55			return err
 56		}
 57
 58		// Load config (lightweight, no full app init).
 59		cfg, err := config.Load(cwd, dataDir, false)
 60		if err != nil {
 61			return fmt.Errorf("failed to load config: %w", err)
 62		}
 63
 64		// Connect to DB.
 65		conn, err := db.Connect(ctx, cfg.Options.DataDirectory)
 66		if err != nil {
 67			return fmt.Errorf("failed to connect to database: %w", err)
 68		}
 69		defer conn.Close()
 70
 71		querier := db.New(conn)
 72		sessionSvc := session.NewService(querier)
 73		messageSvc := message.NewService(querier)
 74
 75		// Load session.
 76		sess, err := sessionSvc.Get(ctx, sessionID)
 77		if err != nil {
 78			return fmt.Errorf("failed to get session: %w", err)
 79		}
 80
 81		// Load messages.
 82		messages, err := messageSvc.List(ctx, sessionID)
 83		if err != nil {
 84			return fmt.Errorf("failed to list messages: %w", err)
 85		}
 86
 87		// Determine model name from first assistant message.
 88		var modelName string
 89		for _, msg := range messages {
 90			if msg.Role == message.Assistant && msg.Model != "" {
 91				modelName = msg.Model
 92				break
 93			}
 94		}
 95
 96		// Export to ATIF.
 97		traj, err := trajectory.ExportSession(sess, messages, "Crush", version.Version, modelName)
 98		if err != nil {
 99			return fmt.Errorf("failed to export trajectory: %w", err)
100		}
101
102		var data []byte
103		switch format {
104		case "html":
105			data, err = trajectory.RenderHTML(traj)
106			if err != nil {
107				return fmt.Errorf("failed to render HTML: %w", err)
108			}
109		case "json":
110			data, err = json.MarshalIndent(traj, "", "  ")
111			if err != nil {
112				return fmt.Errorf("failed to marshal trajectory: %w", err)
113			}
114		default:
115			return fmt.Errorf("unknown format: %s (use 'json' or 'html')", format)
116		}
117
118		// Write output.
119		if outputFile != "" {
120			if err := os.WriteFile(outputFile, data, 0o644); err != nil {
121				return fmt.Errorf("failed to write output file: %w", err)
122			}
123			cmd.Printf("Exported trajectory to %s\n", outputFile)
124		} else {
125			cmd.Println(string(data))
126		}
127
128		return nil
129	},
130}
131
132func init() {
133	trajectoryExportCmd.Flags().StringP("session", "s", "", "Session ID to export (required)")
134	trajectoryExportCmd.Flags().StringP("output", "o", "", "Output file path (defaults to stdout)")
135	trajectoryExportCmd.Flags().StringP("format", "f", "json", "Output format: json or html")
136	_ = trajectoryExportCmd.MarkFlagRequired("session")
137
138	trajectoryCmd.AddCommand(trajectoryExportCmd)
139}