Upload 2 files

aworld/runners/callback/decorator.py

@@ -0,0 +1,170 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.
+
+import functools
+import inspect
+import logging
+from typing import Dict, Any, Callable, Optional, Union, Awaitable
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger("callback_registry")
+
+
+class CallbackRegistry:
+    """Callback function registry, used to manage and execute callback functions"""
+    
+    # Registry for storing decorated callback functions
+    _registry: Dict[str, Callable] = {}
+    
+    @classmethod
+    def register(cls, key_name: str, func: Callable) -> Callable:
+        """Register callback function to the registry
+        
+        Args:
+            key_name: Unique identifier for the callback function
+            func: Callback function to register
+            
+        Returns:
+            Registered callback function
+        """
+        # Check if a callback function with the same key_name already exists
+        if key_name in cls._registry:
+            existing_func = cls._registry[key_name]
+            logger.warning(
+                f"Callback function '{key_name}' already exists and will be overwritten! "
+                f"Original function: {existing_func.__name__ if hasattr(existing_func, '__name__') else str(existing_func)}, "
+                f"New function: {func.__name__ if hasattr(func, '__name__') else str(func)}"
+            )
+        
+        cls._registry[key_name] = func
+        return func
+    
+    @classmethod
+    def get(cls, key_name: str) -> Optional[Callable]:
+        """Get registered callback function by key_name
+        
+        Args:
+            key_name: Unique identifier for the callback function
+            
+        Returns:
+            Registered callback function, or None if not found
+        """
+        return cls._registry.get(key_name)
+    
+    @classmethod
+    async def execute(
+        cls,
+        key_name: str, 
+        tool: Any, 
+        args: Dict[str, Any], 
+        tool_context: Any, 
+        tool_response: Optional[Dict[str, Any]] = None
+    ) -> Optional[Dict[str, Any]]:
+        """Execute registered callback function
+        
+        Args:
+            key_name: Unique identifier for the callback function
+            tool: Tool object
+            args: Tool arguments
+            tool_context: Tool context
+            tool_response: Tool response (for post-callbacks)
+            
+        Returns:
+            Return value of the callback function, or None if the callback function doesn't exist
+        """
+        callback = cls.get(key_name)
+        if not callback:
+            return None
+        
+        # Determine parameters based on callback type
+        if tool_response is not None:
+            # Post-callback
+            result = callback(tool, args, tool_context, tool_response)
+        else:
+            # Pre-callback
+            result = callback(tool, args, tool_context)
+        
+        # Handle asynchronous callbacks
+        if inspect.isawaitable(result):
+            result = await result
+            
+        return result
+    
+    @classmethod
+    def list(cls) -> Dict[str, str]:
+        """List all registered callback functions
+        
+        Returns:
+            Dictionary containing callback function names and descriptions
+        """
+        return {
+            key: func.__name__ if hasattr(func, '__name__') else str(func)
+            for key, func in cls._registry.items()
+        }
+
+
+def reg_callback(key_name: str):
+    """Decorator for registering callback functions
+    
+    Args:
+        key_name: Unique identifier for the callback function
+        
+    Returns:
+        Decorator function
+    """
+    def decorator(func):
+        # Register function to the global registry
+        CallbackRegistry.register(key_name, func)
+        
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+        
+        return wrapper
+    
+    return decorator
+
+
+# For backward compatibility, keep these functions
+def get_callback(key_name: str) -> Optional[Callable]:
+    """Get registered callback function by key_name
+    
+    Args:
+        key_name: Unique identifier for the callback function
+        
+    Returns:
+        Registered callback function, or None if not found
+    """
+    return CallbackRegistry.get(key_name)
+
+
+async def execute_callback(
+    key_name: str, 
+    tool: Any, 
+    args: Dict[str, Any], 
+    tool_context: Any, 
+    tool_response: Optional[Dict[str, Any]] = None
+) -> Optional[Dict[str, Any]]:
+    """Execute registered callback function
+    
+    Args:
+        key_name: Unique identifier for the callback function
+        tool: Tool object
+        args: Tool arguments
+        tool_context: Tool context
+        tool_response: Tool response (for post-callbacks)
+        
+    Returns:
+        Return value of the callback function, or None if the callback function doesn't exist
+    """
+    return await CallbackRegistry.execute(key_name, tool, args, tool_context, tool_response)
+
+
+def list_callbacks() -> Dict[str, str]:
+    """List all registered callback functions
+    
+    Returns:
+        Dictionary containing callback function names and descriptions
+    """
+    return CallbackRegistry.list() 

aworld/runners/callback/tool.py

@@ -0,0 +1,66 @@
+import json
+from typing import AsyncGenerator
+
+from aworld.core.task import TaskResponse
+from aworld.models.model_response import ModelResponse
+from aworld.output import Output, MessageOutput
+from aworld.runners.callback.decorator import CallbackRegistry
+from aworld.runners.handler.base import DefaultHandler
+from aworld.core.common import TaskItem, Observation
+from aworld.core.context.base import Context
+from aworld.core.event.base import Message, Constants, TopicType
+from aworld.logs.util import logger
+
+class ToolCallbackHandler(DefaultHandler):
+    def __init__(self, runner):
+        self.runner = runner
+
+    async def handle(self, message):
+        if message.category != Constants.TOOL_CALLBACK:
+            return
+        logger.info(f"-------ToolCallbackHandler start handle message----: {message}")
+        outputs = self.runner.task.outputs
+        output = None
+        try:
+            payload = message.payload
+            if not payload or not payload[0]:
+                return
+            observation=payload[0]
+            if not isinstance(observation, Observation):
+                return
+            if not observation.action_result:
+                return
+            for res in observation.action_result:
+                if not res or not res.content or not res.tool_name or not res.action_name:
+                    continue
+                callback_func = CallbackRegistry.get(res.tool_name + "__" + res.action_name)
+                if not callback_func:
+                    continue
+                callback_func(res)
+                logger.info(f"-------ToolCallbackHandler callback_func-res: {res}")
+            logger.info(f"-------ToolCallbackHandler end  handle message: {observation}")
+        except Exception as e:
+            # todo
+            logger.warning(f"ToolCallbackHandler Failed to parse payload: {e}")
+            yield Message(
+                category=Constants.TASK,
+                payload=TaskItem(msg="Failed to parse output.", data=payload, stop=True),
+                sender=self.name(),
+                session_id=Context.instance().session_id,
+                topic=TopicType.ERROR
+            )
+        finally:
+            #todo
+            if output:
+                if not output.metadata:
+                    output.metadata = {}
+                output.metadata['sender'] = message.sender
+                output.metadata['receiver'] = message.receiver
+                await outputs.add_output(output)
+            # 1\Update the current message node status
+            # 2\Update the incoming message node status
+
+
+        return
+
+