promptkit/internal/llm/openai_compatible_client.go

package llm

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"time"

	"go.maximumdirect.net/promptkit/internal/defaults"
	"go.maximumdirect.net/promptkit/internal/domain"
	"go.maximumdirect.net/promptkit/internal/jsonvalue"
)

var (
	ErrInvalidRequest    = errors.New("invalid generate request")
	ErrRequestFailed     = errors.New("llm request failed")
	ErrTransport         = errors.New("llm transport failed")
	ErrUnexpectedStatus  = errors.New("llm returned non-success status")
	ErrMalformedResponse = errors.New("malformed llm response")
)

const maxOpenAIChatResponseBytes int64 = 16 << 20

type requestFailedError struct {
	cause error
}

func (e *requestFailedError) Error() string {
	return ErrRequestFailed.Error()
}

func (e *requestFailedError) Unwrap() []error {
	return []error{ErrRequestFailed, ErrTransport, e.cause}
}

type OpenAICompatibleConfig struct {
	Timeout    time.Duration
	HTTPClient *http.Client
}

type OpenAICompatibleClient struct {
	httpClient *http.Client
}

func NewOpenAICompatibleClient(cfg OpenAICompatibleConfig) *OpenAICompatibleClient {
	timeout := cfg.Timeout
	if timeout <= 0 {
		timeout = defaults.LLMRequestTimeoutDefault
	}

	var client *http.Client
	if cfg.HTTPClient != nil {
		cloned := *cfg.HTTPClient
		if cloned.Timeout <= 0 {
			cloned.Timeout = timeout
		}
		client = &cloned
	} else {
		client = &http.Client{Timeout: timeout}
	}

	return &OpenAICompatibleClient{
		httpClient: client,
	}
}

func (c *OpenAICompatibleClient) Generate(ctx context.Context, req domain.GenerateRequest) (*domain.GenerateResponse, error) {
	prepared, err := c.PrepareGeneration(req)
	if err != nil {
		return nil, err
	}
	outcome := prepared.Execute(ctx)
	if outcome.Err != nil {
		return nil, outcome.Err
	}
	return outcome.Response, nil
}

// PreparedGeneration holds one encoded request and resolved authentication for
// repeated attempts. Each execution creates a fresh HTTP request and context.
type PreparedGeneration struct {
	endpoint                  string
	payload                   []byte
	apiKey                    string
	timeoutSeconds            int
	httpClient                *http.Client
	structuredOutputRequested bool
}

// AttemptOutcome describes one HTTP execution without exposing response bodies
// or headers to the orchestration layer.
type AttemptOutcome struct {
	Response                  *domain.GenerateResponse
	Err                       error
	StatusCode                int
	Usage                     domain.TokenUsage
	UsagePresence             *domain.TokenUsagePresence
	RetryNotBefore            time.Time
	Evidence                  HTTPExecutionEvidence
	StructuredOutputRequested *bool
	ResponseDecoded           bool
	GenerationCompleted       bool
}

