internal/gateway/inference_execute.go

Rama main · réplica de lectura del repositorio original.

package gateway

import (
	"context"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/enterpilot/gomodel/internal/core"
	"github.com/enterpilot/gomodel/internal/usage"
)

// ExecuteChatCompletion executes a non-streaming chat request and records usage on success.
func (o *InferenceOrchestrator) ExecuteChatCompletion(ctx context.Context, workflow *core.Workflow, req *core.ChatRequest, requestID, endpoint string) (*ChatCompletionResult, error) {
	if err := o.validateProviderAndRequest(req != nil, "chat request is required"); err != nil {
		return nil, err
	}
	return executeResultWithSlowdown(ctx, workflow, func() (*ChatCompletionResult, error) {
		return executeTranslatedResult(o, ctx, workflow, req, requestID, endpoint, chatExecutionSpec)
	})
}

// DispatchChatCompletion executes a non-streaming chat request without usage side effects.
func (o *InferenceOrchestrator) DispatchChatCompletion(
	ctx context.Context,
	workflow *core.Workflow,
	req *core.ChatRequest,
) (*core.ChatResponse, ExecutionMeta, error) {
	if err := o.validateProviderAndRequest(req != nil, "chat request is required"); err != nil {
		return nil, ExecutionMeta{}, err
	}
	return dispatchTranslatedWithSlowdown(ctx, workflow, func() (*core.ChatResponse, ExecutionMeta, error) {
		return o.executeChatCompletion(ctx, workflow, req)
	})
}

// StreamChatCompletion opens a chat SSE stream. Stream usage is recorded by the caller's stream observer.
func (o *InferenceOrchestrator) StreamChatCompletion(ctx context.Context, workflow *core.Workflow, req *core.ChatRequest) (*StreamResult, error) {
	if err := o.validateProviderAndRequest(req != nil, "chat request is required"); err != nil {
		return nil, err
	}
	started := time.Now()
	streamReq, providerType, providerName, usageModel := o.ResolveChatRoute(workflow, req)
	ctx, generation := core.StartGenerationTrace(ctx)
	stream, meta, err := o.streamChatCompletion(ctx, workflow, streamReq, providerType, providerName, usageModel)
	if err != nil {
		generation.Finish(core.GenerationOutcome{Request: req})
		return nil, err
	}
	return &StreamResult{
		Stream:           stream,
		slowdownFactor:   workflowSlowdown(workflow),
		inferenceStarted: started,
		Meta:             meta,
		generation:       newChatGenerationObserver(generation, req),
	}, nil
}

// ExecuteResponses executes a non-streaming Responses API request and records usage on success.
func (o *InferenceOrchestrator) ExecuteResponses(ctx context.Context, workflow *core.Workflow, req *core.ResponsesRequest, requestID, endpoint string) (*ResponsesResult, error) {
	if err := o.validateProviderAndRequest(req != nil, "responses request is required"); err != nil {
		return nil, err
	}
	return executeResultWithSlowdown(ctx, workflow, func() (*ResponsesResult, error) {
		return executeTranslatedResult(o, ctx, workflow, req, requestID, endpoint, responsesExecutionSpec)
	})
}

// DispatchResponses executes a non-streaming Responses request without usage side effects.
func (o *InferenceOrchestrator) DispatchResponses(
	ctx context.Context,
	workflow *core.Workflow,
	req *core.ResponsesRequest,
) (*core.ResponsesResponse, ExecutionMeta, error) {
	if err := o.validateProviderAndRequest(req != nil, "responses request is required"); err != nil {
		return nil, ExecutionMeta{}, err
	}
	return dispatchTranslatedWithSlowdown(ctx, workflow, func() (*core.ResponsesResponse, ExecutionMeta, error) {
		return o.executeResponses(ctx, workflow, req)
	})
}

