notarius/internal/framework/llm/openai_compatible_client.go

package llm

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"

	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
)

const openAICompatibleProviderName = "openai-compatible"

// OpenAICompatibleClientConfig configures the direct HTTP structured-output adapter.
type OpenAICompatibleClientConfig struct {
	BaseURL        string
	Model          string
	APIKey         string
	MaxRetries     int
	HTTPClient     *http.Client
	RequestTimeout time.Duration
}

// OpenAICompatibleClient sends OpenAI-compatible chat-completion requests with
// response_format.type=json_schema.
type OpenAICompatibleClient struct {
	baseURL        string
	model          string
	apiKey         string
	maxRetries     int
	httpClient     *http.Client
	requestTimeout time.Duration
}

var _ contracts.StructuredLLMClient = (*OpenAICompatibleClient)(nil)

func NewOpenAICompatibleClient(cfg OpenAICompatibleClientConfig) (*OpenAICompatibleClient, error) {
	normalized, err := normalizeOpenAICompatibleConfig(cfg)
	if err != nil {
		return nil, err
	}

	client := normalized.HTTPClient
	if client == nil {
		client = http.DefaultClient
	}

	return &OpenAICompatibleClient{
		baseURL:        normalized.BaseURL,
		model:          normalized.Model,
		apiKey:         normalized.APIKey,
		maxRetries:     normalized.MaxRetries,
		httpClient:     client,
		requestTimeout: normalized.RequestTimeout,
	}, nil
}

func (c *OpenAICompatibleClient) CompleteStructured(
	ctx context.Context,
	req contracts.StructuredCompletionRequest,
	out any,
) (contracts.StructuredCompletionResponse, error) {
	if c == nil {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("openai-compatible client must not be nil")
	}
	if err := validateOutputTarget(out); err != nil {
		return contracts.StructuredCompletionResponse{}, err
	}

	model := strings.TrimSpace(req.Model)
	if model == "" {
		model = c.model
	}
	if model == "" {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("structured completion model must not be empty")
	}

	schemaName := strings.TrimSpace(req.ResponseSchemaName)
	if schemaName == "" {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("structured completion response schema name must not be empty")
	}
	if len(bytes.TrimSpace(req.ResponseSchema)) == 0 || !json.Valid(req.ResponseSchema) {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("structured completion response schema JSON must be valid")
	}

	messages, err := toOpenAICompatibleMessages(req.Messages)
	if err != nil {
		return contracts.StructuredCompletionResponse{}, err
	}

	endpoint := buildChatCompletionsURL(c.baseURL)
	var lastErr error
	for attempt := 0; attempt <= c.maxRetries; attempt++ {
		content, metadata, callErr := c.completeStructuredOnce(ctx, endpoint, model, messages, schemaName, req.ResponseSchema)
		if callErr == nil {
			if decodeErr := json.Unmarshal(content, out); decodeErr != nil {
				callErr = retryableError{err: fmt.Errorf("decode structured output: %w", decodeErr)}
			} else {
				return contracts.StructuredCompletionResponse{
					Content:          content,
					Provider:         openAICompatibleProviderName,
					Model:            firstNonEmpty(metadata.Model, model),
					PromptTokens:     metadata.PromptTokens,
					CompletionTokens: metadata.CompletionTokens,
					TotalTokens:      metadata.TotalTokens,
				}, nil
			}
		}

		if ctx.Err() != nil {
			return contracts.StructuredCompletionResponse{}, ctx.Err()
		}
		lastErr = c.redactError(callErr)
		if !canRetry(ctx, attempt, c.maxRetries, callErr) {
			return contracts.StructuredCompletionResponse{}, lastErr
		}
	}

	if lastErr == nil {
		lastErr = fmt.Errorf("structured completion failed")
	}
	return contracts.StructuredCompletionResponse{}, lastErr
}

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

type openAICompatibleRequest struct {
	Model          string                                `json:"model"`
	Messages       []openAICompatibleMessage             `json:"messages"`
	ResponseFormat openAICompatibleStructuredOutputShape `json:"response_format"`
}

type openAICompatibleStructuredOutputShape struct {
	Type       string                         `json:"type"`
	JSONSchema openAICompatibleSchemaEnvelope `json:"json_schema"`
}

