Add request API key support
This commit is contained in:
@@ -63,6 +63,7 @@ type RunRequest struct {
|
||||
PromptID string
|
||||
PromptVersion string
|
||||
ProfileID string
|
||||
APIKey string `json:"-" yaml:"-"`
|
||||
Inputs map[string]ArtifactRef
|
||||
Vars map[string]string
|
||||
Execution *ExecutionTargetOverride
|
||||
@@ -207,6 +208,7 @@ type ExecutionTarget struct {
|
||||
ServiceTier string `yaml:"service_tier" json:"service_tier"`
|
||||
ReasoningEffort string `yaml:"reasoning_effort" json:"reasoning_effort"`
|
||||
APIKeyEnv string `yaml:"api_key_env" json:"api_key_env"`
|
||||
APIKey string `yaml:"-" json:"-"`
|
||||
ExtraParams map[string]any `yaml:"extra_params" json:"extra_params"`
|
||||
}
|
||||
|
||||
|
||||
@@ -20,6 +20,7 @@ func TestPreparedRunJSONDoesNotIncludeSecretValues(t *testing.T) {
|
||||
Endpoint: "http://llm/v1",
|
||||
Model: "gpt-test",
|
||||
APIKeyEnv: envName,
|
||||
APIKey: secret,
|
||||
},
|
||||
InputHashes: map[string]string{"transcript": "hash-1"},
|
||||
RenderedPromptHash: "rendered-hash",
|
||||
|
||||
@@ -92,6 +92,20 @@ func TestTextFormatterDoesNotIncludeResolvedAPIKeyValue(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestTextFormatterDoesNotIncludeDirectAPIKeyValue(t *testing.T) {
|
||||
const directKey = "direct-format-key"
|
||||
prepared := samplePreparedRun()
|
||||
prepared.EffectiveModelParams.APIKey = directKey
|
||||
|
||||
out, err := FormatPreparedRun(prepared, PreparedRunFormatText)
|
||||
if err != nil {
|
||||
t.Fatalf("expected no error, got %v", err)
|
||||
}
|
||||
if strings.Contains(string(out), directKey) {
|
||||
t.Fatalf("text output should not include direct api key value: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTextFormatterIncludesMessageCacheControlBeforeContent(t *testing.T) {
|
||||
prepared := samplePreparedRun()
|
||||
prepared.Messages = []domain.RenderedMessage{
|
||||
@@ -269,6 +283,20 @@ func TestJSONFormatterDoesNotIncludeResolvedAPIKeyValue(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestJSONFormatterDoesNotIncludeDirectAPIKeyValue(t *testing.T) {
|
||||
const directKey = "direct-format-key"
|
||||
prepared := samplePreparedRun()
|
||||
prepared.EffectiveModelParams.APIKey = directKey
|
||||
|
||||
out, err := FormatPreparedRun(prepared, PreparedRunFormatJSON)
|
||||
if err != nil {
|
||||
t.Fatalf("expected no error, got %v", err)
|
||||
}
|
||||
if strings.Contains(string(out), directKey) {
|
||||
t.Fatalf("json output should not include direct api key value: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParsePreparedRunOutputFormatRecognizesSupportedNames(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
||||
@@ -105,7 +105,9 @@ func (c *OpenAICompatibleClient) Generate(ctx context.Context, req domain.Genera
|
||||
return nil, fmt.Errorf("%w: failed to create request: %v", ErrRequestFailed, err)
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
if envName := strings.TrimSpace(req.Target.APIKeyEnv); envName != "" {
|
||||
if apiKey := strings.TrimSpace(req.Target.APIKey); apiKey != "" {
|
||||
httpReq.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
} else if envName := strings.TrimSpace(req.Target.APIKeyEnv); envName != "" {
|
||||
apiKey := strings.TrimSpace(os.Getenv(envName))
|
||||
if apiKey == "" {
|
||||
return nil, fmt.Errorf("%w: api key environment variable %q is not set", ErrInvalidRequest, envName)
|
||||
|
||||
@@ -148,6 +148,38 @@ func TestOpenAICompatibleClientGenerateSuccess(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAICompatibleClientDirectAPIKeyPreferredOverEnv(t *testing.T) {
|
||||
const directKey = "direct-llm-key"
|
||||
t.Setenv("SCRIPTORIUM_TEST_API_KEY", "env-key")
|
||||
|
||||
var gotAuth string
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotAuth = r.Header.Get("Authorization")
|
||||
_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"ok"}}]}`))
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{BaseURL: ts.URL + "/v1"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
_, err = client.Generate(context.Background(), domain.GenerateRequest{
|
||||
Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hi"}}},
|
||||
Target: domain.ExecutionTarget{
|
||||
Model: "model",
|
||||
APIKeyEnv: "SCRIPTORIUM_TEST_API_KEY",
|
||||
APIKey: directKey,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("expected no error, got %v", err)
|
||||
}
|
||||
if gotAuth != "Bearer "+directKey {
|
||||
t.Fatalf("unexpected Authorization header: %q", gotAuth)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAICompatibleClientSerializesCacheControlledMessageAsContentBlock(t *testing.T) {
|
||||
var observedBody map[string]any
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
@@ -195,13 +195,14 @@ func (r *Runner) Prepare(ctx context.Context, req domain.RunRequest) (*domain.Pr
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %w", ErrInvalidRequest, err)
|
||||
}
|
||||
effectiveModel.APIKey = req.APIKey
|
||||
if strings.TrimSpace(effectiveModel.Endpoint) == "" {
|
||||
return nil, fmt.Errorf("%w: execution endpoint is required", ErrInvalidRequest)
|
||||
}
|
||||
if strings.TrimSpace(effectiveModel.Model) == "" {
|
||||
return nil, fmt.Errorf("%w: execution model is required", ErrInvalidRequest)
|
||||
}
|
||||
if err := validateAPIKeyEnv(effectiveModel.APIKeyEnv); err != nil {
|
||||
if err := validateAPIKey(effectiveModel.APIKeyEnv, effectiveModel.APIKey); err != nil {
|
||||
return nil, fmt.Errorf("%w: %w", ErrInvalidRequest, err)
|
||||
}
|
||||
|
||||
@@ -434,7 +435,10 @@ func resolveExecutionTarget(profileValue *domain.ExecutionProfile, override *dom
|
||||
return out, presence, nil
|
||||
}
|
||||
|
||||
func validateAPIKeyEnv(apiKeyEnv string) error {
|
||||
func validateAPIKey(apiKeyEnv string, apiKey string) error {
|
||||
if strings.TrimSpace(apiKey) != "" {
|
||||
return nil
|
||||
}
|
||||
envName := strings.TrimSpace(apiKeyEnv)
|
||||
if envName == "" {
|
||||
return nil
|
||||
|
||||
@@ -1179,6 +1179,32 @@ func TestRunnerRunAPIKeyEnvMissingEnvironmentValueFailsClearly(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunnerRunDirectAPIKeyBypassesMissingEnvAndReachesLLM(t *testing.T) {
|
||||
const directKey = "direct-runner-key"
|
||||
promptRepo := &fakePromptRepo{def: promptDef(domain.FormatText, domain.ValidationNone, 0)}
|
||||
execRepo := &fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{
|
||||
"exec": {ID: "exec", Endpoint: "http://profile/v1", Model: "profile-model", APIKeyEnv: "SCRIPTORIUM_MISSING_KEY"},
|
||||
}}
|
||||
llmClient := &fakeLLM{resp: &domain.GenerateResponse{Content: "ok"}}
|
||||
runner := NewRunner(promptRepo, execRepo, defaultArtifactReader(), defaultRenderer(), llmClient, nil)
|
||||
|
||||
_, err := runner.Run(context.Background(), domain.RunRequest{
|
||||
PromptID: "p",
|
||||
ProfileID: "exec",
|
||||
APIKey: directKey,
|
||||
Inputs: singleInputRef(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("expected no error, got %v", err)
|
||||
}
|
||||
if llmClient.lastReq.Target.APIKey != directKey {
|
||||
t.Fatalf("expected direct API key to reach LLM request")
|
||||
}
|
||||
if llmClient.lastReq.Target.APIKeyEnv != "SCRIPTORIUM_MISSING_KEY" {
|
||||
t.Fatalf("expected api_key_env name to remain on target, got %q", llmClient.lastReq.Target.APIKeyEnv)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunnerRunRuntimeAPIKeyEnvOverrideWorks(t *testing.T) {
|
||||
const envName = "SCRIPTORIUM_RUNTIME_API_KEY"
|
||||
t.Setenv(envName, "runtime-secret")
|
||||
|
||||
Reference in New Issue
Block a user