// StreamResponses opens a Responses API SSE stream. Stream usage is recorded by the caller's stream observer.
func (o *InferenceOrchestrator) StreamResponses(ctx context.Context, workflow *core.Workflow, req *core.ResponsesRequest) (*StreamResult, error) {
	if err := o.validateProviderAndRequest(req != nil, "responses request is required"); err != nil {
		return nil, err
	}
	started := time.Now()
	providerType, providerName, usageModel := o.routeMetadata(workflow, req.Model)
	if (workflow == nil || workflow.UsageEnabled()) && o.ShouldEnforceReturningUsageData() {
		ctx = core.WithEnforceReturningUsageData(ctx, true)
	}
	ctx, generation := core.StartGenerationTrace(ctx)
	stream, meta, err := o.streamResponses(ctx, workflow, req, providerType, providerName, usageModel)
	if err != nil {
		generation.Finish(core.GenerationOutcome{Request: req})
		return nil, err
	}
	return &StreamResult{
		Stream:           stream,
		slowdownFactor:   workflowSlowdown(workflow),
		inferenceStarted: started,
		Meta:             meta,
		generation:       newResponsesGenerationObserver(generation, req),
	}, nil
}

// ExecuteEmbeddings executes an embeddings request and records usage on success.
func (o *InferenceOrchestrator) ExecuteEmbeddings(ctx context.Context, workflow *core.Workflow, req *core.EmbeddingRequest, requestID, endpoint string) (*EmbeddingResult, error) {
	if err := o.validateProviderAndRequest(req != nil, "embeddings request is required"); err != nil {
		return nil, err
	}
	started := time.Now()
	resp, meta, err := o.executeEmbeddings(ctx, workflow, req)
	if err != nil {
		return nil, err
	}
	meta.Model = resp.Model
	pricingModel := usagePricingModel(workflow, req.Model, meta.FailoverModel, meta.Model)
	o.logUsage(ctx, workflow, pricingModel, meta.ProviderType, meta.ProviderName, func(pricing *core.ModelPricing) *usage.UsageEntry {
		return usage.ExtractFromEmbeddingResponse(resp, requestID, meta.ProviderType, endpoint, pricing)
	})
	if err := waitForInferenceSlowdown(ctx, workflow, time.Since(started)); err != nil {
		return nil, err
	}
	return &EmbeddingResult{
		Response: resp,
		Meta:     meta,
	}, nil
}

// DispatchEmbeddings executes an embeddings request without usage side effects.
func (o *InferenceOrchestrator) DispatchEmbeddings(
	ctx context.Context,
	workflow *core.Workflow,
	req *core.EmbeddingRequest,
) (*core.EmbeddingResponse, string, string, error) {
	if err := o.validateProviderAndRequest(req != nil, "embeddings request is required"); err != nil {
		return nil, "", "", err
	}
	started := time.Now()
	resp, meta, err := o.executeEmbeddings(ctx, workflow, req)
	if err != nil {
		return nil, "", "", err
	}
	if err := waitForInferenceSlowdown(ctx, workflow, time.Since(started)); err != nil {
		return nil, "", "", err
	}
	return resp, meta.ProviderType, meta.ProviderName, nil
}

// ResolveChatRoute returns the provider route and the request to send for chat streams.
func (o *InferenceOrchestrator) ResolveChatRoute(workflow *core.Workflow, req *core.ChatRequest) (*core.ChatRequest, string, string, string) {
	providerType, providerName, usageModel := o.routeMetadata(workflow, "")
	if req != nil {
		usageModel = ResolvedModelFromWorkflow(workflow, req.Model)
	}
	if req == nil || !req.Stream || (workflow != nil && !workflow.UsageEnabled()) || !o.ShouldEnforceReturningUsageData() {
		return req, providerType, providerName, usageModel
	}

	streamReq := CloneChatRequestForStreamUsage(req)
	if streamReq.StreamOptions == nil {
		streamReq.StreamOptions = &core.StreamOptions{}
	}
	streamReq.StreamOptions.IncludeUsage = true
	return streamReq, providerType, providerName, usageModel
}

func (o *InferenceOrchestrator) routeMetadata(workflow *core.Workflow, failoverModel string) (string, string, string) {
	providerType := ProviderTypeFromWorkflow(workflow)
	providerName := ProviderNameFromWorkflow(workflow)
	model := ResolvedModelFromWorkflow(workflow, failoverModel)
	return providerType, providerName, model
}

