112 lines
3.2 KiB
Go
112 lines
3.2 KiB
Go
package runtime
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
|
|
"github.com/krow/krow-backend/go-api/internal/repo"
|
|
)
|
|
|
|
// PostgresSink writes trajectories to agent_runs.
|
|
//
|
|
// Hand-written rather than built through repo.Repo's descriptor machinery, and
|
|
// deliberately so: that layer exists to serve the generic CRUD contract in
|
|
// docs/api-contract.md — filters, sorts, pagination, resource descriptors —
|
|
// and a trajectory has none of those. It is written once, whole, by the
|
|
// runtime, and read back by run id. One INSERT with explicit bind parameters
|
|
// is the honest shape for that, and inventing a resource descriptor to reach
|
|
// it would add a layer that only ever gets used one way.
|
|
type PostgresSink struct {
|
|
db repo.Querier
|
|
}
|
|
|
|
// NewPostgresSink builds a sink over a pool or transaction.
|
|
func NewPostgresSink(db repo.Querier) *PostgresSink {
|
|
return &PostgresSink{db: db}
|
|
}
|
|
|
|
var _ Sink = (*PostgresSink)(nil)
|
|
|
|
const insertRunSQL = `
|
|
INSERT INTO agent_runs (
|
|
run_id, parent_run_id, org_id, user_id,
|
|
agent_id, agent_version, tier, model,
|
|
started_at, ended_at, termination, entries,
|
|
input_tokens, output_tokens, cached_tokens, total_tokens, model_calls
|
|
) VALUES (
|
|
$1, $2, $3::uuid, $4::uuid,
|
|
$5, $6, $7, $8,
|
|
$9, $10, $11, $12::jsonb,
|
|
$13, $14, $15, $16, $17
|
|
)
|
|
ON CONFLICT (run_id) DO NOTHING`
|
|
|
|
// Save writes one finished trajectory.
|
|
//
|
|
// ON CONFLICT DO NOTHING because a run id is generated once and written once:
|
|
// a conflict means a retry of a save that already landed, and the first write
|
|
// is the authoritative one. Failing the second attempt would turn a harmless
|
|
// duplicate into a lost answer, since the caller treats a save error as
|
|
// something to report.
|
|
func (s *PostgresSink) Save(ctx context.Context, t *Trajectory) error {
|
|
if t == nil {
|
|
return fmt.Errorf("runtime: no trajectory to save")
|
|
}
|
|
if t.RunID == "" {
|
|
return fmt.Errorf("runtime: a trajectory needs a run id")
|
|
}
|
|
if !t.Termination.Valid() {
|
|
return fmt.Errorf("runtime: %q is not a termination reason", t.Termination)
|
|
}
|
|
// The table's tenancy column is NOT NULL, and a run with no organization is
|
|
// a bug upstream rather than a row to write. Caught here so the failure
|
|
// names the cause instead of surfacing as a constraint violation.
|
|
if t.OrgID == "" {
|
|
return fmt.Errorf("runtime: a trajectory needs an org id")
|
|
}
|
|
|
|
entries, err := json.Marshal(t.Entries)
|
|
if err != nil {
|
|
return fmt.Errorf("runtime: encoding trajectory entries: %w", err)
|
|
}
|
|
// A nil slice marshals to "null", which the jsonb_typeof CHECK refuses.
|
|
// A run that recorded nothing is still a run worth keeping.
|
|
if len(t.Entries) == 0 {
|
|
entries = []byte("[]")
|
|
}
|
|
|
|
_, err = s.db.Exec(ctx, insertRunSQL,
|
|
t.RunID,
|
|
nullIfEmpty(t.ParentRunID),
|
|
t.OrgID,
|
|
nullIfEmpty(t.UserID),
|
|
t.AgentID,
|
|
t.AgentVersion,
|
|
t.Tier,
|
|
t.Model,
|
|
t.StartedAt,
|
|
t.EndedAt,
|
|
string(t.Termination),
|
|
string(entries),
|
|
t.Usage.InputTokens,
|
|
t.Usage.OutputTokens,
|
|
t.Usage.CachedTokens,
|
|
t.Usage.TotalTokens,
|
|
t.Usage.ModelCalls,
|
|
)
|
|
if err != nil {
|
|
return fmt.Errorf("runtime: saving trajectory %s: %w", t.RunID, err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// nullIfEmpty keeps an empty optional out of a uuid column, where "" is not a
|
|
// value the type accepts.
|
|
func nullIfEmpty(s string) any {
|
|
if s == "" {
|
|
return nil
|
|
}
|
|
return s
|
|
}
|