package store import ( "context" "database/sql" "errors" "fmt" "time" ) // RequestLog is how a project keeps its requests and answers: for how long, // and what share of them. Zero Retain keeps none. type RequestLog struct { Retain time.Duration // Sample is the share kept, above 0 and at most 1. Sample float64 } // SetRequestLog turns keeping requests and answers on for a project, or off // with hours 0. A sample outside (0, 1] keeps every request. func (s *DB) SetRequestLog(ctx context.Context, projectID string, hours int, sample float64) error { var h, sm any if hours > 0 { h = hours if sample <= 0 || sample > 1 { sample = 1 } sm = sample } res, err := s.db.ExecContext(ctx, `UPDATE projects SET request_log_hours = ?, request_log_sample = ? WHERE id = ?`, h, sm, projectID) if err != nil { return err } if n, _ := res.RowsAffected(); n == 0 { return fmt.Errorf("store: no project %q", projectID) } return nil } // RequestLogOf is a project's setting; zero when it keeps none. func (s *DB) RequestLogOf(ctx context.Context, projectID string) (RequestLog, error) { var hours sql.NullInt64 var sample sql.NullFloat64 err := s.db.QueryRowContext(ctx, `SELECT request_log_hours, request_log_sample FROM projects WHERE id = ?`, projectID).Scan(&hours, &sample) if errors.Is(err, sql.ErrNoRows) || err == nil && (!hours.Valid || hours.Int64 <= 0) { return RequestLog{}, nil } if err != nil { return RequestLog{}, err } l := RequestLog{Retain: time.Duration(hours.Int64) * time.Hour, Sample: 1} if sample.Valid && sample.Float64 > 0 && sample.Float64 <= 1 { l.Sample = sample.Float64 } return l, nil } // RequestBody is one kept request and its answer. type RequestBody struct { RequestID string `json:"request_id"` ProjectID string `json:"project_id"` At time.Time `json:"at"` ExpiresAt time.Time `json:"expires_at"` Path string `json:"path"` Model string `json:"model"` ServedModel string `json:"served_model"` Status int `json:"status"` Request string `json:"request"` // Response is the answer's text as the client was sent it; Truncated // says it was longer than what is kept. Response string `json:"response"` Truncated bool `json:"truncated,omitempty"` } // PutRequestBody keeps one. func (s *DB) PutRequestBody(ctx context.Context, b RequestBody) error { tr := 0 if b.Truncated { tr = 1 } _, err := s.db.ExecContext(ctx, `INSERT INTO request_bodies (request_id, project_id, ts, expires_at, path, model, served_model, status, request, response, truncated) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT (request_id) DO NOTHING`, b.RequestID, b.ProjectID, millis(b.At), millis(b.ExpiresAt), b.Path, b.Model, b.ServedModel, b.Status, b.Request, b.Response, tr) return err } // RequestBodyByID finds one, ErrNotFound when none is kept. func (s *DB) RequestBodyByID(ctx context.Context, id string) (RequestBody, error) { var b RequestBody var ts, exp int64 var tr int err := s.db.QueryRowContext(ctx, `SELECT request_id, project_id, ts, expires_at, path, model, served_model, status, request, response, truncated FROM request_bodies WHERE request_id = ?`, id). Scan(&b.RequestID, &b.ProjectID, &ts, &exp, &b.Path, &b.Model, &b.ServedModel, &b.Status, &b.Request, &b.Response, &tr) if errors.Is(err, sql.ErrNoRows) { return RequestBody{}, ErrNotFound } if err != nil { return RequestBody{}, err } b.At, b.ExpiresAt, b.Truncated = time.UnixMilli(ts), time.UnixMilli(exp), tr == 1 return b, nil } // PruneRequestBodies removes those past their retention, and those of // projects that keep none any more. func (s *DB) PruneRequestBodies(ctx context.Context, now time.Time) (int64, error) { res, err := s.db.ExecContext(ctx, `DELETE FROM request_bodies WHERE expires_at < ? OR project_id IN (SELECT id FROM projects WHERE request_log_hours IS NULL OR request_log_hours <= 0)`, millis(now)) if err != nil { return 0, err } return res.RowsAffected() } // SetInCluster says whether a project's requests stay in the cluster: no // model served by an outside provider, asked for or as a fallback. func (s *DB) SetInCluster(ctx context.Context, projectID string, on bool) error { v := 0 if on { v = 1 } res, err := s.db.ExecContext(ctx, `UPDATE projects SET in_cluster = ? WHERE id = ?`, v, projectID) if err != nil { return err } if n, _ := res.RowsAffected(); n == 0 { return fmt.Errorf("store: no project %q", projectID) } return nil } // InCluster is whether a project's requests stay in the cluster. func (s *DB) InCluster(ctx context.Context, projectID string) (bool, error) { var v int err := s.db.QueryRowContext(ctx, `SELECT in_cluster FROM projects WHERE id = ?`, projectID).Scan(&v) if errors.Is(err, sql.ErrNoRows) { return false, nil } return v == 1, err }