func (o *InferenceOrchestrator) executeChatCompletion(
	ctx context.Context,
	workflow *core.Workflow,
	req *core.ChatRequest,
) (*core.ChatResponse, ExecutionMeta, error) {
	if err := o.validateProviderAndRequest(req != nil, "chat request is required"); err != nil {
		return nil, ExecutionMeta{}, err
	}
	call := withRequestOverrides(o, workflow, chatRequestSelectorOf, ChatRequestWithOverrides, rebuildChatAttempt, o.chatCompletionProviderCall)
	return executeTranslatedProviderRequest(o, ctx, workflow, req, req.Model, req.Provider, CloneChatRequestForSelector, call, chatResponseProvider)
}

func (o *InferenceOrchestrator) streamChatCompletion(
	ctx context.Context,
	workflow *core.Workflow,
	req *core.ChatRequest,
	providerType, providerName, usageModel string,
) (io.ReadCloser, ExecutionMeta, error) {
	if err := o.validateProviderAndRequest(req != nil, "chat request is required"); err != nil {
		return nil, ExecutionMeta{}, err
	}
	call := withRequestOverrides(o, workflow, chatRequestSelectorOf, ChatRequestWithOverrides, rebuildChatAttempt, o.streamChatCompletionProviderCall)
	return streamTranslatedProviderRequest(o, ctx, workflow, req, req.Model, req.Provider, providerType, providerName, usageModel, CloneChatRequestForSelector, call)
}

func (o *InferenceOrchestrator) executeResponses(
	ctx context.Context,
	workflow *core.Workflow,
	req *core.ResponsesRequest,
) (*core.ResponsesResponse, ExecutionMeta, error) {
	if err := o.validateProviderAndRequest(req != nil, "responses request is required"); err != nil {
		return nil, ExecutionMeta{}, err
	}
	call := withRequestOverrides(o, workflow, responsesRequestSelectorOf, ResponsesRequestWithOverrides, rebuildResponsesAttempt, patchedResponsesCall(o, workflow, o.responsesProviderCall))
	return executeTranslatedProviderRequest(o, ctx, workflow, req, req.Model, req.Provider, CloneResponsesRequestForSelector, call, responsesResponseProvider)
}

// patchedResponsesCall runs the Responses attempt patcher, when configured,
// in front of call. The attempt's provider type comes from the request's own
// selector, which the primary carries after resolution and a failover clone
// carries from its target, falling back to the workflow's provider type.
func patchedResponsesCall[T any](o *InferenceOrchestrator, workflow *core.Workflow, call func(context.Context, *core.ResponsesRequest, string) (T, error)) func(context.Context, *core.ResponsesRequest, string) (T, error) {
	if o.responsesAttemptPatcher == nil {
		return call
	}
	return func(ctx context.Context, req *core.ResponsesRequest, providerName string) (T, error) {
		providerType := o.ProviderTypeForSelector(core.ModelSelector{Model: req.Model, Provider: req.Provider}, ProviderTypeFromWorkflow(workflow))
		patched, err := o.responsesAttemptPatcher.PatchResponsesAttempt(ctx, req, providerType)
		if err != nil {
			var zero T
			return zero, err
		}
		return call(ctx, patched, providerName)
	}
}

func (o *InferenceOrchestrator) streamResponses(
	ctx context.Context,
	workflow *core.Workflow,
	req *core.ResponsesRequest,
	providerType, providerName, usageModel string,
) (io.ReadCloser, ExecutionMeta, error) {
	if err := o.validateProviderAndRequest(req != nil, "responses request is required"); err != nil {
		return nil, ExecutionMeta{}, err
	}
	call := withRequestOverrides(o, workflow, responsesRequestSelectorOf, ResponsesRequestWithOverrides, rebuildResponsesAttempt, patchedResponsesCall(o, workflow, o.streamResponsesProviderCall))
	return streamTranslatedProviderRequest(o, ctx, workflow, req, req.Model, req.Provider, providerType, providerName, usageModel, CloneResponsesRequestForSelector, call)
}

