xpander-sdk 2.0.161__py3-none-any.whl → 2.0.192__py3-none-any.whl

xpander_sdk/modules/events/decorators/on_tool.py

@@ -0,0 +1,384 @@
+"""
+xpander_sdk.decorators.on_tool
+
+This module provides decorators for tool invocation lifecycle hooks:
+- `@on_tool_before`: Execute before tool invocation
+- `@on_tool_after`: Execute after successful tool invocation
+- `@on_tool_error`: Execute when tool invocation fails
+
+The decorators ensure that registered functions:
+- Accept parameters: Tool, payload, payload_extension, tool_call_id, agent_version
+- Can be either synchronous or asynchronous
+- Are called at the appropriate lifecycle stage during tool execution
+
+Execution Notes:
+- Before-hooks execute before tool invocation and can perform validation or logging
+- After-hooks execute after successful invocation with access to the result
+- Error-hooks execute when an exception occurs during invocation
+- Multiple hooks of the same type can be registered and will execute in registration order
+- Exceptions in hooks are logged but don't prevent tool execution
+
+Example usage:
+--------------
+>>> @on_tool_before
+... async def log_tool_invocation(tool, payload, payload_extension, tool_call_id, agent_version):
+...     logger.info(f"Invoking tool {tool.name} with payload: {payload}")
+
+>>> @on_tool_after
+... def record_tool_result(tool, payload, payload_extension, tool_call_id, agent_version, result):
+...     logger.info(f"Tool {tool.name} completed with result: {result}")
+
+>>> @on_tool_error
+... async def handle_tool_error(tool, payload, payload_extension, tool_call_id, agent_version, error):
+...     logger.error(f"Tool {tool.name} failed: {error}")
+"""
+
+import asyncio
+from functools import wraps
+from inspect import iscoroutinefunction
+from typing import Optional, Callable, Any, Dict, List
+
+from xpander_sdk.models.configuration import Configuration
+
+
+class ToolHooksRegistry:
+    """
+    Registry for tool invocation lifecycle hooks.
+    
+    This class maintains class-level lists of hooks that are executed at different
+    stages of tool invocation: before, after, and on error.
+    """
+    
+    _before_hooks: List[Callable] = []
+    _after_hooks: List[Callable] = []
+    _error_hooks: List[Callable] = []
+    
+    @classmethod
+    def register_before_hook(cls, hook: Callable) -> None:
+        """
+        Register a before-invocation hook.
+        
+        Args:
+            hook (Callable): The hook function to execute before tool invocation.
+        """
+        cls._before_hooks.append(hook)
+    
+    @classmethod
+    def register_after_hook(cls, hook: Callable) -> None:
+        """
+        Register an after-invocation hook.
+        
+        Args:
+            hook (Callable): The hook function to execute after successful tool invocation.
+        """
+        cls._after_hooks.append(hook)
+    
+    @classmethod
+    def register_error_hook(cls, hook: Callable) -> None:
+        """
+        Register an error hook.
+        
+        Args:
+            hook (Callable): The hook function to execute when tool invocation fails.
+        """
+        cls._error_hooks.append(hook)
+    
+    @classmethod
+    async def execute_before_hooks(
+        cls,
+        tool: Any,
+        payload: Any,
+        payload_extension: Optional[Dict[str, Any]] = None,
+        tool_call_id: Optional[str] = None,
+        agent_version: Optional[str] = None
+    ) -> None:
+        """
+        Execute all registered before-invocation hooks.
+        
+        Args:
+            tool: The Tool object being invoked.
+            payload: The payload being sent to the tool.
+            payload_extension: Additional payload data.
+            tool_call_id: Unique ID of the tool call.
+            agent_version: Version of the agent making the call.
+        """
+        from loguru import logger
+        
+        for hook in cls._before_hooks:
+            try:
+                if asyncio.iscoroutinefunction(hook):
+                    await hook(tool, payload, payload_extension, tool_call_id, agent_version)
+                else:
+                    hook(tool, payload, payload_extension, tool_call_id, agent_version)
+            except Exception as e:
+                logger.error(f"Before-hook {hook.__name__} failed: {e}")
+    
+    @classmethod
+    async def execute_after_hooks(
+        cls,
+        tool: Any,
+        payload: Any,
+        payload_extension: Optional[Dict[str, Any]] = None,
+        tool_call_id: Optional[str] = None,
+        agent_version: Optional[str] = None,
+        result: Any = None
+    ) -> None:
+        """
+        Execute all registered after-invocation hooks.
+        
+        Args:
+            tool: The Tool object that was invoked.
+            payload: The payload sent to the tool.
+            payload_extension: Additional payload data.
+            tool_call_id: Unique ID of the tool call.
+            agent_version: Version of the agent that made the call.
+            result: The result returned by the tool invocation.
+        """
+        from loguru import logger
+        
+        for hook in cls._after_hooks:
+            try:
+                if asyncio.iscoroutinefunction(hook):
+                    await hook(tool, payload, payload_extension, tool_call_id, agent_version, result)
+                else:
+                    hook(tool, payload, payload_extension, tool_call_id, agent_version, result)
+            except Exception as e:
+                logger.error(f"After-hook {hook.__name__} failed: {e}")
+    
+    @classmethod
+    async def execute_error_hooks(
+        cls,
+        tool: Any,
+        payload: Any,
+        payload_extension: Optional[Dict[str, Any]] = None,
+        tool_call_id: Optional[str] = None,
+        agent_version: Optional[str] = None,
+        error: Optional[Exception] = None
+    ) -> None:
+        """
+        Execute all registered error hooks.
+        
+        Args:
+            tool: The Tool object that failed.
+            payload: The payload sent to the tool.
+            payload_extension: Additional payload data.
+            tool_call_id: Unique ID of the tool call.
+            agent_version: Version of the agent that made the call.
+            error: The exception that occurred during invocation.
+        """
+        from loguru import logger
+        
+        for hook in cls._error_hooks:
+            try:
+                if asyncio.iscoroutinefunction(hook):
+                    await hook(tool, payload, payload_extension, tool_call_id, agent_version, error)
+                else:
+                    hook(tool, payload, payload_extension, tool_call_id, agent_version, error)
+            except Exception as e:
+                logger.error(f"Error-hook {hook.__name__} failed: {e}")
+
+
+def on_tool_before(
+    _func: Optional[Callable] = None,
+    *,
+    configuration: Optional[Configuration] = None
+):
+    """
+    Decorator to register a handler as a before-invocation hook for tool calls.
+    
+    The decorated function will be executed before any tool invocation. The function:
+    - Must accept parameters: tool, payload, payload_extension, tool_call_id, agent_version
+    - Can be either synchronous or asynchronous
+    - Can perform validation, logging, or modification of invocation context
+    
+    Args:
+        _func (Optional[Callable]):
+            The function to decorate (for direct usage like `@on_tool_before`).
+        configuration (Optional[Configuration]):
+            An optional configuration object (reserved for future use).
+    
+    Example:
+        >>> from typing import Optional, Dict, Any
+        >>> from xpander_sdk import Tool
+        >>> 
+        >>> @on_tool_before
+        ... async def validate_tool_input(
+        ...     tool: Tool,
+        ...     payload: Any,
+        ...     payload_extension: Optional[Dict[str, Any]] = None,
+        ...     tool_call_id: Optional[str] = None,
+        ...     agent_version: Optional[str] = None
+        ... ):
+        ...     logger.info(f"Pre-invoking tool: {tool.name}")
+        ...     if not payload:
+        ...         logger.warning("Empty payload provided")
+        
+        >>> @on_tool_before
+        ... def log_tool_metrics(
+        ...     tool: Tool,
+        ...     payload: Any,
+        ...     payload_extension: Optional[Dict[str, Any]] = None,
+        ...     tool_call_id: Optional[str] = None,
+        ...     agent_version: Optional[str] = None
+        ... ):
+        ...     metrics.record_tool_invocation(tool.name, tool_call_id)
+    """
+    
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            return await func(*args, **kwargs)
+        
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+        
+        wrapped = async_wrapper if iscoroutinefunction(func) else sync_wrapper
+        
+        # Register hook in the registry
+        ToolHooksRegistry.register_before_hook(wrapped)
+        
+        return wrapped
+    
+    if _func and callable(_func):
+        return decorator(_func)
+    
+    return decorator
+
+
+def on_tool_after(
+    _func: Optional[Callable] = None,
+    *,
+    configuration: Optional[Configuration] = None
+):
+    """
+    Decorator to register a handler as an after-invocation hook for tool calls.
+    
+    The decorated function will be executed after successful tool invocation. The function:
+    - Must accept parameters: tool, payload, payload_extension, tool_call_id, agent_version, result
+    - Can be either synchronous or asynchronous
+    - Can perform logging, analytics, or result processing
+    
+    Args:
+        _func (Optional[Callable]):
+            The function to decorate (for direct usage like `@on_tool_after`).
+        configuration (Optional[Configuration]):
+            An optional configuration object (reserved for future use).
+    
+    Example:
+        >>> from typing import Optional, Dict, Any
+        >>> from xpander_sdk import Tool
+        >>> 
+        >>> @on_tool_after
+        ... async def log_tool_success(
+        ...     tool: Tool,
+        ...     payload: Any,
+        ...     payload_extension: Optional[Dict[str, Any]] = None,
+        ...     tool_call_id: Optional[str] = None,
+        ...     agent_version: Optional[str] = None,
+        ...     result: Any = None
+        ... ):
+        ...     logger.info(f"Tool {tool.name} succeeded with result: {result}")
+        ...     analytics.record_success(tool.name, tool_call_id)
+        
+        >>> @on_tool_after
+        ... def cache_tool_result(
+        ...     tool: Tool,
+        ...     payload: Any,
+        ...     payload_extension: Optional[Dict[str, Any]] = None,
+        ...     tool_call_id: Optional[str] = None,
+        ...     agent_version: Optional[str] = None,
+        ...     result: Any = None
+        ... ):
+        ...     cache.set(f"tool_{tool.id}_{hash(payload)}", result)
+    """
+    
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            return await func(*args, **kwargs)
+        
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+        
+        wrapped = async_wrapper if iscoroutinefunction(func) else sync_wrapper
+        
+        # Register hook in the registry
+        ToolHooksRegistry.register_after_hook(wrapped)
+        
+        return wrapped
+    
+    if _func and callable(_func):
+        return decorator(_func)
+    
+    return decorator
+
+
+def on_tool_error(
+    _func: Optional[Callable] = None,
+    *,
+    configuration: Optional[Configuration] = None
+):
+    """
+    Decorator to register a handler as an error hook for tool calls.
+    
+    The decorated function will be executed when tool invocation fails. The function:
+    - Must accept parameters: tool, payload, payload_extension, tool_call_id, agent_version, error
+    - Can be either synchronous or asynchronous
+    - Can perform error logging, alerting, or recovery actions
+    
+    Args:
+        _func (Optional[Callable]):
+            The function to decorate (for direct usage like `@on_tool_error`).
+        configuration (Optional[Configuration]):
+            An optional configuration object (reserved for future use).
+    
+    Example:
+        >>> from typing import Optional, Dict, Any
+        >>> from xpander_sdk import Tool
+        >>> 
+        >>> @on_tool_error
+        ... async def alert_on_tool_failure(
+        ...     tool: Tool,
+        ...     payload: Any,
+        ...     payload_extension: Optional[Dict[str, Any]] = None,
+        ...     tool_call_id: Optional[str] = None,
+        ...     agent_version: Optional[str] = None,
+        ...     error: Optional[Exception] = None
+        ... ):
+        ...     logger.error(f"Tool {tool.name} failed: {error}")
+        ...     await send_alert(f"Tool failure: {tool.name}", str(error))
+        
+        >>> @on_tool_error
+        ... def log_tool_failure(
+        ...     tool: Tool,
+        ...     payload: Any,
+        ...     payload_extension: Optional[Dict[str, Any]] = None,
+        ...     tool_call_id: Optional[str] = None,
+        ...     agent_version: Optional[str] = None,
+        ...     error: Optional[Exception] = None
+        ... ):
+        ...     error_log.write(f"{tool.name},{tool_call_id},{str(error)}\n")
+    """
+    
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            return await func(*args, **kwargs)
+        
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+        
+        wrapped = async_wrapper if iscoroutinefunction(func) else sync_wrapper
+        
+        # Register hook in the registry
+        ToolHooksRegistry.register_error_hook(wrapped)
+        
+        return wrapped
+    
+    if _func and callable(_func):
+        return decorator(_func)
+    
+    return decorator

