base-layer/promptkit
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 5 Wiki Activity

Use resolved targets for client routing

Browse Source
This commit is contained in:
Eric Rakestraw
2026-11-05 19:52:48 +00:00
parent de46935a0e
commit 448e64dc40
15 changed files with 157 additions and 396 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
internal/llm/prepared_generation_test.go
View File

@@ -26,8 +26,7 @@ func TestPreparedGenerationFreezesPayloadAndAuthentication(t *testing.T) {
const envName = "PROMPTKIT_PREPARED_GENERATION_TEST_KEY"
t.Setenv(envName, tc.initialKey)
var bodies, auth, urls []string
client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
BaseURL: "https://example.invalid/v1", Model: "default-model",
client := NewOpenAICompatibleClient(OpenAICompatibleConfig{
HTTPClient: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
body, err := io.ReadAll(req.Body)
if err != nil {
@@ -39,12 +38,9 @@ func TestPreparedGenerationFreezesPayloadAndAuthentication(t *testing.T) {
return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(`{"choices":[{"message":{"content":"ok"}}]}`))}, nil
})},
})
if err != nil {
t.Fatal(err)
}
req := domain.GenerateRequest{
Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "original"}}},
Target: domain.ExecutionTarget{APIKeyEnv: envName},
Target: domain.ExecutionTarget{Endpoint: "https://example.invalid/v1", Model: "default-model", APIKeyEnv: envName},
}
prepared, err := client.PrepareGeneration(req)
if err != nil {
@@ -68,8 +64,7 @@ func TestPreparedGenerationFreezesPayloadAndAuthentication(t *testing.T) {
func TestPreparedGenerationAttemptsHaveIndependentContextsAndCloseBodies(t *testing.T) {
var contexts []context.Context
var closed int
client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
BaseURL: "https://example.invalid/v1", Model: "model",
client := NewOpenAICompatibleClient(OpenAICompatibleConfig{
HTTPClient: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
if err := req.Context().Err(); err != nil {
return nil, err
@@ -81,12 +76,9 @@ func TestPreparedGenerationAttemptsHaveIndependentContextsAndCloseBodies(t *test
}}, nil
})},
})
if err != nil {
t.Fatal(err)
}
prepared, err := client.PrepareGeneration(domain.GenerateRequest{
Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hi"}}},
Target: domain.ExecutionTarget{TimeoutSeconds: 1},
Target: domain.ExecutionTarget{Endpoint: "https://example.invalid/v1", Model: "model", TimeoutSeconds: 1},
})
if err != nil {
t.Fatal(err)
@@ -115,17 +107,14 @@ func TestPreparedGenerationAttemptsHaveIndependentContextsAndCloseBodies(t *test
}
func TestPreparedGenerationRetainsUsageFromMalformedChoice(t *testing.T) {
client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
BaseURL: "https://example.invalid/v1", Model: "model",
client := NewOpenAICompatibleClient(OpenAICompatibleConfig{
HTTPClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(`{"usage":{"prompt_tokens":0},"choices":[]}`))}, nil
})},
})
if err != nil {
t.Fatal(err)
}
prepared, err := client.PrepareGeneration(domain.GenerateRequest{
Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hi"}}},
Target: domain.ExecutionTarget{Endpoint: "https://example.invalid/v1", Model: "model"},
})
if err != nil {
t.Fatal(err)
@@ -137,17 +126,14 @@ func TestPreparedGenerationRetainsUsageFromMalformedChoice(t *testing.T) {
}
func TestPreparedGenerationDoesNotUseInvalidlyFramedUsage(t *testing.T) {
client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
BaseURL: "https://example.invalid/v1", Model: "model",
client := NewOpenAICompatibleClient(OpenAICompatibleConfig{
HTTPClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(`{"usage":{"prompt_tokens":9},"choices":[]}{}`))}, nil
})},
})
if err != nil {
t.Fatal(err)
}
prepared, err := client.PrepareGeneration(domain.GenerateRequest{
Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hi"}}},
Target: domain.ExecutionTarget{Endpoint: "https://example.invalid/v1", Model: "model"},
})
if err != nil {
t.Fatal(err)
@@ -160,17 +146,14 @@ func TestPreparedGenerationDoesNotUseInvalidlyFramedUsage(t *testing.T) {
func TestPreparedGenerationRetainsReadCancellation(t *testing.T) {
closed := false
client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
BaseURL: "https://example.invalid/v1", Model: "model",
client := NewOpenAICompatibleClient(OpenAICompatibleConfig{
HTTPClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
return &http.Response{StatusCode: http.StatusOK, Body: &readErrorBody{err: context.Canceled, closed: &closed}}, nil
})},
})
if err != nil {
t.Fatal(err)
}
prepared, err := client.PrepareGeneration(domain.GenerateRequest{
Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hi"}}},
Target: domain.ExecutionTarget{Endpoint: "https://example.invalid/v1", Model: "model"},
})
if err != nil {
t.Fatal(err)