type translatedExecutionSpec[Req any, Resp any, Result any] struct {
	execute func(*InferenceOrchestrator, context.Context, *core.Workflow, Req) (Resp, ExecutionMeta, error)
	model   func(Resp) string
	request func(Req) string
	usage   func(Resp, string, string, string, *core.ModelPricing) *usage.UsageEntry
	build   func(Resp, ExecutionMeta) Result
}

var chatExecutionSpec = translatedExecutionSpec[*core.ChatRequest, *core.ChatResponse, *ChatCompletionResult]{
	execute: (*InferenceOrchestrator).executeChatCompletion,
	model:   chatResponseModel,
	request: chatRequestModel,
	usage: func(resp *core.ChatResponse, requestID, providerType, endpoint string, pricing *core.ModelPricing) *usage.UsageEntry {
		return usage.ExtractFromChatResponse(resp, requestID, providerType, endpoint, pricing)
	},
	build: func(resp *core.ChatResponse, meta ExecutionMeta) *ChatCompletionResult {
		return &ChatCompletionResult{Response: resp, Meta: meta}
	},
}

var responsesExecutionSpec = translatedExecutionSpec[*core.ResponsesRequest, *core.ResponsesResponse, *ResponsesResult]{
	execute: (*InferenceOrchestrator).executeResponses,
	model:   responsesResponseModel,
	request: responsesRequestModel,
	usage: func(resp *core.ResponsesResponse, requestID, providerType, endpoint string, pricing *core.ModelPricing) *usage.UsageEntry {
		return usage.ExtractFromResponsesResponse(resp, requestID, providerType, endpoint, pricing)
	},
	build: func(resp *core.ResponsesResponse, meta ExecutionMeta) *ResponsesResult {
		return &ResponsesResult{Response: resp, Meta: meta}
	},
}

func executeTranslatedResult[Req any, Resp any, Result any](
	o *InferenceOrchestrator,
	ctx context.Context,
	workflow *core.Workflow,
	req Req,
	requestID, endpoint string,
	spec translatedExecutionSpec[Req, Resp, Result],
) (Result, error) {
	ctx, generation := core.StartGenerationTrace(ctx)
	resp, meta, err := executeWithUsage(o, ctx, workflow,
		func() (Resp, ExecutionMeta, error) {
			return spec.execute(o, ctx, workflow, req)
		},
		requestModel(req, spec.request),
		spec.model,
		func(resp Resp, providerType string, pricing *core.ModelPricing) *usage.UsageEntry {
			return spec.usage(resp, requestID, providerType, endpoint, pricing)
		},
	)
	outcome := core.GenerationOutcome{Request: req}
	if err == nil {
		outcome.Response = resp
	}
	generation.Finish(outcome)
	if err != nil {
		var zero Result
		return zero, err
	}
	return spec.build(resp, meta), nil
}

func executeWithUsage[Resp any](
	o *InferenceOrchestrator,
	ctx context.Context,
	workflow *core.Workflow,
	execute func() (Resp, ExecutionMeta, error),
	requestedModel string,
	modelFromResponse func(Resp) string,
	entry func(Resp, string, *core.ModelPricing) *usage.UsageEntry,
) (Resp, ExecutionMeta, error) {
	resp, meta, err := execute()
	if err != nil {
		var zero Resp
		return zero, ExecutionMeta{}, err
	}
	meta.Model = modelFromResponse(resp)
	pricingModel := usagePricingModel(workflow, requestedModel, meta.FailoverModel, meta.Model)
	o.logUsage(ctx, workflow, pricingModel, meta.ProviderType, meta.ProviderName, func(pricing *core.ModelPricing) *usage.UsageEntry {
		return entry(resp, meta.ProviderType, pricing)
	})
	return resp, meta, nil
}

func requestModel[Req any](req Req, model func(Req) string) string {
	if model == nil {
		return ""
	}
	return strings.TrimSpace(model(req))
}