type openAICompatibleSchemaEnvelope struct {
	Name   string          `json:"name"`
	Strict bool            `json:"strict"`
	Schema json.RawMessage `json:"schema"`
}

type openAICompatibleChatCompletionsResponse struct {
	Model   string `json:"model"`
	Choices []struct {
		Message struct {
			Content json.RawMessage `json:"content"`
		} `json:"message"`
	} `json:"choices"`
	Usage *openAICompatibleUsage `json:"usage,omitempty"`
}

type openAICompatibleUsage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

type openAICompatibleResponseMetadata struct {
	Model            string
	PromptTokens     int
	CompletionTokens int
	TotalTokens      int
}

func normalizeOpenAICompatibleConfig(cfg OpenAICompatibleClientConfig) (OpenAICompatibleClientConfig, error) {
	cfg.BaseURL = strings.TrimSpace(cfg.BaseURL)
	cfg.Model = strings.TrimSpace(cfg.Model)
	cfg.APIKey = strings.TrimSpace(cfg.APIKey)
	if cfg.MaxRetries < 0 {
		return OpenAICompatibleClientConfig{}, fmt.Errorf("max retries must be zero or greater")
	}
	if cfg.BaseURL == "" {
		return OpenAICompatibleClientConfig{}, fmt.Errorf("base URL must not be empty")
	}
	if _, err := url.ParseRequestURI(cfg.BaseURL); err != nil {
		return OpenAICompatibleClientConfig{}, fmt.Errorf("base URL must be valid: %w", err)
	}
	if cfg.Model == "" {
		return OpenAICompatibleClientConfig{}, fmt.Errorf("model must not be empty")
	}
	cfg.BaseURL = strings.TrimRight(cfg.BaseURL, "/")
	return cfg, nil
}

func (c *OpenAICompatibleClient) completeStructuredOnce(
	ctx context.Context,
	endpoint string,
	model string,
	messages []openAICompatibleMessage,
	responseSchemaName string,
	responseSchemaJSON json.RawMessage,
) (json.RawMessage, openAICompatibleResponseMetadata, error) {
	requestCtx := ctx
	var cancel context.CancelFunc
	if c.requestTimeout > 0 {
		requestCtx, cancel = context.WithTimeout(ctx, c.requestTimeout)
		defer cancel()
	}

	requestBody := openAICompatibleRequest{
		Model:    model,
		Messages: messages,
		ResponseFormat: openAICompatibleStructuredOutputShape{
			Type: "json_schema",
			JSONSchema: openAICompatibleSchemaEnvelope{
				Name:   responseSchemaName,
				Strict: true,
				Schema: responseSchemaJSON,
			},
		},
	}

	payload, err := json.Marshal(requestBody)
	if err != nil {
		return nil, openAICompatibleResponseMetadata{}, fmt.Errorf("marshal provider request: %w", err)
	}

	httpReq, err := http.NewRequestWithContext(requestCtx, http.MethodPost, endpoint, bytes.NewReader(payload))
	if err != nil {
		return nil, openAICompatibleResponseMetadata{}, fmt.Errorf("build provider request: %w", err)
	}
	httpReq.Header.Set("Content-Type", "application/json")
	if c.apiKey != "" {
		httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
	}

	httpResp, err := c.httpClient.Do(httpReq)
	if err != nil {
		return nil, openAICompatibleResponseMetadata{}, retryableError{err: fmt.Errorf("provider request failed: %w", err)}
	}
	defer func() {
		_ = httpResp.Body.Close()
	}()

	rawResp, err := io.ReadAll(httpResp.Body)
	if err != nil {
		return nil, openAICompatibleResponseMetadata{}, retryableError{err: fmt.Errorf("read provider response: %w", err)}
	}

	if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
		statusErr := parseProviderErrorBody(httpResp.StatusCode, rawResp)
		if httpResp.StatusCode == http.StatusTooManyRequests || httpResp.StatusCode >= 500 {
			return nil, openAICompatibleResponseMetadata{}, retryableError{err: statusErr}
		}
		return nil, openAICompatibleResponseMetadata{}, statusErr
	}

	return decodeChatCompletionsResponse(rawResp)
}

