Refactor get_team_by_id and improve plan retrieval error handling

Renamed the get_team_by_id method parameter from 'id' to 'team_id' for clarity and updated its usage in CosmosDBClient and DatabaseBase. Refactored the /plan endpoint in router.py to use Optional plan_id and added improved error handling with logging for plan retrieval.

Author: Fr4nc3

src/backend/common/database/cosmosdb.py

@@ -293,7 +293,7 @@ async def get_team(self, team_id: str) -> Optional[TeamConfiguration]:
         teams = await self.query_items(query, parameters, TeamConfiguration)
         return teams[0] if teams else None
 
-    async def get_team_by_id(self, id: str) -> Optional[TeamConfiguration]:
+    async def get_team_by_id(self, team_id: str) -> Optional[TeamConfiguration]:
         """Retrieve a specific team configuration by its document id.
 
         Args:
@@ -302,9 +302,9 @@ async def get_team_by_id(self, id: str) -> Optional[TeamConfiguration]:
         Returns:
             TeamConfiguration object or None if not found
         """
-        query = "SELECT * FROM c WHERE c.id=@id AND c.data_type=@data_type"
+        query = "SELECT * FROM c WHERE c.team_id=@team_id AND c.data_type=@data_type"
         parameters = [
-            {"name": "@id", "value": id},
+            {"name": "@team_id", "value": team_id},
             {"name": "@data_type", "value": DataType.team_config},
         ]
         teams = await self.query_items(query, parameters, TeamConfiguration)

src/backend/common/database/database_base.py

@@ -138,7 +138,7 @@ async def get_team(self, team_id: str) -> Optional[TeamConfiguration]:
         pass
 
     @abstractmethod
-    async def get_team_by_id(self, id: str) -> Optional[TeamConfiguration]:
+    async def get_team_by_id(self, team_id: str) -> Optional[TeamConfiguration]:
         """Retrieve a team configuration by internal id."""
         pass
 

src/backend/v3/api/router.py

@@ -1255,7 +1255,7 @@ async def get_plans(request: Request):
 
 # Get plans is called in the initial side rendering of the frontend
 @app_v3.get("/plan")
-async def get_plan_by_id(request: Request, plan_id: str):
+async def get_plan_by_id(request: Request,  plan_id: Optional[str] = Query(None),):
     """
     Retrieve plans for the current user.
 
@@ -1326,29 +1326,32 @@ async def get_plan_by_id(request: Request, plan_id: str):
 
     # # Initialize memory context
     memory_store = await DatabaseFactory.get_database(user_id=user_id)
+    try:
+        if plan_id:
+            plan = await memory_store.get_plan_by_plan_id(plan_id=plan_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
 
-    if plan_id:
-        plan = await memory_store.get_plan_by_plan_id(plan_id=plan_id)
-        if not plan:
+            team = await memory_store.get_team_by_id(team_id=plan.team_id)
+            agent_messages = await memory_store.get_agent_messages(plan_id=plan.plan_id)
+            m_plan = await memory_store.get_mplan(plan_id=plan.plan_id)
+            return {
+                "plan": plan,
+                "team": team if team else None,
+                "messages": agent_messages,
+                "m_plan": m_plan if m_plan else None,
+            }
+        else:
             track_event_if_configured(
-                "GetPlanBySessionNotFound",
-                {"status_code": 400, "detail": "Plan not found"},
+                "GetPlanId", {"status_code": 400, "detail": "no plan id"}
             )
-            raise HTTPException(status_code=404, detail="Plan not found")
-
-        # Use get_steps_by_plan to match the original implementation
-
-        team = await memory_store.get_team_by_id(team_id=plan.team_id)
-        agent_messages = await memory_store.get_agent_messages(plan_id=plan.plan_id)
-        m_plan = await memory_store.get_mplan(plan_id=plan.plan_id)
-        return {
-            "plan": plan,
-            "team": team if team else None,
-            "messages": agent_messages,
-            "m_plan": m_plan if m_plan else None,
-        }
-    else:
-        track_event_if_configured(
-            "GetPlanId", {"status_code": 400, "detail": "no plan id"}
-        )
-        raise HTTPException(status_code=400, detail="no plan id")
+            raise HTTPException(status_code=400, detail="no plan id")
+    except Exception as e:
+        logging.error(f"Error retrieving plan: {str(e)}")
+        raise HTTPException(status_code=500, detail="Internal server error occurred")