From b061ecb1e36cd43357ab6f7b2f8708b225f78edd Mon Sep 17 00:00:00 2001
From: Aliaksandr Stsiapanay <aliaksandr.stsiapanay@gmail.com>
Date: Fri, 15 Dec 2023 18:38:14 +0200
Subject: [PATCH] fix: fix api key for base assistant (#92)

Co-authored-by: Aliaksandr Stsiapanay <aliaksandr_stsiapanay@epam.com>
---
 .../com/epam/aidial/core/config/Assistants.java   |  3 ++-
 .../java/com/epam/aidial/core/config/Config.java  | 15 +--------------
 .../epam/aidial/core/config/FileConfigStore.java  | 15 +++++++++------
 3 files changed, 12 insertions(+), 21 deletions(-)

diff --git a/src/main/java/com/epam/aidial/core/config/Assistants.java b/src/main/java/com/epam/aidial/core/config/Assistants.java
index 6ad5b7e48..b5248c005 100644
--- a/src/main/java/com/epam/aidial/core/config/Assistants.java
+++ b/src/main/java/com/epam/aidial/core/config/Assistants.java
@@ -2,11 +2,12 @@
 
 import lombok.Data;
 
+import java.util.HashMap;
 import java.util.Map;
 
 @Data
 public class Assistants {
     private String endpoint;
     private Features features;
-    private Map<String, Assistant> assistants = Map.of();
+    private Map<String, Assistant> assistants = new HashMap<>();
 }
\ No newline at end of file
diff --git a/src/main/java/com/epam/aidial/core/config/Config.java b/src/main/java/com/epam/aidial/core/config/Config.java
index 6034701a2..0441a8407 100644
--- a/src/main/java/com/epam/aidial/core/config/Config.java
+++ b/src/main/java/com/epam/aidial/core/config/Config.java
@@ -34,19 +34,6 @@ public Deployment selectDeployment(String deploymentId) {
         }
 
         Assistants assistants = assistant;
-        Assistant assistant = assistants.getAssistants().get(deploymentId);
-        if (assistant != null) {
-            return assistant;
-        }
-
-        if (assistants.getEndpoint() != null && ASSISTANT.equals(deploymentId)) {
-            Assistant baseAssistant = new Assistant();
-            baseAssistant.setName(ASSISTANT);
-            baseAssistant.setEndpoint(assistants.getEndpoint());
-            baseAssistant.setFeatures(assistants.getFeatures());
-            return baseAssistant;
-        }
-
-        return null;
+        return assistants.getAssistants().get(deploymentId);
     }
 }
diff --git a/src/main/java/com/epam/aidial/core/config/FileConfigStore.java b/src/main/java/com/epam/aidial/core/config/FileConfigStore.java
index ad2edd414..9449d42eb 100644
--- a/src/main/java/com/epam/aidial/core/config/FileConfigStore.java
+++ b/src/main/java/com/epam/aidial/core/config/FileConfigStore.java
@@ -65,12 +65,6 @@ private void load(boolean fail) {
             }
 
             Assistants assistants = config.getAssistant();
-            // base assistant
-            if (assistants.getEndpoint() != null) {
-                Assistant baseAssistant = new Assistant();
-                baseAssistant.setName(ASSISTANT);
-                associateDeploymentWithApiKey(config, baseAssistant);
-            }
             for (Map.Entry<String, Assistant> entry : assistants.getAssistants().entrySet()) {
                 String name = entry.getKey();
                 Assistant assistant = entry.getValue();
@@ -83,6 +77,15 @@ private void load(boolean fail) {
 
                 setMissingFeatures(assistant, assistants.getFeatures());
             }
+            // base assistant
+            if (assistants.getEndpoint() != null) {
+                Assistant baseAssistant = new Assistant();
+                baseAssistant.setName(ASSISTANT);
+                baseAssistant.setEndpoint(assistants.getEndpoint());
+                baseAssistant.setFeatures(assistants.getFeatures());
+                associateDeploymentWithApiKey(config, baseAssistant);
+                assistants.getAssistants().put(ASSISTANT, baseAssistant);
+            }
 
             for (Map.Entry<String, Application> entry : config.getApplications().entrySet()) {
                 String name = entry.getKey();