prepare agents

nexus/agents/api/views.py

@@ -33,7 +33,7 @@ def get_agents(agents):
         agents = []
 
         for agent_dto in agents_dto:
-            agent: Agent = AgentUsecase().create_agent(agent_dto, project_uuid)
+            agent: Agent = AgentUsecase().create_agent(request.user, agent_dto, project_uuid)
 
             agents.append({"agent_name": agent.display_name, "agent_external_id": agent.external_id})
 

nexus/task_managers/file_database/bedrock.py

@@ -1,5 +1,6 @@
 import uuid
 import json
+import time
 from typing import Dict, List, Any, Tuple
 from os.path import basename
 
@@ -25,15 +26,20 @@ def __init__(self) -> None:
         self.bedrock_agent_runtime = self.__get_bedrock_agent_runtime()
         self.bedrock_runtime = self.__get_bedrock_runtime()
 
-        self.agent_for_amazon_bedrock = AgentsForAmazonBedrock
+        self.agent_for_amazon_bedrock = AgentsForAmazonBedrock()
         self.agent_foundation_model = [
-            'anthropic.claude-3-sonnet-20nexus/task_managers/file_database/bedrock.py:34240229-v1:0',
+            'anthropic.claude-3-sonnet-2034240229-v1:0',
             'anthropic.claude-3-5-sonnet-20240620-v1:0',
             'anthropic.claude-3-haiku-20240307-v1:0'
         ]
+
+    def prepare_agent(self, agent_id: str):
+        self.bedrock_agent.prepare_agent(agentId=agent_id)
+        time.sleep(5)
+        return
 
     def create_supervisor(self, supervisor_name, supervisor_description, supervisor_instructions):
-        supervisor_id, supervisor_alias, supervisor_arn = self.agent_for_amazon_bedrock().create_agent(
+        supervisor_id, supervisor_alias, supervisor_arn = self.agent_for_amazon_bedrock.create_agent(
             agent_collaboration='SUPERVISOR_ROUTER',
             agent_name=supervisor_name,
             agent_description=supervisor_description,
@@ -43,27 +49,28 @@ def create_supervisor(self, supervisor_name, supervisor_description, supervisor_
         return supervisor_id
 
     def create_agent_alias(self, agent_id: str, alias_name: str):
-        sub_agent_alias_id, sub_agent_alias_arn = self.agent_for_amazon_bedrock().create_agent_alias(
+        sub_agent_alias_id, sub_agent_alias_arn = self.agent_for_amazon_bedrock.create_agent_alias(
             agent_id=agent_id, alias_name=alias_name
         )
         return sub_agent_alias_id, sub_agent_alias_arn
 
     def associate_sub_agents(self, supervisor_id: str, agents_list: str) -> str:
-        supervisor_agent_alias_id, supervisor_agent_alias_arn = self.agent_for_amazon_bedrock().associate_sub_agents(
+        supervisor_agent_alias_id, supervisor_agent_alias_arn = self.agent_for_amazon_bedrock.associate_sub_agents(
             supervisor_agent_id=supervisor_id,
             sub_agents_list=agents_list,
         )
         return supervisor_agent_alias_id
 
     def create_agent(self, agent_name: str, agent_description: str, agent_instructions: str) -> Tuple[str, str, str]:
-        agent_id, agent_alias, agent_arn = self.agent_for_amazon_bedrock().create_agent(
+        agent_id, agent_alias, agent_arn = self.agent_for_amazon_bedrock.create_agent(
             agent_name=agent_name,
             agent_description=agent_description,
             agent_instructions=agent_instructions,
             model_ids=self.agent_foundation_model,
             agent_collaboration="DISABLED",
             code_interpretation=False
         )
+        time.sleep(5)
         return agent_id, agent_alias, agent_arn
 
     def invoke_model(self, prompt: str, config_data: Dict):

nexus/usecases/agents/agents.py

@@ -20,6 +20,10 @@ class AgentUsecase:
     def __init__(self, external_agent_client=BedrockFileDatabase):
         self.external_agent_client = external_agent_client()
 
+    def prepare_agent(self, agent_id: str):
+        self.external_agent_client.prepare_agent(agent_id)
+        return
+
     def create_external_supervisor(
         self,
         supervisor_name: str,
@@ -45,7 +49,7 @@ def create_external_agent_alias(self, agent_id: str, alias_name: str) -> Tuple[s
         )
         return agent_alias_id, agent_alias_arn
 
-    def create_agent(self, agent_dto: AgentDTO, project_uuid: str, alias_name: str = "v1"):
+    def create_agent(self, user, agent_dto: AgentDTO, project_uuid: str, alias_name: str = "v1"):
         def format_instructions(instructions: List[str]):
             return "\n".join(instructions)
 
@@ -55,11 +59,17 @@ def format_instructions(instructions: List[str]):
             agent_description=agent_dto.description,
             agent_instructions=format_instructions(agent_dto.instructions),
         )
+
+        self.prepare_agent(external_id)
+
+        self.external_agent_client.agent_for_amazon_bedrock.wait_agent_status_update(external_id)
+
         sub_agent_alias_id, sub_agent_alias_arn = self.create_external_agent_alias(
             agent_id=external_id, alias_name=alias_name
         )
 
         agent = Agent.objects.create(
+            created_by=user,
             project_id=project_uuid,
             external_id=external_id,
             slug=agent_slug,