From b3772f4930f3420a5eeb034f3c7f5b4e6e33c6e3 Mon Sep 17 00:00:00 2001 From: Christopher Speller Date: Wed, 21 Aug 2024 09:30:16 -0700 Subject: [PATCH] Use a common http client for LLM upstream calls (#236) --- server/ai/anthropic/anthropic.go | 5 +++-- server/ai/anthropic/client.go | 6 +++--- server/ai/asksage/asksage.go | 5 +++-- server/ai/asksage/client.go | 4 ++-- server/ai/openai/openai.go | 7 +++++-- server/plugin.go | 19 +++++++++++++------ 6 files changed, 29 insertions(+), 17 deletions(-) diff --git a/server/ai/anthropic/anthropic.go b/server/ai/anthropic/anthropic.go index 1d08f0cf..b8ad60cc 100644 --- a/server/ai/anthropic/anthropic.go +++ b/server/ai/anthropic/anthropic.go @@ -2,6 +2,7 @@ package anthropic import ( "fmt" + "net/http" "github.com/mattermost/mattermost-plugin-ai/server/ai" "github.com/mattermost/mattermost-plugin-ai/server/metrics" @@ -16,8 +17,8 @@ type Anthropic struct { metricsService metrics.LLMetrics } -func New(llmService ai.ServiceConfig, metricsService metrics.LLMetrics) *Anthropic { - client := NewClient(llmService.APIKey) +func New(llmService ai.ServiceConfig, httpClient *http.Client, metricsService metrics.LLMetrics) *Anthropic { + client := NewClient(llmService.APIKey, httpClient) return &Anthropic{ client: client, diff --git a/server/ai/anthropic/client.go b/server/ai/anthropic/client.go index 2ef3eb63..b1a7fd1a 100644 --- a/server/ai/anthropic/client.go +++ b/server/ai/anthropic/client.go @@ -75,13 +75,13 @@ type MessageStreamEvent struct { type Client struct { apiKey string - httpClient http.Client + httpClient *http.Client } -func NewClient(apiKey string) *Client { +func NewClient(apiKey string, httpClient *http.Client) *Client { return &Client{ apiKey: apiKey, - httpClient: http.Client{}, + httpClient: httpClient, } } diff --git a/server/ai/asksage/asksage.go b/server/ai/asksage/asksage.go index 2301f4cd..5d975d5d 100644 --- a/server/ai/asksage/asksage.go +++ b/server/ai/asksage/asksage.go @@ -1,6 +1,7 @@ package asksage import ( + "net/http" "strings" "github.com/mattermost/mattermost-plugin-ai/server/ai" @@ -14,8 +15,8 @@ type AskSage struct { metric metrics.LLMetrics } -func New(llmService ai.ServiceConfig, metric metrics.LLMetrics) *AskSage { - client := NewClient("") +func New(llmService ai.ServiceConfig, httpClient *http.Client, metric metrics.LLMetrics) *AskSage { + client := NewClient("", httpClient) client.Login(GetTokenParams{ Email: llmService.Username, Password: llmService.Password, diff --git a/server/ai/asksage/client.go b/server/ai/asksage/client.go index 8f650895..4017cb19 100644 --- a/server/ai/asksage/client.go +++ b/server/ai/asksage/client.go @@ -66,10 +66,10 @@ type Persona struct { type Dataset string -func NewClient(authToken string) *Client { +func NewClient(authToken string, httpClient *http.Client) *Client { return &Client{ AuthToken: authToken, - HTTPClient: &http.Client{}, + HTTPClient: httpClient, } } diff --git a/server/ai/openai/openai.go b/server/ai/openai/openai.go index 4a5db9f1..ffffcdfd 100644 --- a/server/ai/openai/openai.go +++ b/server/ai/openai/openai.go @@ -9,6 +9,7 @@ import ( "image" "image/png" "io" + "net/http" "net/url" "strings" "time" @@ -38,12 +39,13 @@ const OpenAIMaxImageSize = 20 * 1024 * 1024 // 20 MB var ErrStreamingTimeout = errors.New("timeout streaming") -func NewCompatible(llmService ai.ServiceConfig, metricsService metrics.LLMetrics) *OpenAI { +func NewCompatible(llmService ai.ServiceConfig, httpClient *http.Client, metricsService metrics.LLMetrics) *OpenAI { apiKey := llmService.APIKey endpointURL := strings.TrimSuffix(llmService.APIURL, "/") defaultModel := llmService.DefaultModel config := openaiClient.DefaultConfig(apiKey) config.BaseURL = endpointURL + config.HTTPClient = httpClient parsedURL, err := url.Parse(endpointURL) if err == nil && strings.HasSuffix(parsedURL.Host, "openai.azure.com") { @@ -64,13 +66,14 @@ func NewCompatible(llmService ai.ServiceConfig, metricsService metrics.LLMetrics } } -func New(llmService ai.ServiceConfig, metricsService metrics.LLMetrics) *OpenAI { +func New(llmService ai.ServiceConfig, httpClient *http.Client, metricsService metrics.LLMetrics) *OpenAI { defaultModel := llmService.DefaultModel if defaultModel == "" { defaultModel = openaiClient.GPT3Dot5Turbo } config := openaiClient.DefaultConfig(llmService.APIKey) config.OrgID = llmService.OrgID + config.HTTPClient = httpClient streamingTimeout := StreamingTimeoutDefault if llmService.StreamingTimeoutSeconds > 0 { diff --git a/server/plugin.go b/server/plugin.go index 04a621af..0f40afc1 100644 --- a/server/plugin.go +++ b/server/plugin.go @@ -7,6 +7,7 @@ import ( "os" "os/exec" "sync" + "time" "errors" @@ -22,6 +23,7 @@ import ( "github.com/mattermost/mattermost/server/public/model" "github.com/mattermost/mattermost/server/public/plugin" "github.com/mattermost/mattermost/server/public/pluginapi" + "github.com/mattermost/mattermost/server/public/shared/httpservice" "github.com/nicksnyder/go-i18n/v2/i18n" ) @@ -72,6 +74,8 @@ type Plugin struct { bots []*Bot i18n *i18n.Bundle + + llmUpstreamHTTPClient *http.Client } func resolveffmpegPath() string { @@ -100,6 +104,9 @@ func (p *Plugin) OnActivate() error { p.i18n = i18nInit() + p.llmUpstreamHTTPClient = httpservice.MakeHTTPServicePlugin(p.API).MakeClient(true) + p.llmUpstreamHTTPClient.Timeout = time.Minute * 10 // LLM requests can be slow + if err := p.MigrateServicesToBots(); err != nil { p.pluginAPI.Log.Error("failed to migrate services to bots", "error", err) // Don't fail on migration errors @@ -144,13 +151,13 @@ func (p *Plugin) getLLM(llmBotConfig ai.BotConfig) ai.LanguageModel { var llm ai.LanguageModel switch llmBotConfig.Service.Type { case "openai": - llm = openai.New(llmBotConfig.Service, llmMetrics) + llm = openai.New(llmBotConfig.Service, p.llmUpstreamHTTPClient, llmMetrics) case "openaicompatible": - llm = openai.NewCompatible(llmBotConfig.Service, llmMetrics) + llm = openai.NewCompatible(llmBotConfig.Service, p.llmUpstreamHTTPClient, llmMetrics) case "anthropic": - llm = anthropic.New(llmBotConfig.Service, llmMetrics) + llm = anthropic.New(llmBotConfig.Service, p.llmUpstreamHTTPClient, llmMetrics) case "asksage": - llm = asksage.New(llmBotConfig.Service, llmMetrics) + llm = asksage.New(llmBotConfig.Service, p.llmUpstreamHTTPClient, llmMetrics) } cfg := p.getConfiguration() @@ -175,9 +182,9 @@ func (p *Plugin) getTranscribe() ai.Transcriber { llmMetrics := p.metricsService.GetMetricsForAIService(botConfig.Name) switch botConfig.Service.Type { case "openai": - return openai.New(botConfig.Service, llmMetrics) + return openai.New(botConfig.Service, p.llmUpstreamHTTPClient, llmMetrics) case "openaicompatible": - return openai.NewCompatible(botConfig.Service, llmMetrics) + return openai.NewCompatible(botConfig.Service, p.llmUpstreamHTTPClient, llmMetrics) } return nil }