audita/internal/framework/llm/instructor_client.go

package llm

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"reflect"
	"strings"
	"time"

	"gitea.maximumdirect.net/eric/audita/internal/framework/contracts"

	"github.com/jxnl/instructor-go/pkg/instructor"
	openai "github.com/sashabaranov/go-openai"
)

const (
	defaultMaxRetries = 3
)

type Mode string

const (
	ModeJSON     Mode = "json"
	ModeToolCall Mode = "tool_call"
)

type InstructorClientConfig struct {
	BaseURL        string
	Model          string
	APIKey         string
	MaxRetries     int
	Mode           Mode
	HTTPClient     *http.Client
	RequestTimeout time.Duration
}

// chatCompletionClient is intentionally narrow to avoid leaking provider types
// outside this adapter package.
type chatCompletionClient interface {
	CreateChatCompletion(ctx context.Context, request openai.ChatCompletionRequest, responseType any) (response openai.ChatCompletionResponse, err error)
}

// InstructorClient adapts instructor-go behind Audita's internal structured
// LLM interface.
type InstructorClient struct {
	cfg    InstructorClientConfig
	client chatCompletionClient
}

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

func NewInstructorClient(cfg InstructorClientConfig) (*InstructorClient, error) {
	normalized, err := normalizeConfig(cfg)
	if err != nil {
		return nil, err
	}

	openaiConfig := openai.DefaultConfig(normalized.APIKey)
	openaiConfig.BaseURL = normalized.BaseURL
	openaiConfig.HTTPClient = resolvedHTTPClient(normalized.HTTPClient, normalized.RequestTimeout)

	mode, err := toInstructorMode(normalized.Mode)
	if err != nil {
		return nil, err
	}

	client := instructor.FromOpenAI(
		openai.NewClientWithConfig(openaiConfig),
		instructor.WithMode(mode),
		instructor.WithMaxRetries(normalized.MaxRetries),
	)

	return &InstructorClient{
		cfg:    normalized,
		client: client,
	}, nil
}

func (c *InstructorClient) CompleteStructured(
	ctx context.Context,
	req contracts.StructuredCompletionRequest,
	out any,
) (contracts.StructuredCompletionResponse, error) {
	if err := validateOutputTarget(out); err != nil {
		return contracts.StructuredCompletionResponse{}, err
	}

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

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

	chatRequest := openai.ChatCompletionRequest{
		Model:    model,
		Messages: messages,
	}

	resp, err := c.client.CreateChatCompletion(ctx, chatRequest, out)
	if err != nil {
		return contracts.StructuredCompletionResponse{}, sanitizeError(err, c.cfg.APIKey)
	}

	content, err := json.Marshal(out)
	if err != nil {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("marshal structured completion output: %w", err)
	}

	return contracts.StructuredCompletionResponse{
		Content:          content,
		Provider:         "openai-compatible",
		Model:            model,
		PromptTokens:     resp.Usage.PromptTokens,
		CompletionTokens: resp.Usage.CompletionTokens,
		TotalTokens:      resp.Usage.TotalTokens,
	}, nil
}

func normalizeConfig(cfg InstructorClientConfig) (InstructorClientConfig, error) {
	cfg.BaseURL = strings.TrimSpace(cfg.BaseURL)
	cfg.Model = strings.TrimSpace(cfg.Model)
	cfg.APIKey = strings.TrimSpace(cfg.APIKey)
	if cfg.MaxRetries == 0 {
		cfg.MaxRetries = defaultMaxRetries
	}
	if cfg.MaxRetries < 0 {
		return InstructorClientConfig{}, fmt.Errorf("max retries must be zero or greater")
	}
	if cfg.BaseURL == "" {
		return InstructorClientConfig{}, fmt.Errorf("base URL must not be empty")
	}
	if cfg.Model == "" {
		return InstructorClientConfig{}, fmt.Errorf("model must not be empty")
	}
	if cfg.Mode == "" {
		cfg.Mode = ModeJSON
	}
	return cfg, nil
}

func toInstructorMode(mode Mode) (instructor.Mode, error) {
	switch mode {
	case ModeJSON:
		return instructor.ModeJSON, nil
	case ModeToolCall:
		return instructor.ModeToolCall, nil
	default:
		return "", fmt.Errorf("unsupported LLM mode %q", mode)
	}
}

func resolvedHTTPClient(base *http.Client, timeout time.Duration) *http.Client {
	if base == nil {
		if timeout <= 0 {
			return http.DefaultClient
		}
		return &http.Client{Timeout: timeout}
	}

	if timeout <= 0 {
		return base
	}

	cloned := *base
	cloned.Timeout = timeout
	return &cloned
}

func toOpenAIMessages(messages []contracts.LLMMessage) ([]openai.ChatCompletionMessage, error) {
	result := make([]openai.ChatCompletionMessage, 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] = openai.ChatCompletionMessage{
			Role:    role,
			Content: content,
		}
	}
	return result, nil
}

func validateOutputTarget(out any) error {
	if out == nil {
		return fmt.Errorf("output target must not be nil")
	}
	value := reflect.ValueOf(out)
	if value.Kind() != reflect.Ptr || value.IsNil() {
		return fmt.Errorf("output target must be a non-nil pointer")
	}
	return nil
}

func sanitizeError(err error, apiKey string) error {
	if err == nil {
		return nil
	}
	msg := err.Error()
	key := strings.TrimSpace(apiKey)
	if key != "" {
		msg = strings.ReplaceAll(msg, key, "[REDACTED]")
		msg = strings.ReplaceAll(msg, "Bearer "+key, "Bearer [REDACTED]")
	}
	return fmt.Errorf("%s", msg)
}