easyai-ai-gateway/apps/api/internal/store/response_chains.go

package store

import (
	"context"
	"encoding/json"
	"errors"
	"strings"
	"time"

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

var ErrInvalidPreviousResponseID = errors.New("invalid previous response id")
var ErrResponseChainTooDeep = errors.New("response chain exceeds the maximum depth")

const maxResponseChainHistory = 100

type ResponseChain struct {
	PublicResponseID   string
	ParentResponseID   string
	GatewayTaskID      string
	RequestedModel     string
	GatewayUserID      string
	UserID             string
	UserSource         string
	GatewayTenantID    string
	TenantID           string
	TenantKey          string
	APIKeyID           string
	PlatformID         string
	PlatformModelID    string
	ClientID           string
	UpstreamProtocol   string
	UpstreamEndpoint   string
	UpstreamResponseID string
	RequestSnapshot    map[string]any
	ResponseSnapshot   map[string]any
	InternalSnapshot   map[string]any
	CreatedAt          time.Time
	ExpiresAt          time.Time
}

type CreateResponseChainInput struct {
	ResponseChain
}

func (s *Store) CreateResponseChain(ctx context.Context, input CreateResponseChainInput) error {
	requestJSON, err := json.Marshal(sanitizeJSONForStorage(input.RequestSnapshot))
	if err != nil {
		return err
	}
	responseJSON, err := json.Marshal(sanitizeJSONForStorage(input.ResponseSnapshot))
	if err != nil {
		return err
	}
	internalJSON, err := json.Marshal(sanitizeJSONForStorage(input.InternalSnapshot))
	if err != nil {
		return err
	}
	_, err = s.pool.Exec(ctx, `
INSERT INTO gateway_response_chains (
  public_response_id, parent_response_id, gateway_task_id,
  gateway_user_id, user_id, user_source, gateway_tenant_id, tenant_id, tenant_key, api_key_id,
  platform_id, platform_model_id, client_id, upstream_protocol, upstream_endpoint, upstream_response_id,
  request_snapshot, response_snapshot, internal_snapshot, expires_at
) VALUES (
  $1, NULLIF($2, ''), $3::uuid,
  NULLIF($4, '')::uuid, $5, $6, NULLIF($7, '')::uuid, NULLIF($8, ''), NULLIF($9, ''), NULLIF($10, ''),
  $11::uuid, $12::uuid, $13, $14, $15, NULLIF($16, ''),
  $17::jsonb, $18::jsonb, $19::jsonb, $20
)
ON CONFLICT (public_response_id) DO NOTHING`,
		input.PublicResponseID, input.ParentResponseID, input.GatewayTaskID,
		input.GatewayUserID, input.UserID, input.UserSource, input.GatewayTenantID, input.TenantID, input.TenantKey, input.APIKeyID,
		input.PlatformID, input.PlatformModelID, input.ClientID, input.UpstreamProtocol, input.UpstreamEndpoint, input.UpstreamResponseID,
		requestJSON, responseJSON, internalJSON, input.ExpiresAt,
	)
	return err
}

func (s *Store) GetResponseChain(ctx context.Context, publicResponseID string, owner GatewayTask) (ResponseChain, error) {
	chain, err := s.scanResponseChain(s.pool.QueryRow(ctx, `
	SELECT chain.public_response_id, COALESCE(chain.parent_response_id, ''), chain.gateway_task_id::text,
	       task.requested_model,
	       COALESCE(chain.gateway_user_id::text, ''), chain.user_id, chain.user_source,
	       COALESCE(chain.gateway_tenant_id::text, ''), COALESCE(chain.tenant_id, ''), COALESCE(chain.tenant_key, ''), COALESCE(chain.api_key_id, ''),
	       chain.platform_id::text, chain.platform_model_id::text, chain.client_id, chain.upstream_protocol, chain.upstream_endpoint,
	       COALESCE(chain.upstream_response_id, ''), chain.request_snapshot, chain.response_snapshot, chain.internal_snapshot,
	       chain.created_at, chain.expires_at
FROM gateway_response_chains chain
JOIN gateway_tasks task ON task.id = chain.gateway_task_id
WHERE chain.public_response_id = $1`, strings.TrimSpace(publicResponseID)))
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return ResponseChain{}, ErrInvalidPreviousResponseID
		}
		return ResponseChain{}, err
	}
	if chain.ExpiresAt.Before(time.Now()) || !sameResponseChainOwner(chain, owner) {
		return ResponseChain{}, ErrInvalidPreviousResponseID
	}
	return chain, nil
}