xpander_sdk/modules/events/events_module.py

@@ -81,6 +81,7 @@ class Events(ModuleBase):
 
     worker: Optional[DeployedAsset] = None
     test_task: Optional[LocalTaskTest] = None
+    _heartbeat_task: Optional[asyncio.Task] = None
     
     # Class-level registries for boot and shutdown handlers
     _boot_handlers: List[BootHandler] = []
@@ -358,12 +359,12 @@ class Events(ModuleBase):
             plan_following_status = await task.aget_plan_following_status()
             if not plan_following_status.can_finish:
                 # Check if we've exceeded max retries
-                if retry_count >= 2:  # 0, 1, 2 = 3 total attempts
+                if retry_count >= 50:  # 0, 1, 2 = 50 total attempts
                     logger.warning(f"Failed to complete plan after {retry_count + 1} attempts. Remaining incomplete tasks.")
                     return
                 
                 # Recursively call with incremented retry count
-                logger.info(f"Plan not complete, retrying (attempt {retry_count + 2}/3)")
+                logger.info(f"Plan not complete, retrying (attempt {retry_count + 2})")
                 await self.handle_task_execution_request(
                     agent_worker,
                     task,
@@ -500,7 +501,12 @@ class Events(ModuleBase):
                     )
                 )
 
-                self.track(asyncio.create_task(self.heartbeat_loop(agent_worker.id)))
+                # Cancel previous heartbeat task if it exists and start a new one
+                if self._heartbeat_task and not self._heartbeat_task.done():
+                    logger.debug(f"Canceling previous heartbeat task for worker {agent_worker.id}")
+                    self._heartbeat_task.cancel()
+                self._heartbeat_task = asyncio.create_task(self.heartbeat_loop(agent_worker.id))
+                self.track(self._heartbeat_task)
 
             elif event.event == EventType.AgentExecution:
                 task = Task(**json.loads(event.data), configuration=self.configuration)

