Add prod/stgn AuthType for JetBrains AI service

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## 2.2.7
+
+### Internal changes: 
+- Allow prod/stgn auth-type for JetBrains AI service.
+
 ## 2.2.6
 
 ### Public changes

package.json

@@ -8,7 +8,7 @@
     "url": "https://github.com/JetBrains-Research/coqpilot"
   },
   "publisher": "JetBrains-Research",
-  "version": "2.2.6",
+  "version": "2.2.7",
   "engines": {
     "vscode": "^1.82.0"
   },
@@ -214,6 +214,11 @@
                   "description": "Api key to communicate with the Grazie api. Now available for JetBrains employees only.",
                   "default": "None"
                 },
+                "authType": {
+                  "type": "string",
+                  "description": "Use stgn if you are an internal JetBrains AI user and use prod otherwise.",
+                  "default": "stgn"
+                },
                 "choices": {
                   "type": "number",
                   "description": "Number of attempts to generate proof for one hole with this model.",

src/llm/llmServices/grazie/grazieApi.ts

@@ -8,7 +8,8 @@ export type GrazieChatRole = "User" | "System" | "Assistant";
 export type GrazieFormattedHistory = { role: GrazieChatRole; text: string }[];
 
 interface GrazieConfig {
-    gateawayUrl: string;
+    gateawayUrlStgn: string;
+    gateawayUrlProd: string;
     chatUrl: string;
     quotaUrl: string;
 }
@@ -17,8 +18,10 @@ export class GrazieApi {
     private readonly config: GrazieConfig = {
         chatUrl: "v5/llm/chat/stream/v3",
         quotaUrl: "v5/quota/get",
-        gateawayUrl:
+        gateawayUrlStgn:
             "https://api.app.stgn.grazie.aws.intellij.net/application/",
+        gateawayUrlProd:
+            "https://api.app.prod.grazie.aws.intellij.net/application/",
     };
 
     constructor(private readonly debug: DebugWrappers) {}
@@ -28,7 +31,12 @@ export class GrazieApi {
         history: GrazieFormattedHistory
     ): Promise<string> {
         const body = this.createRequestBody(history, params);
-        return this.post(this.config.chatUrl, body, params.apiKey);
+        return this.post(
+            this.config.chatUrl,
+            body,
+            params.apiKey,
+            params.authType
+        );
     }
 
     async checkQuota(apiToken: string): Promise<any> {
@@ -55,7 +63,8 @@ export class GrazieApi {
     private async post(
         url: string,
         body: string,
-        apiToken: string
+        apiToken: string,
+        authType: "stgn" | "prod"
     ): Promise<string> {
         const headers = await this.createHeaders(apiToken);
 
@@ -65,8 +74,13 @@ export class GrazieApi {
             headers: headers,
         });
 
+        const baseUrl =
+            authType === "stgn"
+                ? this.config.gateawayUrlStgn
+                : this.config.gateawayUrlProd;
+
         const response = await this.fetchAndProcessEvents(
-            this.config.gateawayUrl + url,
+            baseUrl + url,
             body,
             headers
         );

src/llm/llmServices/grazie/grazieModelParamsResolver.ts

@@ -23,6 +23,13 @@ export class GrazieModelParamsResolver
         .requiredToBeConfigured()
         .validateAtRuntimeOnly();
 
+    readonly authType = this.resolveParam<"stgn" | "prod">("authType")
+        .requiredToBeConfigured()
+        .validate([
+            (value) => value === "stgn" || value === "prod",
+            "be either 'stgn' or 'prod'",
+        ]);
+
     readonly maxTokensToGenerate = this.resolveParam<number>(
         "maxTokensToGenerate"
     )

src/llm/llmServices/modelParams.ts

@@ -42,6 +42,7 @@ export interface OpenAiModelParams extends ModelParams {
 export interface GrazieModelParams extends ModelParams {
     modelName: string;
     apiKey: string;
+    authType: "stgn" | "prod";
 }
 
 export interface LMStudioModelParams extends ModelParams {
@@ -134,9 +135,18 @@ export const grazieModelParamsSchema: JSONSchemaType<GrazieModelParams> = {
     properties: {
         modelName: { type: "string" },
         apiKey: { type: "string" },
+        authType: {
+            type: "string",
+            enum: ["stgn", "prod"],
+        },
         ...(modelParamsSchema.properties as PropertiesSchema<ModelParams>),
     },
-    required: ["modelName", "apiKey", ...modelParamsSchema.required],
+    required: [
+        "modelName",
+        "apiKey",
+        "authType",
+        ...modelParamsSchema.required,
+    ],
     additionalProperties: false,
 };
 

src/llm/userModelParams.ts

@@ -59,6 +59,7 @@ export interface OpenAiUserModelParams extends UserModelParams {
 export interface GrazieUserModelParams extends UserModelParams {
     modelName: string;
     apiKey: string;
+    authType: string;
 }
 
 export interface LMStudioUserModelParams extends UserModelParams {
@@ -135,9 +136,10 @@ export const grazieUserModelParamsSchema: JSONSchemaType<GrazieUserModelParams>
         properties: {
             modelName: { type: "string" },
             apiKey: { type: "string" },
+            authType: { type: "string" },
             ...(userModelParamsSchema.properties as PropertiesSchema<UserModelParams>),
         },
-        required: ["modelId", "modelName", "apiKey"],
+        required: ["modelId", "modelName", "apiKey", "authType"],
         additionalProperties: false,
     };
 

src/test/benchmark/report.json

@@ -0,0 +1 @@
+{}

src/test/llm/llmServices/grazieService.test.ts

@@ -25,6 +25,7 @@ import {
 
 suite("[LLMService] Test `GrazieService`", function () {
     const apiKey = process.env.GRAZIE_API_KEY;
+    const authType = process.env.GRAZIE_AUTH_TYPE;
     const choices = 15;
     const inputFile = ["small_document.v"];
 
@@ -37,14 +38,15 @@ suite("[LLMService] Test `GrazieService`", function () {
     };
 
     testIf(
-        apiKey !== undefined,
-        "`GRAZIE_API_KEY` is not specified",
+        apiKey !== undefined && authType !== undefined,
+        "`GRAZIE_API_KEY` or `GRAZIE_AUTH_TYPE` is not specified",
         this.title,
         `Simple generation: 1 request, ${choices} choices`,
         async () => {
             const inputParams: GrazieUserModelParams = {
                 ...requiredInputParamsTemplate,
                 apiKey: apiKey!,
+                authType: authType!,
             };
             const grazieService = new GrazieService();
             await testLLMServiceCompletesAdmitFromFile(
@@ -60,6 +62,7 @@ suite("[LLMService] Test `GrazieService`", function () {
         const inputParams: GrazieUserModelParams = {
             ...requiredInputParamsTemplate,
             apiKey: "undefined",
+            authType: "stgn",
         };
         await withLLMService(new GrazieService(), async (grazieService) => {
             testResolveValidCompleteParameters(grazieService, inputParams);
@@ -79,6 +82,7 @@ suite("[LLMService] Test `GrazieService`", function () {
         const inputParams: GrazieUserModelParams = {
             ...requiredInputParamsTemplate,
             apiKey: "undefined",
+            authType: "stgn",
             maxTokensToGenerate: 6666, // should be overriden by GrazieService
         };
         withLLMService(new GrazieService(), async (grazieService) => {
@@ -96,6 +100,7 @@ suite("[LLMService] Test `GrazieService`", function () {
         const inputParams: GrazieUserModelParams = {
             ...requiredInputParamsTemplate,
             apiKey: "undefined",
+            authType: "stgn",
         };
         await withLLMServiceAndParams(
             new GrazieService(),

src/test/llm/parseUserModelParams.test.ts

@@ -72,6 +72,7 @@ suite("Parse `UserModelParams` from JSON test", () => {
         ...validUserModelParamsCompelete,
         modelName: "gpt-model",
         apiKey: "api-key",
+        authType: "stgn",
     };
     const validLMStudioUserModelParamsComplete = {
         ...validUserModelParamsCompelete,