diff --git a/src/backend/agents/base_agent.py b/src/backend/agents/base_agent.py index 46b34960f..59a6bc1ec 100644 --- a/src/backend/agents/base_agent.py +++ b/src/backend/agents/base_agent.py @@ -21,6 +21,7 @@ Step, StepStatus, ) +from models.messages_kernel import AgentType from src.backend.event_utils import track_event_if_configured @@ -76,7 +77,7 @@ async def handle_action_request( AssistantMessage(content=message.action, source="GroupChatManager"), UserMessage( content=f"{step.human_feedback}. Now make the function call", - source="HumanAgent", + source=AgentType.HUMAN.value, ), ] ) diff --git a/src/backend/agents/group_chat_manager.py b/src/backend/agents/group_chat_manager.py index 32d7f2386..13c67f115 100644 --- a/src/backend/agents/group_chat_manager.py +++ b/src/backend/agents/group_chat_manager.py @@ -21,7 +21,7 @@ Step, StepStatus, ) - +from models.messages_kernel import AgentType from src.backend.event_utils import track_event_if_configured @@ -57,7 +57,7 @@ async def handle_input_task( user_id=self._user_id, plan_id="", content=f"{message.description}", - source="HumanAgent", + source=AgentType.HUMAN.value, step_id="", ) ) @@ -68,7 +68,7 @@ async def handle_input_task( "session_id": message.session_id, "user_id": self._user_id, "content": message.description, - "source": "HumanAgent", + "source": AgentType.HUMAN.value, }, ) diff --git a/src/backend/agents/human.py b/src/backend/agents/human.py index 5d1a72d81..d43e9f50a 100644 --- a/src/backend/agents/human.py +++ b/src/backend/agents/human.py @@ -13,6 +13,7 @@ Step, ) from src.backend.event_utils import track_event_if_configured +from models.messages_kernel import AgentType @default_subscription @@ -52,7 +53,7 @@ async def handle_step_feedback( user_id=self.user_id, plan_id=step.plan_id, content=f"Received feedback for step: {step.action}", - source="HumanAgent", + source=AgentType.HUMAN.value, step_id=message.step_id, ) ) @@ -64,7 +65,7 @@ async def handle_step_feedback( "user_id": self.user_id, "plan_id": step.plan_id, "content": f"Received feedback for step: {step.action}", - "source": "HumanAgent", + "source": AgentType.HUMAN.value, "step_id": message.step_id, }, ) diff --git a/src/backend/agents/planner.py b/src/backend/agents/planner.py index e7975be3f..940397a12 100644 --- a/src/backend/agents/planner.py +++ b/src/backend/agents/planner.py @@ -25,7 +25,7 @@ StepStatus, HumanFeedbackStatus, ) - +from models.messages_kernel import AgentType from src.backend.event_utils import track_event_if_configured @@ -133,7 +133,7 @@ async def handle_plan_clarification( user_id=self._user_id, plan_id="", content=f"{message.human_clarification}", - source="HumanAgent", + source=AgentType.HUMAN.value, step_id="", ) ) @@ -144,7 +144,7 @@ async def handle_plan_clarification( "session_id": message.session_id, "user_id": self._user_id, "content": f"{message.human_clarification}", - "source": "HumanAgent", + "source": AgentType.HUMAN.value, }, ) diff --git a/src/backend/app_kernel.py b/src/backend/app_kernel.py index 30f19e551..145b8d552 100644 --- a/src/backend/app_kernel.py +++ b/src/backend/app_kernel.py @@ -413,7 +413,7 @@ async def approve_step_endpoint( agents = await get_agents(human_feedback.session_id, user_id) # Send the approval to the group chat manager - group_chat_manager = agents["GroupChatManager"] + group_chat_manager = agents[AgentType.GROUP_CHAT_MANAGER.value] # Handle the approval human_feedback_json = human_feedback.json() diff --git a/src/backend/models/messages.py b/src/backend/models/messages.py index 60453cb57..ebfd83aa6 100644 --- a/src/backend/models/messages.py +++ b/src/backend/models/messages.py @@ -2,13 +2,7 @@ from enum import Enum from typing import Literal, Optional -from autogen_core.components.models import ( - AssistantMessage, - FunctionExecutionResultMessage, - LLMMessage, - SystemMessage, - UserMessage, -) +from .messages_kernel import AgentType from pydantic import BaseModel, Field @@ -23,15 +17,15 @@ class DataType(str, Enum): class BAgentType(str, Enum): """Enumeration of agent types.""" - human_agent = "HumanAgent" - hr_agent = "HrAgent" - marketing_agent = "MarketingAgent" - procurement_agent = "ProcurementAgent" - product_agent = "ProductAgent" - generic_agent = "GenericAgent" - tech_support_agent = "TechSupportAgent" - group_chat_manager = "GroupChatManager" - planner_agent = "PlannerAgent" + AgentType.HUMAN.value = "HumanAgent" + AgentType.HR.value = "HrAgent" + AgentType.MARKETING.value = "MarketingAgent" + AgentType.PROCUREMENT.value = "ProcurementAgent" + AgentType.PRODUCT.value = "ProductAgent" + AgentType.GENERIC.value = "GenericAgent" + AgentType.TECH_SUPPORT.value = "TechSupportAgent" + AgentType.GROUP_CHAT_MANAGER.value = "GroupChatManager" + AgentType.PLANNER.value = "PlannerAgent" # Add other agents as needed diff --git a/src/backend/tests/test_group_chat_manager_integration.py b/src/backend/tests/test_group_chat_manager_integration.py index 0a05b5c20..6068cf5c9 100644 --- a/src/backend/tests/test_group_chat_manager_integration.py +++ b/src/backend/tests/test_group_chat_manager_integration.py @@ -222,9 +222,9 @@ async def initialize_group_chat_manager(self): # Create agent dictionary for the group chat manager available_agents = { - "PlannerAgent": planner_agent, - "HumanAgent": human_agent, - "GenericAgent": generic_agent + "planner_agent": planner_agent, + "human_agent": human_agent, + "generic_agent": generic_agent } # Create the group chat manager diff --git a/src/backend/utils_kernel.py b/src/backend/utils_kernel.py index 20f9e8a0a..dfe637a6a 100644 --- a/src/backend/utils_kernel.py +++ b/src/backend/utils_kernel.py @@ -95,7 +95,7 @@ async def get_agents(session_id: str, user_id: str) -> Dict[str, Any]: } # Convert to the agent name dictionary format used by the rest of the app - agents = {agent_classes[agent_type]: agent for agent_type, agent in raw_agents.items()} + agents = {agent_type.value: agent for agent_type, agent in raw_agents.items()} # Cache the agents agent_instances[cache_key] = agents