Add team agent management to database layer

Introduces methods to add, delete, and retrieve team agents in both the CosmosDBClient and DatabaseBase classes. Also renames the model class CurrentAgent to CurrentTeamAgent for clarity and updates references accordingly.

Author: Fr4nc3

src/backend/common/database/cosmosdb.py

@@ -12,6 +12,7 @@
     AgentMessage,
     AgentMessageData,
     BaseDataModel,
+    CurrentTeamAgent,
     DataType,
     Plan,
     Step,
@@ -491,3 +492,45 @@ async def get_agent_messages(self, plan_id: str) -> List[AgentMessageData]:
         ]
 
         return await self.query_items(query, parameters, AgentMessageData)
+
+    async def add_team_agent(self, team_agent: CurrentTeamAgent) -> None:
+        """Add an agent message to the database."""
+        await self.add_item(team_agent)
+
+    async def delete_team_agent(self, team_id: str, agent_name: str) -> None:
+        """Delete the current team for a user."""
+        query = "SELECT c.id, c.session_id FROM c WHERE c.team_id=@team_id AND c.data_type=@data_type AND c.agent_name=@agent_name"
+
+        params = [
+            {"name": "@team_id", "value": team_id},
+            {"name": "@agent_name", "value": agent_name},
+            {"name": "@data_type", "value": DataType.current_team_agent},
+        ]
+        items = self.container.query_items(query=query, parameters=params)
+        print("Items to delete:", items)
+        if items:
+            async for doc in items:
+                try:
+                    await self.container.delete_item(
+                        doc["id"], partition_key=doc["session_id"]
+                    )
+                except Exception as e:
+                    self.logger.warning(
+                        "Failed deleting current team doc %s: %s", doc.get("id"), e
+                    )
+
+        return True
+
+    async def get_team_agent(
+        self, team_id: str, agent_name: str
+    ) -> Optional[CurrentTeamAgent]:
+        """Retrieve a team agent by team_id and agent_name."""
+        query = "SELECT c.id, c.session_id FROM c WHERE c.team_id=@team_id AND c.data_type=@data_type AND c.agent_name=@agent_name"
+        params = [
+            {"name": "@team_id", "value": team_id},
+            {"name": "@agent_name", "value": agent_name},
+            {"name": "@data_type", "value": DataType.current_team_agent},
+        ]
+
+        results = await self.query_items(query, params, CurrentTeamAgent)
+        return results[0] if results else None

src/backend/common/database/database_base.py

@@ -8,6 +8,7 @@
 from ..models.messages_af import (
     AgentMessageData,
     BaseDataModel,
+    CurrentTeamAgent,
     Plan,
     Step,
     TeamConfiguration,
@@ -231,3 +232,20 @@ async def update_agent_message(self, message: AgentMessageData) -> None:
     async def get_agent_messages(self, plan_id: str) -> Optional[AgentMessageData]:
         """Retrieve agent messages by plan_id."""
         pass
+
+    @abstractmethod
+    async def add_team_agent(self, team_agent: CurrentTeamAgent) -> None:
+        """Add an agent message to the database."""
+        pass
+
+    @abstractmethod
+    async def delete_team_agent(self, team_id: str, agent_name: str) -> None:
+        """Delete a team agent from the database."""
+        pass
+
+    @abstractmethod
+    async def get_team_agent(
+        self, team_id: str, agent_name: str
+    ) -> Optional[CurrentTeamAgent]:
+        """Retrieve a team agent by team_id and agent_name."""
+        pass

src/backend/common/models/messages_af.py

@@ -112,7 +112,7 @@ class UserCurrentTeam(BaseDataModel):
     user_id: str
     team_id: str
 
-class CurrentAgent(BaseDataModel):
+class CurrentTeamAgent(BaseDataModel):
     """Represents the current agent of a user."""
     data_type: Literal[DataType.current_team_agent] = DataType.current_team_agent
     team_id: str