feat: add Responses API compatibility
This commit is contained in:
@@ -0,0 +1,183 @@
|
||||
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(input.RequestSnapshot)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
responseJSON, err := json.Marshal(input.ResponseSnapshot)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
internalJSON, err := json.Marshal(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
|
||||
}
|
||||
Reference in New Issue
Block a user