Add team-specific methods to CosmosMemoryContext

Introduces dedicated CRUD methods for TeamConfiguration objects in CosmosMemoryContext, including add, update, retrieve, and delete operations. Refactors JsonService to use these new methods for team configuration management, improving clarity and encapsulation. Adds a test script to verify the correct behavior of all team-specific methods.

Author: Fr4nc3

src/backend/context/cosmos_memory_kernel.py

@@ -17,7 +17,14 @@
 
 # Import the AppConfig instance
 from app_config import config
-from models.messages_kernel import BaseDataModel, Plan, Session, Step, AgentMessage
+from models.messages_kernel import (
+    BaseDataModel,
+    Plan,
+    Session,
+    Step,
+    AgentMessage,
+    TeamConfiguration,
+)
 
 
 # Add custom JSON encoder class for datetime objects
@@ -38,6 +45,7 @@ class CosmosMemoryContext(MemoryStoreBase):
         "plan": Plan,
         "step": Step,
         "agent_message": AgentMessage,
+        "team_config": TeamConfiguration,
         # Messages are handled separately
     }
 
@@ -168,65 +176,6 @@ async def get_item_by_id(
             logging.exception(f"Failed to retrieve item from Cosmos DB: {e}")
             return None
 
-    async def query_items_with_parameters(
-        self, query: str, parameters: List[Dict[str, Any]], limit: int = 1000
-    ) -> List[Dict[str, Any]]:
-        """Query items from Cosmos DB with parameters and return raw dictionaries."""
-        await self.ensure_initialized()
-
-        try:
-            items = self._container.query_items(
-                query=query, 
-                parameters=parameters,
-                enable_cross_partition_query=True
-            )
-            result_list = []
-            count = 0
-            async for item in items:
-                if count >= limit:
-                    break
-                result_list.append(item)
-                count += 1
-            return result_list
-        except Exception as e:
-            logging.exception(f"Failed to query items from Cosmos DB: {e}")
-            return []
-
-    async def get_async(self, key: str) -> Optional[Dict[str, Any]]:
-        """Get an item by its key/ID."""
-        await self.ensure_initialized()
-        
-        try:
-            # Query by ID across all partitions since we don't know the partition key
-            query = "SELECT * FROM c WHERE c.id=@id"
-            parameters = [{"name": "@id", "value": key}]
-            
-            items = self._container.query_items(
-                query=query, 
-                parameters=parameters,
-                enable_cross_partition_query=True
-            )
-            
-            async for item in items:
-                return item
-            return None
-        except Exception as e:
-            logging.exception(f"Failed to get item from Cosmos DB: {e}")
-            return None
-
-    async def delete_async(self, key: str) -> None:
-        """Delete an item by its key/ID."""
-        await self.ensure_initialized()
-        
-        try:
-            # First get the item to find its partition key
-            item = await self.get_async(key)
-            if item:
-                partition_key = item.get("session_id", item.get("user_id", key))
-                await self._container.delete_item(item=key, partition_key=partition_key)
-        except Exception as e:
-            logging.exception(f"Failed to delete item from Cosmos DB: {e}")
-
     async def query_items(
         self,
         query: str,
@@ -400,6 +349,131 @@ async def get_agent_messages_by_session(
         messages = await self.query_items(query, parameters, AgentMessage)
         return messages
 
+    async def add_team(self, team: TeamConfiguration) -> None:
+        """Add a team configuration to Cosmos DB.
+
+        Args:
+            team: The TeamConfiguration to add
+        """
+        await self.add_item(team)
+
+    async def update_team(self, team: TeamConfiguration) -> None:
+        """Update an existing team configuration in Cosmos DB.
+
+        Args:
+            team: The TeamConfiguration to update
+        """
+        await self.update_item(team)
+
+    async def get_team(self, team_id: str) -> Optional[TeamConfiguration]:
+        """Retrieve a specific team configuration by team_id.
+
+        Args:
+            team_id: The team_id of the team configuration to retrieve
+
+        Returns:
+            TeamConfiguration object or None if not found
+        """
+        query = "SELECT * FROM c WHERE c.team_id=@team_id AND c.data_type=@data_type"
+        parameters = [
+            {"name": "@team_id", "value": team_id},
+            {"name": "@data_type", "value": "team_config"},
+        ]
+        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]:
+        """Retrieve a specific team configuration by its document id.
+
+        Args:
+            id: The document id of the team configuration to retrieve
+
+        Returns:
+            TeamConfiguration object or None if not found
+        """
+        query = "SELECT * FROM c WHERE c.id=@id AND c.data_type=@data_type"
+        parameters = [
+            {"name": "@id", "value": id},
+            {"name": "@data_type", "value": "team_config"},
+        ]
+        teams = await self.query_items(query, parameters, TeamConfiguration)
+        return teams[0] if teams else None
+
+    async def get_all_teams_by_user(self, user_id: str) -> List[TeamConfiguration]:
+        """Retrieve all team configurations for a specific user.
+
+        Args:
+            user_id: The user_id to get team configurations for
+
+        Returns:
+            List of TeamConfiguration objects
+        """
+        query = "SELECT * FROM c WHERE c.user_id=@user_id AND c.data_type=@data_type ORDER BY c.created DESC"
+        parameters = [
+            {"name": "@user_id", "value": user_id},
+            {"name": "@data_type", "value": "team_config"},
+        ]
+        teams = await self.query_items(query, parameters, TeamConfiguration)
+        return teams
+
+    async def delete_team(self, team_id: str) -> bool:
+        """Delete a team configuration by team_id.
+
+        Args:
+            team_id: The team_id of the team configuration to delete
+
+        Returns:
+            True if team was found and deleted, False otherwise
+        """
+        await self.ensure_initialized()
+
+        try:
+            # First find the team to get its document id and partition key
+            team = await self.get_team(team_id)
+            if team:
+                # Use the session_id as partition key, or fall back to user_id if no session_id
+                partition_key = (
+                    team.session_id
+                    if hasattr(team, "session_id") and team.session_id
+                    else team.user_id
+                )
+                await self._container.delete_item(
+                    item=team.id, partition_key=partition_key
+                )
+                return True
+            return False
+        except Exception as e:
+            logging.exception(f"Failed to delete team from Cosmos DB: {e}")
+            return False
+
+    async def delete_team_by_id(self, id: str) -> bool:
+        """Delete a team configuration by its document id.
+
+        Args:
+            id: The document id of the team configuration to delete
+
+        Returns:
+            True if team was found and deleted, False otherwise
+        """
+        await self.ensure_initialized()
+
+        try:
+            # First find the team to get its partition key
+            team = await self.get_team_by_id(id)
+            if team:
+                # Use the session_id as partition key, or fall back to user_id if no session_id
+                partition_key = (
+                    team.session_id
+                    if hasattr(team, "session_id") and team.session_id
+                    else team.user_id
+                )
+                await self._container.delete_item(item=id, partition_key=partition_key)
+                return True
+            return False
+        except Exception as e:
+            logging.exception(f"Failed to delete team from Cosmos DB: {e}")
+            return False
+
     async def add_message(self, message: ChatMessageContent) -> None:
         """Add a message to the memory and save to Cosmos DB."""
         await self.ensure_initialized()

src/backend/services/json_service.py

@@ -149,15 +149,8 @@ async def save_team_configuration(self, team_config: TeamConfiguration) -> str:
             The unique ID of the saved configuration
         """
         try:
-            # Convert to dictionary and add data_type for proper querying
-            config_dict = team_config.model_dump()
-            config_dict["data_type"] = "team_config"
-            
-            # Use the cosmos memory context to save the team configuration
-            await self.memory_context.upsert_async(
-                collection_name=team_config.id,
-                record=config_dict
-            )
+            # Use the specific add_team method from cosmos memory context
+            await self.memory_context.add_team(team_config)
 
             self.logger.info(
                 "Successfully saved team configuration with ID: %s", team_config.id
@@ -182,22 +175,22 @@ async def get_team_configuration(
             TeamConfiguration object or None if not found
         """
         try:
-            # Get the specific configuration using cosmos memory context
-            config_dict = await self.memory_context.get_async(config_id)
-            
-            if config_dict is None:
+            # Get the specific configuration using the team-specific method
+            team_config = await self.memory_context.get_team_by_id(config_id)
+
+            if team_config is None:
                 return None
-            
+
             # Verify the configuration belongs to the user
-            if config_dict.get("user_id") != user_id:
+            if team_config.user_id != user_id:
                 self.logger.warning(
                     "Access denied: config %s does not belong to user %s",
                     config_id,
                     user_id,
                 )
                 return None
 
-            return TeamConfiguration.model_validate(config_dict)
+            return team_config
 
         except (KeyError, TypeError, ValueError) as e:
             self.logger.error("Error retrieving team configuration: %s", str(e))
@@ -216,28 +209,8 @@ async def get_all_team_configurations(
             List of TeamConfiguration objects
         """
         try:
-            # Query configurations using SQL with parameters through cosmos memory context
-            query = "SELECT * FROM c WHERE c.user_id=@user_id AND c.data_type=@data_type"
-            parameters = [
-                {"name": "@user_id", "value": user_id},
-                {"name": "@data_type", "value": "team_config"},
-            ]
-            
-            configs = await self.memory_context.query_items_with_parameters(
-                query, parameters, limit=1000
-            )
-
-            team_configs = []
-            for config_dict in configs:
-                try:
-                    team_config = TeamConfiguration.model_validate(config_dict)
-                    team_configs.append(team_config)
-                except (ValueError, TypeError) as e:
-                    self.logger.warning(
-                        "Failed to parse team configuration: %s", str(e)
-                    )
-                    continue
-
+            # Use the specific get_all_teams_by_user method
+            team_configs = await self.memory_context.get_all_teams_by_user(user_id)
             return team_configs
 
         except (KeyError, TypeError, ValueError) as e:
@@ -257,27 +230,32 @@ async def delete_team_configuration(self, config_id: str, user_id: str) -> bool:
         """
         try:
             # First, verify the configuration exists and belongs to the user
-            config_dict = await self.memory_context.get_async(config_id)
-            
-            if config_dict is None:
+            team_config = await self.memory_context.get_team_by_id(config_id)
+
+            if team_config is None:
                 self.logger.warning(
                     "Team configuration not found for deletion: %s", config_id
                 )
                 return False
-            
+
             # Verify the configuration belongs to the user
-            if config_dict.get("user_id") != user_id:
+            if team_config.user_id != user_id:
                 self.logger.warning(
-                    "Access denied: cannot delete config %s for user %s", 
-                    config_id, user_id
+                    "Access denied: cannot delete config %s for user %s",
+                    config_id,
+                    user_id,
                 )
                 return False
 
-            # Delete the configuration using cosmos memory context
-            await self.memory_context.delete_async(config_id)
+            # Delete the configuration using the specific delete_team_by_id method
+            success = await self.memory_context.delete_team_by_id(config_id)
+
+            if success:
+                self.logger.info(
+                    "Successfully deleted team configuration: %s", config_id
+                )
 
-            self.logger.info("Successfully deleted team configuration: %s", config_id)
-            return True
+            return success
 
         except (KeyError, TypeError, ValueError) as e:
             self.logger.error("Error deleting team configuration: %s", str(e))