Fixing object id find

redhat-composer-ai · Nov 11, 2024 · 6daa682 · 6daa682
1 parent e134739
commit 6daa682
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/src/main/java/com/redhat/composer/services/AssistantInfoService.java b/src/main/java/com/redhat/composer/services/AssistantInfoService.java
@@ -3,6 +3,8 @@
 import java.util.List;
 import java.util.stream.Stream;
 
+import org.bson.types.ObjectId;
+
 import com.redhat.composer.model.mongo.AssistantEntity;
 import com.redhat.composer.model.mongo.LLMConnectionEntity;
 import com.redhat.composer.model.mongo.RetrieverConnectionEntity;
@@ -25,16 +27,17 @@ public class AssistantInfoService {
   MapperUtil mapperUtil;
 
   public AssistantEntity createAssistant(AssistantCreationRequest request) {
-
-    LLMConnectionEntity llm = (LLMConnectionEntity) LLMConnectionEntity.findByIdOptional(request.getLlmConnectionId()).orElseThrow(() -> new IllegalArgumentException("LLM Connection not found"));
-    RetrieverConnectionEntity retriever = (RetrieverConnectionEntity) RetrieverConnectionEntity.findByIdOptional(request.getRetrieverConnectionId()).orElseThrow(() -> new IllegalArgumentException("Retriever Connection not found"));
-
     AssistantEntity assistant = new AssistantEntity();
+
+    LLMConnectionEntity llm = (LLMConnectionEntity) LLMConnectionEntity.findByIdOptional(new ObjectId(request.getLlmConnectionId())).orElseThrow(() -> new IllegalArgumentException("LLM Connection not found"));
+    if (request.getRetrieverConnectionId() != null) {    
+      RetrieverConnectionEntity retriever = (RetrieverConnectionEntity) RetrieverConnectionEntity.findByIdOptional(new ObjectId(request.getRetrieverConnectionId())).orElseThrow(() -> new IllegalArgumentException("Retriever Connection not found"));
+      assistant.setRetrieverConnectionId(retriever.id);
+    }
     assistant.setName(request.getName());
     assistant.setDisplayName(request.getDisplayName());
     assistant.setDescription(request.getDescription());
     assistant.setLlmConnectionId(llm.id);
-    assistant.setRetrieverConnectionId(retriever.id);
     assistant.persist();
     return assistant;
   }