// PrepareGeneration validates and freezes one built-in request before dispatch.
func (c *OpenAICompatibleClient) PrepareGeneration(req domain.GenerateRequest) (*PreparedGeneration, error) {
	if err := domain.ValidateExecutionTargetSettings(req.Target); err != nil {
		return nil, fmt.Errorf("%w: %v", ErrInvalidRequest, err)
	}

	endpoint, err := domain.NormalizeOpenAICompatibleBaseEndpoint(req.Target.Endpoint)
	if err != nil {
		return nil, fmt.Errorf("%w: invalid endpoint: %v", ErrInvalidRequest, err)
	}
	endpoint, err = url.JoinPath(endpoint, defaults.OpenAIChatCompletionsPath)
	if err != nil {
		return nil, fmt.Errorf("%w: invalid endpoint path: %v", ErrInvalidRequest, err)
	}

	// Custom JSON encoders may mutate their receivers. Both payload encoding
	// and encoding-failure diagnostics must use values detached from metadata.
	req.Target.ExtraParams = jsonvalue.CloneTrusted(req.Target.ExtraParams).(map[string]any)
	wirePayload, err := openAIChatRequestPayload(req)
	if err != nil {
		return nil, fmt.Errorf("%w: %v", ErrInvalidRequest, err)
	}

	payload, err := json.Marshal(wirePayload)
	if err != nil {
		for key, value := range req.Target.ExtraParams {
			if _, extraErr := json.Marshal(value); extraErr != nil {
				return nil, fmt.Errorf("%w: extra_params.%s must be JSON-serializable: %v", ErrInvalidRequest, key, extraErr)
			}
		}
		return nil, fmt.Errorf("%w: failed to encode request: %v", ErrRequestFailed, err)
	}
	apiKey := resolveAPIKey(req.Target)
	envName := strings.TrimSpace(req.Target.APIKeyEnv)
	if apiKey == "" && req.Target.APIKeyRequired {
		if envName != "" {
			return nil, fmt.Errorf("%w: api key environment variable %q is not set", ErrInvalidRequest, envName)
		}
		return nil, fmt.Errorf("%w: api key is required", ErrInvalidRequest)
	}
	httpClient := c.httpClient
	if httpClient == nil {
		httpClient = &http.Client{Timeout: defaults.LLMRequestTimeoutDefault}
	}
	return &PreparedGeneration{
		endpoint: endpoint, payload: payload, apiKey: apiKey,
		timeoutSeconds: req.Target.TimeoutSeconds, httpClient: httpClient,
		structuredOutputRequested: wirePayload["response_format"] != nil,
	}, nil
}

func resolveAPIKey(target domain.ExecutionTarget) string {
	apiKey := strings.TrimSpace(target.APIKey)
	if apiKey == "" {
		if envName := strings.TrimSpace(target.APIKeyEnv); envName != "" {
			apiKey = strings.TrimSpace(os.Getenv(envName))
		}
	}
	return apiKey
}

// InitialEndpoint returns the frozen request URL before redirects.
func (p *PreparedGeneration) InitialEndpoint() string { return p.endpoint }

// Execute dispatches one attempt. The returned response is non-nil only when
// the successful envelope contains a usable first choice.
func (p *PreparedGeneration) Execute(ctx context.Context) AttemptOutcome {
	return p.execute(ctx, false)
}

// ExecuteWithEvidence performs one prepared attempt and retains HTTP and
// decoded-generation facts even when a later context check returns an error.
func (p *PreparedGeneration) ExecuteWithEvidence(ctx context.Context) AttemptOutcome {
	return p.execute(ctx, true)
}

