camel-ai 0.2.69a7__tar.gz → 0.2.70__tar.gz

{camel_ai-0.2.69a7 → camel_ai-0.2.70}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: camel-ai
-Version: 0.2.69a7
+Version: 0.2.70
 Summary: Communicative Agents for AI Society Study
 Project-URL: Homepage, https://www.camel-ai.org/
 Project-URL: Repository, https://github.com/camel-ai/camel
@@ -87,10 +87,12 @@ Requires-Dist: openapi-spec-validator<0.8,>=0.7.1; extra == 'all'
 Requires-Dist: openpyxl>=3.1.5; extra == 'all'
 Requires-Dist: pandas<2,>=1.5.3; extra == 'all'
 Requires-Dist: pandasai<3,>=2.3.0; extra == 'all'
+Requires-Dist: pgvector<0.3,>=0.2.4; extra == 'all'
 Requires-Dist: playwright>=1.50.0; extra == 'all'
 Requires-Dist: prance<24,>=23.6.21.0; extra == 'all'
 Requires-Dist: praw<8,>=7.7.1; extra == 'all'
 Requires-Dist: pre-commit<4,>=3; extra == 'all'
+Requires-Dist: psycopg[binary]<4,>=3.1.18; extra == 'all'
 Requires-Dist: pyautogui<0.10,>=0.9.54; extra == 'all'
 Requires-Dist: pydub<0.26,>=0.25.1; extra == 'all'
 Requires-Dist: pygithub<3,>=2.6.0; extra == 'all'
@@ -319,6 +321,8 @@ Requires-Dist: google-cloud-storage<3,>=2.18.0; extra == 'storage'
 Requires-Dist: mem0ai>=0.1.73; extra == 'storage'
 Requires-Dist: nebula3-python==3.8.2; extra == 'storage'
 Requires-Dist: neo4j<6,>=5.18.0; extra == 'storage'
+Requires-Dist: pgvector<0.3,>=0.2.4; extra == 'storage'
+Requires-Dist: psycopg[binary]<4,>=3.1.18; extra == 'storage'
 Requires-Dist: pymilvus<3,>=2.4.0; extra == 'storage'
 Requires-Dist: pyobvector>=0.1.18; extra == 'storage'
 Requires-Dist: pytidb-experimental==0.0.1.dev4; extra == 'storage'

{camel_ai-0.2.69a7 → camel_ai-0.2.70}/camel/__init__.py

@@ -14,7 +14,7 @@
 
 from camel.logger import disable_logging, enable_logging, set_log_level
 
-__version__ = '0.2.69a7'
+__version__ = '0.2.70'
 
 __all__ = [
     '__version__',

{camel_ai-0.2.69a7 → camel_ai-0.2.70}/camel/societies/role_playing.py

@@ -556,13 +556,12 @@ class RolePlaying:
             )
         user_msg = self._reduce_message_options(user_response.msgs)
 
-        # To prevent recording the same memory more than once (once in chat
-        # step and once in role play), and the model generates only one
-        # response when multi-response support is enabled.
-        if (
-            'n' in self.user_agent.model_backend.model_config_dict.keys()
-            and self.user_agent.model_backend.model_config_dict['n'] > 1
-        ):
+        # To prevent recording missing messages: ChatAgent.step automatically
+        # saves the response to memory only when a single message is returned.
+        # When multi-response support is enabled (n > 1), it is the caller's
+        # responsibility to record the selected message. Therefore, we record
+        # it here after choosing one message via `_reduce_message_options()`.
+        if self._is_multi_response(self.user_agent):
             self.user_agent.record_message(user_msg)
 
         assistant_response = self.assistant_agent.step(user_msg)
@@ -579,13 +578,7 @@ class RolePlaying:
             )
         assistant_msg = self._reduce_message_options(assistant_response.msgs)
 
