diff --git a/internal/domain/domain.go b/internal/domain/domain.go index 94b98dd..694f29d 100644 --- a/internal/domain/domain.go +++ b/internal/domain/domain.go @@ -237,6 +237,8 @@ type TokenUsage struct { PromptTokens int CompletionTokens int TotalTokens int + CachedTokens int + CacheWriteTokens int } // ValidationResult represents the outcome of an output validation. diff --git a/internal/llm/openai_compatible_client.go b/internal/llm/openai_compatible_client.go index 8e32b74..a1a7790 100644 --- a/internal/llm/openai_compatible_client.go +++ b/internal/llm/openai_compatible_client.go @@ -151,6 +151,8 @@ func (c *OpenAICompatibleClient) Generate(ctx context.Context, req domain.Genera PromptTokens: wireResp.Usage.PromptTokens, CompletionTokens: wireResp.Usage.CompletionTokens, TotalTokens: wireResp.Usage.TotalTokens, + CachedTokens: wireResp.Usage.PromptTokensDetails.CachedTokens, + CacheWriteTokens: wireResp.Usage.CacheWriteTokens, }, }, nil } @@ -168,12 +170,9 @@ func openAIChatRequestFromGenerateRequest(req domain.GenerateRequest, defaultMod Model: model, } - wireReq.Messages = make([]openAIChatMessage, 0, len(req.Prompt.Messages)) + wireReq.Messages = make([]openAIChatRequestMessage, 0, len(req.Prompt.Messages)) for _, msg := range req.Prompt.Messages { - wireReq.Messages = append(wireReq.Messages, openAIChatMessage{ - Role: msg.Role, - Content: msg.Content, - }) + wireReq.Messages = append(wireReq.Messages, openAIChatRequestMessageFromRenderedMessage(msg)) } if req.Target.Temperature != 0 { @@ -200,28 +199,48 @@ func openAIChatRequestFromGenerateRequest(req domain.GenerateRequest, defaultMod } type openAIChatRequest struct { - Model string `json:"model"` - Messages []openAIChatMessage `json:"messages"` - Temperature *float64 `json:"temperature,omitempty"` - MaxTokens *int `json:"max_tokens,omitempty"` - TopP *float64 `json:"top_p,omitempty"` - ServiceTier string `json:"service_tier,omitempty"` - ResponseFormat *openAIResponseFormat `json:"response_format,omitempty"` + Model string `json:"model"` + Messages []openAIChatRequestMessage `json:"messages"` + Temperature *float64 `json:"temperature,omitempty"` + MaxTokens *int `json:"max_tokens,omitempty"` + TopP *float64 `json:"top_p,omitempty"` + ServiceTier string `json:"service_tier,omitempty"` + ResponseFormat *openAIResponseFormat `json:"response_format,omitempty"` } -type openAIChatMessage struct { +type openAIChatRequestMessage struct { + Role string `json:"role"` + Content any `json:"content"` +} + +type openAIChatTextContentBlock struct { + Type string `json:"type"` + Text string `json:"text"` + CacheControl *openAICacheControl `json:"cache_control,omitempty"` +} + +type openAICacheControl struct { + Type string `json:"type"` + TTL string `json:"ttl,omitempty"` +} + +type openAIChatResponseMessage struct { Role string `json:"role"` Content string `json:"content"` } type openAIChatResponse struct { Choices []struct { - Message openAIChatMessage `json:"message"` + Message openAIChatResponseMessage `json:"message"` } `json:"choices"` Usage struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - TotalTokens int `json:"total_tokens"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + PromptTokensDetails struct { + CachedTokens int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` + CacheWriteTokens int `json:"cache_write_tokens"` } `json:"usage"` } @@ -236,6 +255,28 @@ type openAIJSONSchemaEnvelope struct { Schema any `json:"schema"` } +func openAIChatRequestMessageFromRenderedMessage(msg domain.RenderedMessage) openAIChatRequestMessage { + wireMsg := openAIChatRequestMessage{ + Role: msg.Role, + Content: msg.Content, + } + if msg.CacheControl == nil { + return wireMsg + } + + wireMsg.Content = []openAIChatTextContentBlock{ + { + Type: "text", + Text: msg.Content, + CacheControl: &openAICacheControl{ + Type: string(msg.CacheControl.Type), + TTL: msg.CacheControl.TTL, + }, + }, + } + return wireMsg +} + func toOpenAIResponseFormat(spec *domain.StructuredOutputSpec) (*openAIResponseFormat, error) { if spec == nil { return nil, nil diff --git a/internal/llm/openai_compatible_client_test.go b/internal/llm/openai_compatible_client_test.go index 702b08d..f7371f2 100644 --- a/internal/llm/openai_compatible_client_test.go +++ b/internal/llm/openai_compatible_client_test.go @@ -89,6 +89,9 @@ func TestOpenAICompatibleClientGenerateSuccess(t *testing.T) { if resp.Usage.PromptTokens != 11 || resp.Usage.CompletionTokens != 22 || resp.Usage.TotalTokens != 33 { t.Fatalf("unexpected usage: %+v", resp.Usage) } + if resp.Usage.CachedTokens != 0 || resp.Usage.CacheWriteTokens != 0 { + t.Fatalf("expected absent cache usage fields to remain zero, got %+v", resp.Usage) + } if obs.Authorization != "Bearer secret-key" { t.Fatalf("unexpected Authorization header: %q", obs.Authorization) @@ -144,6 +147,155 @@ func TestOpenAICompatibleClientGenerateSuccess(t *testing.T) { } } +func TestOpenAICompatibleClientSerializesCacheControlledMessageAsContentBlock(t *testing.T) { + var observedBody map[string]any + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer r.Body.Close() + if err := json.NewDecoder(r.Body).Decode(&observedBody); err != nil { + t.Fatalf("failed to decode request body: %v", err) + } + _, _ = 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: "system", + Content: "Stable instructions.", + CacheControl: &domain.CacheControl{ + Type: domain.CacheControlEphemeral, + TTL: "1h", + }, + }, + {Role: "user", Content: "Dynamic request."}, + }}, + Target: domain.ExecutionTarget{Model: "model"}, + }) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + + for _, forbidden := range []string{"cache_control", "session_id", "extra_params"} { + if _, exists := observedBody[forbidden]; exists { + t.Fatalf("expected top-level %s to be omitted, got %#v", forbidden, observedBody[forbidden]) + } + } + + msgs, ok := observedBody["messages"].([]any) + if !ok || len(msgs) != 2 { + t.Fatalf("unexpected messages payload: %#v", observedBody["messages"]) + } + msg0 := msgs[0].(map[string]any) + if msg0["role"] != "system" { + t.Fatalf("unexpected first message role: %#v", msg0["role"]) + } + contentBlocks, ok := msg0["content"].([]any) + if !ok || len(contentBlocks) != 1 { + t.Fatalf("expected first message content block array, got %#v", msg0["content"]) + } + block := contentBlocks[0].(map[string]any) + if block["type"] != "text" || block["text"] != "Stable instructions." { + t.Fatalf("unexpected text content block: %#v", block) + } + cacheControl, ok := block["cache_control"].(map[string]any) + if !ok { + t.Fatalf("expected cache_control on content block, got %#v", block) + } + if cacheControl["type"] != string(domain.CacheControlEphemeral) || cacheControl["ttl"] != "1h" { + t.Fatalf("unexpected cache_control payload: %#v", cacheControl) + } + + msg1 := msgs[1].(map[string]any) + if msg1["role"] != "user" || msg1["content"] != "Dynamic request." { + t.Fatalf("expected uncached message to keep string content, got %#v", msg1) + } +} + +func TestOpenAICompatibleClientOmitsEmptyCacheControlTTL(t *testing.T) { + var observedBody map[string]any + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer r.Body.Close() + if err := json.NewDecoder(r.Body).Decode(&observedBody); err != nil { + t.Fatalf("failed to decode request body: %v", err) + } + _, _ = 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: "system", + Content: "Stable instructions.", + CacheControl: &domain.CacheControl{ + Type: domain.CacheControlEphemeral, + }, + }, + }}, + Target: domain.ExecutionTarget{Model: "model"}, + }) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + + msgs := observedBody["messages"].([]any) + msg0 := msgs[0].(map[string]any) + contentBlocks := msg0["content"].([]any) + block := contentBlocks[0].(map[string]any) + cacheControl := block["cache_control"].(map[string]any) + if cacheControl["type"] != string(domain.CacheControlEphemeral) { + t.Fatalf("unexpected cache_control type: %#v", cacheControl) + } + if _, exists := cacheControl["ttl"]; exists { + t.Fatalf("expected empty ttl to be omitted, got %#v", cacheControl) + } +} + +func TestOpenAICompatibleClientParsesCacheUsage(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(`{ + "choices": [{"message": {"role": "assistant", "content": "ok"}}], + "usage": { + "prompt_tokens": 100, + "completion_tokens": 20, + "total_tokens": 120, + "prompt_tokens_details": {"cached_tokens": 80}, + "cache_write_tokens": 60 + } +}`)) + })) + defer ts.Close() + + client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{BaseURL: ts.URL + "/v1", Model: "model"}) + if err != nil { + t.Fatal(err) + } + + resp, err := client.Generate(context.Background(), domain.GenerateRequest{ + Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hi"}}}, + }) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + if resp.Usage.PromptTokens != 100 || resp.Usage.CompletionTokens != 20 || resp.Usage.TotalTokens != 120 { + t.Fatalf("unexpected base usage fields: %+v", resp.Usage) + } + if resp.Usage.CachedTokens != 80 || resp.Usage.CacheWriteTokens != 60 { + t.Fatalf("unexpected cache usage fields: %+v", resp.Usage) + } +} + func TestOpenAICompatibleClientOmitsResponseFormatWhenNoStructuredOutput(t *testing.T) { var observedBody map[string]any ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {