Feature/multi agents collaboration

nexus/agents/api/routers.py

@@ -0,0 +1,7 @@
+from django.urls import path
+from nexus.agents.api.views import PushAgents
+
+
+urlpatterns = [
+    path('agents/push', PushAgents.as_view(), name="push-agents"),
+]

nexus/agents/api/views.py

@@ -0,0 +1,43 @@
+from rest_framework.views import APIView
+from rest_framework.response import Response
+from rest_framework.permissions import IsAuthenticated
+
+from nexus.usecases.projects import get_project_by_uuid
+from nexus.usecases.agents import AgentDTO, AgentUsecase
+from nexus.agents.models import Agent
+
+from nexus.projects.api.permissions import ProjectPermission
+
+
+class PushAgents(APIView):
+
+    permission_classes = [IsAuthenticated, ProjectPermission]
+
+    def post(self, request, *args, **kwargs):
+        def get_agents(agents):
+            dto_list = []
+            for agent_slug in agents:
+                agent = agents.get(agent_slug)
+                agent.update({"slug": agent_slug})
+                dto_list.append(AgentDTO(**agent))
+            return dto_list
+
+        project_uuid = request.data.get("project")
+        project = get_project_by_uuid(project_uuid)
+
+        agents = request.data.get("agents")
+        agents_dto = get_agents(agents)
+
+        print(agents_dto)
+
+        agents = []
+
+        for agent_dto in agents_dto:
+            agent: Agent = AgentUsecase().create_agent(request.user, agent_dto, project_uuid)
+
+            agents.append({"agent_name": agent.display_name, "agent_external_id": agent.external_id})
+
+        return Response({
+            "project": str(project.uuid),
+            "agents": agents
+        })

nexus/agents/apps.py

@@ -0,0 +1,6 @@
+from django.apps import AppConfig
+
+
+class AgentsConfig(AppConfig):
+    default_auto_field = 'django.db.models.BigAutoField'
+    name = 'nexus.agents'

nexus/agents/migrations/0001_initial.py

@@ -0,0 +1,63 @@
+# Generated by Django 4.2.10 on 2024-12-27 20:15
+
+from django.conf import settings
+from django.db import migrations, models
+import django.db.models.deletion
+import uuid
+
+
+class Migration(migrations.Migration):
+
+    initial = True
+
+    dependencies = [
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        ('projects', '0008_merge_20241010_2217'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='Team',
+            fields=[
+                ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('external_id', models.CharField(help_text='Supervisor ID', max_length=255)),
+                ('project', models.OneToOneField(on_delete=django.db.models.deletion.CASCADE, to='projects.project')),
+            ],
+        ),
+        migrations.CreateModel(
+            name='Agent',
+            fields=[
+                ('uuid', models.UUIDField(default=uuid.uuid4, primary_key=True, serialize=False)),
+                ('created_at', models.DateTimeField(auto_now_add=True)),
+                ('modified_at', models.DateTimeField(blank=True, null=True)),
+                ('external_id', models.CharField(max_length=255)),
+                ('slug', models.SlugField(max_length=255)),
+                ('display_name', models.CharField(max_length=255)),
+                ('model', models.CharField(max_length=255)),
+                ('is_official', models.BooleanField(default=False)),
+                ('metadata', models.JSONField(default=dict)),
+                ('created_by', models.ForeignKey(on_delete=django.db.models.deletion.PROTECT, related_name='created_%(class)ss', to=settings.AUTH_USER_MODEL)),
+                ('modified_by', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.PROTECT, related_name='modified_%(class)ss', to=settings.AUTH_USER_MODEL)),
+                ('project', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='projects.project')),
+            ],
+            options={
+                'abstract': False,
+            },
+        ),
+        migrations.CreateModel(
+            name='ActiveAgent',
+            fields=[
+                ('uuid', models.UUIDField(default=uuid.uuid4, primary_key=True, serialize=False)),
+                ('created_at', models.DateTimeField(auto_now_add=True)),
+                ('modified_at', models.DateTimeField(blank=True, null=True)),
+                ('is_official', models.BooleanField(default=False)),
+                ('agent', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='agents.agent')),
+                ('created_by', models.ForeignKey(on_delete=django.db.models.deletion.PROTECT, related_name='created_%(class)ss', to=settings.AUTH_USER_MODEL)),
+                ('modified_by', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.PROTECT, related_name='modified_%(class)ss', to=settings.AUTH_USER_MODEL)),
+                ('team', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='agents.team')),
+            ],
+            options={
+                'abstract': False,
+            },
+        ),
+    ]

nexus/agents/models.py

@@ -0,0 +1,24 @@
+from django.db import models
+from nexus.projects.models import Project
+from nexus.db.models import BaseModel
+
+
+class Agent(BaseModel):
+    external_id = models.CharField(max_length=255)
+    slug = models.SlugField(max_length=255)
+    display_name = models.CharField(max_length=255)
+    model = models.CharField(max_length=255)
+    is_official = models.BooleanField(default=False)
+    project = models.ForeignKey(Project, on_delete=models.CASCADE)
+    metadata = models.JSONField(default=dict)
+
+
+class Team(models.Model):
+    external_id = models.CharField(max_length=255, help_text="Supervisor ID")
+    project = models.OneToOneField(Project, on_delete=models.CASCADE)
+
+
+class ActiveAgent(BaseModel):
+    agent = models.ForeignKey(Agent, on_delete=models.CASCADE)
+    team = models.ForeignKey(Team, on_delete=models.CASCADE)
+    is_official = models.BooleanField(default=False)

nexus/projects/api/permissions.py

@@ -9,7 +9,16 @@
 class ProjectPermission(permissions.BasePermission):
     def has_permission(self, request, view):
         try:
-            project_uuid = view.kwargs["project_uuid"]
+            uuids = [
+                view.kwargs.get("project_uuid"),
+                request.data.get("project"),
+                request.data.get("project_uuid")
+            ]
+
+            uuids = set(uuids)
+            uuids.remove(None)
+            project_uuid = list(uuids)[0]
+
             project = ProjectsUseCase().get_by_uuid(project_uuid)
             return has_project_permission(
                 user=request.user,

nexus/settings.py

@@ -82,6 +82,7 @@
     'nexus.logs',
     'nexus.sentry',
     'nexus.zeroshot',
+    'nexus.agents',
 ]
 
 MIDDLEWARE = [

nexus/task_managers/file_database/bedrock.py

@@ -1,12 +1,14 @@
 import uuid
 import json
-from typing import Dict, List, Any
+import time
+from typing import Dict, List, Any, Tuple
 from os.path import basename
 
 import boto3
 from django.conf import settings
 
 from nexus.task_managers.file_database.file_database import FileDataBase, FileResponseDTO
+from nexus.agents.src.utils.bedrock_agent_helper import AgentsForAmazonBedrock
 
 
 class BedrockFileDatabase(FileDataBase):
@@ -24,6 +26,53 @@ 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_foundation_model = [
+            '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(
+            agent_collaboration='SUPERVISOR_ROUTER',
+            agent_name=supervisor_name,
+            agent_description=supervisor_description,
+            agent_instructions=supervisor_instructions,
+            model_ids=self.agent_foundation_model,
+        )
+        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(
+            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_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_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):
         data = {
             "top_p": config_data.get("top_p"),

nexus/urls.py

@@ -11,6 +11,8 @@
 from nexus.projects.api.routers import urlpatterns as projects_routes
 from nexus.logs.api.routers import urlpatterns as logs_routes
 from nexus.zeroshot.api.routers import urlpatterns as zeroshot_routes
+from nexus.agents.api.routers import urlpatterns as agent_routes
+
 
 schema_view = get_schema_view(
     openapi.Info(
@@ -28,6 +30,7 @@
 url_api += actions_routes
 url_api += projects_routes
 url_api += logs_routes
+url_api += agent_routes
 
 urlpatterns = [
     path("", lambda _: HttpResponse()),

nexus/usecases/agents/__init__.py

@@ -0,0 +1 @@
+from .agents import AgentUsecase, AgentDTO

nexus/usecases/agents/agents.py

@@ -0,0 +1,105 @@
+from typing import List, Dict, Tuple
+from dataclasses import dataclass
+
+from nexus.agents.models import Agent, Team
+from nexus.task_managers.file_database.bedrock import BedrockFileDatabase
+
+
+@dataclass
+class AgentDTO:
+    slug: str
+    name: str
+    description: str
+    instructions: List[str]
+    guardrails: List[str]
+    skills: List[Dict]
+    model: str
+
+
+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,
+        supervisor_description: str,
+        supervisor_instructions: str,
+    ) -> str:
+        return self.external_agent_client.create_supervisor(
+            supervisor_name,
+            supervisor_description,
+            supervisor_instructions,
+        )
+
+    def create_external_agent(self, agent_name: str, agent_description: str, agent_instructions: str):
+        return self.external_agent_client.create_agent(
+            agent_name,
+            agent_description,
+            agent_instructions
+        )
+
+    def create_external_agent_alias(self, agent_id: str, alias_name: str) -> Tuple[str, str]:
+        agent_alias_id, agent_alias_arn = self.external_agent_client.create_agent_alias(
+            agent_id=agent_id, alias_name=alias_name
+        )
+        return agent_alias_id, agent_alias_arn
+
+    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)
+
+        agent_slug = f"{agent_dto.slug}-{project_uuid}"
+        external_id, _agent_alias_id, _agent_alias_arn = self.create_external_agent(
+            agent_name=agent_slug,
+            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,
+            display_name=agent_dto.name,
+            model=agent_dto.model,
+            metadata={
+                "engine": "BEDROCK",
+                "_agent_alias_id": _agent_alias_id,
+                "_agent_alias_arn": _agent_alias_arn,
+                "agent_alias_id": sub_agent_alias_id,
+                "agent_alias_arn": sub_agent_alias_arn,
+            }
+        )
+        return agent
+
+    def create_supervisor(
+        self,
+        project_uuid: str,
+        supervisor_name: str,
+        supervisor_description: str,
+        supervisor_instructions: str
+    ):
+        external_id = self.create_external_supervisor(
+            supervisor_name,
+            supervisor_description,
+            supervisor_instructions,
+        )
+
+        Team.objects.create(
+            project_id=project_uuid,
+            external_id=external_id
+        )
+        return