Prevent having to pass posts to services (#207)

server/ai/anthropic/anthropic.go

@@ -13,19 +13,17 @@ type Anthropic struct {
 	client         *Client
 	defaultModel   string
 	tokenLimit     int
-	metricsService metrics.Metrics
-	name           string
+	metricsService metrics.LLMetrics
 }
 
-func New(botConfig ai.BotConfig, metricsService metrics.Metrics) *Anthropic {
-	client := NewClient(botConfig.Service.APIKey)
+func New(llmService ai.ServiceConfig, metricsService metrics.LLMetrics) *Anthropic {
+	client := NewClient(llmService.APIKey)
 
 	return &Anthropic{
 		client:         client,
-		defaultModel:   botConfig.Service.DefaultModel,
-		tokenLimit:     botConfig.Service.TokenLimit,
+		defaultModel:   llmService.DefaultModel,
+		tokenLimit:     llmService.TokenLimit,
 		metricsService: metricsService,
-		name:           botConfig.Name,
 	}
 }
 
@@ -84,7 +82,7 @@ func (a *Anthropic) createCompletionRequest(conversation ai.BotConversation, opt
 }
 
 func (a *Anthropic) ChatCompletion(conversation ai.BotConversation, opts ...ai.LanguageModelOption) (*ai.TextStreamResult, error) {
-	a.metricsService.IncrementLLMRequests(a.name)
+	a.metricsService.IncrementLLMRequests()
 
 	request := a.createCompletionRequest(conversation, opts)
 	request.Stream = true
@@ -97,7 +95,7 @@ func (a *Anthropic) ChatCompletion(conversation ai.BotConversation, opts ...ai.L
 }
 
 func (a *Anthropic) ChatCompletionNoStream(conversation ai.BotConversation, opts ...ai.LanguageModelOption) (string, error) {
-	a.metricsService.IncrementLLMRequests(a.name)
+	a.metricsService.IncrementLLMRequests()
 
 	request := a.createCompletionRequest(conversation, opts)
 	request.Stream = false

server/ai/asksage/asksage.go

@@ -8,25 +8,23 @@ import (
 )
 
 type AskSage struct {
-	client         *Client
-	defaultModel   string
-	maxTokens      int
-	metricsService metrics.Metrics
-	name           string
+	client       *Client
+	defaultModel string
+	maxTokens    int
+	metric       metrics.LLMetrics
 }
 
-func New(botConfig ai.BotConfig, metricsService metrics.Metrics) *AskSage {
+func New(llmService ai.ServiceConfig, metric metrics.LLMetrics) *AskSage {
 	client := NewClient("")
 	client.Login(GetTokenParams{
-		Email:    botConfig.Service.Username,
-		Password: botConfig.Service.Password,
+		Email:    llmService.Username,
+		Password: llmService.Password,
 	})
 	return &AskSage{
-		client:         client,
-		defaultModel:   botConfig.Service.DefaultModel,
-		maxTokens:      botConfig.Service.TokenLimit,
-		metricsService: metricsService,
-		name:           botConfig.Name,
+		client:       client,
+		defaultModel: llmService.DefaultModel,
+		maxTokens:    llmService.TokenLimit,
+		metric:       metric,
 	}
 }
 
@@ -80,7 +78,7 @@ func (s *AskSage) ChatCompletion(conversation ai.BotConversation, opts ...ai.Lan
 }
 
 func (s *AskSage) ChatCompletionNoStream(conversation ai.BotConversation, opts ...ai.LanguageModelOption) (string, error) {
-	s.metricsService.IncrementLLMRequests(s.name)
+	s.metric.IncrementLLMRequests()
 
 	params := s.queryParamsFromConfig(s.createConfig(opts))
 	params.Message = conversationToMessagesList(conversation)

server/ai/openai/openai.go

@@ -27,8 +27,7 @@ type OpenAI struct {
 	defaultModel     string
 	tokenLimit       int
 	streamingTimeout time.Duration
-	metricsService   metrics.Metrics
-	name             string
+	metricsService   metrics.LLMetrics
 }
 
 const StreamingTimeoutDefault = 10 * time.Second
@@ -39,10 +38,10 @@ const OpenAIMaxImageSize = 20 * 1024 * 1024 // 20 MB
 
 var ErrStreamingTimeout = errors.New("timeout streaming")
 
-func NewCompatible(botConfig ai.BotConfig, metricsService metrics.Metrics) *OpenAI {
-	apiKey := botConfig.Service.APIKey
-	endpointURL := strings.TrimSuffix(botConfig.Service.APIURL, "/")
-	defaultModel := botConfig.Service.DefaultModel
+func NewCompatible(llmService ai.ServiceConfig, metricsService metrics.LLMetrics) *OpenAI {
+	apiKey := llmService.APIKey
+	endpointURL := strings.TrimSuffix(llmService.APIURL, "/")
+	defaultModel := llmService.DefaultModel
 	config := openaiClient.DefaultConfig(apiKey)
 	config.BaseURL = endpointURL
 
@@ -53,39 +52,37 @@ func NewCompatible(botConfig ai.BotConfig, metricsService metrics.Metrics) *Open
 	}
 
 	streamingTimeout := StreamingTimeoutDefault
-	if botConfig.Service.StreamingTimeoutSeconds > 0 {
-		streamingTimeout = time.Duration(botConfig.Service.StreamingTimeoutSeconds) * time.Second
+	if llmService.StreamingTimeoutSeconds > 0 {
+		streamingTimeout = time.Duration(llmService.StreamingTimeoutSeconds) * time.Second
 	}
 	return &OpenAI{
 		client:           openaiClient.NewClientWithConfig(config),
 		defaultModel:     defaultModel,
-		tokenLimit:       botConfig.Service.TokenLimit,
+		tokenLimit:       llmService.TokenLimit,
 		streamingTimeout: streamingTimeout,
 		metricsService:   metricsService,
-		name:             botConfig.Name,
 	}
 }
 
-func New(botConfig ai.BotConfig, metricsService metrics.Metrics) *OpenAI {
-	defaultModel := botConfig.Service.DefaultModel
+func New(llmService ai.ServiceConfig, metricsService metrics.LLMetrics) *OpenAI {
+	defaultModel := llmService.DefaultModel
 	if defaultModel == "" {
 		defaultModel = openaiClient.GPT3Dot5Turbo
 	}
-	config := openaiClient.DefaultConfig(botConfig.Service.APIKey)
-	config.OrgID = botConfig.Service.OrgID
+	config := openaiClient.DefaultConfig(llmService.APIKey)
+	config.OrgID = llmService.OrgID
 
 	streamingTimeout := StreamingTimeoutDefault
-	if botConfig.Service.StreamingTimeoutSeconds > 0 {
-		streamingTimeout = time.Duration(botConfig.Service.StreamingTimeoutSeconds) * time.Second
+	if llmService.StreamingTimeoutSeconds > 0 {
+		streamingTimeout = time.Duration(llmService.StreamingTimeoutSeconds) * time.Second
 	}
 
 	return &OpenAI{
 		client:           openaiClient.NewClientWithConfig(config),
 		defaultModel:     defaultModel,
-		tokenLimit:       botConfig.Service.TokenLimit,
+		tokenLimit:       llmService.TokenLimit,
 		streamingTimeout: streamingTimeout,
 		metricsService:   metricsService,
-		name:             botConfig.Name,
 	}
 }
 
@@ -351,7 +348,7 @@ func (s *OpenAI) completionRequestFromConfig(cfg ai.LLMConfig) openaiClient.Chat
 }
 
 func (s *OpenAI) ChatCompletion(conversation ai.BotConversation, opts ...ai.LanguageModelOption) (*ai.TextStreamResult, error) {
-	s.metricsService.IncrementLLMRequests(s.name)
+	s.metricsService.IncrementLLMRequests()
 
 	request := s.completionRequestFromConfig(s.createConfig(opts))
 	request = modifyCompletionRequestWithConversation(request, conversation)

server/metrics/metrics.go

@@ -25,7 +25,7 @@ type Metrics interface {
 	IncrementHTTPRequests()
 	IncrementHTTPErrors()
 
-	IncrementLLMRequests(llmName string)
+	GetMetricsForAIService(llmName string) *llmMetrics
 }
 
 type InstanceInfo struct {
@@ -152,8 +152,26 @@ func (m *metrics) IncrementHTTPErrors() {
 	}
 }
 
-func (m *metrics) IncrementLLMRequests(llmName string) {
+func (m *metrics) GetMetricsForAIService(llmName string) *llmMetrics {
+	if m == nil {
+		return nil
+	}
+
+	return &llmMetrics{
+		llmRequestsTotal: m.llmRequestsTotal.MustCurryWith(prometheus.Labels{"llm_name": llmName}),
+	}
+}
+
+type LLMetrics interface {
+	IncrementLLMRequests()
+}
+
+type llmMetrics struct {
+	llmRequestsTotal *prometheus.CounterVec
+}
+
+func (m *llmMetrics) IncrementLLMRequests() {
 	if m != nil {
-		m.llmRequestsTotal.With(prometheus.Labels{"llm_name": llmName}).Inc()
+		m.llmRequestsTotal.With(prometheus.Labels{}).Inc()
 	}
 }

server/plugin.go

@@ -128,16 +128,18 @@ func (p *Plugin) OnActivate() error {
 }
 
 func (p *Plugin) getLLM(llmBotConfig ai.BotConfig) ai.LanguageModel {
+	metrics := p.metricsService.GetMetricsForAIService(llmBotConfig.Name)
+
 	var llm ai.LanguageModel
 	switch llmBotConfig.Service.Type {
 	case "openai":
-		llm = openai.New(llmBotConfig, p.metricsService)
+		llm = openai.New(llmBotConfig.Service, metrics)
 	case "openaicompatible":
-		llm = openai.NewCompatible(llmBotConfig, p.metricsService)
+		llm = openai.NewCompatible(llmBotConfig.Service, metrics)
 	case "anthropic":
-		llm = anthropic.New(llmBotConfig, p.metricsService)
+		llm = anthropic.New(llmBotConfig.Service, metrics)
 	case "asksage":
-		llm = asksage.New(llmBotConfig, p.metricsService)
+		llm = asksage.New(llmBotConfig.Service, metrics)
 	}
 
 	cfg := p.getConfiguration()
@@ -159,11 +161,12 @@ func (p *Plugin) getTranscribe() ai.Transcriber {
 			break
 		}
 	}
+	metrics := p.metricsService.GetMetricsForAIService(botConfig.Name)
 	switch botConfig.Service.Type {
 	case "openai":
-		return openai.New(botConfig, p.metricsService)
+		return openai.New(botConfig.Service, metrics)
 	case "openaicompatible":
-		return openai.NewCompatible(botConfig, p.metricsService)
+		return openai.NewCompatible(botConfig.Service, metrics)
 	}
 	return nil
 }