Refactor agent message handling and plan retrieval

Updated database interfaces and implementations to use AgentMessageData instead of AgentMessageResponse for agent message operations. Added ordering to CosmosDB queries for plans and messages. Modified API route for plan retrieval to return plan, team, messages, and m_plan as structured objects instead of a list, improving response consistency and context.

Author: Fr4nc3

src/backend/common/database/cosmosdb.py

@@ -237,7 +237,7 @@ async def get_all_plans_by_team_id(self, team_id: str) -> List[Plan]:
 
     async def get_all_plans_by_team_id_status(self, team_id: str, status: str) -> List[Plan]:
         """Retrieve all plans for a specific team."""
-        query = "SELECT * FROM c WHERE c.team_id=@team_id AND c.data_type=@data_type and c.user_id=@user_id and c.overall_status=@status"
+        query = "SELECT * FROM c WHERE c.team_id=@team_id AND c.data_type=@data_type and c.user_id=@user_id and c.overall_status=@status ORDER BY c._ts DESC"
         parameters = [
             {"name": "@user_id", "value": self.user_id},
             {"name": "@team_id", "value": team_id},
@@ -474,7 +474,7 @@ async def update_mplan(self, mplan: messages.MPlan) -> None:
 
     async def get_mplan(self, plan_id: str) -> Optional[messages.MPlan]:
         """Retrieve a mplan configuration by mplan_id."""
-        query = "SELECT * FROM c WHERE c.plan_id=@plan_id AND c.data_type=@data_type"
+        query = "SELECT * FROM c WHERE c.plan_id=@plan_id AND c.data_type=@data_type ORDER BY c._ts ASC"
         parameters = [
             {"name": "@plan_id", "value": plan_id},
             {"name": "@data_type", "value": DataType.m_plan},
@@ -483,20 +483,20 @@ async def get_mplan(self, plan_id: str) -> Optional[messages.MPlan]:
         return results[0] if results else None
 
 
-    async def add_agent_message(self, message: messages.AgentMessageResponse) -> None:
+    async def add_agent_message(self, message: AgentMessageData) -> None:
         """Add an agent message to the database."""
         await self.add_item(message)
 
-    async def update_agent_message(self, message: messages.AgentMessageResponse) -> None:
+    async def update_agent_message(self, message: AgentMessageData) -> None:
         """Update an agent message in the database."""
         await self.update_item(message)
 
-    async def get_agent_messages(self, plan_id: str) -> List[messages.AgentMessageResponse]:
+    async def get_agent_messages(self, plan_id: str) -> List[AgentMessageData]:
         """Retrieve an agent message by message_id."""
-        query = "SELECT * FROM c WHERE c.plan_id=@plan_id AND c.data_type=@data_type"
+        query = "SELECT * FROM c WHERE c.plan_id=@plan_id AND c.data_type=@data_type ORDER BY c._ts ASC"
         parameters = [
             {"name": "@plan_id", "value": plan_id},
             {"name": "@data_type", "value": DataType.m_plan_message},
         ]
 
-        return await self.query_items(query, parameters, messages.AgentMessageResponse)
+        return await self.query_items(query, parameters, AgentMessageData)

src/backend/common/database/database_base.py

@@ -222,11 +222,11 @@ async def add_agent_message(self, message: AgentMessageData) -> None:
         pass
 
     @abstractmethod
-    async def update_agent_message(self, message: AgentMessageResponse) -> None:
+    async def update_agent_message(self, message: AgentMessageData) -> None:
         """Update an agent message in the database."""
         pass
 
     @abstractmethod
-    async def get_agent_messages(self, plan_id: str) -> Optional[AgentMessageResponse]:
+    async def get_agent_messages(self, plan_id: str) -> Optional[AgentMessageData]:
         """Retrieve an agent message by message_id."""
         pass

src/backend/v3/api/router.py

@@ -559,7 +559,8 @@ async def user_clarification(
 
             try:
                 result = await PlanService.handle_human_clarification(
-                    human_feedback, user_id)
+                    human_feedback, user_id
+                )
                 print("Human clarification processed:", result)
             except ValueError as ve:
                 print(f"ValueError processing human clarification: {ve}")
@@ -573,7 +574,9 @@ async def user_clarification(
                     "user_id": user_id,
                 },
             )
-            return {"status": "clarification recorded",}
+            return {
+                "status": "clarification recorded",
+            }
         else:
             logging.warning(
                 f"No orchestration or plan found for request_id: {human_feedback.request_id}"
@@ -582,6 +585,7 @@ async def user_clarification(
                 status_code=404, detail="No active plan found for clarification"
             )
 
+
 @app_v3.post("/agent_message")
 async def agent_message_user(
     agent_message: messages.AgentMessageResponse, request: Request
@@ -642,14 +646,14 @@ async def agent_message_user(
     # Set the approval in the orchestration config
 
     try:
-    
+
         result = await PlanService.handle_agent_messages(agent_message, user_id)
         print("Agent message processed:", result)
     except ValueError as ve:
         print(f"ValueError processing agent message: {ve}")
     except Exception as e:
         print(f"Error processing agent message: {e}")
-   
+
     track_event_if_configured(
         "AgentMessageReceived",
         {
@@ -658,7 +662,9 @@ async def agent_message_user(
             "user_id": user_id,
         },
     )
-    return {"status": "message recorded",}
+    return {
+        "status": "message recorded",
+    }
 
 
 @app_v3.post("/upload_team_config")
@@ -726,19 +732,19 @@ async def upload_team_config(
                 track_event_if_configured(
                     "Team configuration RAI validation failed",
                     {
-                    "status": "failed",
-                    "user_id": user_id,
-                    "filename": file.filename,
-                    "reason": rai_error,
-                },
-            )
-            
+                        "status": "failed",
+                        "user_id": user_id,
+                        "filename": file.filename,
+                        "reason": rai_error,
+                    },
+                )
+
             raise HTTPException(status_code=400, detail=rai_error)
 
         track_event_if_configured(
-                "Team configuration RAI validation passed",
-                {"status": "passed", "user_id": user_id, "filename": file.filename},
-    )
+            "Team configuration RAI validation passed",
+            {"status": "passed", "user_id": user_id, "filename": file.filename},
+        )
         # Initialize memory store and service
         memory_store = await DatabaseFactory.get_database(user_id=user_id)
         team_service = TeamService(memory_store)
@@ -1340,16 +1346,16 @@ async def get_plan_by_id(request: Request, plan_id: str):
             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)
-        messages = []
 
-        plan_with_steps = PlanWithSteps(**plan.model_dump(), steps=steps)
-        plan_with_steps.update_step_counts()
-
-        # Format dates in messages according to locale
-        formatted_messages = []
-
-        return [plan_with_steps, formatted_messages]
+        team = await memory_store.get_team_by_id(team_id=plan.team_id)
+        messages = await memory_store.get_agent_messages(plan_id=plan.plan_id)
+        m_plan = await memory_store.get_m_plan_by_plan_id(plan_id=plan.plan_id)
+        return {
+            "plan": plan.model_dump(),
+            "team": team.model_dump() if team else None,
+            "messages": [msg.model_dump() for msg in messages],
+            "m_plan": m_plan.model_dump() if m_plan else None,
+        }
     else:
         track_event_if_configured(
             "GetPlanId", {"status_code": 400, "detail": "no plan id"}