Handle Tuple return type

Author: jnkrmg

application_sdk/decorators/automation_activity.py

@@ -4,7 +4,7 @@
 
 import inspect
 import json
-from typing import Any, Callable, Dict, List, Union, get_args, get_origin
+from typing import Any, Callable, Dict, List, Optional, Union, get_args, get_origin
 
 from pydantic import BaseModel
 import requests
@@ -14,107 +14,180 @@
 
 logger = get_logger(__name__)
 
-
 # Global registry to collect decorated activities
 ACTIVITY_SPECS: List[dict[str, Any]] = []
 
 
 def _type_to_json_schema(annotation: Any) -> Dict[str, Any]:
     """Convert Python type annotation to JSON Schema using reflection."""
+    # Pydantic models
     if inspect.isclass(annotation) and issubclass(annotation, BaseModel):
         schema = annotation.model_json_schema(mode="serialization")
         if "$defs" in schema:
-            defs = schema.pop("$defs")
-            schema = {**schema, **defs}
+            schema = {**schema, **schema.pop("$defs")}
         return schema
 
     origin = get_origin(annotation)
-    if origin is Union:
-        args = get_args(annotation)
-        if type(None) in args:
-            non_none_types = [arg for arg in args if arg is not type(None)]
-            if non_none_types:
-                schema = _type_to_json_schema(non_none_types[0])
-                schema["nullable"] = True
-                return schema
+    args = get_args(annotation)
+
+    # Optional types (Union with None)
+    if origin is Union and type(None) in args:
+        non_none = [a for a in args if a is not type(None)]
+        if non_none:
+            schema = _type_to_json_schema(non_none[0])
+            schema["nullable"] = True
+            return schema
 
+    # Dict types
     if origin is dict or annotation is dict:
-        args = get_args(annotation)
-        if args:
-            return {
-                "type": "object",
-                "additionalProperties": _type_to_json_schema(args[1]),
-            }
-        return {"type": "object"}
+        return {
+            "type": "object",
+            "additionalProperties": _type_to_json_schema(args[1]) if args else {},
+        }
 
+    # List types
     if origin is list or annotation is list:
-        args = get_args(annotation)
+        return {"type": "array", "items": _type_to_json_schema(args[0]) if args else {}}
+
+    # Tuple types
+    if origin is tuple or annotation is tuple:
         if args:
-            return {"type": "array", "items": _type_to_json_schema(args[0])}
+            item_schemas = [_type_to_json_schema(arg) for arg in args]
+            return {
+                "type": "array",
+                "items": item_schemas,
+                "minItems": len(item_schemas),
+                "maxItems": len(item_schemas),
+            }
         return {"type": "array"}
 
-    type_mapping = {
-        str: {"type": "string"},
-        int: {"type": "integer"},
-        float: {"type": "number"},
-        bool: {"type": "boolean"},
-        Any: {},
-    }
-
-    if annotation in type_mapping:
-        return type_mapping[annotation]
+    # Primitive types
+    type_map = {str: "string", int: "integer", float: "number", bool: "boolean"}
+    if annotation in type_map:
+        return {"type": type_map[annotation]}
 
     return {}
 
 
+def _create_object_schema(
+    properties: Dict[str, Any], required: Optional[List[str]] = None
+) -> Dict[str, Any]:
+    """Create a JSON Schema object with properties and required fields."""
+    schema = {"type": "object", "properties": properties}
+    if required:
+        schema["required"] = required
+    return schema
+
+
 def _generate_input_schema(func: Any) -> Dict[str, Any]:
     """Generate JSON Schema for function inputs using reflection."""
-    sig = inspect.signature(func)
-    properties = {}
-    required: list[str] = []
-
-    for param_name, param in sig.parameters.items():
-        if param_name == "self":
+    properties: Dict[str, Any] = {}
+    required: List[str] = []
+    for name, param in inspect.signature(func).parameters.items():
+        if name == "self":
             continue
