Refactor user team selection and management logic

Introduces new methods for getting, setting, and updating a user's current team in the database layer. Refactors the team selection endpoint and service to use these methods, improving clarity and separation of concerns. Updates abstract base and service interfaces to support the new workflow.

Author: Fr4nc3

src/backend/common/database/cosmosdb.py

@@ -43,6 +43,7 @@ class CosmosDBClient(DatabaseBase):
         "step": Step,
         "agent_message": AgentMessage,
         "team_config": TeamConfiguration,
+        "user_current_team": UserCurrentTeam,
     }
 
     def __init__(
@@ -432,7 +433,33 @@ async def update_team(self, team: TeamConfiguration) -> None:
             team: The TeamConfiguration to update
         """
         await self.update_item(team)
-    
-    async def set_user_default_team(current_team: UserCurrentTeam) -> bool:
+
+    async def get_current_team(self, user_id: str, team_id: str) -> UserCurrentTeam:
+        """Retrieve the current team for a user."""
+        await self._ensure_initialized()
+        if self.container is None:
+            return None
+
+        query = "SELECT * FROM c WHERE c.user_id=@user_id AND c.is_default=true"
+        parameters = [
+            {"name": "@user_id", "value": user_id},
+            {"name": "@team_id", "value": team_id},
+        ]
+
+        items = self.container.query_items(query=query, parameters=parameters)
+        async for item in items:
+            return UserCurrentTeam(**item)
+
+        return None
+    async def set_current_team(self, current_team: UserCurrentTeam) -> None:
+        """Set the current team for a user."""
+        await self._ensure_initialized()
+
+
+        await self.container.upsert_item(current_team.dict())
+    async def update_current_team(self, current_team: UserCurrentTeam) -> None:
+        """Update the current team for a user."""
+        await self._ensure_initialized()
 
 
+        await self.container.upsert_item(current_team.dict())

src/backend/common/database/database_base.py

@@ -192,7 +192,15 @@ async def get_steps_for_plan(self, plan_id: str) -> List[Step]:
         pass
 
     @abstractmethod
-    async def set_user_default_team(self, UserCurrentTeam: UserCurrentTeam) -> bool:
-        """Set the default team for a user."""
+    async def get_current_team(self, user_id: str, team_id: str) -> UserCurrentTeam:
+        """Retrieve the current team for a user."""
         pass
 
+    @abstractmethod
+    async def set_current_team(self, current_team: UserCurrentTeam) -> None:
+        pass
+
+    @abstractmethod
+    async def update_current_team(self, current_team: UserCurrentTeam) -> None:
+        """Update the current team for a user."""
+        pass    

src/backend/v3/api/router.py

@@ -791,14 +791,20 @@ async def select_team_endpoint(selection: TeamSelectionRequest, request: Request
                 status_code=404, 
                 detail=f"Team configuration '{selection.team_id}' not found or access denied"
             )
-        current_team = await team_service.get_user_current_team(user_id=user_id)
-        if current_team and current_team.team_id != selection.team_id:
-            current_team.team_id = selection.team_id
-            await team_service.set_user_current_team(user_id=user_id, current_team=current_team, )
-        else:
-            current_team = UserCurrentTeam(
-                user_id=user_id,
-                team_id=selection.team_id
+        set_team = await team_service.handle_team_selection(user_id=user_id, team_id=selection.team_id)
+        if not set_team:
+            track_event_if_configured(
+            "Team selected",
+            {
+                "status": "failed",
+                "team_id": selection.team_id,
+                "team_name": team_configuration.name,
+                "user_id": user_id
+            },
+        )
+            raise HTTPException(
+                status_code=404,
+                detail=f"Team configuration '{selection.team_id}' failed to set"
             )
 
         # save to in-memory config for current user

src/backend/v3/common/services/team_service.py

@@ -5,15 +5,21 @@
 from datetime import datetime, timezone
 from typing import Any, Dict, List, Optional, Tuple
 
-from azure.core.credentials import AzureKeyCredential
-from azure.core.exceptions import (ClientAuthenticationError,
-                                   HttpResponseError, ResourceNotFoundError)
-from azure.identity import DefaultAzureCredential
+from azure.core.exceptions import (
+    ClientAuthenticationError,
+    HttpResponseError,
+    ResourceNotFoundError,
+)
+
 from azure.search.documents.indexes import SearchIndexClient
 from common.config.app_config import config
 from common.database.database_base import DatabaseBase
-from common.models.messages_kernel import (StartingTask, TeamAgent,
-                                           TeamConfiguration)
+from common.models.messages_kernel import (
+    StartingTask,
+    TeamAgent,
+    TeamConfiguration,
+    UserCurrentTeam,
+)
 from v3.common.services.foundry_service import FoundryService
 
 
@@ -237,7 +243,7 @@ async def delete_user_current_team(self, user_id: str) -> bool:
             self.logger.error("Error deleting current team: %s", str(e))
             return False
 
-    async def set_default_team(self, user_id: str, team_id: str) -> bool:
+    async def handle_team_selection(self, user_id: str, team_id: str) -> bool:
         """
         Set a default team for a user.
 
@@ -249,9 +255,16 @@ async def set_default_team(self, user_id: str, team_id: str) -> bool:
             True if successful, False otherwise
         """
         try:
-            await self.memory_context.set_user_default_team(user_id, team_id)
-            self.logger.info("Successfully set default team for user %s: %s", user_id, team_id)
-            return True
+            current_team = await self.memory_context.get_current_team(user_id, team_id)
+
+            if current_team is None:
+                current_team = UserCurrentTeam(user_id=user_id, team_id=team_id)
+                await self.memory_context.set_current_team(current_team)
+                return True
+            else:
+                current_team.team_id = team_id
+                await self.memory_context.update_current_team(current_team)
+                return True
 
         except Exception as e:
             self.logger.error("Error setting default team: %s", str(e))
@@ -384,7 +397,7 @@ async def validate_team_models(
             missing_models: List[str] = []
             for model in required_models:
                 # Temporary bypass for known deployed models
-                if model.lower() in ['gpt-4o', 'o3', 'gpt-4', 'gpt-35-turbo']:
+                if model.lower() in ["gpt-4o", "o3", "gpt-4", "gpt-35-turbo"]:
                     continue
                 if model not in available_models:
                     missing_models.append(model)