From 5005066fd1126461e39ee73b9fcd09bcd1131afc Mon Sep 17 00:00:00 2001
From: Aleksandr Bortnikov <albortnikov@ozon.ru>
Date: Sun, 1 Mar 2026 09:32:12 +0300
Subject: [PATCH] fix: added WithUsePromptCaching,
 NewHTTPProviderWithMaxTokensFieldAndRequestTimeoutAndUsePromptCaching, used
 usePromptCaching for openai protocol only

---
 pkg/providers/factory_provider.go       | 10 +++++++---
 pkg/providers/http_provider.go          | 17 +++++++++++++++++
 pkg/providers/openai_compat/provider.go | 19 ++++++++++++++-----
 3 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/pkg/providers/factory_provider.go b/pkg/providers/factory_provider.go
index 53f7a08a0..dd795c1e6 100644
--- a/pkg/providers/factory_provider.go
+++ b/pkg/providers/factory_provider.go
@@ -84,12 +84,14 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
 		if apiBase == "" {
 			apiBase = getDefaultAPIBase(protocol)
 		}
-		return NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(
+
+		return NewHTTPProviderWithMaxTokensFieldAndRequestTimeoutAndUsePromptCaching(
 			cfg.APIKey,
 			apiBase,
 			cfg.Proxy,
 			cfg.MaxTokensField,
 			cfg.RequestTimeout,
+			true,
 		), modelID, nil
 
 	case "openrouter", "groq", "zhipu", "gemini", "nvidia",
@@ -103,12 +105,13 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
 		if apiBase == "" {
 			apiBase = getDefaultAPIBase(protocol)
 		}
-		return NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(
+		return NewHTTPProviderWithMaxTokensFieldAndRequestTimeoutAndUsePromptCaching(
 			cfg.APIKey,
 			apiBase,
 			cfg.Proxy,
 			cfg.MaxTokensField,
 			cfg.RequestTimeout,
+			false,
 		), modelID, nil
 
 	case "anthropic":
@@ -128,12 +131,13 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
 		if cfg.APIKey == "" {
 			return nil, "", fmt.Errorf("api_key is required for anthropic protocol (model: %s)", cfg.Model)
 		}
-		return NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(
+		return NewHTTPProviderWithMaxTokensFieldAndRequestTimeoutAndUsePromptCaching(
 			cfg.APIKey,
 			apiBase,
 			cfg.Proxy,
 			cfg.MaxTokensField,
 			cfg.RequestTimeout,
+			false,
 		), modelID, nil
 
 	case "antigravity":
diff --git a/pkg/providers/http_provider.go b/pkg/providers/http_provider.go
index 5c328f418..71ff94506 100644
--- a/pkg/providers/http_provider.go
+++ b/pkg/providers/http_provider.go
@@ -42,6 +42,23 @@ func NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(
 	}
 }
 
+func NewHTTPProviderWithMaxTokensFieldAndRequestTimeoutAndUsePromptCaching(
+	apiKey, apiBase, proxy, maxTokensField string,
+	requestTimeoutSeconds int,
+	usePromptCaching bool,
+) *HTTPProvider {
+	return &HTTPProvider{
+		delegate: openai_compat.NewProvider(
+			apiKey,
+			apiBase,
+			proxy,
+			openai_compat.WithMaxTokensField(maxTokensField),
+			openai_compat.WithRequestTimeout(time.Duration(requestTimeoutSeconds)*time.Second),
+			openai_compat.WithUsePromptCaching(usePromptCaching),
+		),
+	}
+}
+
 func (p *HTTPProvider) Chat(
 	ctx context.Context,
 	messages []Message,
diff --git a/pkg/providers/openai_compat/provider.go b/pkg/providers/openai_compat/provider.go
index 5dab9b03e..be0c799bc 100644
--- a/pkg/providers/openai_compat/provider.go
+++ b/pkg/providers/openai_compat/provider.go
@@ -29,10 +29,11 @@ type (
 )
 
 type Provider struct {
-	apiKey         string
-	apiBase        string
-	maxTokensField string // Field name for max tokens (e.g., "max_completion_tokens" for o1/glm models)
-	httpClient     *http.Client
+	apiKey           string
+	apiBase          string
+	maxTokensField   string // Field name for max tokens (e.g., "max_completion_tokens" for o1/glm models)
+	usePromptCaching bool
+	httpClient       *http.Client
 }
 
 type Option func(*Provider)
@@ -53,6 +54,14 @@ func WithRequestTimeout(timeout time.Duration) Option {
 	}
 }
 
+// WithUsePromptCaching pass in true to enable prompt caching.
+// Prompt caching is only supported by OpenAI-native endpoints.
+func WithUsePromptCaching(usePromptCaching bool) Option {
+	return func(p *Provider) {
+		p.usePromptCaching = usePromptCaching
+	}
+}
+
 func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider {
 	client := &http.Client{
 		Timeout: defaultRequestTimeout,
@@ -157,7 +166,7 @@ func (p *Provider) Chat(
 	// Prompt caching is only supported by OpenAI-native endpoints.
 	// Gemini and other providers reject unknown fields, so skip for non-OpenAI APIs.
 	if cacheKey, ok := options["prompt_cache_key"].(string); ok && cacheKey != "" {
-		if !strings.Contains(p.apiBase, "generativelanguage.googleapis.com") {
+		if p.usePromptCaching {
 			requestBody["prompt_cache_key"] = cacheKey
 		}
 	}