🔧 #164 Make client connection pool configurable across all providers (#…

…251)
EinStack · May 8, 2024 · bd03442 · bd03442
1 parent ab244b2
commit bd03442
Show file tree

Hide file tree

Showing 12 changed files with 60 additions and 23 deletions.
diff --git 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
@@ -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
@@ -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

diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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,