Refactor team selection and plan retrieval endpoints

Simplifies the /api/plans endpoint by removing session_id and team_id query parameters and related logic. Refactors get_current_team to only require user_id, updating its usage across database and service layers. Enhances team initialization and configuration upload endpoints to better handle current team selection and configuration updates.

Author: Fr4nc3

src/backend/app_kernel.py

@@ -598,9 +598,7 @@ async def approve_step_endpoint(
 @app.get("/api/plans")
 async def get_plans(
     request: Request,
-    session_id: Optional[str] = Query(None),
     plan_id: Optional[str] = Query(None),
-    team_id: Optional[str] = Query(None),
 ):
     """
     Retrieve plans for the current user.
@@ -673,20 +671,7 @@ async def get_plans(
 
     # # Initialize memory context
     memory_store = await DatabaseFactory.get_database(user_id=user_id)
-    if session_id:
-        plan = await memory_store.get_plan_by_session(session_id=session_id)
-        if not plan:
-            track_event_if_configured(
-                "GetPlanBySessionNotFound",
-                {"status_code": 400, "detail": "Plan not found"},
-            )
-            raise HTTPException(status_code=404, detail="Plan not found")
 
-        # Use get_steps_by_plan to match the original implementation
-        steps = await memory_store.get_steps_by_plan(plan_id=plan.id)
-        plan_with_steps = PlanWithSteps(**plan.model_dump(), steps=steps)
-        plan_with_steps.update_step_counts()
-        return [plan_with_steps]
     if plan_id:
         plan = await memory_store.get_plan_by_plan_id(plan_id=plan_id)
         if not plan:

src/backend/common/database/cosmosdb.py

@@ -434,23 +434,22 @@ async def update_team(self, team: TeamConfiguration) -> None:
         """
         await self.update_item(team)
 
-    async def get_current_team(self, user_id: str, team_id: str) -> UserCurrentTeam:
+    async def get_current_team(self, user_id: str) -> Optional[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"
+        query = "SELECT * FROM c WHERE c.data_type=@data_type AND c.user_id=@user_id"
         parameters = [
+            {"name": "@data_type", "value": "user_current_team"},
             {"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)
+        # Get the appropriate model class
+        teams = await self.query_items(query, parameters, UserCurrentTeam)
+        return teams[0] if teams else None
 
-        return None
     async def set_current_team(self, current_team: UserCurrentTeam) -> None:
         """Set the current team for a user."""
         await self._ensure_initialized()

src/backend/common/database/database_base.py

@@ -192,7 +192,7 @@ async def get_steps_for_plan(self, plan_id: str) -> List[Step]:
         pass
 
     @abstractmethod
-    async def get_current_team(self, user_id: str, team_id: str) -> UserCurrentTeam:
+    async def get_current_team(self, user_id: str) -> Optional[UserCurrentTeam]:
         """Retrieve the current team for a user."""
         pass
 

src/backend/v3/api/router.py

@@ -13,7 +13,7 @@
                                            TeamSelectionRequest, UserCurrentTeam)
 from common.utils.event_utils import track_event_if_configured
 from common.utils.utils_kernel import rai_success, rai_validate_team_config
-from fastapi import (APIRouter, BackgroundTasks, Depends, FastAPI, File,
+from fastapi import (APIRouter, BackgroundTasks, Depends, FastAPI, File, Query,
                      HTTPException, Request, UploadFile, WebSocket,
                      WebSocketDisconnect)
 from kernel_agents.agent_factory import AgentFactory
@@ -103,7 +103,11 @@ async def init_team(
       # Initialize memory store and service
       memory_store = await DatabaseFactory.get_database(user_id=user_id)
       team_service = TeamService(memory_store)
-
+      user_current_team = await memory_store.get_current_team(user_id=user_id)
+      if not user_current_team:
+          await team_service.handle_team_selection(user_id=user_id, team_id=init_team_id)
+      else:
+          init_team_id = user_current_team.team_id
       # Verify the team exists and user has access to it
       team_configuration = await team_service.get_team_configuration(init_team_id, user_id)
       if team_configuration is None:
@@ -120,7 +124,8 @@ async def init_team(
 
       return {
           "status": "Request started successfully",
-          "team_id": init_team_id
+          "team_id": init_team_id,
+          "team": team_configuration
       }
 
     except Exception as e:
@@ -351,7 +356,7 @@ async def user_clarification(human_feedback: messages.UserClarificationResponse,
 
 
 @app_v3.post("/upload_team_config")
-async def upload_team_config_endpoint(request: Request, file: UploadFile = File(...)):
+async def upload_team_config_endpoint(request: Request, file: UploadFile = File(...),     team_id: Optional[str] = Query(None),):
     """
     Upload and save a team configuration JSON file.
 
@@ -487,6 +492,10 @@ async def upload_team_config_endpoint(request: Request, file: UploadFile = File(
 
         # Save the configuration
         try:
+            print("Saving team configuration...", team_id)
+            if team_id:
+                team_config.team_id = team_id
+                team_config.id = team_id  # Ensure id is also set for updates
             team_id = await team_service.save_team_configuration(team_config)
         except ValueError as e:
             raise HTTPException(

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

@@ -255,7 +255,7 @@ async def handle_team_selection(self, user_id: str, team_id: str) -> bool:
             True if successful, False otherwise
         """
         try:
-            current_team = await self.memory_context.get_current_team(user_id, team_id)
+            current_team = await self.memory_context.get_current_team(user_id)
 
             if current_team is None:
                 current_team = UserCurrentTeam(user_id=user_id, team_id=team_id)