Implement Phase 9 structured LLM adapter spike
This commit is contained in:
221
internal/framework/llm/instructor_client.go
Normal file
221
internal/framework/llm/instructor_client.go
Normal file
@@ -0,0 +1,221 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"reflect"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gitea.maximumdirect.net/eric/audita/internal/framework/contracts"
|
||||
|
||||
"github.com/jxnl/instructor-go/pkg/instructor"
|
||||
openai "github.com/sashabaranov/go-openai"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultMaxRetries = 3
|
||||
)
|
||||
|
||||
type Mode string
|
||||
|
||||
const (
|
||||
ModeJSON Mode = "json"
|
||||
ModeToolCall Mode = "tool_call"
|
||||
)
|
||||
|
||||
type InstructorClientConfig struct {
|
||||
BaseURL string
|
||||
Model string
|
||||
APIKey string
|
||||
MaxRetries int
|
||||
Mode Mode
|
||||
HTTPClient *http.Client
|
||||
RequestTimeout time.Duration
|
||||
}
|
||||
|
||||
// chatCompletionClient is intentionally narrow to avoid leaking provider types
|
||||
// outside this adapter package.
|
||||
type chatCompletionClient interface {
|
||||
CreateChatCompletion(ctx context.Context, request openai.ChatCompletionRequest, responseType any) (response openai.ChatCompletionResponse, err error)
|
||||
}
|
||||
|
||||
// InstructorClient adapts instructor-go behind Audita's internal structured
|
||||
// LLM interface.
|
||||
type InstructorClient struct {
|
||||
cfg InstructorClientConfig
|
||||
client chatCompletionClient
|
||||
}
|
||||
|
||||
var _ contracts.StructuredLLMClient = (*InstructorClient)(nil)
|
||||
|
||||
func NewInstructorClient(cfg InstructorClientConfig) (*InstructorClient, error) {
|
||||
normalized, err := normalizeConfig(cfg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
openaiConfig := openai.DefaultConfig(normalized.APIKey)
|
||||
openaiConfig.BaseURL = normalized.BaseURL
|
||||
openaiConfig.HTTPClient = resolvedHTTPClient(normalized.HTTPClient, normalized.RequestTimeout)
|
||||
|
||||
mode, err := toInstructorMode(normalized.Mode)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
client := instructor.FromOpenAI(
|
||||
openai.NewClientWithConfig(openaiConfig),
|
||||
instructor.WithMode(mode),
|
||||
instructor.WithMaxRetries(normalized.MaxRetries),
|
||||
)
|
||||
|
||||
return &InstructorClient{
|
||||
cfg: normalized,
|
||||
client: client,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c *InstructorClient) CompleteStructured(
|
||||
ctx context.Context,
|
||||
req contracts.StructuredCompletionRequest,
|
||||
out any,
|
||||
) (contracts.StructuredCompletionResponse, error) {
|
||||
if err := validateOutputTarget(out); err != nil {
|
||||
return contracts.StructuredCompletionResponse{}, err
|
||||
}
|
||||
|
||||
model := strings.TrimSpace(req.Model)
|
||||
if model == "" {
|
||||
model = c.cfg.Model
|
||||
}
|
||||
if model == "" {
|
||||
return contracts.StructuredCompletionResponse{}, fmt.Errorf("structured completion model must not be empty")
|
||||
}
|
||||
|
||||
messages, err := toOpenAIMessages(req.Messages)
|
||||
if err != nil {
|
||||
return contracts.StructuredCompletionResponse{}, err
|
||||
}
|
||||
|
||||
chatRequest := openai.ChatCompletionRequest{
|
||||
Model: model,
|
||||
Messages: messages,
|
||||
}
|
||||
|
||||
resp, err := c.client.CreateChatCompletion(ctx, chatRequest, out)
|
||||
if err != nil {
|
||||
return contracts.StructuredCompletionResponse{}, sanitizeError(err, c.cfg.APIKey)
|
||||
}
|
||||
|
||||
content, err := json.Marshal(out)
|
||||
if err != nil {
|
||||
return contracts.StructuredCompletionResponse{}, fmt.Errorf("marshal structured completion output: %w", err)
|
||||
}
|
||||
|
||||
return contracts.StructuredCompletionResponse{
|
||||
Content: content,
|
||||
Provider: "openai-compatible",
|
||||
Model: model,
|
||||
PromptTokens: resp.Usage.PromptTokens,
|
||||
CompletionTokens: resp.Usage.CompletionTokens,
|
||||
TotalTokens: resp.Usage.TotalTokens,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// TODO(phase9): integrate bounded scheduler/semaphore in the next Phase 9 prompt.
|
||||
|
||||
func normalizeConfig(cfg InstructorClientConfig) (InstructorClientConfig, error) {
|
||||
cfg.BaseURL = strings.TrimSpace(cfg.BaseURL)
|
||||
cfg.Model = strings.TrimSpace(cfg.Model)
|
||||
cfg.APIKey = strings.TrimSpace(cfg.APIKey)
|
||||
if cfg.MaxRetries == 0 {
|
||||
cfg.MaxRetries = defaultMaxRetries
|
||||
}
|
||||
if cfg.MaxRetries < 0 {
|
||||
return InstructorClientConfig{}, fmt.Errorf("max retries must be zero or greater")
|
||||
}
|
||||
if cfg.BaseURL == "" {
|
||||
return InstructorClientConfig{}, fmt.Errorf("base URL must not be empty")
|
||||
}
|
||||
if cfg.Model == "" {
|
||||
return InstructorClientConfig{}, fmt.Errorf("model must not be empty")
|
||||
}
|
||||
if cfg.Mode == "" {
|
||||
cfg.Mode = ModeJSON
|
||||
}
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
func toInstructorMode(mode Mode) (instructor.Mode, error) {
|
||||
switch mode {
|
||||
case ModeJSON:
|
||||
return instructor.ModeJSON, nil
|
||||
case ModeToolCall:
|
||||
return instructor.ModeToolCall, nil
|
||||
default:
|
||||
return "", fmt.Errorf("unsupported LLM mode %q", mode)
|
||||
}
|
||||
}
|
||||
|
||||
func resolvedHTTPClient(base *http.Client, timeout time.Duration) *http.Client {
|
||||
if base == nil {
|
||||
if timeout <= 0 {
|
||||
return http.DefaultClient
|
||||
}
|
||||
return &http.Client{Timeout: timeout}
|
||||
}
|
||||
|
||||
if timeout <= 0 {
|
||||
return base
|
||||
}
|
||||
|
||||
cloned := *base
|
||||
cloned.Timeout = timeout
|
||||
return &cloned
|
||||
}
|
||||
|
||||
func toOpenAIMessages(messages []contracts.LLMMessage) ([]openai.ChatCompletionMessage, error) {
|
||||
result := make([]openai.ChatCompletionMessage, len(messages))
|
||||
for i, message := range messages {
|
||||
role := strings.TrimSpace(message.Role)
|
||||
content := strings.TrimSpace(message.Content)
|
||||
if role == "" {
|
||||
return nil, fmt.Errorf("message[%d] role must not be empty", i)
|
||||
}
|
||||
if content == "" {
|
||||
return nil, fmt.Errorf("message[%d] content must not be empty", i)
|
||||
}
|
||||
result[i] = openai.ChatCompletionMessage{
|
||||
Role: role,
|
||||
Content: content,
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func validateOutputTarget(out any) error {
|
||||
if out == nil {
|
||||
return fmt.Errorf("output target must not be nil")
|
||||
}
|
||||
value := reflect.ValueOf(out)
|
||||
if value.Kind() != reflect.Ptr || value.IsNil() {
|
||||
return fmt.Errorf("output target must be a non-nil pointer")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func sanitizeError(err error, apiKey string) error {
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
msg := err.Error()
|
||||
key := strings.TrimSpace(apiKey)
|
||||
if key != "" {
|
||||
msg = strings.ReplaceAll(msg, key, "[REDACTED]")
|
||||
msg = strings.ReplaceAll(msg, "Bearer "+key, "Bearer [REDACTED]")
|
||||
}
|
||||
return fmt.Errorf("%s", msg)
|
||||
}
|
||||
282
internal/framework/llm/instructor_client_test.go
Normal file
282
internal/framework/llm/instructor_client_test.go
Normal file
@@ -0,0 +1,282 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.maximumdirect.net/eric/audita/internal/framework/contracts"
|
||||
|
||||
openai "github.com/sashabaranov/go-openai"
|
||||
)
|
||||
|
||||
func TestNewInstructorClientValidation(t *testing.T) {
|
||||
_, err := NewInstructorClient(InstructorClientConfig{
|
||||
BaseURL: " ",
|
||||
Model: "x",
|
||||
})
|
||||
if err == nil || !strings.Contains(err.Error(), "base URL") {
|
||||
t.Fatalf("expected base URL validation error, got %v", err)
|
||||
}
|
||||
|
||||
_, err = NewInstructorClient(InstructorClientConfig{
|
||||
BaseURL: "http://localhost:1234/v1",
|
||||
Model: " ",
|
||||
})
|
||||
if err == nil || !strings.Contains(err.Error(), "model") {
|
||||
t.Fatalf("expected model validation error, got %v", err)
|
||||
}
|
||||
|
||||
_, err = NewInstructorClient(InstructorClientConfig{
|
||||
BaseURL: "http://localhost:1234/v1",
|
||||
Model: "test-model",
|
||||
MaxRetries: -1,
|
||||
})
|
||||
if err == nil || !strings.Contains(err.Error(), "max retries") {
|
||||
t.Fatalf("expected retries validation error, got %v", err)
|
||||
}
|
||||
|
||||
_, err = NewInstructorClient(InstructorClientConfig{
|
||||
BaseURL: "http://localhost:1234/v1",
|
||||
Model: "test-model",
|
||||
Mode: "unsupported",
|
||||
})
|
||||
if err == nil || !strings.Contains(err.Error(), "unsupported LLM mode") {
|
||||
t.Fatalf("expected mode validation error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInstructorClientCompleteStructuredSuccessNoAPIKey(t *testing.T) {
|
||||
var seenPath string
|
||||
var seenHost string
|
||||
var seenAuth string
|
||||
var seenModel string
|
||||
client, err := NewInstructorClient(InstructorClientConfig{
|
||||
BaseURL: "https://local-compat.example/v1",
|
||||
Model: "test-model",
|
||||
HTTPClient: &http.Client{
|
||||
Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
seenPath = r.URL.Path
|
||||
seenHost = r.URL.Host
|
||||
seenAuth = r.Header.Get("Authorization")
|
||||
|
||||
var req map[string]any
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
t.Fatalf("decode request: %v", err)
|
||||
}
|
||||
model, _ := req["model"].(string)
|
||||
seenModel = model
|
||||
|
||||
return newJSONHTTPResponse(http.StatusOK, chatCompletionResponseBody(`{"name":"Robby","age":22}`)), nil
|
||||
}),
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewInstructorClient: %v", err)
|
||||
}
|
||||
|
||||
type person struct {
|
||||
Name string `json:"name"`
|
||||
Age int `json:"age"`
|
||||
}
|
||||
var out person
|
||||
resp, err := client.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{
|
||||
StageName: "proposal:test",
|
||||
Messages: []contracts.LLMMessage{
|
||||
{Role: openai.ChatMessageRoleUser, Content: "extract person"},
|
||||
},
|
||||
}, &out)
|
||||
if err != nil {
|
||||
t.Fatalf("CompleteStructured: %v", err)
|
||||
}
|
||||
|
||||
if seenPath != "/v1/chat/completions" {
|
||||
t.Fatalf("unexpected request path: %q", seenPath)
|
||||
}
|
||||
if seenHost != "local-compat.example" {
|
||||
t.Fatalf("unexpected request host: %q", seenHost)
|
||||
}
|
||||
if seenModel != "test-model" {
|
||||
t.Fatalf("unexpected model: %q", seenModel)
|
||||
}
|
||||
if seenAuth != "" {
|
||||
t.Fatalf("expected empty Authorization header for empty API key, got %q", seenAuth)
|
||||
}
|
||||
if out.Name != "Robby" || out.Age != 22 {
|
||||
t.Fatalf("unexpected output: %+v", out)
|
||||
}
|
||||
if resp.Provider != "openai-compatible" || resp.Model != "test-model" {
|
||||
t.Fatalf("unexpected response metadata: %+v", resp)
|
||||
}
|
||||
if resp.TotalTokens != 18 || resp.PromptTokens != 11 || resp.CompletionTokens != 7 {
|
||||
t.Fatalf("unexpected usage metadata: %+v", resp)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInstructorClientRetriesOnMalformedJSON(t *testing.T) {
|
||||
var attempts int32
|
||||
client, err := NewInstructorClient(InstructorClientConfig{
|
||||
BaseURL: "https://retry.example/v1",
|
||||
Model: "test-model",
|
||||
MaxRetries: 1,
|
||||
HTTPClient: &http.Client{
|
||||
Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
current := atomic.AddInt32(&attempts, 1)
|
||||
if current == 1 {
|
||||
return newJSONHTTPResponse(http.StatusOK, chatCompletionResponseBody(`{"name":"broken"`)), nil
|
||||
}
|
||||
return newJSONHTTPResponse(http.StatusOK, chatCompletionResponseBody(`{"name":"Recovered","age":30}`)), nil
|
||||
}),
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewInstructorClient: %v", err)
|
||||
}
|
||||
|
||||
type person struct {
|
||||
Name string `json:"name"`
|
||||
Age int `json:"age"`
|
||||
}
|
||||
var out person
|
||||
_, err = client.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{
|
||||
Messages: []contracts.LLMMessage{
|
||||
{Role: openai.ChatMessageRoleUser, Content: "extract person"},
|
||||
},
|
||||
}, &out)
|
||||
if err != nil {
|
||||
t.Fatalf("CompleteStructured: %v", err)
|
||||
}
|
||||
if out.Name != "Recovered" || out.Age != 30 {
|
||||
t.Fatalf("unexpected output after retry: %+v", out)
|
||||
}
|
||||
if got := atomic.LoadInt32(&attempts); got != 2 {
|
||||
t.Fatalf("expected 2 attempts, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInstructorClientContextDeadlinePropagates(t *testing.T) {
|
||||
client, err := NewInstructorClient(InstructorClientConfig{
|
||||
BaseURL: "https://slow.example/v1",
|
||||
Model: "test-model",
|
||||
HTTPClient: &http.Client{
|
||||
Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
select {
|
||||
case <-r.Context().Done():
|
||||
return nil, r.Context().Err()
|
||||
case <-time.After(250 * time.Millisecond):
|
||||
return newJSONHTTPResponse(http.StatusOK, chatCompletionResponseBody(`{"name":"slow","age":1}`)), nil
|
||||
}
|
||||
}),
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewInstructorClient: %v", err)
|
||||
}
|
||||
|
||||
type person struct {
|
||||
Name string `json:"name"`
|
||||
}
|
||||
var out person
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
_, err = client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
|
||||
Messages: []contracts.LLMMessage{
|
||||
{Role: openai.ChatMessageRoleUser, Content: "extract person"},
|
||||
},
|
||||
}, &out)
|
||||
if err == nil {
|
||||
t.Fatalf("expected context deadline error")
|
||||
}
|
||||
if !strings.Contains(strings.ToLower(err.Error()), "context deadline") {
|
||||
t.Fatalf("expected deadline-related error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInstructorClientSanitizesAPIKeyInErrors(t *testing.T) {
|
||||
apiKey := "super-secret-key"
|
||||
client := &InstructorClient{
|
||||
cfg: InstructorClientConfig{
|
||||
BaseURL: "http://localhost:1234/v1",
|
||||
Model: "test-model",
|
||||
APIKey: apiKey,
|
||||
Mode: ModeJSON,
|
||||
},
|
||||
client: fakeChatCompletionClient{
|
||||
err: fmt.Errorf("provider failed with Authorization: Bearer %s", apiKey),
|
||||
},
|
||||
}
|
||||
|
||||
var out map[string]any
|
||||
_, err := client.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{
|
||||
Messages: []contracts.LLMMessage{
|
||||
{Role: openai.ChatMessageRoleUser, Content: "extract"},
|
||||
},
|
||||
}, &out)
|
||||
if err == nil {
|
||||
t.Fatalf("expected error")
|
||||
}
|
||||
if strings.Contains(err.Error(), apiKey) {
|
||||
t.Fatalf("error leaked API key: %v", err)
|
||||
}
|
||||
if !strings.Contains(err.Error(), "[REDACTED]") {
|
||||
t.Fatalf("expected redaction marker in error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
type fakeChatCompletionClient struct {
|
||||
err error
|
||||
}
|
||||
|
||||
func (f fakeChatCompletionClient) CreateChatCompletion(ctx context.Context, request openai.ChatCompletionRequest, responseType any) (response openai.ChatCompletionResponse, err error) {
|
||||
_ = ctx
|
||||
_ = request
|
||||
_ = responseType
|
||||
return openai.ChatCompletionResponse{}, f.err
|
||||
}
|
||||
|
||||
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) {
|
||||
return f(r)
|
||||
}
|
||||
|
||||
func newJSONHTTPResponse(status int, body string) *http.Response {
|
||||
return &http.Response{
|
||||
StatusCode: status,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
}
|
||||
}
|
||||
|
||||
func chatCompletionResponseBody(content string) string {
|
||||
body, _ := json.Marshal(map[string]any{
|
||||
"id": "chatcmpl-test",
|
||||
"object": "chat.completion",
|
||||
"created": 12345,
|
||||
"model": "test-model",
|
||||
"choices": []map[string]any{
|
||||
{
|
||||
"index": 0,
|
||||
"message": map[string]any{
|
||||
"role": "assistant",
|
||||
"content": content,
|
||||
},
|
||||
"finish_reason": "stop",
|
||||
},
|
||||
},
|
||||
"usage": map[string]any{
|
||||
"prompt_tokens": 11,
|
||||
"completion_tokens": 7,
|
||||
"total_tokens": 18,
|
||||
},
|
||||
})
|
||||
return string(body)
|
||||
}
|
||||
Reference in New Issue
Block a user