xpander_sdk/modules/tasks/models/task.py

@@ -39,26 +39,15 @@ class AgentExecutionStatus(str, Enum):
     Stopped = "stopped"
 
 
-class HumanInTheLoop(BaseModel):
+class HumanInTheLoopRequest(BaseModel):
     """
     Model representing human-in-the-loop approval records for tasks.
     
     Attributes:
-        operation_id (str): Unique identifier of the operation requiring approval.
-        approved_by (Optional[str]): User who approved the operation.
-        rejected_by (Optional[str]): User who rejected the operation.
-        title (Optional[str]): Title/subject of the approval request.
-        description (Optional[str]): Detailed description of the approval.
-        content (str): Content or action that requires approval.
+        wait_node_id (str): The id of the node that triggered this HITL.
     """
     
-    operation_id: str
-    approved_by: Optional[str] = None
-    rejected_by: Optional[str] = None
-    title: Optional[str] = None
-    description: Optional[str] = None
-    content: str
-
+    wait_node_id: str
 
 class AgentExecutionInput(BaseModel):
     """

xpander_sdk/modules/tasks/sub_modules/task.py

@@ -43,7 +43,7 @@ from httpx import HTTPStatusError
 import httpx
 import json
 from httpx_sse import aconnect_sse
-from pydantic import Field, computed_field
+from pydantic import Field
 
 from xpander_sdk.consts.api_routes import APIRoute
 from xpander_sdk.core.xpander_api_client import APIClient
@@ -67,7 +67,7 @@ from xpander_sdk.modules.events.utils.generic import get_events_base, get_events
 from xpander_sdk.modules.tasks.models.task import (
     AgentExecutionInput,
     AgentExecutionStatus,
-    HumanInTheLoop,
+    HumanInTheLoopRequest,
     ExecutionMetricsReport,
     PendingECARequest,
     TaskReportRequest,
@@ -82,12 +82,13 @@ from xpander_sdk.modules.tools_repository.models.mcp import (
     MCPServerDetails,
 )
 from xpander_sdk.utils.event_loop import run_sync
+from xpander_sdk.models.compactization import TaskCompactizationEvent
 
 # Type variable for Task class methods
 T = TypeVar("T", bound="Task")
 
 TaskUpdateEventData = Union[
-    T, ToolCallRequest, ToolCallResult, MCPOAuthGetTokenResponse, DeepPlanning
+    TaskCompactizationEvent, T, ToolCallRequest, ToolCallResult, MCPOAuthGetTokenResponse, DeepPlanning
 ]
 
 
@@ -96,7 +97,7 @@ class TaskUpdateEvent(XPanderSharedModel):
     task_id: str
     organization_id: str
     time: datetime
-    data: TaskUpdateEventData
+    data: Any
 
 
 class Task(XPanderSharedModel):
@@ -124,7 +125,7 @@ class Task(XPanderSharedModel):
         sub_executions (Optional[List[str]]): List of sub-execution IDs.
         is_manually_stopped (Optional[bool]): Flag indicating if the task was manually stopped.
         payload_extension (Optional[dict]): Additional data for the task.
-        hitl_request (Optional[HumanInTheLoop]): Human-in-the-loop request state.
+        hitl_request (Optional[HumanInTheLoopRequest]): Human-in-the-loop request state.
         pending_eca_request (Optional[PendingECARequest]): Pending ECA request, if any.
         source (Optional[str]): Source information of the task.
         output_format (Optional[OutputFormat]): Desired output format of the task.
@@ -136,6 +137,7 @@ class Task(XPanderSharedModel):
         triggering_agent_id (Optional[str]): Optional triggering agent id.
         title (Optional[str]): Optional task title.
         deep_planning: Optional[DeepPlanning] = Field(default_factory=DeepPlanning)
+        execution_attempts: Optional[int] = 1
 
     Example:
         >>> task = Task.load(task_id="task_123")
@@ -170,12 +172,13 @@ class Task(XPanderSharedModel):
     sub_executions: Optional[List[str]] = []
     is_manually_stopped: Optional[bool] = False
     payload_extension: Optional[dict] = None
-    hitl_request: Optional[HumanInTheLoop] = None
+    hitl_request: Optional[HumanInTheLoopRequest] = None
     pending_eca_request: Optional[PendingECARequest] = None
     source: Optional[str] = None
     output_format: Optional[OutputFormat] = None
     output_schema: Optional[Dict] = None
     events_streaming: Optional[bool] = False
+    is_orchestration: Optional[bool] = False
     additional_context: Optional[str] = None
     expected_output: Optional[str] = (None,)
     mcp_servers: Optional[List[MCPServerDetails]] = ([],)
@@ -184,6 +187,7 @@ class Task(XPanderSharedModel):
     think_mode: Optional[ThinkMode] = ThinkMode.Default
     disable_attachment_injection: Optional[bool] = False
     deep_planning: Optional[DeepPlanning] = Field(default_factory=DeepPlanning)
+    execution_attempts: Optional[int] = 1
 
     # metrics
     tokens: Optional[Tokens] = None
@@ -337,10 +341,13 @@ class Task(XPanderSharedModel):
         """
         return run_sync(self.aset_status(status=status, result=result))
 