func (s *Store) ListResponseChainHistory(ctx context.Context, current ResponseChain) ([]ResponseChain, error) {
	rows, err := s.pool.Query(ctx, `
WITH RECURSIVE history AS (
  SELECT chain.*, 0 AS depth
  FROM gateway_response_chains chain
  WHERE chain.public_response_id = $1
  UNION ALL
  SELECT parent.*, history.depth + 1
  FROM gateway_response_chains parent
  JOIN history ON parent.public_response_id = history.parent_response_id
	  WHERE history.depth < $2
)
	SELECT history.public_response_id, COALESCE(history.parent_response_id, ''), history.gateway_task_id::text,
	       task.requested_model,
	       COALESCE(history.gateway_user_id::text, ''), history.user_id, history.user_source,
	       COALESCE(history.gateway_tenant_id::text, ''), COALESCE(history.tenant_id, ''), COALESCE(history.tenant_key, ''), COALESCE(history.api_key_id, ''),
	       history.platform_id::text, history.platform_model_id::text, history.client_id, history.upstream_protocol, history.upstream_endpoint,
	       COALESCE(history.upstream_response_id, ''), history.request_snapshot, history.response_snapshot, history.internal_snapshot,
	       history.created_at, history.expires_at
FROM history
JOIN gateway_tasks task ON task.id = history.gateway_task_id
ORDER BY depth DESC`, current.PublicResponseID, maxResponseChainHistory)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := make([]ResponseChain, 0)
	for rows.Next() {
		item, scanErr := s.scanResponseChain(rows)
		if scanErr != nil {
			return nil, scanErr
		}
		out = append(out, item)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	if len(out) > maxResponseChainHistory {
		return nil, ErrResponseChainTooDeep
	}
	return out, nil
}

type responseChainScanner interface {
	Scan(dest ...any) error
}

func (s *Store) scanResponseChain(scanner responseChainScanner) (ResponseChain, error) {
	var item ResponseChain
	var requestJSON, responseJSON, internalJSON []byte
	err := scanner.Scan(
		&item.PublicResponseID, &item.ParentResponseID, &item.GatewayTaskID, &item.RequestedModel,
		&item.GatewayUserID, &item.UserID, &item.UserSource,
		&item.GatewayTenantID, &item.TenantID, &item.TenantKey, &item.APIKeyID,
		&item.PlatformID, &item.PlatformModelID, &item.ClientID,
		&item.UpstreamProtocol, &item.UpstreamEndpoint, &item.UpstreamResponseID,
		&requestJSON, &responseJSON, &internalJSON, &item.CreatedAt, &item.ExpiresAt,
	)
	if err != nil {
		return ResponseChain{}, err
	}
	item.RequestSnapshot = decodeObject(requestJSON)
	item.ResponseSnapshot = decodeObject(responseJSON)
	item.InternalSnapshot = decodeObject(internalJSON)
	return item, nil
}

func sameResponseChainOwner(chain ResponseChain, task GatewayTask) bool {
	if chain.GatewayUserID != "" || task.GatewayUserID != "" {
		if chain.GatewayUserID == "" || chain.GatewayUserID != task.GatewayUserID {
			return false
		}
	} else if chain.UserID != task.UserID || chain.UserSource != task.UserSource {
		return false
	}
	return chain.GatewayTenantID == task.GatewayTenantID &&
		chain.TenantID == task.TenantID &&
		chain.TenantKey == task.TenantKey
}