func (p *PreparedGeneration) execute(ctx context.Context, captureEvidence bool) AttemptOutcome {
	requestContext := ctx
	if p.timeoutSeconds > 0 {
		var cancel context.CancelFunc
		requestContext, cancel = context.WithTimeout(ctx, time.Duration(p.timeoutSeconds)*time.Second)
		defer cancel()
	}
	httpReq, err := http.NewRequestWithContext(requestContext, http.MethodPost, p.endpoint, bytes.NewReader(p.payload))
	if err != nil {
		return AttemptOutcome{Err: fmt.Errorf("%w: failed to create request: %v", ErrRequestFailed, err)}
	}
	httpReq.Header.Set("Content-Type", "application/json")
	if p.apiKey != "" {
		httpReq.Header.Set("Authorization", "Bearer "+p.apiKey)
	}

	var (
		httpResp *http.Response
		outcome  AttemptOutcome
	)
	if captureEvidence {
		httpResp, err, outcome.Evidence = doWithEvidence(p.httpClient, httpReq)
		if outcome.Evidence.DispatchInvoked {
			requested := p.structuredOutputRequested
			outcome.StructuredOutputRequested = &requested
		}
	} else {
		httpResp, err = p.httpClient.Do(httpReq)
	}
	if httpResp != nil && (captureEvidence || err == nil) {
		outcome.StatusCode = httpResp.StatusCode
	}
	if err != nil {
		outcome.Err = &requestFailedError{cause: err}
		return outcome
	}
	receivedAt := time.Now()
	defer httpResp.Body.Close()
	outcome.StatusCode = httpResp.StatusCode

	if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
		outcome.RetryNotBefore = retryNotBefore(httpResp.Header.Values("Retry-After"), receivedAt)
		outcome.Err = providerHTTPErrorFromBody(
			httpResp.StatusCode,
			httpResp.ContentLength,
			httpResp.Body,
		)
		if err := observedContextError(requestContext, ctx, outcome.Err); err != nil {
			outcome.Err = err
		}
		return outcome
	}
	if httpResp.ContentLength > maxOpenAIChatResponseBytes {
		outcome.Err = openAIChatResponseTooLargeError()
		if captureEvidence {
			if contextErr := observedContextError(requestContext, ctx, outcome.Err); contextErr != nil {
				outcome.Err = contextErr
			}
		}
		return outcome
	}

	wireResp, err := decodeOpenAIChatResponse(httpResp.Body)
	if err != nil {
		outcome.Err = err
		if contextErr := observedContextError(requestContext, ctx, err); contextErr != nil {
			outcome.Err = contextErr
		}
		return outcome
	}
	outcome.ResponseDecoded = true
	outcome.Usage = domain.TokenUsage{
		PromptTokens: tokenCount(wireResp.Usage.PromptTokens), CompletionTokens: tokenCount(wireResp.Usage.CompletionTokens),
		TotalTokens: tokenCount(wireResp.Usage.TotalTokens), CachedTokens: tokenCount(wireResp.Usage.PromptTokensDetails.CachedTokens),
		CacheWriteTokens: tokenCount(wireResp.Usage.CacheWriteTokens),
	}
	outcome.UsagePresence = &domain.TokenUsagePresence{
		PromptTokens: wireResp.Usage.PromptTokens != nil, CompletionTokens: wireResp.Usage.CompletionTokens != nil,
		TotalTokens: wireResp.Usage.TotalTokens != nil, CachedTokens: wireResp.Usage.PromptTokensDetails.CachedTokens != nil,
		CacheWriteTokens: wireResp.Usage.CacheWriteTokens != nil,
	}

	if len(wireResp.Choices) == 0 {
		outcome.Err = fmt.Errorf("%w: no choices returned", ErrMalformedResponse)
		if captureEvidence {
			if contextErr := observedContextError(requestContext, ctx, outcome.Err); contextErr != nil {
				outcome.Err = contextErr
			}
		}
		return outcome
	}
	content := wireResp.Choices[0].Message.Content
	if content == nil {
		outcome.Err = fmt.Errorf("%w: first choice has missing message content", ErrMalformedResponse)
		if captureEvidence {
			if contextErr := observedContextError(requestContext, ctx, outcome.Err); contextErr != nil {
				outcome.Err = contextErr
			}
		}
		return outcome
	}
	outcome.Response = &domain.GenerateResponse{Content: *content, Usage: outcome.Usage, UsagePresence: outcome.UsagePresence}
	outcome.GenerationCompleted = true
	if captureEvidence {
		if contextErr := observedContextError(requestContext, ctx, nil); contextErr != nil {
			outcome.Err = contextErr
		}
	}
	return outcome
}

func observedContextError(attempt, caller context.Context, cause error) error {
	if err := caller.Err(); err != nil {
		return &requestFailedError{cause: err}
	}
	if err := attempt.Err(); err != nil {
		return &requestFailedError{cause: err}
	}
	if errors.Is(cause, context.Canceled) || errors.Is(cause, context.DeadlineExceeded) {
		return &requestFailedError{cause: cause}
	}
	return nil
}

func retryNotBefore(values []string, receivedAt time.Time) time.Time {
	if len(values) == 0 {
		return time.Time{}
	}
	value := strings.TrimSpace(values[0])
	if value == "" {
		return time.Time{}
	}
	digits := true
	for _, r := range value {
		if r < '0' || r > '9' {
			digits = false
			break
		}
	}
	if digits {
		seconds, err := strconv.ParseUint(value, 10, 64)
		if err != nil || seconds > uint64(math.MaxInt64/int64(time.Second)) {
			return receivedAt.Add(time.Duration(math.MaxInt64))
		}
		return receivedAt.Add(time.Duration(seconds) * time.Second)
	}
	if when, err := http.ParseTime(value); err == nil {
		if when.Before(receivedAt) {
			return receivedAt
		}
		return when
	}
	return time.Time{}
}

