added azureModelMapping Logic

grafana · Sep 22, 2023 · f6970d8 · f6970d8
1 parent a0d5195
commit f6970d8
Show file tree

Hide file tree

Showing 2 changed files with 111 additions and 6 deletions.
diff --git a/pkg/plugin/settings.go b/pkg/plugin/settings.go
@@ -10,10 +10,10 @@ import (
 const openAIKey = "openAIKey"
 
 type OpenAISettings struct {
-	URL            string            `json:"url"`
-	OrganizationID string            `json:"organizationId"`
-	UseAzure       bool              `json:"useAzure"`
-	AzureMapping   map[string]string `json:"azureOpenAIModelMapping"`
+	URL            string     `json:"url"`
+	OrganizationID string     `json:"organizationId"`
+	UseAzure       bool       `json:"useAzure"`
+	AzureMapping   [][]string `json:"azureModelMapping"`
 	apiKey         string
 }
 

diff --git a/src/components/AppConfig/AppConfig.tsx b/src/components/AppConfig/AppConfig.tsx
@@ -3,15 +3,18 @@ import { lastValueFrom } from 'rxjs';
 import { css } from '@emotion/css';
 import { AppPluginMeta, GrafanaTheme2, PluginConfigPageProps, PluginMeta } from '@grafana/data';
 import { getBackendSrv } from '@grafana/runtime';
-import { Button, Field, FieldSet, Input, SecretInput, useStyles2, Switch } from '@grafana/ui';
+import { Button, Field, FieldSet, Input, SecretInput, useStyles2, Switch, InlineField, IconButton, Select, InlineFieldRow } from '@grafana/ui';
 import { testIds } from '../testIds';
 
 type OpenAISettings = {
   url?: string;
   organizationId?: string;
   useAzure?: boolean;
+  azureModelMapping?: ModelDeploymentMap;
 }
 
+type ModelDeploymentMap = Array<[string, string]>;
+
 export type AppPluginSettings = {
   openAI?: OpenAISettings;
 };
@@ -27,12 +30,95 @@ type State = {
   openAIKey: string;
   // A flag to tell us if we should use Azure OpenAI.
   useAzureOpenAI: boolean;
+  // A mapping of Azure models to OpenAI models.
+  azureModelMapping: ModelDeploymentMap;
   // A flag to tell us if state was updated
   updated: boolean;
 };
 
 export interface AppConfigProps extends PluginConfigPageProps<AppPluginMeta<AppPluginSettings>> { }
 
+function ModelMappingConfig({ modelMapping, modelNames, onChange }: {
+  modelMapping: Array<[string, string]>;
+  modelNames: string[];
+  onChange: (modelMapping: Array<[string, string]>) => void;
+}) {
+  return (
+    <>
+      <IconButton
+        name="plus"
+        aria-label="Add model mapping"
+        onClick={e => {
+          e.preventDefault();
+          onChange([...modelMapping, ['', '']]);
+        }}
+      />
+      {modelMapping.map(([model, deployment], i) => (
+        <ModelMappingField
+          key={i}
+          model={model}
+          deployment={deployment}
+          modelNames={modelNames}
+          onChange={(model, deployment) => {
+            onChange([
+              ...modelMapping.slice(0, i),
+              [model, deployment],
+              ...modelMapping.slice(i + 1),
+            ]);
+          }}
+          onRemove={() => onChange([
+            ...modelMapping.slice(0, i),
+            ...modelMapping.slice(i + 1),
+          ])}
+        />
+      )
+      )}
+    </>
+  );
+}
+
+function ModelMappingField({ model, deployment, modelNames, onChange, onRemove }: {
+  model: string;
+  deployment: string;
+  modelNames: string[],
+  onChange: (model: string, deployment: string) => void;
+  onRemove: () => void;
+}): JSX.Element {
+  return (
+    <InlineFieldRow>
+      <InlineField label="Model">
+        <Select
+          placeholder="model label"
+          options={modelNames
+            .filter(n => n !== deployment && n !== '')
+            .map(value => ({ label: value, value }))
+          }
+          value={model}
+          onChange={event => event.value !== undefined && onChange(event.value, deployment)}
+        />
+      </InlineField>
+      <InlineField label="Deployment">
+        <Input
+          width={40} 
+          name="AzureDeployment"
+          placeholder="deployment name"
+          value={deployment}
+          onChange={event => event.currentTarget.value !== undefined && onChange(model, event.currentTarget.value)}
+        />
+      </InlineField>
+      <IconButton
+        name="trash-alt"
+        aria-label="Remove model mapping"
+        onClick={e => {
+          e.preventDefault();
+          onRemove()
+        }}
+      />
+    </InlineFieldRow>
+  );
+}
+
+
 export const AppConfig = ({ plugin }: AppConfigProps) => {
   const s = useStyles2(getStyles);
   const { enabled, pinned, jsonData, secureJsonFields } = plugin.meta;
@@ -42,6 +128,7 @@ export const AppConfig = ({ plugin }: AppConfigProps) => {
     openAIKey: '',
     isOpenAIKeySet: Boolean(secureJsonFields?.openAIKey),
     useAzureOpenAI: jsonData?.openAI?.useAzure || false,
+    azureModelMapping: jsonData?.openAI?.azureModelMapping || [],
     updated: false,
   });
 
@@ -53,6 +140,7 @@ export const AppConfig = ({ plugin }: AppConfigProps) => {
       openAIOrganizationID: '',
       isOpenAIKeySet: false,
       useAzureOpenAI: state.useAzureOpenAI,
+      azureModelMapping: [],
       updated: true,
     });
 
@@ -89,7 +177,7 @@ export const AppConfig = ({ plugin }: AppConfigProps) => {
 
         <Field label="OpenAI API Organization ID" description="Your OpenAI API Organization ID">
           <Input
-            width={60}
+            width={60} 
             name="openAIOrganizationID"
             data-testid={testIds.appConfig.openAIOrganizationID}
             value={state.openAIOrganizationID}
@@ -112,6 +200,22 @@ export const AppConfig = ({ plugin }: AppConfigProps) => {
           />
         </Field>
 
+        {state.useAzureOpenAI && (
+          <Field label="Azure OpenAI Model Mapping" description="">
+            <ModelMappingConfig
+              modelMapping={state.azureModelMapping}
+              modelNames={["gpt-3.5-turbo", "gpt-4"]}
+              onChange={
+                azureModelMapping => setState({
+                  ...state,
+                  azureModelMapping,
+                  updated: true,
+                })
+              }
+            />
+          </Field>
+        )}
+
         <div className={s.marginTop}>
           <Button
             type="submit"
@@ -125,6 +229,7 @@ export const AppConfig = ({ plugin }: AppConfigProps) => {
                     url: state.openAIUrl,
                     organizationId: state.openAIOrganizationID,
                     useAzure: state.useAzureOpenAI,
+                    azureModelMapping: state.azureModelMapping,
                   },
                 },
                 // This cannot be queried later by the frontend.