Add OpenAI-compatible LLM runtime
This commit is contained in:
55
internal/framework/llm/client_common.go
Normal file
55
internal/framework/llm/client_common.go
Normal file
@@ -0,0 +1,55 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"reflect"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
type retryableError struct {
|
||||||
|
err error
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e retryableError) Error() string {
|
||||||
|
if e.err == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return e.err.Error()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e retryableError) Unwrap() error {
|
||||||
|
return e.err
|
||||||
|
}
|
||||||
|
|
||||||
|
func validateOutputTarget(out any) error {
|
||||||
|
if out == nil {
|
||||||
|
return fmt.Errorf("structured completion output target must be a non-nil pointer")
|
||||||
|
}
|
||||||
|
value := reflect.ValueOf(out)
|
||||||
|
if value.Kind() != reflect.Pointer || value.IsNil() {
|
||||||
|
return fmt.Errorf("structured completion output target must be a non-nil pointer")
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func canRetry(ctx context.Context, attempt int, maxRetries int, err error) bool {
|
||||||
|
if attempt >= maxRetries {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if ctx.Err() != nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
var retryable retryableError
|
||||||
|
return errors.As(err, &retryable)
|
||||||
|
}
|
||||||
|
|
||||||
|
func firstNonEmpty(values ...string) string {
|
||||||
|
for _, value := range values {
|
||||||
|
if trimmed := strings.TrimSpace(value); trimmed != "" {
|
||||||
|
return trimmed
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ""
|
||||||
|
}
|
||||||
364
internal/framework/llm/openai_compatible_client.go
Normal file
364
internal/framework/llm/openai_compatible_client.go
Normal file
@@ -0,0 +1,364 @@
|
|||||||
|
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)
|
||||||
|
}
|
||||||
494
internal/framework/llm/openai_compatible_client_test.go
Normal file
494
internal/framework/llm/openai_compatible_client_test.go
Normal file
@@ -0,0 +1,494 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"sync/atomic"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
||||||
|
)
|
||||||
|
|
||||||
|
type testArtifact struct {
|
||||||
|
Value string `json:"value"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestNewOpenAICompatibleClientValidation(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
cfg OpenAICompatibleClientConfig
|
||||||
|
want string
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "empty base URL",
|
||||||
|
cfg: OpenAICompatibleClientConfig{
|
||||||
|
BaseURL: " ",
|
||||||
|
Model: "model",
|
||||||
|
},
|
||||||
|
want: "base URL",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "invalid base URL",
|
||||||
|
cfg: OpenAICompatibleClientConfig{
|
||||||
|
BaseURL: "://bad",
|
||||||
|
Model: "model",
|
||||||
|
},
|
||||||
|
want: "base URL",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "empty model",
|
||||||
|
cfg: OpenAICompatibleClientConfig{
|
||||||
|
BaseURL: "https://example.test/v1",
|
||||||
|
Model: " ",
|
||||||
|
},
|
||||||
|
want: "model",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "negative retries",
|
||||||
|
cfg: OpenAICompatibleClientConfig{
|
||||||
|
BaseURL: "https://example.test/v1",
|
||||||
|
Model: "model",
|
||||||
|
MaxRetries: -1,
|
||||||
|
},
|
||||||
|
want: "max retries",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range tests {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
_, err := NewOpenAICompatibleClient(tc.cfg)
|
||||||
|
if err == nil || !strings.Contains(err.Error(), tc.want) {
|
||||||
|
t.Fatalf("expected error containing %q, got %v", tc.want, err)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAICompatibleClientSuccessfulStructuredCompletion(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
_, _ = io.WriteString(w, `{
|
||||||
|
"model":"provider-model",
|
||||||
|
"choices":[{"message":{"content":"{\"value\":\"ok\"}"}}],
|
||||||
|
"usage":{"prompt_tokens":11,"completion_tokens":7,"total_tokens":18}
|
||||||
|
}`)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client := newTestClient(t, server.URL, "default-model", 0)
|
||||||
|
var out testArtifact
|
||||||
|
resp, err := client.CompleteStructured(context.Background(), validStructuredRequest(""), &out)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CompleteStructured: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if out.Value != "ok" {
|
||||||
|
t.Fatalf("unexpected decoded output: %+v", out)
|
||||||
|
}
|
||||||
|
if string(resp.Content) != `{"value":"ok"}` {
|
||||||
|
t.Fatalf("unexpected raw content: %s", resp.Content)
|
||||||
|
}
|
||||||
|
if resp.Provider != openAICompatibleProviderName {
|
||||||
|
t.Fatalf("unexpected provider: %q", resp.Provider)
|
||||||
|
}
|
||||||
|
if resp.Model != "provider-model" {
|
||||||
|
t.Fatalf("unexpected model: %q", resp.Model)
|
||||||
|
}
|
||||||
|
if resp.PromptTokens != 11 || resp.CompletionTokens != 7 || resp.TotalTokens != 18 {
|
||||||
|
t.Fatalf("unexpected token metadata: %+v", resp)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAICompatibleClientRequestBodyIncludesStructuredOutputShape(t *testing.T) {
|
||||||
|
var seenPath string
|
||||||
|
var seenAuthorization string
|
||||||
|
var seenReq map[string]any
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
seenPath = r.URL.Path
|
||||||
|
seenAuthorization = r.Header.Get("Authorization")
|
||||||
|
if err := json.NewDecoder(r.Body).Decode(&seenReq); err != nil {
|
||||||
|
t.Fatalf("decode request: %v", err)
|
||||||
|
}
|
||||||
|
_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"{\"value\":\"ok\"}"}}]}`)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client, err := NewOpenAICompatibleClient(OpenAICompatibleClientConfig{
|
||||||
|
BaseURL: server.URL + "/v1",
|
||||||
|
Model: "default-model",
|
||||||
|
APIKey: "secret-key",
|
||||||
|
MaxRetries: 0,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewOpenAICompatibleClient: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var out testArtifact
|
||||||
|
_, err = client.CompleteStructured(context.Background(), validStructuredRequest("request-model"), &out)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CompleteStructured: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if seenPath != "/v1/chat/completions" {
|
||||||
|
t.Fatalf("unexpected request path: %q", seenPath)
|
||||||
|
}
|
||||||
|
if seenAuthorization != "Bearer secret-key" {
|
||||||
|
t.Fatalf("unexpected authorization header: %q", seenAuthorization)
|
||||||
|
}
|
||||||
|
if seenReq["model"] != "request-model" {
|
||||||
|
t.Fatalf("unexpected model: %v", seenReq["model"])
|
||||||
|
}
|
||||||
|
|
||||||
|
messages, ok := seenReq["messages"].([]any)
|
||||||
|
if !ok || len(messages) != 1 {
|
||||||
|
t.Fatalf("unexpected messages: %#v", seenReq["messages"])
|
||||||
|
}
|
||||||
|
message, ok := messages[0].(map[string]any)
|
||||||
|
if !ok {
|
||||||
|
t.Fatalf("unexpected message shape: %#v", messages[0])
|
||||||
|
}
|
||||||
|
if message["role"] != "user" || message["content"] != "extract this" {
|
||||||
|
t.Fatalf("unexpected message: %#v", message)
|
||||||
|
}
|
||||||
|
|
||||||
|
responseFormat, ok := seenReq["response_format"].(map[string]any)
|
||||||
|
if !ok {
|
||||||
|
t.Fatalf("expected response_format object, got %T", seenReq["response_format"])
|
||||||
|
}
|
||||||
|
if responseFormat["type"] != "json_schema" {
|
||||||
|
t.Fatalf("unexpected response_format.type: %v", responseFormat["type"])
|
||||||
|
}
|
||||||
|
jsonSchema, ok := responseFormat["json_schema"].(map[string]any)
|
||||||
|
if !ok {
|
||||||
|
t.Fatalf("expected response_format.json_schema object, got %T", responseFormat["json_schema"])
|
||||||
|
}
|
||||||
|
if jsonSchema["name"] != "test_artifact" {
|
||||||
|
t.Fatalf("unexpected schema name: %v", jsonSchema["name"])
|
||||||
|
}
|
||||||
|
if jsonSchema["strict"] != true {
|
||||||
|
t.Fatalf("expected strict=true, got %v", jsonSchema["strict"])
|
||||||
|
}
|
||||||
|
schema, ok := jsonSchema["schema"].(map[string]any)
|
||||||
|
if !ok {
|
||||||
|
t.Fatalf("expected schema object, got %T", jsonSchema["schema"])
|
||||||
|
}
|
||||||
|
if schema["type"] != "object" {
|
||||||
|
t.Fatalf("unexpected schema: %#v", schema)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAICompatibleClientDefaultModelFallbackAndOverride(t *testing.T) {
|
||||||
|
var seenModels []string
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
var req map[string]any
|
||||||
|
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||||
|
t.Fatalf("decode request: %v", err)
|
||||||
|
}
|
||||||
|
seenModels = append(seenModels, fmt.Sprint(req["model"]))
|
||||||
|
_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"{\"value\":\"ok\"}"}}]}`)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client := newTestClient(t, server.URL, "default-model", 0)
|
||||||
|
var first testArtifact
|
||||||
|
if _, err := client.CompleteStructured(context.Background(), validStructuredRequest(""), &first); err != nil {
|
||||||
|
t.Fatalf("first CompleteStructured: %v", err)
|
||||||
|
}
|
||||||
|
var second testArtifact
|
||||||
|
if _, err := client.CompleteStructured(context.Background(), validStructuredRequest("override-model"), &second); err != nil {
|
||||||
|
t.Fatalf("second CompleteStructured: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(seenModels) != 2 || seenModels[0] != "default-model" || seenModels[1] != "override-model" {
|
||||||
|
t.Fatalf("unexpected models: %v", seenModels)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAICompatibleClientInvalidOutputTarget(t *testing.T) {
|
||||||
|
client := newTestClient(t, "https://example.test/v1", "default-model", 0)
|
||||||
|
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
out any
|
||||||
|
}{
|
||||||
|
{name: "nil", out: nil},
|
||||||
|
{name: "non-pointer", out: testArtifact{}},
|
||||||
|
{name: "nil pointer", out: (*testArtifact)(nil)},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range tests {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
_, err := client.CompleteStructured(context.Background(), validStructuredRequest(""), tc.out)
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "output target") {
|
||||||
|
t.Fatalf("expected output target error, got %v", err)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAICompatibleClientMissingAndInvalidSchema(t *testing.T) {
|
||||||
|
client := newTestClient(t, "https://example.test/v1", "default-model", 0)
|
||||||
|
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
mutate func(*contracts.StructuredCompletionRequest)
|
||||||
|
want string
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "missing schema name",
|
||||||
|
mutate: func(req *contracts.StructuredCompletionRequest) {
|
||||||
|
req.ResponseSchemaName = " "
|
||||||
|
},
|
||||||
|
want: "schema name",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "missing schema JSON",
|
||||||
|
mutate: func(req *contracts.StructuredCompletionRequest) {
|
||||||
|
req.ResponseSchema = nil
|
||||||
|
},
|
||||||
|
want: "schema JSON",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "invalid schema JSON",
|
||||||
|
mutate: func(req *contracts.StructuredCompletionRequest) {
|
||||||
|
req.ResponseSchema = json.RawMessage(`{"type":`)
|
||||||
|
},
|
||||||
|
want: "schema JSON",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range tests {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
req := validStructuredRequest("")
|
||||||
|
tc.mutate(&req)
|
||||||
|
var out testArtifact
|
||||||
|
_, err := client.CompleteStructured(context.Background(), req, &out)
|
||||||
|
if err == nil || !strings.Contains(err.Error(), tc.want) {
|
||||||
|
t.Fatalf("expected error containing %q, got %v", tc.want, err)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAICompatibleClientRejectsEmptyMessages(t *testing.T) {
|
||||||
|
client := newTestClient(t, "https://example.test/v1", "default-model", 0)
|
||||||
|
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
mutate func(*contracts.StructuredCompletionRequest)
|
||||||
|
want string
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "no messages",
|
||||||
|
mutate: func(req *contracts.StructuredCompletionRequest) {
|
||||||
|
req.Messages = nil
|
||||||
|
},
|
||||||
|
want: "messages",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "empty role",
|
||||||
|
mutate: func(req *contracts.StructuredCompletionRequest) {
|
||||||
|
req.Messages[0].Role = " "
|
||||||
|
},
|
||||||
|
want: "role",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "empty content",
|
||||||
|
mutate: func(req *contracts.StructuredCompletionRequest) {
|
||||||
|
req.Messages[0].Content = " "
|
||||||
|
},
|
||||||
|
want: "content",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range tests {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
req := validStructuredRequest("")
|
||||||
|
tc.mutate(&req)
|
||||||
|
var out testArtifact
|
||||||
|
_, err := client.CompleteStructured(context.Background(), req, &out)
|
||||||
|
if err == nil || !strings.Contains(err.Error(), tc.want) {
|
||||||
|
t.Fatalf("expected error containing %q, got %v", tc.want, err)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAICompatibleClientProviderNon2xxBehavior(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusBadRequest)
|
||||||
|
_, _ = io.WriteString(w, `{"error":{"message":"bad request"}}`)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client := newTestClient(t, server.URL, "default-model", 0)
|
||||||
|
var out testArtifact
|
||||||
|
_, err := client.CompleteStructured(context.Background(), validStructuredRequest(""), &out)
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "status 400: bad request") {
|
||||||
|
t.Fatalf("expected provider status error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAICompatibleClientRetries429And5xx(t *testing.T) {
|
||||||
|
var attempts atomic.Int32
|
||||||
|
statuses := []int{http.StatusTooManyRequests, http.StatusInternalServerError, http.StatusOK}
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
attempt := int(attempts.Add(1)) - 1
|
||||||
|
if statuses[attempt] != http.StatusOK {
|
||||||
|
w.WriteHeader(statuses[attempt])
|
||||||
|
_, _ = io.WriteString(w, `{"error":{"message":"try again"}}`)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"{\"value\":\"ok\"}"}}]}`)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client := newTestClient(t, server.URL, "default-model", 2)
|
||||||
|
var out testArtifact
|
||||||
|
if _, err := client.CompleteStructured(context.Background(), validStructuredRequest(""), &out); err != nil {
|
||||||
|
t.Fatalf("CompleteStructured: %v", err)
|
||||||
|
}
|
||||||
|
if attempts.Load() != 3 {
|
||||||
|
t.Fatalf("expected 3 attempts, got %d", attempts.Load())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAICompatibleClientRetriesMalformedResponses(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
firstBody string
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "malformed provider envelope",
|
||||||
|
firstBody: `{"choices":[]}`,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "malformed assistant JSON",
|
||||||
|
firstBody: `{"choices":[{"message":{"content":"{"}}]}`,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range tests {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
var attempts atomic.Int32
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if attempts.Add(1) == 1 {
|
||||||
|
_, _ = io.WriteString(w, tc.firstBody)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"{\"value\":\"ok\"}"}}]}`)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client := newTestClient(t, server.URL, "default-model", 1)
|
||||||
|
var out testArtifact
|
||||||
|
if _, err := client.CompleteStructured(context.Background(), validStructuredRequest(""), &out); err != nil {
|
||||||
|
t.Fatalf("CompleteStructured: %v", err)
|
||||||
|
}
|
||||||
|
if attempts.Load() != 2 {
|
||||||
|
t.Fatalf("expected 2 attempts, got %d", attempts.Load())
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAICompatibleClientNoRetryForNonRetryable4xx(t *testing.T) {
|
||||||
|
var attempts atomic.Int32
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
attempts.Add(1)
|
||||||
|
w.WriteHeader(http.StatusForbidden)
|
||||||
|
_, _ = io.WriteString(w, `{"error":{"message":"forbidden"}}`)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client := newTestClient(t, server.URL, "default-model", 3)
|
||||||
|
var out testArtifact
|
||||||
|
_, err := client.CompleteStructured(context.Background(), validStructuredRequest(""), &out)
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "status 403") {
|
||||||
|
t.Fatalf("expected forbidden error, got %v", err)
|
||||||
|
}
|
||||||
|
if attempts.Load() != 1 {
|
||||||
|
t.Fatalf("expected 1 attempt, got %d", attempts.Load())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAICompatibleClientProviderErrorRedactsAPIKey(t *testing.T) {
|
||||||
|
const apiKey = "secret-api-key"
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusInternalServerError)
|
||||||
|
_, _ = io.WriteString(w, `{"error":{"message":"Bearer secret-api-key failed for secret-api-key"}}`)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client, err := NewOpenAICompatibleClient(OpenAICompatibleClientConfig{
|
||||||
|
BaseURL: server.URL,
|
||||||
|
Model: "default-model",
|
||||||
|
APIKey: apiKey,
|
||||||
|
MaxRetries: 0,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewOpenAICompatibleClient: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var out testArtifact
|
||||||
|
_, err = client.CompleteStructured(context.Background(), validStructuredRequest(""), &out)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatalf("expected provider error")
|
||||||
|
}
|
||||||
|
if strings.Contains(err.Error(), apiKey) || strings.Contains(err.Error(), "Bearer "+apiKey) {
|
||||||
|
t.Fatalf("expected API key to be redacted, got %q", err.Error())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAICompatibleClientRespectsContextCancellation(t *testing.T) {
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
cancel()
|
||||||
|
|
||||||
|
client := newTestClient(t, "https://example.test/v1", "default-model", 1)
|
||||||
|
var out testArtifact
|
||||||
|
_, err := client.CompleteStructured(ctx, validStructuredRequest(""), &out)
|
||||||
|
if !errors.Is(err, context.Canceled) {
|
||||||
|
t.Fatalf("expected context canceled, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func newTestClient(t *testing.T, baseURL string, model string, maxRetries int) *OpenAICompatibleClient {
|
||||||
|
t.Helper()
|
||||||
|
client, err := NewOpenAICompatibleClient(OpenAICompatibleClientConfig{
|
||||||
|
BaseURL: baseURL,
|
||||||
|
Model: model,
|
||||||
|
MaxRetries: maxRetries,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewOpenAICompatibleClient: %v", err)
|
||||||
|
}
|
||||||
|
return client
|
||||||
|
}
|
||||||
|
|
||||||
|
func validStructuredRequest(model string) contracts.StructuredCompletionRequest {
|
||||||
|
return contracts.StructuredCompletionRequest{
|
||||||
|
Messages: []contracts.LLMMessage{
|
||||||
|
{Role: " user ", Content: " extract this "},
|
||||||
|
},
|
||||||
|
Model: model,
|
||||||
|
ResponseSchemaName: " test_artifact ",
|
||||||
|
ResponseSchema: testResponseSchema(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func testResponseSchema() json.RawMessage {
|
||||||
|
return json.RawMessage(`{
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"value": {"type": "string"}
|
||||||
|
},
|
||||||
|
"required": ["value"],
|
||||||
|
"additionalProperties": false
|
||||||
|
}`)
|
||||||
|
}
|
||||||
122
internal/framework/llm/scheduler.go
Normal file
122
internal/framework/llm/scheduler.go
Normal file
@@ -0,0 +1,122 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"sync"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Scheduler bounds concurrent LLM backend calls.
|
||||||
|
type Scheduler struct {
|
||||||
|
maxConcurrency int
|
||||||
|
mu sync.Mutex
|
||||||
|
inFlight int
|
||||||
|
queue []*waiter
|
||||||
|
}
|
||||||
|
|
||||||
|
type waiter struct {
|
||||||
|
ready chan struct{}
|
||||||
|
queued bool
|
||||||
|
granted bool
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewScheduler creates a scheduler with a fixed concurrency limit.
|
||||||
|
func NewScheduler(maxConcurrency int) (*Scheduler, error) {
|
||||||
|
if maxConcurrency <= 0 {
|
||||||
|
return nil, fmt.Errorf("max concurrency must be greater than zero")
|
||||||
|
}
|
||||||
|
return &Scheduler{
|
||||||
|
maxConcurrency: maxConcurrency,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Acquire blocks until a permit is available or the context is canceled.
|
||||||
|
// The returned release function is safe to call multiple times.
|
||||||
|
func (s *Scheduler) Acquire(ctx context.Context) (func(), error) {
|
||||||
|
if s == nil {
|
||||||
|
return nil, fmt.Errorf("scheduler must not be nil")
|
||||||
|
}
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
s.mu.Lock()
|
||||||
|
if s.inFlight < s.maxConcurrency && len(s.queue) == 0 {
|
||||||
|
s.inFlight++
|
||||||
|
s.mu.Unlock()
|
||||||
|
return s.releaseFunc(), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
w := &waiter{
|
||||||
|
ready: make(chan struct{}),
|
||||||
|
queued: true,
|
||||||
|
}
|
||||||
|
s.queue = append(s.queue, w)
|
||||||
|
s.mu.Unlock()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case <-w.ready:
|
||||||
|
return s.releaseFunc(), nil
|
||||||
|
case <-ctx.Done():
|
||||||
|
s.mu.Lock()
|
||||||
|
if w.queued {
|
||||||
|
s.removeQueuedWaiterLocked(w)
|
||||||
|
s.mu.Unlock()
|
||||||
|
return nil, ctx.Err()
|
||||||
|
}
|
||||||
|
if w.granted {
|
||||||
|
s.inFlight--
|
||||||
|
s.grantQueuedLocked()
|
||||||
|
}
|
||||||
|
s.mu.Unlock()
|
||||||
|
return nil, ctx.Err()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Run acquires a permit, executes fn, and releases the permit.
|
||||||
|
func (s *Scheduler) Run(ctx context.Context, fn func(context.Context) error) error {
|
||||||
|
if fn == nil {
|
||||||
|
return fmt.Errorf("scheduler function must not be nil")
|
||||||
|
}
|
||||||
|
release, err := s.Acquire(ctx)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
defer release()
|
||||||
|
return fn(ctx)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Scheduler) releaseFunc() func() {
|
||||||
|
var once sync.Once
|
||||||
|
return func() {
|
||||||
|
once.Do(func() {
|
||||||
|
s.mu.Lock()
|
||||||
|
if s.inFlight > 0 {
|
||||||
|
s.inFlight--
|
||||||
|
s.grantQueuedLocked()
|
||||||
|
}
|
||||||
|
s.mu.Unlock()
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Scheduler) grantQueuedLocked() {
|
||||||
|
for s.inFlight < s.maxConcurrency && len(s.queue) > 0 {
|
||||||
|
w := s.queue[0]
|
||||||
|
s.queue = s.queue[1:]
|
||||||
|
w.queued = false
|
||||||
|
w.granted = true
|
||||||
|
s.inFlight++
|
||||||
|
close(w.ready)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Scheduler) removeQueuedWaiterLocked(target *waiter) {
|
||||||
|
for i, w := range s.queue {
|
||||||
|
if w == target {
|
||||||
|
w.queued = false
|
||||||
|
s.queue = append(s.queue[:i], s.queue[i+1:]...)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
164
internal/framework/llm/scheduler_test.go
Normal file
164
internal/framework/llm/scheduler_test.go
Normal file
@@ -0,0 +1,164 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"runtime"
|
||||||
|
"sync"
|
||||||
|
"sync/atomic"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestNewSchedulerValidation(t *testing.T) {
|
||||||
|
if _, err := NewScheduler(0); err == nil {
|
||||||
|
t.Fatalf("expected validation error")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSchedulerMaxConcurrency(t *testing.T) {
|
||||||
|
s, err := NewScheduler(2)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewScheduler: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var inFlight int32
|
||||||
|
var maxInFlight int32
|
||||||
|
release := make(chan struct{})
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
|
||||||
|
for i := 0; i < 12; i++ {
|
||||||
|
wg.Add(1)
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
runErr := s.Run(context.Background(), func(context.Context) error {
|
||||||
|
current := atomic.AddInt32(&inFlight, 1)
|
||||||
|
for {
|
||||||
|
seen := atomic.LoadInt32(&maxInFlight)
|
||||||
|
if current <= seen || atomic.CompareAndSwapInt32(&maxInFlight, seen, current) {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
<-release
|
||||||
|
atomic.AddInt32(&inFlight, -1)
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
if runErr != nil {
|
||||||
|
t.Errorf("Run error: %v", runErr)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
|
waitForAtomicAtLeast(t, &maxInFlight, 2)
|
||||||
|
close(release)
|
||||||
|
wg.Wait()
|
||||||
|
|
||||||
|
if got := atomic.LoadInt32(&maxInFlight); got > 2 {
|
||||||
|
t.Fatalf("expected max in-flight <= 2, got %d", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSchedulerCancellationWhileQueued(t *testing.T) {
|
||||||
|
s, err := NewScheduler(1)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewScheduler: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
release, err := s.Acquire(context.Background())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Acquire: %v", err)
|
||||||
|
}
|
||||||
|
defer release()
|
||||||
|
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
errCh := make(chan error, 1)
|
||||||
|
go func() {
|
||||||
|
_, acquireErr := s.Acquire(ctx)
|
||||||
|
errCh <- acquireErr
|
||||||
|
}()
|
||||||
|
|
||||||
|
waitForQueueDepth(t, s, 1)
|
||||||
|
cancel()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case acquireErr := <-errCh:
|
||||||
|
if !errors.Is(acquireErr, context.Canceled) {
|
||||||
|
t.Fatalf("expected context canceled, got %v", acquireErr)
|
||||||
|
}
|
||||||
|
case <-time.After(time.Second):
|
||||||
|
t.Fatalf("timed out waiting for queued acquire to cancel")
|
||||||
|
}
|
||||||
|
|
||||||
|
release()
|
||||||
|
if err := s.Run(context.Background(), func(context.Context) error { return nil }); err != nil {
|
||||||
|
t.Fatalf("expected scheduler to accept work after cancellation, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSchedulerReleaseFunctionIsIdempotent(t *testing.T) {
|
||||||
|
s, err := NewScheduler(1)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewScheduler: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
release, err := s.Acquire(context.Background())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Acquire: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
release()
|
||||||
|
release()
|
||||||
|
|
||||||
|
if err := s.Run(context.Background(), func(context.Context) error { return nil }); err != nil {
|
||||||
|
t.Fatalf("expected permit to be released once, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSchedulerRunReleasesPermitAfterError(t *testing.T) {
|
||||||
|
s, err := NewScheduler(1)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewScheduler: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
expected := errors.New("failed")
|
||||||
|
err = s.Run(context.Background(), func(context.Context) error {
|
||||||
|
return expected
|
||||||
|
})
|
||||||
|
if !errors.Is(err, expected) {
|
||||||
|
t.Fatalf("expected %v, got %v", expected, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := s.Run(context.Background(), func(context.Context) error { return nil }); err != nil {
|
||||||
|
t.Fatalf("expected permit to be released after error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func waitForAtomicAtLeast(t *testing.T, value *int32, want int32) {
|
||||||
|
t.Helper()
|
||||||
|
deadline := time.Now().Add(time.Second)
|
||||||
|
for time.Now().Before(deadline) {
|
||||||
|
if atomic.LoadInt32(value) >= want {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
runtime.Gosched()
|
||||||
|
}
|
||||||
|
t.Fatalf("timed out waiting for value >= %d; got %d", want, atomic.LoadInt32(value))
|
||||||
|
}
|
||||||
|
|
||||||
|
func waitForQueueDepth(t *testing.T, s *Scheduler, want int) {
|
||||||
|
t.Helper()
|
||||||
|
deadline := time.Now().Add(time.Second)
|
||||||
|
for time.Now().Before(deadline) {
|
||||||
|
s.mu.Lock()
|
||||||
|
depth := len(s.queue)
|
||||||
|
s.mu.Unlock()
|
||||||
|
if depth >= want {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
runtime.Gosched()
|
||||||
|
}
|
||||||
|
s.mu.Lock()
|
||||||
|
depth := len(s.queue)
|
||||||
|
s.mu.Unlock()
|
||||||
|
t.Fatalf("timed out waiting for queue depth >= %d; got %d", want, depth)
|
||||||
|
}
|
||||||
52
internal/framework/llm/secrets.go
Normal file
52
internal/framework/llm/secrets.go
Normal file
@@ -0,0 +1,52 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"errors"
|
||||||
|
"sort"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
const secretReplacement = "[REDACTED]"
|
||||||
|
|
||||||
|
// RedactSecrets replaces configured secret values in diagnostics.
|
||||||
|
func RedactSecrets(message string, secrets []string) string {
|
||||||
|
if message == "" || len(secrets) == 0 {
|
||||||
|
return message
|
||||||
|
}
|
||||||
|
|
||||||
|
normalized := normalizeSecrets(secrets)
|
||||||
|
for _, secret := range normalized {
|
||||||
|
message = strings.ReplaceAll(message, secret, secretReplacement)
|
||||||
|
}
|
||||||
|
return message
|
||||||
|
}
|
||||||
|
|
||||||
|
// ErrorWithSecretsRedacted returns an error with known secret values removed
|
||||||
|
// from its message.
|
||||||
|
func ErrorWithSecretsRedacted(err error, secrets []string) error {
|
||||||
|
if err == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return errors.New(RedactSecrets(err.Error(), secrets))
|
||||||
|
}
|
||||||
|
|
||||||
|
func normalizeSecrets(secrets []string) []string {
|
||||||
|
seen := make(map[string]struct{}, len(secrets))
|
||||||
|
result := make([]string, 0, len(secrets))
|
||||||
|
for _, secret := range secrets {
|
||||||
|
secret = strings.TrimSpace(secret)
|
||||||
|
if secret == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, ok := seen[secret]; ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
seen[secret] = struct{}{}
|
||||||
|
result = append(result, secret)
|
||||||
|
}
|
||||||
|
|
||||||
|
sort.Slice(result, func(i, j int) bool {
|
||||||
|
return len(result[i]) > len(result[j])
|
||||||
|
})
|
||||||
|
return result
|
||||||
|
}
|
||||||
43
internal/framework/llm/secrets_test.go
Normal file
43
internal/framework/llm/secrets_test.go
Normal file
@@ -0,0 +1,43 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"errors"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestRedactSecrets(t *testing.T) {
|
||||||
|
got := RedactSecrets("api key secret-token and Bearer secret-token failed", []string{"", "secret-token", "secret-token"})
|
||||||
|
|
||||||
|
if strings.Contains(got, "secret-token") {
|
||||||
|
t.Fatalf("expected secret to be redacted, got %q", got)
|
||||||
|
}
|
||||||
|
if strings.Count(got, secretReplacement) != 2 {
|
||||||
|
t.Fatalf("expected two redactions, got %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRedactSecretsPrefersLongerSecrets(t *testing.T) {
|
||||||
|
got := RedactSecrets("token token-extra", []string{"token", "token-extra"})
|
||||||
|
|
||||||
|
if strings.Contains(got, "token") {
|
||||||
|
t.Fatalf("expected overlapping secrets to be redacted, got %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestErrorWithSecretsRedacted(t *testing.T) {
|
||||||
|
err := ErrorWithSecretsRedacted(errors.New("secret-value failed"), []string{"secret-value"})
|
||||||
|
|
||||||
|
if err == nil {
|
||||||
|
t.Fatalf("expected redacted error")
|
||||||
|
}
|
||||||
|
if strings.Contains(err.Error(), "secret-value") {
|
||||||
|
t.Fatalf("expected secret to be redacted, got %q", err.Error())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestErrorWithSecretsRedactedNil(t *testing.T) {
|
||||||
|
if err := ErrorWithSecretsRedacted(nil, []string{"secret"}); err != nil {
|
||||||
|
t.Fatalf("expected nil error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user