func tokenCount(value *int) int {
	if value == nil {
		return 0
	}
	return *value
}

type responseErrorReader struct {
	reader io.Reader
	err    error
}

func (r *responseErrorReader) Read(p []byte) (int, error) {
	n, err := r.reader.Read(p)
	if err != nil && !errors.Is(err, io.EOF) {
		r.err = err
	}
	return n, err
}

func decodeOpenAIChatResponse(body io.Reader) (openAIChatResponse, error) {
	observed := &responseErrorReader{reader: body}
	limited := &io.LimitedReader{
		R: observed,
		N: maxOpenAIChatResponseBytes + 1,
	}
	decoder := json.NewDecoder(limited)

	var response openAIChatResponse
	if err := decoder.Decode(&response); err != nil {
		if errors.Is(observed.err, context.Canceled) || errors.Is(observed.err, context.DeadlineExceeded) {
			return openAIChatResponse{}, &requestFailedError{cause: observed.err}
		}
		if limited.N == 0 {
			return openAIChatResponse{}, openAIChatResponseTooLargeError()
		}
		if observed.err != nil {
			return openAIChatResponse{}, &requestFailedError{cause: observed.err}
		}
		return openAIChatResponse{}, fmt.Errorf("%w: failed to decode response", ErrMalformedResponse)
	}
	if errors.Is(observed.err, context.Canceled) || errors.Is(observed.err, context.DeadlineExceeded) {
		return openAIChatResponse{}, &requestFailedError{cause: observed.err}
	}
	if limited.N == 0 {
		return openAIChatResponse{}, openAIChatResponseTooLargeError()
	}

	var trailing any
	if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) {
		if errors.Is(observed.err, context.Canceled) || errors.Is(observed.err, context.DeadlineExceeded) {
			return openAIChatResponse{}, &requestFailedError{cause: observed.err}
		}
		if limited.N == 0 {
			return openAIChatResponse{}, openAIChatResponseTooLargeError()
		}
		if observed.err != nil {
			return openAIChatResponse{}, &requestFailedError{cause: observed.err}
		}
		return openAIChatResponse{}, fmt.Errorf("%w: response contains trailing data", ErrMalformedResponse)
	}
	if errors.Is(observed.err, context.Canceled) || errors.Is(observed.err, context.DeadlineExceeded) {
		return openAIChatResponse{}, &requestFailedError{cause: observed.err}
	}
	if limited.N == 0 {
		return openAIChatResponse{}, openAIChatResponseTooLargeError()
	}
	if observed.err != nil {
		return openAIChatResponse{}, &requestFailedError{cause: observed.err}
	}

	return response, nil
}

func openAIChatResponseTooLargeError() error {
	return fmt.Errorf(
		"%w: response exceeds %d-byte limit",
		ErrMalformedResponse,
		maxOpenAIChatResponseBytes,
	)
}

func openAIChatRequestPayload(req domain.GenerateRequest) (map[string]any, error) {
	model := strings.TrimSpace(req.Target.Model)
	if model == "" {
		return nil, errors.New("model is required")
	}

	sessionID, err := domain.NormalizeSessionID(req.Prompt.SessionID)
	if err != nil {
		return nil, err
	}

	messages := make([]openAIChatRequestMessage, 0, len(req.Prompt.Messages))
	for _, msg := range req.Prompt.Messages {
		messages = append(messages, openAIChatRequestMessageFromRenderedMessage(msg))
	}

	out := map[string]any{"model": model, "messages": messages}
	if sessionID != "" {
		out["session_id"] = sessionID
	}
	if req.Target.Temperature != 0 || req.TargetPresence.Temperature {
		out["temperature"] = req.Target.Temperature
	}
	if req.Target.MaxTokens != 0 || req.TargetPresence.MaxTokens {
		out["max_tokens"] = req.Target.MaxTokens
	}
	if req.Target.TopP != 0 || req.TargetPresence.TopP {
		out["top_p"] = req.Target.TopP
	}
	if strings.TrimSpace(req.Target.ServiceTier) != "" {
		out["service_tier"] = req.Target.ServiceTier
	}
	if strings.TrimSpace(req.Target.ReasoningEffort) != "" {
		out["reasoning_effort"] = req.Target.ReasoningEffort
	}
	if req.StructuredOutput != nil {
		responseFormat, err := toOpenAIResponseFormat(req.StructuredOutput)
		if err != nil {
			return nil, err
		}
		if responseFormat != nil {
			out["response_format"] = responseFormat
		}
	}

	for key, value := range req.Target.ExtraParams {
		if key == "" {
			return nil, errors.New("extra_params key must not be empty")
		}
		if IsReservedOpenAIChatRequestField(key) {
			return nil, fmt.Errorf("extra_params key %q collides with reserved request field", key)
		}
		out[key] = value
	}

	return out, nil
}

