Merge pull request #338 from grafana/backend-model-mapping

Author: sd2k

packages/grafana-llm-app/pkg/plugin/app.go

@@ -53,15 +53,20 @@ func NewApp(ctx context.Context, appSettings backend.AppInstanceSettings) (insta
 		return nil, err
 	}
 
+	if app.settings.Models == nil {
+		// backwards-compat: if Model settings is nil, use the default one
+		app.settings.Models = DEFAULT_MODEL_SETTINGS
+	}
+
 	switch app.settings.OpenAI.Provider {
 	case openAIProviderOpenAI:
-		p, err := NewOpenAIProvider(app.settings.OpenAI)
+		p, err := NewOpenAIProvider(app.settings.OpenAI, app.settings.Models)
 		if err != nil {
 			return nil, err
 		}
 		app.llmProvider = p
 	case openAIProviderAzure:
-		p, err := NewAzureProvider(app.settings.OpenAI)
+		p, err := NewAzureProvider(app.settings.OpenAI, app.settings.Models.Default)
 		if err != nil {
 			return nil, err
 		}

packages/grafana-llm-app/pkg/plugin/azure_provider.go

@@ -11,16 +11,18 @@ import (
 )
 
 type azure struct {
-	settings OpenAISettings
-	oc       *openai.Client
+	settings     OpenAISettings
+	defaultModel Model
+	oc           *openai.Client
 }
 
-func NewAzureProvider(settings OpenAISettings) (LLMProvider, error) {
+func NewAzureProvider(settings OpenAISettings, defaultModel Model) (LLMProvider, error) {
 	client := &http.Client{
 		Timeout: 2 * time.Minute,
 	}
 	p := &azure{
-		settings: settings,
+		settings:     settings,
+		defaultModel: defaultModel,
 	}
 
 	// go-openai expects the URL without the '/openai' suffix, which is
@@ -48,14 +50,25 @@ func (p *azure) Models(ctx context.Context) (ModelResponse, error) {
 	return ModelResponse{Data: models}, nil
 }
 
-func (p *azure) ChatCompletion(ctx context.Context, req ChatCompletionRequest) (openai.ChatCompletionResponse, error) {
+func (p *azure) getDeployment(model Model) (string, error) {
 	mapping, err := p.getAzureMapping()
 	if err != nil {
-		return openai.ChatCompletionResponse{}, err
+		return "", err
+	}
+	if model == "" {
+		model = p.defaultModel
 	}
-	deployment := mapping[req.Model]
+	deployment := mapping[model]
 	if deployment == "" {
-		return openai.ChatCompletionResponse{}, fmt.Errorf("%w: no deployment found for model: %s", errBadRequest, req.Model)
+		return "", fmt.Errorf("%w: no deployment found for model: %s", errBadRequest, model)
+	}
+	return deployment, nil
+}
+
+func (p *azure) ChatCompletion(ctx context.Context, req ChatCompletionRequest) (openai.ChatCompletionResponse, error) {
+	deployment, err := p.getDeployment(req.Model)
+	if err != nil {
+		return openai.ChatCompletionResponse{}, err
 	}
 
 	r := req.ChatCompletionRequest
@@ -69,14 +82,10 @@ func (p *azure) ChatCompletion(ctx context.Context, req ChatCompletionRequest) (
 }
 
 func (p *azure) ChatCompletionStream(ctx context.Context, req ChatCompletionRequest) (<-chan ChatCompletionStreamResponse, error) {
-	mapping, err := p.getAzureMapping()
+	deployment, err := p.getDeployment(req.Model)
 	if err != nil {
 		return nil, err
 	}
-	deployment := mapping[req.Model]
-	if deployment == "" {
-		return nil, fmt.Errorf("%w: no deployment found for model: %s", errBadRequest, req.Model)
-	}
 
 	r := req.ChatCompletionRequest
 	// For the Azure mapping we want to use the name of the mapped deployment as the model.

packages/grafana-llm-app/pkg/plugin/grafana_provider.go

@@ -59,7 +59,7 @@ func (p *grafanaProvider) Models(ctx context.Context) (ModelResponse, error) {
 
 func (p *grafanaProvider) ChatCompletion(ctx context.Context, req ChatCompletionRequest) (openai.ChatCompletionResponse, error) {
 	r := req.ChatCompletionRequest
-	r.Model = req.Model.toOpenAI()
+	r.Model = req.Model.toOpenAI(DEFAULT_MODEL_SETTINGS)
 	resp, err := p.oc.CreateChatCompletion(ctx, r)
 	if err != nil {
 		log.DefaultLogger.Error("error creating grafana chat completion", "err", err)
@@ -70,6 +70,6 @@ func (p *grafanaProvider) ChatCompletion(ctx context.Context, req ChatCompletion
 
 func (p *grafanaProvider) ChatCompletionStream(ctx context.Context, req ChatCompletionRequest) (<-chan ChatCompletionStreamResponse, error) {
 	r := req.ChatCompletionRequest
-	r.Model = req.Model.toOpenAI()
+	r.Model = req.Model.toOpenAI(DEFAULT_MODEL_SETTINGS)
 	return streamOpenAIRequest(ctx, r, p.oc)
 }

packages/grafana-llm-app/pkg/plugin/llm_provider.go

@@ -49,15 +49,17 @@ func (m *Model) UnmarshalJSON(data []byte) error {
 	return fmt.Errorf("unrecognized model: %s", dataString)
 }
 
-func (m Model) toOpenAI() string {
-	// TODO: Add ability to change which model is used for each abstraction in settings.
-	switch m {
-	case ModelBase:
-		return "gpt-3.5-turbo"
-	case ModelLarge:
-		return "gpt-4-turbo"
+func (m Model) toOpenAI(modelSettings *ModelSettings) string {
+	if modelSettings == nil || len(modelSettings.Mapping) == 0 {
+		switch m {
+		case ModelBase:
+			return "gpt-3.5-turbo"
+		case ModelLarge:
+			return "gpt-4-turbo"
+		}
+		panic(fmt.Sprintf("unrecognized model: %s", m))
 	}
-	panic("unknown model: " + m)
+	return modelSettings.getModel(m)
 }
 
 type ChatCompletionRequest struct {

packages/grafana-llm-app/pkg/plugin/openai_provider.go

@@ -15,10 +15,11 @@ import (
 
 type openAI struct {
 	settings OpenAISettings
+	models   *ModelSettings
 	oc       *openai.Client
 }
 
-func NewOpenAIProvider(settings OpenAISettings) (LLMProvider, error) {
+func NewOpenAIProvider(settings OpenAISettings, models *ModelSettings) (LLMProvider, error) {
 	client := &http.Client{
 		Timeout: 2 * time.Minute,
 	}
@@ -32,6 +33,7 @@ func NewOpenAIProvider(settings OpenAISettings) (LLMProvider, error) {
 	cfg.OrgID = settings.OrganizationID
 	return &openAI{
 		settings: settings,
+		models:   models,
 		oc:       openai.NewClientWithConfig(cfg),
 	}, nil
 }
@@ -53,7 +55,7 @@ type openAIChatCompletionRequest struct {
 
 func (p *openAI) ChatCompletion(ctx context.Context, req ChatCompletionRequest) (openai.ChatCompletionResponse, error) {
 	r := req.ChatCompletionRequest
-	r.Model = req.Model.toOpenAI()
+	r.Model = req.Model.toOpenAI(p.models)
 	resp, err := p.oc.CreateChatCompletion(ctx, r)
 	if err != nil {
 		log.DefaultLogger.Error("error creating openai chat completion", "err", err)
@@ -64,7 +66,7 @@ func (p *openAI) ChatCompletion(ctx context.Context, req ChatCompletionRequest)
 
 func (p *openAI) ChatCompletionStream(ctx context.Context, req ChatCompletionRequest) (<-chan ChatCompletionStreamResponse, error) {
 	r := req.ChatCompletionRequest
-	r.Model = req.Model.toOpenAI()
+	r.Model = req.Model.toOpenAI(p.models)
 	return streamOpenAIRequest(ctx, r, p.oc)
 }
 

packages/grafana-llm-app/pkg/plugin/resources_test.go

@@ -237,6 +237,30 @@ func TestCallOpenAIProxy(t *testing.T) {
 
 			expStatus: http.StatusOK,
 		},
+		{
+			name: "openai - empty model",
+
+			settings: Settings{
+				OpenAI: OpenAISettings{
+					OrganizationID: "myOrg",
+					Provider:       openAIProviderOpenAI,
+				},
+			},
+			apiKey: "abcd1234",
+
+			method: http.MethodPost,
+			path:   "/openai/v1/chat/completions",
+			body:   []byte(`{"messages": [{"content":"some stuff"}]}`),
+
+			expReqHeaders: http.Header{
+				"Authorization":       {"Bearer abcd1234"},
+				"OpenAI-Organization": {"myOrg"},
+			},
+			expReqPath: "/v1/chat/completions",
+			expReqBody: []byte(`{"model": "gpt-3.5-turbo", "messages": [{"content":"some stuff"}]}`),
+
+			expStatus: http.StatusOK,
+		},
 		{
 			name: "openai - streaming",
 
@@ -265,6 +289,34 @@ func TestCallOpenAIProxy(t *testing.T) {
 			// newlines (required by the SSE spec) are escaped.
 			expBody: []byte("data: {\"id\":\"\",\"object\":\"\",\"created\":0,\"model\":\"\",\"choices\":null,\"system_fingerprint\":\"\"}\n\ndata: [DONE]\n\n"),
 		},
+		{
+			name: "openai - streaming - empty model",
+
+			settings: Settings{
+				OpenAI: OpenAISettings{
+					OrganizationID: "myOrg",
+					Provider:       openAIProviderOpenAI,
+				},
+			},
+			apiKey: "abcd1234",
+
+			method: http.MethodPost,
+			path:   "/openai/v1/chat/completions",
+			body:   []byte(`{"stream": true, "messages": [{"content":"some stuff"}]}`),
+
+			expReqHeaders: http.Header{
+				"Authorization":       {"Bearer abcd1234"},
+				"OpenAI-Organization": {"myOrg"},
+			},
+			expReqPath: "/v1/chat/completions",
+			expReqBody: []byte(`{"model": "gpt-3.5-turbo", "stream": true, "messages": [{"content":"some stuff"}]}`),
+
+			expStatus: http.StatusOK,
+
+			// We need to use regular strings rather than raw strings here otherwise the double
+			// newlines (required by the SSE spec) are escaped.
+			expBody: []byte("data: {\"id\":\"\",\"object\":\"\",\"created\":0,\"model\":\"\",\"choices\":null,\"system_fingerprint\":\"\"}\n\ndata: [DONE]\n\n"),
+		},
 		{
 			name: "azure",
 
@@ -293,6 +345,62 @@ func TestCallOpenAIProxy(t *testing.T) {
 
 			expStatus: http.StatusOK,
 		},
+		{
+			name: "azure - abstract model",
+
+			settings: Settings{
+				OpenAI: OpenAISettings{
+					OrganizationID: "myOrg",
+					Provider:       openAIProviderAzure,
+					AzureMapping: [][]string{
+						{"gpt-3.5-turbo", "gpt-35-turbo"},
+					},
+				},
+			},
+
+			apiKey: "abcd1234",
+
+			method: http.MethodPost,
+			path:   "/openai/v1/chat/completions",
+			body:   []byte(`{"model": "base", "messages": [{"content":"some stuff"}]}`),
+
+			expReqHeaders: http.Header{
+				"api-key": {"abcd1234"},
+			},
+			expReqPath: "/openai/deployments/gpt-35-turbo/chat/completions",
+			// the 'model' field should have been removed.
+			expReqBody: []byte(`{"messages":[{"content":"some stuff"}]}`),
+
+			expStatus: http.StatusOK,
+		},
+		{
+			name: "azure - empty model",
+
+			settings: Settings{
+				OpenAI: OpenAISettings{
+					OrganizationID: "myOrg",
+					Provider:       openAIProviderAzure,
+					AzureMapping: [][]string{
+						{"gpt-3.5-turbo", "gpt-35-turbo"},
+					},
+				},
+			},
+
+			apiKey: "abcd1234",
+
+			method: http.MethodPost,
+			path:   "/openai/v1/chat/completions",
+			body:   []byte(`{"messages": [{"content":"some stuff"}]}`),
+
+			expReqHeaders: http.Header{
+				"api-key": {"abcd1234"},
+			},
+			expReqPath: "/openai/deployments/gpt-35-turbo/chat/completions",
+			// the 'model' field should have been removed.
+			expReqBody: []byte(`{"messages":[{"content":"some stuff"}]}`),
+
+			expStatus: http.StatusOK,
+		},
 		{
 			name: "azure invalid deployment",
 
@@ -310,7 +418,7 @@ func TestCallOpenAIProxy(t *testing.T) {
 			method: http.MethodPost,
 			path:   "/openai/v1/chat/completions",
 			// note no gpt-4 in AzureMapping.
-			body: []byte(`{"model": "gpt-4", "messages": [{"content":"some stuff"}]}`),
+			body: []byte(`{"model": "gpt-4-turbo", "messages": [{"content":"some stuff"}]}`),
 
 			expNilRequest: true,
 
@@ -364,6 +472,31 @@ func TestCallOpenAIProxy(t *testing.T) {
 			expReqPath: "/llm/openai/v1/chat/completions",
 			expReqBody: []byte(`{"model": "gpt-3.5-turbo", "messages": [{"content":"some stuff"]}}`),
 
+			expStatus: http.StatusOK,
+		},
+		{
+			name: "grafana-managed llm gateway - empty model",
+
+			settings: Settings{
+				Tenant:           "123",
+				GrafanaComAPIKey: "abcd1234",
+				OpenAI: OpenAISettings{
+					Provider: openAIProviderGrafana,
+				},
+			},
+			apiKey: "abcd1234",
+
+			method: http.MethodPost,
+			path:   "/openai/v1/chat/completions",
+			body:   []byte(`{"messages": [{"content":"some stuff"}]}`),
+
+			expReqHeaders: http.Header{
+				"Authorization": {"Bearer 123:abcd1234"},
+				"X-Scope-OrgID": {"123"},
+			},
+			expReqPath: "/llm/openai/v1/chat/completions",
+			expReqBody: []byte(`{"model": "gpt-3.5-turbo", "messages": [{"content":"some stuff"]}}`),
+
 			expStatus: http.StatusOK,
 		},
 	} {

packages/grafana-llm-app/pkg/plugin/settings.go

@@ -42,6 +42,36 @@ type OpenAISettings struct {
 	apiKey string
 }
 
+type ModelMapping struct {
+	Model Model  `json:"model"`
+	Name  string `json:"name"`
+}
+
+type ModelSettings struct {
+	// Default model to use when no model is defined, or the model is not found.
+	Default Model `json:"default"`
+
+	// Mapping is mapping from our abstract model names to the provider's model names.
+	Mapping map[Model]string `json:"mapping"`
+}
+
+func (c ModelSettings) getModel(model Model) string {
+	// Helper function to get the name of a model.
+	if name, ok := c.Mapping[model]; ok {
+		return name
+	}
+	// If the model is not found, return the default model.
+	return c.getModel(c.Default)
+}
+
+var DEFAULT_MODEL_SETTINGS = &ModelSettings{
+	Default: ModelBase,
+	Mapping: map[Model]string{
+		ModelBase:  "gpt-3.5-turbo",
+		ModelLarge: "gpt-4-turbo",
+	},
+}
+
 // LLMGatewaySettings contains the configuration for the Grafana Managed Key LLM solution.
 type LLMGatewaySettings struct {
 	// This is the URL of the LLM endpoint of the machine learning backend which proxies
@@ -72,6 +102,9 @@ type Settings struct {
 	// VectorDB settings. May rely on OpenAI settings.
 	Vector vector.VectorSettings `json:"vector"`
 
+	// Models contains the user-specified models.
+	Models *ModelSettings `json:"models"`
+
 	// LLMGateway provides Grafana-managed OpenAI.
 	LLMGateway LLMGatewaySettings `json:"llmGateway"`
 }