Atlas/internal/store/generations.go

package store

import "context"

type Generation struct {
	ID          int64  `json:"id"`
	SessionID   string `json:"session_id"`
	ServiceType string `json:"service_type"`
	ProviderID  *int64 `json:"provider_id"`
	Model       string `json:"model"`
	InputJSON   string `json:"input_json"`
	OutputJSON  string `json:"output_json"`
	TaskID      *int64 `json:"task_id"`
	CreatedAt   string `json:"created_at"`
}

func (s *Store) ListGenerations(ctx context.Context, sessionID, serviceType string, limit, offset int) ([]Generation, int, error) {
	countQ := `SELECT COUNT(*) FROM generations WHERE session_id = ?`
	args := []any{sessionID}
	if serviceType != "" {
		countQ += ` AND service_type = ?`
		args = append(args, serviceType)
	}
	var total int
	if err := s.db.QueryRowContext(ctx, countQ, args...).Scan(&total); err != nil {
		return nil, 0, err
	}
	q := `SELECT id, session_id, service_type, provider_id, model, input_json, output_json, task_id, created_at FROM generations WHERE session_id = ?`
	qArgs := []any{sessionID}
	if serviceType != "" {
		q += ` AND service_type = ?`
		qArgs = append(qArgs, serviceType)
	}
	q += ` ORDER BY id DESC LIMIT ? OFFSET ?`
	qArgs = append(qArgs, limit, offset)
	rows, err := s.db.QueryContext(ctx, q, qArgs...)
	if err != nil {
		return nil, 0, err
	}
	defer rows.Close()
	var items []Generation
	for rows.Next() {
		var g Generation
		if err := rows.Scan(&g.ID, &g.SessionID, &g.ServiceType, &g.ProviderID, &g.Model, &g.InputJSON, &g.OutputJSON, &g.TaskID, &g.CreatedAt); err != nil {
			return nil, 0, err
		}
		items = append(items, g)
	}
	return items, total, rows.Err()
}

func (s *Store) GetGeneration(ctx context.Context, sessionID string, id int64) (*Generation, error) {
	var g Generation
	err := s.db.QueryRowContext(ctx,
		`SELECT id, session_id, service_type, provider_id, model, input_json, output_json, task_id, created_at FROM generations WHERE id = ? AND session_id = ?`, id, sessionID,
	).Scan(&g.ID, &g.SessionID, &g.ServiceType, &g.ProviderID, &g.Model, &g.InputJSON, &g.OutputJSON, &g.TaskID, &g.CreatedAt)
	if err != nil {
		return nil, err
	}
	return &g, nil
}

func (s *Store) InsertGeneration(ctx context.Context, g *Generation) (int64, error) {
	res, err := s.db.ExecContext(ctx,
		`INSERT INTO generations (session_id, service_type, provider_id, model, input_json, output_json, task_id) VALUES (?, ?, ?, ?, ?, ?, ?)`,
		g.SessionID, g.ServiceType, g.ProviderID, g.Model, g.InputJSON, g.OutputJSON, g.TaskID,
	)
	if err != nil {
		return 0, err
	}
	return res.LastInsertId()
}

func (s *Store) DeleteGeneration(ctx context.Context, sessionID string, id int64) error {
	_, err := s.db.ExecContext(ctx, `DELETE FROM generations WHERE id = ? AND session_id = ?`, id, sessionID)
	return err
}

func (s *Store) ListAllGenerations(ctx context.Context, serviceType string, limit, offset int) ([]Generation, int, error) {
	countQ := `SELECT COUNT(*) FROM generations WHERE 1=1`
	args := []any{}
	if serviceType != "" {
		countQ += ` AND service_type = ?`
		args = append(args, serviceType)
	}
	var total int
	if err := s.db.QueryRowContext(ctx, countQ, args...).Scan(&total); err != nil {
		return nil, 0, err
	}
	q := `SELECT id, session_id, service_type, provider_id, model, input_json, output_json, task_id, created_at FROM generations WHERE 1=1`
	qArgs := []any{}
	if serviceType != "" {
		q += ` AND service_type = ?`
		qArgs = append(qArgs, serviceType)
	}
	q += ` ORDER BY id DESC LIMIT ? OFFSET ?`
	qArgs = append(qArgs, limit, offset)
	rows, err := s.db.QueryContext(ctx, q, qArgs...)
	if err != nil {
		return nil, 0, err
	}
	defer rows.Close()
	var items []Generation
	for rows.Next() {
		var g Generation
		if err := rows.Scan(&g.ID, &g.SessionID, &g.ServiceType, &g.ProviderID, &g.Model, &g.InputJSON, &g.OutputJSON, &g.TaskID, &g.CreatedAt); err != nil {
			return nil, 0, err
		}
		items = append(items, g)
	}
	return items, total, rows.Err()
}

func (s *Store) AdminGetGeneration(ctx context.Context, id int64) (*Generation, error) {
	var g Generation
	err := s.db.QueryRowContext(ctx,
		`SELECT id, session_id, service_type, provider_id, model, input_json, output_json, task_id, created_at FROM generations WHERE id = ?`, id,
	).Scan(&g.ID, &g.SessionID, &g.ServiceType, &g.ProviderID, &g.Model, &g.InputJSON, &g.OutputJSON, &g.TaskID, &g.CreatedAt)
	if err != nil {
		return nil, err
	}
	return &g, nil
}

func (s *Store) AdminDeleteGeneration(ctx context.Context, id int64) error {
	_, err := s.db.ExecContext(ctx, `DELETE FROM generations WHERE id = ?`, id)
	return err
}