308 lines
8.8 KiB
Go
308 lines
8.8 KiB
Go
package llm
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/url"
|
|
"os"
|
|
"strings"
|
|
"time"
|
|
|
|
"gitea.maximumdirect.net/eric/scriptorium/internal/defaults"
|
|
"gitea.maximumdirect.net/eric/scriptorium/internal/domain"
|
|
)
|
|
|
|
var (
|
|
ErrInvalidConfig = errors.New("invalid llm client configuration")
|
|
ErrInvalidRequest = errors.New("invalid generate request")
|
|
ErrRequestFailed = errors.New("llm request failed")
|
|
ErrUnexpectedStatus = errors.New("llm returned non-success status")
|
|
ErrMalformedResponse = errors.New("malformed llm response")
|
|
)
|
|
|
|
type OpenAICompatibleConfig struct {
|
|
BaseURL string
|
|
Model string
|
|
Timeout time.Duration
|
|
HTTPClient *http.Client
|
|
}
|
|
|
|
type OpenAICompatibleClient struct {
|
|
baseURL string
|
|
defaultModel string
|
|
timeout time.Duration
|
|
httpClient *http.Client
|
|
}
|
|
|
|
func NewOpenAICompatibleClient(cfg OpenAICompatibleConfig) (*OpenAICompatibleClient, error) {
|
|
baseURL := strings.TrimSpace(cfg.BaseURL)
|
|
if baseURL != "" {
|
|
if _, err := url.ParseRequestURI(baseURL); err != nil {
|
|
return nil, fmt.Errorf("%w: invalid base URL: %v", ErrInvalidConfig, err)
|
|
}
|
|
}
|
|
|
|
timeout := cfg.Timeout
|
|
if timeout <= 0 {
|
|
timeout = defaults.LLMRequestTimeoutDefault
|
|
}
|
|
|
|
var client *http.Client
|
|
if cfg.HTTPClient != nil {
|
|
client = cfg.HTTPClient
|
|
if client.Timeout == 0 {
|
|
client.Timeout = timeout
|
|
}
|
|
} else {
|
|
client = &http.Client{Timeout: timeout}
|
|
}
|
|
|
|
return &OpenAICompatibleClient{
|
|
baseURL: strings.TrimRight(baseURL, "/"),
|
|
defaultModel: cfg.Model,
|
|
timeout: timeout,
|
|
httpClient: client,
|
|
}, nil
|
|
}
|
|
|
|
func (c *OpenAICompatibleClient) Generate(ctx context.Context, req domain.GenerateRequest) (*domain.GenerateResponse, error) {
|
|
if req.Target.TimeoutSeconds < 0 {
|
|
return nil, fmt.Errorf("%w: timeout_seconds must be greater than or equal to 0", ErrInvalidRequest)
|
|
}
|
|
|
|
endpoint := strings.TrimSpace(req.Target.Endpoint)
|
|
if endpoint == "" {
|
|
endpoint = c.baseURL
|
|
}
|
|
if endpoint == "" {
|
|
return nil, fmt.Errorf("%w: endpoint is required", ErrInvalidRequest)
|
|
}
|
|
endpoint = strings.TrimRight(endpoint, "/") + defaults.OpenAIChatCompletionsPath
|
|
|
|
wireReq, err := openAIChatRequestFromGenerateRequest(req, c.defaultModel)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("%w: %v", ErrInvalidRequest, err)
|
|
}
|
|
|
|
payload, err := json.Marshal(wireReq)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("%w: failed to encode request: %v", ErrRequestFailed, err)
|
|
}
|
|
|
|
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(payload))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("%w: failed to create request: %v", ErrRequestFailed, err)
|
|
}
|
|
httpReq.Header.Set("Content-Type", "application/json")
|
|
if envName := strings.TrimSpace(req.Target.APIKeyEnv); envName != "" {
|
|
apiKey := strings.TrimSpace(os.Getenv(envName))
|
|
if apiKey == "" {
|
|
return nil, fmt.Errorf("%w: api key environment variable %q is not set", ErrInvalidRequest, envName)
|
|
}
|
|
httpReq.Header.Set("Authorization", "Bearer "+apiKey)
|
|
}
|
|
|
|
effectiveTimeout := c.timeout
|
|
if req.Target.TimeoutSeconds > 0 {
|
|
effectiveTimeout = time.Duration(req.Target.TimeoutSeconds) * time.Second
|
|
}
|
|
|
|
httpClient := c.httpClient
|
|
if httpClient == nil {
|
|
httpClient = &http.Client{Timeout: effectiveTimeout}
|
|
} else if httpClient.Timeout != effectiveTimeout {
|
|
cloned := *httpClient
|
|
cloned.Timeout = effectiveTimeout
|
|
httpClient = &cloned
|
|
}
|
|
|
|
httpResp, err := httpClient.Do(httpReq)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("%w: %v", ErrRequestFailed, err)
|
|
}
|
|
defer httpResp.Body.Close()
|
|
|
|
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
|
body, _ := io.ReadAll(io.LimitReader(httpResp.Body, 4096))
|
|
return nil, fmt.Errorf("%w: status=%d body=%q", ErrUnexpectedStatus, httpResp.StatusCode, strings.TrimSpace(string(body)))
|
|
}
|
|
|
|
var wireResp openAIChatResponse
|
|
if err := json.NewDecoder(httpResp.Body).Decode(&wireResp); err != nil {
|
|
return nil, fmt.Errorf("%w: failed to decode response: %v", ErrMalformedResponse, err)
|
|
}
|
|
|
|
if len(wireResp.Choices) == 0 {
|
|
return nil, fmt.Errorf("%w: no choices returned", ErrMalformedResponse)
|
|
}
|
|
content := wireResp.Choices[0].Message.Content
|
|
if content == "" {
|
|
return nil, fmt.Errorf("%w: first choice has empty message content", ErrMalformedResponse)
|
|
}
|
|
|
|
return &domain.GenerateResponse{
|
|
Content: content,
|
|
Usage: domain.TokenUsage{
|
|
PromptTokens: wireResp.Usage.PromptTokens,
|
|
CompletionTokens: wireResp.Usage.CompletionTokens,
|
|
TotalTokens: wireResp.Usage.TotalTokens,
|
|
CachedTokens: wireResp.Usage.PromptTokensDetails.CachedTokens,
|
|
CacheWriteTokens: wireResp.Usage.CacheWriteTokens,
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
func openAIChatRequestFromGenerateRequest(req domain.GenerateRequest, defaultModel string) (openAIChatRequest, error) {
|
|
model := strings.TrimSpace(req.Target.Model)
|
|
if model == "" {
|
|
model = strings.TrimSpace(defaultModel)
|
|
}
|
|
if model == "" {
|
|
return openAIChatRequest{}, errors.New("model is required")
|
|
}
|
|
|
|
wireReq := openAIChatRequest{
|
|
Model: model,
|
|
}
|
|
|
|
wireReq.Messages = make([]openAIChatRequestMessage, 0, len(req.Prompt.Messages))
|
|
for _, msg := range req.Prompt.Messages {
|
|
wireReq.Messages = append(wireReq.Messages, openAIChatRequestMessageFromRenderedMessage(msg))
|
|
}
|
|
|
|
if req.Target.Temperature != 0 {
|
|
wireReq.Temperature = &req.Target.Temperature
|
|
}
|
|
if req.Target.MaxTokens != 0 {
|
|
wireReq.MaxTokens = &req.Target.MaxTokens
|
|
}
|
|
if req.Target.TopP != 0 {
|
|
wireReq.TopP = &req.Target.TopP
|
|
}
|
|
if strings.TrimSpace(req.Target.ServiceTier) != "" {
|
|
wireReq.ServiceTier = req.Target.ServiceTier
|
|
}
|
|
if req.StructuredOutput != nil {
|
|
responseFormat, err := toOpenAIResponseFormat(req.StructuredOutput)
|
|
if err != nil {
|
|
return openAIChatRequest{}, err
|
|
}
|
|
wireReq.ResponseFormat = responseFormat
|
|
}
|
|
|
|
return wireReq, nil
|
|
}
|
|
|
|
type openAIChatRequest struct {
|
|
Model string `json:"model"`
|
|
Messages []openAIChatRequestMessage `json:"messages"`
|
|
Temperature *float64 `json:"temperature,omitempty"`
|
|
MaxTokens *int `json:"max_tokens,omitempty"`
|
|
TopP *float64 `json:"top_p,omitempty"`
|
|
ServiceTier string `json:"service_tier,omitempty"`
|
|
ResponseFormat *openAIResponseFormat `json:"response_format,omitempty"`
|
|
}
|
|
|
|
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)
|
|
}
|
|
}
|