add OpenAIOrganizationID

grafana · Jul 27, 2023 · 2a2e4e4 · 2a2e4e4
1 parent 59c8c02
commit 2a2e4e4
Show file tree

Hide file tree

Showing 6 changed files with 27 additions and 2 deletions.
diff --git a/pkg/plugin/app.go b/pkg/plugin/app.go
@@ -22,7 +22,8 @@ var (
 )
 
 type Settings struct {
-	OpenAIURL string `json:"openAIUrl"`
+	OpenAIURL            string `json:"openAIUrl"`
+	OpenAIOrganizationID string `json:"openAIOrganizationId"`
 
 	openAIKey string
 }

diff --git a/pkg/plugin/resources.go b/pkg/plugin/resources.go
@@ -16,9 +16,11 @@ func newOpenAIProxy() http.Handler {
 			config := httpadapter.PluginConfigFromContext(r.In.Context())
 			settings := loadSettings(*config.AppInstanceSettings)
 			apiKey := config.AppInstanceSettings.DecryptedSecureJSONData["apiKey"]
+			organizationID := settings.OpenAIOrganizationID
 			u, _ := url.Parse(settings.OpenAIURL)
 			r.SetURL(u)
 			r.Out.Header.Set("Authorization", "Bearer "+apiKey)
+			r.Out.Header.Set("OpenAI-Organization", organizationID)
 			r.Out.URL.Path = strings.TrimPrefix(r.In.URL.Path, "/openai")
 			log.DefaultLogger.Info("proxying to url", "url", r.Out.URL.String())
 		},

diff --git a/pkg/plugin/stream.go b/pkg/plugin/stream.go
@@ -57,6 +57,7 @@ func (a *App) runOpenAIChatCompletionsStream(ctx context.Context, req *backend.R
 	path := strings.TrimPrefix(req.Path, "/openai")
 	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, settings.OpenAIURL+path, bytes.NewReader(outgoingBody))
 	httpReq.Header.Set("Authorization", fmt.Sprintf("Bearer %s", settings.openAIKey))
+	httpReq.Header.Set("OpenAI-Organization", settings.OpenAIOrganizationID)
 	httpReq.Header.Set("Content-Type", "application/json")
 	lastEventID := "" // no last event id
 	eventStream, err := eventsource.SubscribeWithRequest(lastEventID, httpReq)

diff --git a/src/components/AppConfig/AppConfig.test.tsx b/src/components/AppConfig/AppConfig.test.tsx
@@ -32,6 +32,7 @@ describe('Components/AppConfig', () => {
 
     expect(screen.queryByRole('group', { name: /openai settings/i })).toBeInTheDocument();
     expect(screen.queryByTestId(testIds.appConfig.openAIKey)).toBeInTheDocument();
+    expect(screen.queryByTestId(testIds.appConfig.openAIOrganizationID)).toBeInTheDocument();
     expect(screen.queryByTestId(testIds.appConfig.openAIUrl)).toBeInTheDocument();
     expect(screen.queryByRole('button', { name: /save api settings/i })).toBeInTheDocument();
   });

diff --git a/src/components/AppConfig/AppConfig.tsx b/src/components/AppConfig/AppConfig.tsx
@@ -8,11 +8,14 @@ import { testIds } from '../testIds';
 
 export type AppPluginSettings = {
   openAIUrl?: string;
+  openAIOrganizationID?: string;
 };
 
 type State = {
   // The URL to reach our custom API.
   openAIUrl: string;
+  // The Organization ID for our custom API.
+  openAIOrganizationID: string;
   // Tells us if the API key secret is set.
   isOpenAIKeySet: boolean;
   // A secret key for our custom API.
@@ -26,6 +29,7 @@ export const AppConfig = ({ plugin }: AppConfigProps) => {
   const { enabled, pinned, jsonData, secureJsonFields } = plugin.meta;
   const [state, setState] = useState<State>({
     openAIUrl: jsonData?.openAIUrl || 'https://api.openai.com',
+    openAIOrganizationID: jsonData?.openAIOrganizationID || '',
     openAIKey: '',
     isOpenAIKeySet: Boolean(secureJsonFields?.openAIKey),
   });
@@ -34,6 +38,7 @@ export const AppConfig = ({ plugin }: AppConfigProps) => {
     setState({
       ...state,
       openAIKey: '',
+      openAIOrganizationID: '',
       isOpenAIKeySet: false,
     });
 
@@ -58,6 +63,17 @@ export const AppConfig = ({ plugin }: AppConfigProps) => {
           />
         </Field>
 
+        <Field label="OpenAI API Organization ID" description="Your OpenAI API Organization ID">
+          <Input
+            width={60}
+            name="openAIOrganizationID"
+            data-testid={testIds.appConfig.openAIOrganizationID}
+            value={state.openAIOrganizationID}
+            placeholder={'org-...'}
+            onChange={onChange}
+          />
+        </Field>
+
         <Field label="OpenAI API Key" description="Your OpenAI API Key">
           <SecretInput
             width={60}
@@ -81,6 +97,7 @@ export const AppConfig = ({ plugin }: AppConfigProps) => {
                 pinned,
                 jsonData: {
                   openAIUrl: state.openAIUrl,
+                  openAIOrganizationID: state.openAIOrganizationID,
                 },
                 // This cannot be queried later by the frontend.
                 // We don't want to override it in case it was set previously and left untouched now.
@@ -91,7 +108,9 @@ export const AppConfig = ({ plugin }: AppConfigProps) => {
                     },
               })
             }
-            disabled={Boolean(!state.openAIUrl || (!state.isOpenAIKeySet && !state.openAIKey))}
+            disabled={Boolean(
+              !state.openAIUrl || !state.openAIOrganizationID || (!state.isOpenAIKeySet && !state.openAIKey)
+            )}
           >
             Save API settings
           </Button>

diff --git a/src/components/testIds.ts b/src/components/testIds.ts
@@ -2,6 +2,7 @@ export const testIds = {
   appConfig: {
     container: 'data-testid ac-container',
     openAIKey: 'data-testid ac-openai-api-key',
+    openAIOrganizationID: 'data-testid ac-openai-api-organization-id',
     openAIUrl: 'data-testid ac-openai-api-url',
     submit: 'data-testid ac-submit-form',
   },