func toOpenAICompatibleMessages(messages []contracts.LLMMessage) ([]openAICompatibleMessage, error) {
	if len(messages) == 0 {
		return nil, fmt.Errorf("structured completion messages must not be empty")
	}

	result := make([]openAICompatibleMessage, len(messages))
	for i, message := range messages {
		role := strings.TrimSpace(message.Role)
		content := strings.TrimSpace(message.Content)
		if role == "" {
			return nil, fmt.Errorf("message[%d] role must not be empty", i)
		}
		if content == "" {
			return nil, fmt.Errorf("message[%d] content must not be empty", i)
		}
		result[i] = openAICompatibleMessage{
			Role:    role,
			Content: content,
		}
	}
	return result, nil
}

func buildChatCompletionsURL(baseURL string) string {
	return strings.TrimRight(baseURL, "/") + "/chat/completions"
}

func decodeChatCompletionsResponse(raw []byte) (json.RawMessage, openAICompatibleResponseMetadata, error) {
	var parsed openAICompatibleChatCompletionsResponse
	if err := json.Unmarshal(raw, &parsed); err != nil {
		return nil, openAICompatibleResponseMetadata{}, retryableError{err: fmt.Errorf("decode provider response envelope: %w", err)}
	}
	if len(parsed.Choices) == 0 {
		return nil, openAICompatibleResponseMetadata{}, retryableError{err: fmt.Errorf("provider response missing choices")}
	}

	content, err := extractAssistantContentJSON(parsed.Choices[0].Message.Content)
	if err != nil {
		return nil, openAICompatibleResponseMetadata{}, retryableError{err: err}
	}

	metadata := openAICompatibleResponseMetadata{
		Model: parsed.Model,
	}
	if parsed.Usage != nil {
		metadata.PromptTokens = parsed.Usage.PromptTokens
		metadata.CompletionTokens = parsed.Usage.CompletionTokens
		metadata.TotalTokens = parsed.Usage.TotalTokens
	}
	return content, metadata, nil
}

func extractAssistantContentJSON(raw json.RawMessage) (json.RawMessage, error) {
	trimmedRaw := bytes.TrimSpace(raw)
	if len(trimmedRaw) == 0 || bytes.Equal(trimmedRaw, []byte("null")) {
		return nil, fmt.Errorf("provider response missing assistant message content")
	}

	var textContent string
	if err := json.Unmarshal(trimmedRaw, &textContent); err == nil {
		textContent = strings.TrimSpace(textContent)
		if textContent == "" {
			return nil, fmt.Errorf("provider response assistant message content is empty")
		}
		if !json.Valid([]byte(textContent)) {
			return nil, fmt.Errorf("provider response assistant message content is not valid JSON")
		}
		return json.RawMessage(textContent), nil
	}

	if json.Valid(trimmedRaw) {
		return append(json.RawMessage(nil), trimmedRaw...), nil
	}
	return nil, fmt.Errorf("provider response assistant message content is not valid JSON")
}

func parseProviderErrorBody(status int, body []byte) error {
	trimmed := strings.TrimSpace(string(body))
	if trimmed == "" {
		return fmt.Errorf("provider returned status %d", status)
	}

	var payload map[string]any
	if err := json.Unmarshal(body, &payload); err == nil {
		if nested, ok := payload["error"].(map[string]any); ok {
			if msg, ok := nested["message"].(string); ok && strings.TrimSpace(msg) != "" {
				return fmt.Errorf("provider returned status %d: %s", status, strings.TrimSpace(msg))
			}
		}
		if msg, ok := payload["message"].(string); ok && strings.TrimSpace(msg) != "" {
			return fmt.Errorf("provider returned status %d: %s", status, strings.TrimSpace(msg))
		}
	}

	return fmt.Errorf("provider returned status %d: %s", status, trimmed)
}

func (c *OpenAICompatibleClient) redactError(err error) error {
	secrets := []string{c.apiKey}
	if c.apiKey != "" {
		secrets = append(secrets, "Bearer "+c.apiKey)
	}
	return ErrorWithSecretsRedacted(err, secrets)
}