func usagePricingModel(workflow *core.Workflow, requestedModel, failoverModel, responseModel string) string {
	requestedModel = strings.TrimSpace(requestedModel)
	failoverModel = strings.TrimSpace(failoverModel)
	if failoverModel != "" {
		return failoverModel
	}
	if model := ResolvedModelFromWorkflow(workflow, requestedModel); model != "" {
		return model
	}
	return strings.TrimSpace(responseModel)
}

func executeTranslatedProviderRequest[Req any, Resp any](
	o *InferenceOrchestrator,
	ctx context.Context,
	workflow *core.Workflow,
	req Req,
	model, provider string,
	cloneForSelector func(Req, core.ModelSelector) Req,
	call func(context.Context, Req, string) (Resp, error),
	responseProvider func(Resp) string,
) (Resp, ExecutionMeta, error) {
	resp, meta, err := executeTranslatedWithFailover(ctx, o, workflow, req, model, provider, cloneForSelector,
		func(ctx context.Context, req Req, providerName string) (Resp, string, error) {
			resp, err := call(ctx, req, providerName)
			if err != nil {
				var zero Resp
				return zero, "", err
			}
			return resp, responseProvider(resp), nil
		},
	)
	if err != nil {
		return resp, meta, err
	}
	// The provider type drives routing and audit filters, so it is what the
	// attempt reported above; the client-facing field names the configured
	// instance that served the request, as errors and breaker messages do.
	return nameProviderInstance(resp, meta.ProviderName), meta, nil
}

// nameProviderInstance rewrites the response's public provider field to the
// configured provider instance name. The router stamps the provider type; an
// instance name is more precise whenever two instances share a type, and it
// matches how the gateway names providers everywhere else. A request served
// without a configured name keeps the type.
func nameProviderInstance[Resp any](resp Resp, providerName string) Resp {
	providerName = strings.TrimSpace(providerName)
	if providerName == "" {
		return resp
	}
	switch typed := any(resp).(type) {
	case *core.ChatResponse:
		if typed != nil {
			typed.Provider = providerName
		}
	case *core.ResponsesResponse:
		if typed != nil {
			typed.Provider = providerName
		}
	case *core.EmbeddingResponse:
		if typed != nil {
			typed.Provider = providerName
		}
	}
	return resp
}

func streamTranslatedProviderRequest[Req any](
	o *InferenceOrchestrator,
	ctx context.Context,
	workflow *core.Workflow,
	req Req,
	model, provider string,
	providerType, providerName, usageModel string,
	cloneForSelector func(Req, core.ModelSelector) Req,
	call func(context.Context, Req, string) (io.ReadCloser, error),
) (io.ReadCloser, ExecutionMeta, error) {
	started := time.Now()
	var stream io.ReadCloser
	// See executeTranslatedWithFailover: a rate-saturated primary route skips
	// the provider call and enters the failover sweep with its stored 429.
	err := core.PrimaryRouteSaturated(ctx)
	if err == nil {
		stream, err = call(ctx, req, providerName)
		if err == nil && stream != nil {
			recordProviderAttempt(ctx, providerAttemptFromResult(AttemptKindPrimary, providerType, providerName, currentSelectorForWorkflow(workflow, model, provider), started, nil))
			return stream, ExecutionMeta{ProviderType: providerType, ProviderName: providerName, Model: usageModel}, nil
		}
		if err == nil {
			err = emptyProviderStreamError(providerName)
		}
	}
	recordProviderAttempt(ctx, providerAttemptFromResult(AttemptKindPrimary, providerType, providerName, currentSelectorForWorkflow(workflow, model, provider), started, err))

	return tryFailoverStream(ctx, o, workflow, model, provider, err,
		func(selector core.ModelSelector, providerType, providerName string) (io.ReadCloser, string, string, error) {
			// See executeTranslatedWithFailover: no idempotency key on failover.
			stream, err := call(core.WithIdempotencyKey(ctx, ""), cloneForSelector(req, selector), providerName)
			if err != nil {
				return nil, "", "", err
			}
			if stream == nil {
				return nil, "", "", emptyProviderStreamError(providerName)
			}
			return stream, providerType, selector.Model, nil
		},
	)
}