-    async def asave(self):
+    async def asave(self, with_deep_plan_update: Optional[bool] = False):
         """
         Asynchronously saves the current task state to the backend.
 
+        Args:
+            with_deep_plan_update (Optional[bool]): should update deep plan as well? default false.
+        
         Raises:
             ModuleException: Error related to HTTP requests or task saving.
 
@@ -349,10 +356,15 @@ class Task(XPanderSharedModel):
         """
         client = APIClient(configuration=self.configuration)
         try:
+            exclude = {"configuration"}
+            
+            if not with_deep_plan_update:
+                exclude.add("deep_planning")
+                
             response = await client.make_request(
                 path=APIRoute.UpdateTask.format(task_id=self.id),
                 method="PATCH",
-                payload=self.model_dump_safe(exclude={"configuration","deep_planning"}),
+                payload=self.model_dump_safe(exclude=exclude),
             )
             updated_task = Task(**response, configuration=self.configuration)
             for field, value in updated_task.__dict__.items():
@@ -362,16 +374,19 @@ class Task(XPanderSharedModel):
         except Exception as e:
             raise ModuleException(500, f"Failed to save task: {str(e)}")
 
-    def save(self):
+    def save(self, with_deep_plan_update: Optional[bool] = False):
         """
         Saves the current task state synchronously.
 
         This function wraps the asynchronous asave method.
+        
+        Args:
+            with_deep_plan_update (Optional[bool]): should update deep plan as well? default false.
 
         Example:
             >>> task.save()
         """