-        # To prevent recording the same memory more than once (once in chat
-        # step and once in role play), and the model generates only one
-        # response when multi-response support is enabled.
-        if (
-            'n' in self.assistant_agent.model_backend.model_config_dict.keys()
-            and self.assistant_agent.model_backend.model_config_dict['n'] > 1
-        ):
+        if self._is_multi_response(self.assistant_agent):
             self.assistant_agent.record_message(assistant_msg)
 
         return (
@@ -639,13 +632,7 @@ class RolePlaying:
             )
         user_msg = self._reduce_message_options(user_response.msgs)
 
-        # To prevent recording the same memory more than once (once in chat
-        # step and once in role play), and the model generates only one
-        # response when multi-response support is enabled.
-        if (
-            'n' in self.user_agent.model_backend.model_config_dict.keys()
-            and self.user_agent.model_backend.model_config_dict['n'] > 1
-        ):
+        if self._is_multi_response(self.user_agent):
             self.user_agent.record_message(user_msg)
 
         assistant_response = await self.assistant_agent.astep(user_msg)
@@ -662,13 +649,7 @@ class RolePlaying:
             )
         assistant_msg = self._reduce_message_options(assistant_response.msgs)
 
-        # To prevent recording the same memory more than once (once in chat
-        # step and once in role play), and the model generates only one
-        # response when multi-response support is enabled.
-        if (
-            'n' in self.assistant_agent.model_backend.model_config_dict.keys()
-            and self.assistant_agent.model_backend.model_config_dict['n'] > 1
-        ):
+        if self._is_multi_response(self.assistant_agent):
             self.assistant_agent.record_message(assistant_msg)
 
         return (
@@ -730,3 +711,20 @@ class RolePlaying:
             new_instance.critic = self.critic.clone(with_memory)
 
         return new_instance
+
+    def _is_multi_response(self, agent: ChatAgent) -> bool:
+        r"""Checks if the given agent supports multi-response.
+
+        Args:
+            agent (ChatAgent): The agent to check for multi-response support.
+
+        Returns:
+            bool: True if the agent supports multi-response, False otherwise.
+        """
+        if (
+            'n' in agent.model_backend.model_config_dict.keys()
+            and agent.model_backend.model_config_dict['n'] is not None
+            and agent.model_backend.model_config_dict['n'] > 1
+        ):
+            return True
+        return False

{camel_ai-0.2.69a7 → camel_ai-0.2.70}/camel/societies/workforce/workforce.py

@@ -111,27 +111,24 @@ class Workforce(BaseNode):
         children (Optional[List[BaseNode]], optional): List of child nodes
             under this node. Each child node can be a worker node or
             another workforce node. (default: :obj:`None`)
-        coordinator_agent_kwargs (Optional[Dict], optional): Keyword
-            arguments passed directly to the coordinator :obj:`ChatAgent`
-            constructor. The coordinator manages task assignment and failure
-            handling strategies. See :obj:`ChatAgent` documentation
-            for all available parameters.
-            (default: :obj:`None` - uses ModelPlatformType.DEFAULT,
-            ModelType.DEFAULT)
-        task_agent_kwargs (Optional[Dict], optional): Keyword arguments
-            passed directly to the task planning :obj:`ChatAgent` constructor.
-            The task agent handles task decomposition into subtasks and result
-            composition. See :obj:`ChatAgent` documentation for all
-            available parameters.
-            (default: :obj:`None` - uses ModelPlatformType.DEFAULT,
-            ModelType.DEFAULT)
-        new_worker_agent_kwargs (Optional[Dict], optional): Default keyword
-            arguments passed to :obj:`ChatAgent` constructor for workers
-            created dynamically at runtime when existing workers cannot handle
-            failed tasks. See :obj:`ChatAgent` documentation for all
-            available parameters.
-            (default: :obj:`None` - creates workers with SearchToolkit,
-            CodeExecutionToolkit, and ThinkingToolkit)
+        coordinator_agent (Optional[ChatAgent], optional): A custom coordinator
+            agent instance for task assignment and worker creation. If
+            provided, the workforce will create a new agent using this agent's
+            model configuration but with the required system message and
+            functionality.
+            If None, a default agent will be created using DEFAULT model
+            settings. (default: :obj:`None`)
+        task_agent (Optional[ChatAgent], optional): A custom task planning
+            agent instance for task decomposition and composition. If
+            provided, the workforce will create a new agent using this agent's
+            model configuration but with the required system message and tools
+            (TaskPlanningToolkit). If None, a default agent will be created
+            using DEFAULT model settings. (default: :obj:`None`)
+        new_worker_agent (Optional[ChatAgent], optional): A template agent for
+            workers created dynamically at runtime when existing workers cannot
+            handle failed tasks. If None, workers will be created with default
+            settings including SearchToolkit, CodeExecutionToolkit, and
+            ThinkingToolkit. (default: :obj:`None`)
         graceful_shutdown_timeout (float, optional): The timeout in seconds
             for graceful shutdown when a task fails 3 times. During this
             period, the workforce remains active for debugging.
@@ -147,40 +144,59 @@ class Workforce(BaseNode):
             (default: :obj:`False`)
 
     Example:
-        >>> # Configure with custom model and shared memory
         >>> import asyncio
+        >>> from camel.agents import ChatAgent
+        >>> from camel.models import ModelFactory
+        >>> from camel.types import ModelPlatformType, ModelType
+        >>> from camel.tasks import Task
+        >>>
+        >>> # Simple workforce with default agents
+        >>> workforce = Workforce("Research Team")
+        >>>
+        >>> # Workforce with custom model configuration
         >>> model = ModelFactory.create(
-        ...     ModelPlatformType.OPENAI, ModelType.GPT_4O
+        ...     ModelPlatformType.OPENAI, model_type=ModelType.GPT_4O
         ... )
+        >>> coordinator_agent = ChatAgent(model=model)
+        >>> task_agent = ChatAgent(model=model)
+        >>>
         >>> workforce = Workforce(
         ...     "Research Team",
-        ...     coordinator_agent_kwargs={"model": model, "token_limit": 4000},
-        ...     task_agent_kwargs={"model": model, "token_limit": 8000},
-        ...     share_memory=True  # Enable shared memory
+        ...     coordinator_agent=coordinator_agent,
+        ...     task_agent=task_agent,
         ... )
         >>>
         >>> # Process a task
         >>> async def main():
         ...     task = Task(content="Research AI trends", id="1")
-        ...     result = workforce.process_task(task)
+        ...     result = await workforce.process_task_async(task)
         ...     return result
-        >>> asyncio.run(main())
+        >>>
+        >>> result_task = asyncio.run(main())
+
+    Note:
+        When custom coordinator_agent or task_agent are provided, the workforce
+        will preserve the user's system message and append the required
+        workforce coordination or task planning instructions to it. This
+        ensures both the user's intent is preserved and proper workforce
+        functionality is maintained. All other agent configurations (model,
+        memory, tools, etc.) will also be preserved.
     """
 
     def __init__(
         self,
         description: str,
         children: Optional[List[BaseNode]] = None,
-        coordinator_agent_kwargs: Optional[Dict] = None,
-        task_agent_kwargs: Optional[Dict] = None,
-        new_worker_agent_kwargs: Optional[Dict] = None,
+        coordinator_agent: Optional[ChatAgent] = None,
+        task_agent: Optional[ChatAgent] = None,
+        new_worker_agent: Optional[ChatAgent] = None,  # TODO: use MCP Agent
         graceful_shutdown_timeout: float = 15.0,
         share_memory: bool = False,
     ) -> None:
         super().__init__(description)
         self._child_listening_tasks: Deque[asyncio.Task] = deque()
         self._children = children or []
-        self.new_worker_agent_kwargs = new_worker_agent_kwargs
+        self.new_worker_agent = new_worker_agent
         self.graceful_shutdown_timeout = graceful_shutdown_timeout
         self.share_memory = share_memory
         self.metrics_logger = WorkforceLogger(workforce_id=self.node_id)
@@ -214,58 +230,72 @@ class Workforce(BaseNode):
                     role=role_or_desc,
                 )
 
-        # Warning messages for default model usage
-        if coordinator_agent_kwargs is None:
-            logger.warning(
-                "No coordinator_agent_kwargs provided. Using default "
-                "ChatAgent settings (ModelPlatformType.DEFAULT, "
-                "ModelType.DEFAULT). To customize the coordinator agent "
-                "that assigns tasks and handles failures, pass a dictionary "
-                "with ChatAgent parameters, e.g.: {'model': your_model, "
-                "'tools': your_tools, 'token_limit': 8000}. See ChatAgent "
-                "documentation for all available options."
-            )
-        if task_agent_kwargs is None:
-            logger.warning(
-                "No task_agent_kwargs provided. Using default ChatAgent "
-                "settings (ModelPlatformType.DEFAULT, ModelType.DEFAULT). "
-                "To customize the task planning agent that "
-                "decomposes/composes tasks, pass a dictionary with "
-                "ChatAgent parameters, e.g.: {'model': your_model, "
-                "'token_limit': 16000}. See ChatAgent documentation for "
-                "all available options."
-            )
-        if new_worker_agent_kwargs is None:
-            logger.warning(
-                "No new_worker_agent_kwargs provided. Workers created at "
-                "runtime will use default ChatAgent settings with "
-                "SearchToolkit, CodeExecutionToolkit, and ThinkingToolkit. "
-                "To customize runtime worker creation, pass a dictionary "
-                "with ChatAgent parameters, e.g.: {'model': your_model, "
-                "'tools': your_tools}. See ChatAgent documentation for all "
-                "available options."
-            )
-
-        if self.share_memory:
-            logger.info(
-                "Shared memory enabled. All agents will share their complete "
-                "conversation history and function-calling trajectory for "
-                "better context continuity during task handoffs."
-            )
-
+        # Set up coordinator agent with default system message
         coord_agent_sys_msg = BaseMessage.make_assistant_message(
             role_name="Workforce Manager",
-            content="You are coordinating a group of workers. A worker can be "
-            "a group of agents or a single agent. Each worker is "
+            content="You are coordinating a group of workers. A worker "
+            "can be a group of agents or a single agent. Each worker is "
             "created to solve a specific kind of task. Your job "
             "includes assigning tasks to a existing worker, creating "
             "a new worker for a task, etc.",
         )
-        self.coordinator_agent = ChatAgent(
-            coord_agent_sys_msg,
-            **(coordinator_agent_kwargs or {}),
-        )
 
+        if coordinator_agent is None:
+            logger.warning(
+                "No coordinator_agent provided. Using default "
+                "ChatAgent settings (ModelPlatformType.DEFAULT, "
+                "ModelType.DEFAULT) with default system message."
+            )
+            self.coordinator_agent = ChatAgent(coord_agent_sys_msg)
+        else:
+            logger.info(
+                "Custom coordinator_agent provided. Preserving user's "
+                "system message and appending workforce coordination "
+                "instructions to ensure proper functionality."
+            )
+
+            if coordinator_agent.system_message is not None:
+                user_sys_msg_content = coordinator_agent.system_message.content
+                combined_content = (
+                    f"{user_sys_msg_content}\n\n"
+                    f"{coord_agent_sys_msg.content}"
+                )
+                combined_sys_msg = BaseMessage.make_assistant_message(
+                    role_name=coordinator_agent.system_message.role_name,
+                    content=combined_content,
+                )
+            else:
+                combined_sys_msg = coord_agent_sys_msg
+
+            # Create a new agent with the provided agent's configuration
+            # but with the combined system message
+            self.coordinator_agent = ChatAgent(
+                system_message=combined_sys_msg,
+                model=coordinator_agent.model_backend,
+                memory=coordinator_agent.memory,
+                message_window_size=getattr(
+                    coordinator_agent.memory, "window_size", None
+                ),
+                token_limit=getattr(
+                    coordinator_agent.memory.get_context_creator(),
+                    "token_limit",
+                    None,
+                ),
+                output_language=coordinator_agent.output_language,
+                tools=[
+                    tool.func
+                    for tool in coordinator_agent._internal_tools.values()
+                ],
+                external_tools=[
+                    schema
+                    for schema in coordinator_agent._external_tool_schemas.values()  # noqa: E501
+                ],
+                response_terminators=coordinator_agent.response_terminators,
+                max_iteration=coordinator_agent.max_iteration,
+                stop_event=coordinator_agent.stop_event,
+            )
+
+        # Set up task agent with default system message and required tools
         task_sys_msg = BaseMessage.make_assistant_message(
             role_name="Task Planner",
             content="You are going to compose and decompose tasks. Keep "
@@ -275,13 +305,83 @@ class Workforce(BaseNode):
             "of agents. This ensures efficient execution by minimizing "
             "context switching between agents.",
         )
-        _task_agent_kwargs = dict(task_agent_kwargs or {})
-        extra_tools = TaskPlanningToolkit().get_tools()
-        _task_agent_kwargs["tools"] = [
-            *_task_agent_kwargs.get("tools", []),
-            *extra_tools,
-        ]
-        self.task_agent = ChatAgent(task_sys_msg, **_task_agent_kwargs)
+        task_planning_tools = TaskPlanningToolkit().get_tools()
+
+        if task_agent is None:
+            logger.warning(
+                "No task_agent provided. Using default ChatAgent "
+                "settings (ModelPlatformType.DEFAULT, ModelType.DEFAULT) "
+                "with default system message and TaskPlanningToolkit."
+            )
+            self.task_agent = ChatAgent(
+                task_sys_msg,
+                tools=TaskPlanningToolkit().get_tools(),  # type: ignore[arg-type]
+            )
+        else:
+            logger.info(
+                "Custom task_agent provided. Preserving user's "
+                "system message and appending task planning "
+                "instructions to ensure proper functionality."
+            )
+
+            if task_agent.system_message is not None:
+                user_task_sys_msg_content = task_agent.system_message.content
+                combined_task_content = (
+                    f"{user_task_sys_msg_content}\n\n"
+                    f"{task_sys_msg.content}"
+                )
+                combined_task_sys_msg = BaseMessage.make_assistant_message(
+                    role_name=task_agent.system_message.role_name,
+                    content=combined_task_content,
+                )
+            else:
+                combined_task_sys_msg = task_sys_msg
+
+            # Since ChatAgent constructor uses a dictionary with
+            # function names as keys, we don't need to manually deduplicate.
+            combined_tools = [
+                tool.func for tool in task_agent._internal_tools.values()
+            ] + [tool.func for tool in task_planning_tools]
+
+            # Create a new agent with the provided agent's configuration
+            # but with the combined system message and tools
+            self.task_agent = ChatAgent(
+                system_message=combined_task_sys_msg,
+                model=task_agent.model_backend,
+                memory=task_agent.memory,
+                message_window_size=getattr(
+                    task_agent.memory, "window_size", None
+                ),
+                token_limit=getattr(
+                    task_agent.memory.get_context_creator(),
+                    "token_limit",
+                    None,
+                ),
+                output_language=task_agent.output_language,
+                tools=combined_tools,
+                external_tools=[
+                    schema
+                    for schema in task_agent._external_tool_schemas.values()
+                ],
+                response_terminators=task_agent.response_terminators,
+                max_iteration=task_agent.max_iteration,
+                stop_event=task_agent.stop_event,
+            )
+
+        if new_worker_agent is None:
+            logger.info(
+                "No new_worker_agent provided. Workers created at runtime "
+                "will use default ChatAgent settings with SearchToolkit, "
+                "CodeExecutionToolkit, and ThinkingToolkit. To customize "
+                "runtime worker creation, pass a ChatAgent instance."
+            )
+
+        if self.share_memory:
+            logger.info(
+                "Shared memory enabled. All agents will share their complete "
+                "conversation history and function-calling trajectory for "
+                "better context continuity during task handoffs."
+            )
 
     def __repr__(self):
         return (
@@ -1225,8 +1325,16 @@ class Workforce(BaseNode):
             )
             return TaskAssignResult(assignments=[])
 
-        result_dict = json.loads(response.msg.content, parse_int=str)
-        return TaskAssignResult(**result_dict)
+        try:
+            result_dict = json.loads(response.msg.content, parse_int=str)
+            return TaskAssignResult(**result_dict)
+        except json.JSONDecodeError as e:
+            logger.error(
+                f"JSON parsing error in task assignment: Invalid response "
+                f"format - {e}. Response content: "
+                f"{response.msg.content[:50]}..."
+            )
+            return TaskAssignResult(assignments=[])
 
     def _validate_assignments(
         self, assignments: List[TaskAssignment], valid_ids: Set[str]
@@ -1450,8 +1558,19 @@ class Workforce(BaseNode):
                 "with various tasks.",
             )
         else:
-            result_dict = json.loads(response.msg.content)
-            new_node_conf = WorkerConf(**result_dict)
+            try:
+                result_dict = json.loads(response.msg.content)
+                new_node_conf = WorkerConf(**result_dict)
+            except json.JSONDecodeError as e:
+                logger.error(
+                    f"JSON parsing error in worker creation: Invalid response "
+                    f"format - {e}. Response content: "
+                    f"{response.msg.content[:100]}..."
+                )
+                raise RuntimeError(
+                    f"Failed to create worker for task {task.id}: "
+                    f"Coordinator agent returned malformed JSON response. "
+                )
 
         new_agent = self._create_new_agent(
             new_node_conf.role,
@@ -1486,23 +1605,23 @@ class Workforce(BaseNode):
             content=sys_msg,
         )
 
-        if self.new_worker_agent_kwargs is not None:
-            return ChatAgent(worker_sys_msg, **self.new_worker_agent_kwargs)
-
-        # Default tools for a new agent
-        function_list = [
-            SearchToolkit().search_duckduckgo,
-            *CodeExecutionToolkit().get_tools(),
-            *ThinkingToolkit().get_tools(),
-        ]
+        if self.new_worker_agent is not None:
+            return self.new_worker_agent
+        else:
+            # Default tools for a new agent
+            function_list = [
+                SearchToolkit().search_duckduckgo,
+                *CodeExecutionToolkit().get_tools(),
+                *ThinkingToolkit().get_tools(),
+            ]
 
-        model = ModelFactory.create(
-            model_platform=ModelPlatformType.DEFAULT,
-            model_type=ModelType.DEFAULT,
-            model_config_dict={"temperature": 0},
-        )
+            model = ModelFactory.create(
+                model_platform=ModelPlatformType.DEFAULT,
+                model_type=ModelType.DEFAULT,
+                model_config_dict={"temperature": 0},
+            )
 
-        return ChatAgent(worker_sys_msg, model=model, tools=function_list)  # type: ignore[arg-type]
+            return ChatAgent(worker_sys_msg, model=model, tools=function_list)  # type: ignore[arg-type]
 
     async def _get_returned_task(self) -> Task:
         r"""Get the task that's published by this node and just get returned
@@ -1988,28 +2107,17 @@ class Workforce(BaseNode):
         """
 
         # Create a new instance with the same configuration
-        # Extract the original kwargs from the agents to properly clone them
-        coordinator_kwargs = (
-            getattr(self.coordinator_agent, 'init_kwargs', {}) or {}
-        )
-        task_kwargs = getattr(self.task_agent, 'init_kwargs', {}) or {}
-
         new_instance = Workforce(
             description=self.description,
-            coordinator_agent_kwargs=coordinator_kwargs.copy(),
-            task_agent_kwargs=task_kwargs.copy(),
-            new_worker_agent_kwargs=self.new_worker_agent_kwargs.copy()
-            if self.new_worker_agent_kwargs
+            coordinator_agent=self.coordinator_agent.clone(with_memory),
+            task_agent=self.task_agent.clone(with_memory),
+            new_worker_agent=self.new_worker_agent.clone(with_memory)
+            if self.new_worker_agent
             else None,
             graceful_shutdown_timeout=self.graceful_shutdown_timeout,
             share_memory=self.share_memory,
         )
 
-        new_instance.task_agent = self.task_agent.clone(with_memory)
-        new_instance.coordinator_agent = self.coordinator_agent.clone(
-            with_memory
-        )
-
         for child in self._children:
             if isinstance(child, SingleAgentWorker):
                 cloned_worker = child.worker.clone(with_memory)

{camel_ai-0.2.69a7 → camel_ai-0.2.70}/camel/societies/workforce/workforce_logger.py

@@ -499,54 +499,68 @@ class WorkforceLogger:
 
         tasks_handled_by_worker: Dict[str, int] = {}
 
+        # Helper function to check if a task is the main task (has no parent)
+        def is_main_task(task_id: str) -> bool:
+            return (
+                task_id in self._task_hierarchy
+                and self._task_hierarchy[task_id].get('parent') is None
+            )
+
         for entry in self.log_entries:
             event_type = entry['event_type']
             timestamp = datetime.fromisoformat(entry['timestamp'])
+            task_id = entry.get('task_id', '')
+
             if first_timestamp is None or timestamp < first_timestamp:
                 first_timestamp = timestamp
             if last_timestamp is None or timestamp > last_timestamp:
                 last_timestamp = timestamp
 
             if event_type == 'task_created':
-                kpis['total_tasks_created'] += 1
-                task_creation_timestamps[entry['task_id']] = timestamp
+                # Exclude main task from total count
+                if not is_main_task(task_id):
+                    kpis['total_tasks_created'] += 1
+                task_creation_timestamps[task_id] = timestamp
             elif event_type == 'task_assigned':
-                task_assignment_timestamps[entry['task_id']] = timestamp
+                task_assignment_timestamps[task_id] = timestamp
                 # Queue time tracking has been removed
 
             elif event_type == 'task_started':
                 # Store start time for processing time calculation
-                task_start_times[entry['task_id']] = timestamp.timestamp()
+                task_start_times[task_id] = timestamp.timestamp()
 
             elif event_type == 'task_completed':
-                kpis['total_tasks_completed'] += 1
-                # Count tasks handled by worker
-                if 'worker_id' in entry and entry['worker_id'] is not None:
-                    worker_id = entry['worker_id']
-                    tasks_handled_by_worker[worker_id] = (
-                        tasks_handled_by_worker.get(worker_id, 0) + 1
-                    )
+                # Exclude main task from total count
+                if not is_main_task(task_id):
+                    kpis['total_tasks_completed'] += 1
+                    # Count tasks handled by worker (only for non-main tasks)
+                    if 'worker_id' in entry and entry['worker_id'] is not None:
+                        worker_id = entry['worker_id']
+                        tasks_handled_by_worker[worker_id] = (
+                            tasks_handled_by_worker.get(worker_id, 0) + 1
+                        )
 
-                if entry['task_id'] in task_assignment_timestamps:
+                if task_id in task_assignment_timestamps:
                     completion_time = (
-                        timestamp
-                        - task_assignment_timestamps[entry['task_id']]
+                        timestamp - task_assignment_timestamps[task_id]
                     ).total_seconds()
                     # Store completion time in task hierarchy instead of KPIs
                     # array
-                    if entry['task_id'] in self._task_hierarchy:
-                        self._task_hierarchy[entry['task_id']][
+                    if task_id in self._task_hierarchy:
+                        self._task_hierarchy[task_id][
                             'completion_time_seconds'
                         ] = completion_time
 
             elif event_type == 'task_failed':
-                kpis['total_tasks_failed'] += 1
-                # Count tasks handled by worker (also for failed tasks)
-                if 'worker_id' in entry and entry['worker_id'] is not None:
-                    worker_id = entry['worker_id']
-                    tasks_handled_by_worker[worker_id] = (
-                        tasks_handled_by_worker.get(worker_id, 0) + 1
-                    )
+                # Exclude main task from total count
+                if not is_main_task(task_id):
+                    kpis['total_tasks_failed'] += 1
+                    # Count tasks handled by worker (only for non-main tasks)
+                    if 'worker_id' in entry and entry['worker_id'] is not None:
+                        worker_id = entry['worker_id']
+                        tasks_handled_by_worker[worker_id] = (
+                            tasks_handled_by_worker.get(worker_id, 0) + 1
+                        )
                 error_type = entry['error_type']
                 kpis['error_types_count'][error_type] = (
                     kpis['error_types_count'].get(error_type, 0) + 1