func (o *InferenceOrchestrator) validateProviderAndRequest(requestPresent bool, requiredMessage string) *core.GatewayError {
	if o == nil || o.provider == nil {
		return core.NewInvalidRequestError("provider is not configured", nil)
	}
	if !requestPresent {
		return core.NewInvalidRequestError(requiredMessage, nil)
	}
	return nil
}

// The provider name of the attempt — not the provider type the response body
// carries — names these gateway-raised errors, so an operator running two
// instances of one type can tell which instance answered (see
// ResolvedProviderName).
func (o *InferenceOrchestrator) chatCompletionProviderCall(ctx context.Context, req *core.ChatRequest, providerName string) (*core.ChatResponse, error) {
	resp, err := o.provider.ChatCompletion(ctx, req)
	if err != nil {
		return nil, err
	}
	if resp == nil {
		return nil, emptyProviderResponseError(providerName)
	}
	if len(resp.Choices) == 0 {
		return nil, core.NewNoChoicesProviderError(providerName)
	}
	return resp, nil
}

func (o *InferenceOrchestrator) responsesProviderCall(ctx context.Context, req *core.ResponsesRequest, providerName string) (*core.ResponsesResponse, error) {
	resp, err := o.provider.Responses(ctx, req)
	if err != nil {
		return nil, err
	}
	if resp == nil {
		return nil, emptyProviderResponseError(providerName)
	}
	return resp, nil
}

func (o *InferenceOrchestrator) streamChatCompletionProviderCall(ctx context.Context, req *core.ChatRequest, _ string) (io.ReadCloser, error) {
	return o.provider.StreamChatCompletion(ctx, req)
}

func (o *InferenceOrchestrator) streamResponsesProviderCall(ctx context.Context, req *core.ResponsesRequest, _ string) (io.ReadCloser, error) {
	return o.provider.StreamResponses(ctx, req)
}

func emptyProviderResponseError(provider string) *core.GatewayError {
	return core.NewEmptyProviderResponseError(provider)
}

func emptyProviderStreamError(provider string) *core.GatewayError {
	return core.NewProviderError(provider, http.StatusBadGateway, "provider returned empty stream", nil)
}

func chatResponseModel(resp *core.ChatResponse) string {
	if resp == nil {
		return ""
	}
	return resp.Model
}

func chatRequestModel(req *core.ChatRequest) string {
	if req == nil {
		return ""
	}
	return req.Model
}

func chatResponseProvider(resp *core.ChatResponse) string {
	if resp == nil {
		return ""
	}
	return resp.Provider
}

func embeddingResponseProvider(resp *core.EmbeddingResponse) string {
	if resp == nil {
		return ""
	}
	return resp.Provider
}

func responsesResponseModel(resp *core.ResponsesResponse) string {
	if resp == nil {
		return ""
	}
	return resp.Model
}

func responsesRequestModel(req *core.ResponsesRequest) string {
	if req == nil {
		return ""
	}
	return req.Model
}

func responsesResponseProvider(resp *core.ResponsesResponse) string {
	if resp == nil {
		return ""
	}
	return resp.Provider
}

func (o *InferenceOrchestrator) executeEmbeddings(
	ctx context.Context,
	workflow *core.Workflow,
	req *core.EmbeddingRequest,
) (*core.EmbeddingResponse, ExecutionMeta, error) {
	if err := o.validateProviderAndRequest(req != nil, "embeddings request is required"); err != nil {
		return nil, ExecutionMeta{}, err
	}
	return executeTranslatedProviderRequest(o, ctx, workflow, req, req.Model, req.Provider,
		CloneEmbeddingRequestForSelector, o.embeddingsProviderCall, embeddingResponseProvider)
}

func (o *InferenceOrchestrator) embeddingsProviderCall(ctx context.Context, req *core.EmbeddingRequest, providerName string) (*core.EmbeddingResponse, error) {
	resp, err := o.provider.Embeddings(ctx, req)
	if err != nil {
		return nil, err
	}
	if resp == nil {
		return nil, emptyProviderResponseError(providerName)
	}
	return resp, nil
}