-        return run_sync(self.asave())
+        return run_sync(self.asave(with_deep_plan_update=with_deep_plan_update))
 
     async def astop(self):
         """
@@ -540,16 +555,28 @@ class Task(XPanderSharedModel):
             for f in readable_files:
                 message += f"\n{json.dumps(f)}"
 
-        if self.deep_planning and self.deep_planning.enabled == True:
+        if self.deep_planning and self.deep_planning.enabled == True and self.deep_planning.started:
+            task_backup = self.model_copy() # backup result and status
+            
             self.reload()
-            uncompleted_tasks = [task for task in self.deep_planning.tasks if not task.completed]
-            if len(uncompleted_tasks) != 0:
-                message = "\n".join([
-                    "Task not finished, uncompleted tasks detected:",
-                    f"Uncompleted tasks: {[task.model_dump_json() for task in uncompleted_tasks]}",
-                    "You must complete tasks if fulfilled",
-                    f"User's original request: \"{message}\""
-                ])
+            
+            # restore result and status
+            self.result = task_backup.result
+            self.status = task_backup.status
+            self.tokens = task_backup.tokens
+            
+            if not self.deep_planning.question_raised:
+                uncompleted_tasks = [task for task in self.deep_planning.tasks if not task.completed]
+                if len(uncompleted_tasks) != 0: # make a retry with compactization
+                    from xpander_sdk.utils.agents.compactization_agent import run_task_compactization
+                    compactization_result = run_task_compactization(message=message, task=self, uncompleted_tasks=uncompleted_tasks)
+                    if isinstance(compactization_result, str):
+                        message = compactization_result
+                    else:
+                        message = f"<user_input>{compactization_result.new_task_prompt}</user_input><task_context>{compactization_result.task_context}</task_context>"
+            else:
+                self.deep_planning.question_raised = False # reset question raised indicator
+                self.save(with_deep_plan_update=True)
         
         return message
 
@@ -785,8 +812,15 @@ class Task(XPanderSharedModel):
             ...     print(f"Remaining tasks: {len(status.uncompleted_tasks)}")
         """
         try:
+            task_backup = self.model_copy() # backup result and status
+            await self.areload() # reload
+            
+            # restore result and status
+            self.result = task_backup.result
+            self.status = task_backup.status
+            self.tokens = task_backup.tokens
+            
             if self.deep_planning and self.deep_planning.enabled and self.deep_planning.started and self.deep_planning.enforce:
-                await self.areload()
                 
                 # allow early exit to ask question
                 if self.deep_planning.question_raised: