From e36caf342d3d10edb480d0d57f485803dbdd9157 Mon Sep 17 00:00:00 2001 From: daesu Date: Tue, 7 May 2024 10:19:13 +0100 Subject: [PATCH 1/2] Make client connection pool configurable across all providers --- pkg/providers/anthropic/client.go | 5 ++--- pkg/providers/azureopenai/client.go | 5 ++--- pkg/providers/bedrock/client.go | 5 ++--- pkg/providers/clients/config.go | 10 ++++++++-- pkg/providers/clients/config_test.go | 22 ++++++++++++++++++++++ pkg/providers/cohere/client.go | 5 ++--- pkg/providers/octoml/client.go | 5 ++--- pkg/providers/ollama/client.go | 5 ++--- pkg/providers/openai/client.go | 5 ++--- 9 files changed, 44 insertions(+), 23 deletions(-) diff --git a/pkg/providers/anthropic/client.go b/pkg/providers/anthropic/client.go index 2ddd6099..16df6974 100644 --- a/pkg/providers/anthropic/client.go +++ b/pkg/providers/anthropic/client.go @@ -47,10 +47,9 @@ func NewClient(providerConfig *Config, clientConfig *clients.ClientConfig, tel * errMapper: NewErrorMapper(tel), httpClient: &http.Client{ Timeout: *clientConfig.Timeout, - // TODO: use values from the config Transport: &http.Transport{ - MaxIdleConns: 100, - MaxIdleConnsPerHost: 2, + MaxIdleConns: *clientConfig.MaxIdleConns, + MaxIdleConnsPerHost: *clientConfig.MaxIdleConnsPerHost, }, }, tel: tel, diff --git a/pkg/providers/azureopenai/client.go b/pkg/providers/azureopenai/client.go index 05e59b81..3df92c28 100644 --- a/pkg/providers/azureopenai/client.go +++ b/pkg/providers/azureopenai/client.go @@ -50,11 +50,10 @@ func NewClient(providerConfig *Config, clientConfig *clients.ClientConfig, tel * finishReasonMapper: openai.NewFinishReasonMapper(tel), errMapper: NewErrorMapper(tel), httpClient: &http.Client{ - // TODO: use values from the config Timeout: *clientConfig.Timeout, Transport: &http.Transport{ - MaxIdleConns: 100, - MaxIdleConnsPerHost: 2, + MaxIdleConns: *clientConfig.MaxIdleConns, + MaxIdleConnsPerHost: *clientConfig.MaxIdleConnsPerHost, }, }, tel: tel, diff --git a/pkg/providers/bedrock/client.go b/pkg/providers/bedrock/client.go index f47dc027..44feb197 100644 --- a/pkg/providers/bedrock/client.go +++ b/pkg/providers/bedrock/client.go @@ -57,10 +57,9 @@ func NewClient(providerConfig *Config, clientConfig *clients.ClientConfig, tel * chatRequestTemplate: NewChatRequestFromConfig(providerConfig), httpClient: &http.Client{ Timeout: *clientConfig.Timeout, - // TODO: use values from the config Transport: &http.Transport{ - MaxIdleConns: 100, - MaxIdleConnsPerHost: 2, + MaxIdleConns: *clientConfig.MaxIdleConns, + MaxIdleConnsPerHost: *clientConfig.MaxIdleConnsPerHost, }, }, telemetry: tel, diff --git a/pkg/providers/clients/config.go b/pkg/providers/clients/config.go index d01a2ab1..5ed23291 100644 --- a/pkg/providers/clients/config.go +++ b/pkg/providers/clients/config.go @@ -3,13 +3,19 @@ package clients import "time" type ClientConfig struct { - Timeout *time.Duration `yaml:"timeout,omitempty" json:"timeout" swaggertype:"primitive,string"` + Timeout *time.Duration `yaml:"timeout,omitempty" json:"timeout" swaggertype:"primitive,string"` + MaxIdleConns *int `yaml:"max_idle_connections,omitempty" json:"max_idle_connections" swaggertype:"primitive,integer"` + MaxIdleConnsPerHost *int `yaml:"max_idle_connections_per_host,omitempty" json:"max_idle_connections_per_host" swaggertype:"primitive,integer"` } func DefaultClientConfig() *ClientConfig { defaultTimeout := 10 * time.Second + maxIdleConns := 100 + maxIdleConnsPerHost := 2 return &ClientConfig{ - Timeout: &defaultTimeout, + Timeout: &defaultTimeout, + MaxIdleConns: &maxIdleConns, + MaxIdleConnsPerHost: &maxIdleConnsPerHost, } } diff --git a/pkg/providers/clients/config_test.go b/pkg/providers/clients/config_test.go index 0e725201..bc28444b 100644 --- a/pkg/providers/clients/config_test.go +++ b/pkg/providers/clients/config_test.go @@ -2,6 +2,7 @@ package clients import ( "testing" + "time" "github.com/stretchr/testify/require" ) @@ -11,3 +12,24 @@ func TestClientConfig_DefaultConfig(t *testing.T) { require.NotEmpty(t, config.Timeout) } + +func TestDefaultClientConfig(t *testing.T) { + config := DefaultClientConfig() + + require.NotNil(t, config, "Config must not be nil") + require.NotNil(t, config.Timeout, "Timeout must not be nil") + require.NotNil(t, config.MaxIdleConns, "MaxIdleConns must not be nil") + require.NotNil(t, config.MaxIdleConnsPerHost, "MaxIdleConnsPerHost must not be nil") + + // Check default timeout + expectedTimeout := 10 * time.Second + require.Equal(t, expectedTimeout, *config.Timeout) + + // Check MaxIdleConns + expectedMaxIdleConns := 100 + require.Equal(t, expectedMaxIdleConns, *config.MaxIdleConns) + + // Check MaxIdleConnsPerHost + expectedMaxIdleConnsPerHost := 2 + require.Equal(t, expectedMaxIdleConnsPerHost, *config.MaxIdleConnsPerHost) +} diff --git a/pkg/providers/cohere/client.go b/pkg/providers/cohere/client.go index b4e23567..c355a57c 100644 --- a/pkg/providers/cohere/client.go +++ b/pkg/providers/cohere/client.go @@ -45,10 +45,9 @@ func NewClient(providerConfig *Config, clientConfig *clients.ClientConfig, tel * chatRequestTemplate: NewChatRequestFromConfig(providerConfig), httpClient: &http.Client{ Timeout: *clientConfig.Timeout, - // TODO: use values from the config Transport: &http.Transport{ - MaxIdleConns: 100, - MaxIdleConnsPerHost: 2, + MaxIdleConns: *clientConfig.MaxIdleConns, + MaxIdleConnsPerHost: *clientConfig.MaxIdleConnsPerHost, }, }, errMapper: NewErrorMapper(tel), diff --git a/pkg/providers/octoml/client.go b/pkg/providers/octoml/client.go index 89a1094e..d9092780 100644 --- a/pkg/providers/octoml/client.go +++ b/pkg/providers/octoml/client.go @@ -45,10 +45,9 @@ func NewClient(providerConfig *Config, clientConfig *clients.ClientConfig, tel * errMapper: NewErrorMapper(tel), httpClient: &http.Client{ Timeout: *clientConfig.Timeout, - // TODO: use values from the config Transport: &http.Transport{ - MaxIdleConns: 100, - MaxIdleConnsPerHost: 2, + MaxIdleConns: *clientConfig.MaxIdleConns, + MaxIdleConnsPerHost: *clientConfig.MaxIdleConnsPerHost, }, }, telemetry: tel, diff --git a/pkg/providers/ollama/client.go b/pkg/providers/ollama/client.go index 7dda3f90..177e0bab 100644 --- a/pkg/providers/ollama/client.go +++ b/pkg/providers/ollama/client.go @@ -43,10 +43,9 @@ func NewClient(providerConfig *Config, clientConfig *clients.ClientConfig, tel * chatRequestTemplate: NewChatRequestFromConfig(providerConfig), httpClient: &http.Client{ Timeout: *clientConfig.Timeout, - // TODO: use values from the config Transport: &http.Transport{ - MaxIdleConns: 100, - MaxIdleConnsPerHost: 2, + MaxIdleConns: *clientConfig.MaxIdleConns, + MaxIdleConnsPerHost: *clientConfig.MaxIdleConnsPerHost, }, }, telemetry: tel, diff --git a/pkg/providers/openai/client.go b/pkg/providers/openai/client.go index c39f2ac5..a27093fd 100644 --- a/pkg/providers/openai/client.go +++ b/pkg/providers/openai/client.go @@ -54,10 +54,9 @@ func NewClient(providerConfig *Config, clientConfig *clients.ClientConfig, tel * errMapper: NewErrorMapper(tel), httpClient: &http.Client{ Timeout: *clientConfig.Timeout, - // TODO: use values from the config Transport: &http.Transport{ - MaxIdleConns: 100, - MaxIdleConnsPerHost: 2, + MaxIdleConns: *clientConfig.MaxIdleConns, + MaxIdleConnsPerHost: *clientConfig.MaxIdleConnsPerHost, }, }, tel: tel, From 76cb75e8bcda0241b99d48b2817ea145568ddd89 Mon Sep 17 00:00:00 2001 From: daesu Date: Wed, 8 May 2024 09:53:04 +0100 Subject: [PATCH 2/2] regenerated swagger docs --- docs/docs.go | 6 ++++++ docs/swagger.json | 6 ++++++ docs/swagger.yaml | 4 ++++ 3 files changed, 16 insertions(+) diff --git a/docs/docs.go b/docs/docs.go index 65d48212..a236eac8 100644 --- a/docs/docs.go +++ b/docs/docs.go @@ -373,6 +373,12 @@ const docTemplate = `{ "clients.ClientConfig": { "type": "object", "properties": { + "max_idle_connections": { + "type": "integer" + }, + "max_idle_connections_per_host": { + "type": "integer" + }, "timeout": { "type": "string" } diff --git a/docs/swagger.json b/docs/swagger.json index 1588b4a5..f626c88d 100644 --- a/docs/swagger.json +++ b/docs/swagger.json @@ -370,6 +370,12 @@ "clients.ClientConfig": { "type": "object", "properties": { + "max_idle_connections": { + "type": "integer" + }, + "max_idle_connections_per_host": { + "type": "integer" + }, "timeout": { "type": "string" } diff --git a/docs/swagger.yaml b/docs/swagger.yaml index 9078f891..c4608100 100644 --- a/docs/swagger.yaml +++ b/docs/swagger.yaml @@ -126,6 +126,10 @@ definitions: type: object clients.ClientConfig: properties: + max_idle_connections: + type: integer + max_idle_connections_per_host: + type: integer timeout: type: string type: object