-
-        param_schema = (
+        schema = (
             _type_to_json_schema(param.annotation)
             if param.annotation != inspect.Parameter.empty
             else {}
         )
-
         if param.default != inspect.Parameter.empty:
             if isinstance(param.default, (str, int, float, bool, type(None))):
-                param_schema["default"] = param.default
+                schema["default"] = param.default
         else:
-            required.append(param_name)
-
-        properties[param_name] = param_schema
+            required.append(name)
+        properties[name] = schema
+    return _create_object_schema(properties, required if required else None)
 
-    schema = {"type": "object", "properties": properties}
-    if required:
-        schema["required"] = required
 
-    return schema
+def _validate_output_field_names(
+    func: Any, output_field_names: List[str], return_annotation: Any
+) -> None:
+    """Validate output_field_names against function return type."""
+    origin = get_origin(return_annotation)
+    is_tuple: bool = origin is tuple or return_annotation is tuple
+    if not is_tuple:
+        raise ValueError(
+            f"output_field_names provided for '{func.__name__}', "
+            f"but return type is not a tuple (got {return_annotation})."
+        )
+    args = get_args(return_annotation)
+    if not args:
+        raise ValueError(
+            f"output_field_names provided for '{func.__name__}', but tuple has no type arguments."
+        )
+    if len(output_field_names) != len(args):
+        raise ValueError(
+            f"output_field_names length ({len(output_field_names)}) doesn't match "
+            f"tuple length ({len(args)}) for '{func.__name__}'. Expected: {output_field_names}"
+        )
 
 
-def _generate_output_schema(func: Any) -> Dict[str, Any]:
+def _generate_output_schema(
+    func: Any, output_field_names: Optional[List[str]] = None
+) -> Dict[str, Any]:
     """Generate JSON Schema for function outputs using reflection."""
     return_annotation = inspect.signature(func).return_annotation
-    if return_annotation == inspect.Signature.empty:
+    has_return = return_annotation != inspect.Signature.empty
+
+    if output_field_names:
+        if not has_return:
+            raise ValueError(
+                f"output_field_names provided for '{func.__name__}', but function has no return annotation."
+            )
+        _validate_output_field_names(func, output_field_names, return_annotation)
+
+    if not has_return:
         return {}
-    return _type_to_json_schema(return_annotation)
+
+    schema: dict[str, Any] = _type_to_json_schema(return_annotation)
+
+    # Convert tuple to object if field names provided
+    if output_field_names and schema.get("type") == "array" and "items" in schema:
+        items = schema.get("items", [])
+        if isinstance(items, list) and len(items) == len(output_field_names):
+            properties = {name: item for name, item in zip(output_field_names, items)}
+            return _create_object_schema(properties, output_field_names)
+
+    return schema
 
 
 def automation_activity(
     name: str,
     description: str,
+    output_field_names: Optional[List[str]] = None,
 ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
-    """Decorator to mark an activity for automatic registration."""
+    """Decorator to mark an activity for automatic registration.
+
+    Args:
+        name: Name of the activity/tool.
+        description: Description of what the activity does.
+        output_field_names: Optional list of field names for tuple return types.
+            If provided, tuple outputs will be converted to objects with named properties.
+            Length must match the number of elements in the tuple return type.
+
+    Returns:
+        Decorator function.
+
+    Raises:
+        ValueError: If output_field_names is provided but function has no return annotation,
+            or if return type is not a tuple, or if lengths don't match.
+
+    Example:
+        >>> @automation_activity(
+        ...     name="fetch_entities",
+        ...     description="Fetch entities by DSL",
+        ...     output_field_names=["entities_path", "total_count", "chunk_count"]
+        ... )
+        ... def fetch_entities(self, dsl_query: dict) -> tuple[str, int, int]:
+        ...     return ("path/to/entities", 1000, 10)
+    """
 
     def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
         input_schema: dict[str, Any] = _generate_input_schema(func)
-        output_schema: dict[str, Any] = _generate_output_schema(func)
+        output_schema: dict[str, Any] = _generate_output_schema(
+            func, output_field_names
+        )
 
         logger.info(f"Collected automation activity: {name}")
         ACTIVITY_SPECS.append(
@@ -161,9 +234,7 @@ def flush_activity_registrations(
         )
         return
 
-    logger.info(
-        f"Registering {len(ACTIVITY_SPECS)} activities with automation engine"
-    )
+    logger.info(f"Registering {len(ACTIVITY_SPECS)} activities with automation engine")
 
     # Generate app qualified name
     app_qualified_name: str = f"default/apps/{app_name}"