Refactor RAI agent creation to use team and memory store

Updated RAI agent creation and compliance check functions to require explicit team configuration and memory store parameters. Refactored API endpoints to retrieve team and memory store before RAI checks, improving context handling and error reporting.

Author: Fr4nc3

src/backend/common/utils/utils_af.py

@@ -3,6 +3,8 @@
 import logging
 
 # Converted import path (agent_framework version of FoundryAgentTemplate)
+from common.database.database_base import DatabaseBase
+from common.models.messages_af import TeamConfiguration
 from v4.common.services.team_service import TeamService
 from v4.magentic_agents.foundry_agent import FoundryAgentTemplate  # formerly v4.magentic_agents.foundry_agent
 from v4.config.agent_registry import agent_registry
@@ -34,7 +36,7 @@ async def find_first_available_team(team_service: TeamService, user_id: str) ->
     print("No teams found in priority order")
     return None
 
-async def create_RAI_agent() -> FoundryAgentTemplate:
+async def create_RAI_agent(team: TeamConfiguration, memory_store: DatabaseBase) -> FoundryAgentTemplate:
     """Create and initialize a FoundryAgentTemplate for Responsible AI (RAI) checks."""
     agent_name = "RAIAgent"
     agent_description = "A comprehensive research assistant for integration testing"
@@ -53,6 +55,9 @@ async def create_RAI_agent() -> FoundryAgentTemplate:
     )
 
     model_deployment_name = config.AZURE_OPENAI_DEPLOYMENT_NAME
+    team.team_id = "rai_team"  # Use a fixed team ID for RAI agent
+    team.name = "RAI Team"
+    team.description = "Team responsible for Responsible AI checks"
     agent = FoundryAgentTemplate(
         agent_name=agent_name,
         agent_description=agent_description,
@@ -62,6 +67,8 @@ async def create_RAI_agent() -> FoundryAgentTemplate:
         project_endpoint=config.AZURE_AI_PROJECT_ENDPOINT,
         mcp_config=None,
         search_config=None,
+        team_config=team,
+        memory_store=memory_store,
     )
 
     await agent.open()
@@ -104,14 +111,14 @@ async def _get_agent_response(agent: FoundryAgentTemplate, query: str) -> str:
         return "TRUE"  # Default to blocking on error
 
 
-async def rai_success(description: str) -> bool:
+async def rai_success(description: str, team_config: TeamConfiguration,  memory_store: DatabaseBase) -> bool:
     """
     Run a RAI compliance check on the provided description using the RAIAgent.
     Returns True if content is safe (should proceed), False if it should be blocked.
     """
     agent: FoundryAgentTemplate | None = None
     try:
-        agent = await create_RAI_agent()
+        agent = await create_RAI_agent(team_config, memory_store)
         if not agent:
             logging.error("Failed to instantiate RAIAgent.")
             return False

src/backend/v4/api/router.py

@@ -229,8 +229,25 @@ async def process_request(
               type: string
               description: Error message
     """
-
-    if not await rai_success(input_task.description):
+    try:
+        memory_store = await DatabaseFactory.get_database(user_id=user_id)
+        user_current_team = await memory_store.get_current_team(user_id=user_id)
+        team_id = None
+        if user_current_team:
+            team_id = user_current_team.team_id
+        team = await memory_store.get_team_by_id(team_id=team_id)
+        if not team:
+            raise HTTPException(
+                status_code=404,
+                detail=f"Team configuration '{team_id}' not found or access denied",
+            )
+    except Exception as e:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Error retrieving team configuration: {e}",
+        ) from e
+    
+    if not await rai_success(input_task.description, team, memory_store):
         track_event_if_configured(
             "RAI failed",
             {
@@ -264,17 +281,6 @@ async def process_request(
     try:
         plan_id = str(uuid.uuid4())
         # Initialize memory store and service
-        memory_store = await DatabaseFactory.get_database(user_id=user_id)
-        user_current_team = await memory_store.get_current_team(user_id=user_id)
-        team_id = None
-        if user_current_team:
-            team_id = user_current_team.team_id
-        team = await memory_store.get_team_by_id(team_id=team_id)
-        if not team:
-            raise HTTPException(
-                status_code=404,
-                detail=f"Team configuration '{team_id}' not found or access denied",
-            )
         plan = Plan(
             id=plan_id,
             plan_id=plan_id,
@@ -507,11 +513,28 @@ async def user_clarification(
         raise HTTPException(
             status_code=401, detail="Missing or invalid user information"
         )
+    try:
+        memory_store = await DatabaseFactory.get_database(user_id=user_id)
+        user_current_team = await memory_store.get_current_team(user_id=user_id)
+        team_id = None
+        if user_current_team:
+            team_id = user_current_team.team_id
+        team = await memory_store.get_team_by_id(team_id=team_id)
+        if not team:
+            raise HTTPException(
+                status_code=404,
+                detail=f"Team configuration '{team_id}' not found or access denied",
+            )
+    except Exception as e:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Error retrieving team configuration: {e}",
+        ) from e
     # Set the approval in the orchestration config
     if user_id and human_feedback.request_id:
         # validate rai
         if human_feedback.answer is not None or human_feedback.answer != "":
-            if not await rai_success(human_feedback.answer):
+            if not await rai_success(human_feedback.answer, team, memory_store):
                 track_event_if_configured(
                     "RAI failed",
                     {