// IsReservedOpenAIChatRequestField reports whether name is owned by the
// standard OpenAI-compatible chat request rather than extra parameters.
func IsReservedOpenAIChatRequestField(name string) bool {
	switch name {
	case "model",
		"session_id",
		"messages",
		"temperature",
		"max_tokens",
		"top_p",
		"service_tier",
		"reasoning_effort",
		"response_format":
		return true
	default:
		return false
	}
}

type openAIChatRequestMessage struct {
	Role    string `json:"role"`
	Content any    `json:"content"`
}

type openAIChatTextContentBlock struct {
	Type         string              `json:"type"`
	Text         string              `json:"text"`
	CacheControl *openAICacheControl `json:"cache_control,omitempty"`
}

type openAICacheControl struct {
	Type string `json:"type"`
	TTL  string `json:"ttl,omitempty"`
}

type openAIChatResponseMessage struct {
	Role    string  `json:"role"`
	Content *string `json:"content"`
}

type openAIChatResponse struct {
	Choices []struct {
		Message openAIChatResponseMessage `json:"message"`
	} `json:"choices"`
	Usage struct {
		PromptTokens        *int `json:"prompt_tokens"`
		CompletionTokens    *int `json:"completion_tokens"`
		TotalTokens         *int `json:"total_tokens"`
		PromptTokensDetails struct {
			CachedTokens *int `json:"cached_tokens"`
		} `json:"prompt_tokens_details"`
		CacheWriteTokens *int `json:"cache_write_tokens"`
	} `json:"usage"`
}

type openAIResponseFormat struct {
	Type       string                    `json:"type"`
	JSONSchema *openAIJSONSchemaEnvelope `json:"json_schema,omitempty"`
}

type openAIJSONSchemaEnvelope struct {
	Name   string `json:"name"`
	Strict bool   `json:"strict"`
	Schema any    `json:"schema"`
}

func openAIChatRequestMessageFromRenderedMessage(msg domain.RenderedMessage) openAIChatRequestMessage {
	wireMsg := openAIChatRequestMessage{
		Role:    msg.Role,
		Content: msg.Content,
	}
	if msg.CacheControl == nil {
		return wireMsg
	}

	wireMsg.Content = []openAIChatTextContentBlock{
		{
			Type: "text",
			Text: msg.Content,
			CacheControl: &openAICacheControl{
				Type: string(msg.CacheControl.Type),
				TTL:  msg.CacheControl.TTL,
			},
		},
	}
	return wireMsg
}

func toOpenAIResponseFormat(spec *domain.StructuredOutputSpec) (*openAIResponseFormat, error) {
	if spec == nil {
		return nil, nil
	}

	switch spec.Type {
	case domain.StructuredOutputJSONSchema:
		if spec.JSONSchema == nil {
			return nil, errors.New("json_schema structured output requires schema payload")
		}
		if strings.TrimSpace(spec.JSONSchema.Name) == "" {
			return nil, errors.New("json_schema structured output requires non-empty schema name")
		}
		if spec.JSONSchema.Schema == nil {
			return nil, errors.New("json_schema structured output requires schema document")
		}
		return &openAIResponseFormat{
			Type: "json_schema",
			JSONSchema: &openAIJSONSchemaEnvelope{
				Name:   spec.JSONSchema.Name,
				Strict: spec.JSONSchema.Strict,
				Schema: spec.JSONSchema.Schema,
			},
		}, nil
	default:
		return nil, fmt.Errorf("unsupported structured output type %q", spec.Type)
	}
}