diff --git a/internal/framework/llm/client_common.go b/internal/framework/llm/client_common.go new file mode 100644 index 0000000..a36e439 --- /dev/null +++ b/internal/framework/llm/client_common.go @@ -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 "" +} diff --git a/internal/framework/llm/openai_compatible_client.go b/internal/framework/llm/openai_compatible_client.go new file mode 100644 index 0000000..d72f361 --- /dev/null +++ b/internal/framework/llm/openai_compatible_client.go @@ -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) +} diff --git a/internal/framework/llm/openai_compatible_client_test.go b/internal/framework/llm/openai_compatible_client_test.go new file mode 100644 index 0000000..398645b --- /dev/null +++ b/internal/framework/llm/openai_compatible_client_test.go @@ -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 +}`) +} diff --git a/internal/framework/llm/scheduler.go b/internal/framework/llm/scheduler.go new file mode 100644 index 0000000..9d30223 --- /dev/null +++ b/internal/framework/llm/scheduler.go @@ -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 + } + } +} diff --git a/internal/framework/llm/scheduler_test.go b/internal/framework/llm/scheduler_test.go new file mode 100644 index 0000000..d06ea00 --- /dev/null +++ b/internal/framework/llm/scheduler_test.go @@ -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) +} diff --git a/internal/framework/llm/secrets.go b/internal/framework/llm/secrets.go new file mode 100644 index 0000000..f10da7d --- /dev/null +++ b/internal/framework/llm/secrets.go @@ -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 +} diff --git a/internal/framework/llm/secrets_test.go b/internal/framework/llm/secrets_test.go new file mode 100644 index 0000000..9b5ea47 --- /dev/null +++ b/internal/framework/llm/secrets_test.go @@ -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) + } +}