agentfield 0.1.22rc2__py3-none-any.whl

agentfield/decorators.py

@@ -0,0 +1,527 @@
+"""
+Enhanced decorators for AgentField SDK with automatic workflow tracking.
+Provides always-on workflow tracking for reasoner calls.
+"""
+
+import asyncio
+import functools
+import inspect
+import time
+from typing import Any, Callable, Dict, List, Optional, Union
+
+from agentfield.logger import log_warn
+
+from .execution_context import (
+    ExecutionContext,
+    get_current_context,
+    set_execution_context,
+    reset_execution_context,
+)
+from .agent_registry import get_current_agent_instance
+from .types import ReasonerDefinition
+from .pydantic_utils import convert_function_args, should_convert_args
+from pydantic import ValidationError
+
+
+def reasoner(
+    func=None,
+    *,
+    path: Optional[str] = None,
+    tags: Optional[List[str]] = None,
+    description: Optional[str] = None,
+    track_workflow: bool = True,
+    **kwargs,
+):
+    """
+    Enhanced reasoner decorator with automatic workflow tracking and full feature support.
+
+    Supports both:
+    @reasoner                           # Default: track_workflow=True
+    @reasoner(track_workflow=False)     # Explicit: disable tracking
+    @reasoner(path="/custom/path")      # Custom endpoint path
+    @reasoner(tags=["ai", "nlp"])       # Tags for organization
+    @reasoner(description="...")        # Custom description
+
+    Args:
+        func: The function to decorate (when used without parentheses)
+        path: Custom API endpoint path for this reasoner
+        tags: List of tags for organizing and categorizing reasoners
+        description: Description of what this reasoner does
+        track_workflow: Whether to enable automatic workflow tracking (default: True)
+        **kwargs: Additional metadata to store with the reasoner
+
+    Returns:
+        Decorated function with workflow tracking capabilities and full metadata support
+    """
+
+    def decorator(f: Callable) -> Callable:
+        @functools.wraps(f)
+        async def wrapper(*args, **kwargs):
+            if track_workflow:
+                # Execute with automatic workflow tracking
+                return await _execute_with_tracking(f, *args, **kwargs)
+            else:
+                # Execute without tracking
+                if asyncio.iscoroutinefunction(f):
+                    return await f(*args, **kwargs)
+                else:
+                    return f(*args, **kwargs)
+
+        # Store comprehensive metadata on the function
+        wrapper._is_reasoner = True
+        wrapper._track_workflow = track_workflow
+        wrapper._reasoner_name = f.__name__
+        wrapper._original_func = f
+        wrapper._reasoner_path = path
+        wrapper._reasoner_tags = tags or []
+        wrapper._reasoner_description = (
+            description or f.__doc__ or f"Reasoner: {f.__name__}"
+        )
+
+        # Store any additional metadata
+        for key, value in kwargs.items():
+            setattr(wrapper, f"_reasoner_{key}", value)
+
+        return wrapper
+
+    # Handle both @reasoner and @reasoner(...) syntax
+    if func is None:
+        # Called as @reasoner(track_workflow=False) or @reasoner(path="/custom")
+        return decorator
+    else:
+        # Called as @reasoner (no parentheses)
+        return decorator(func)
+
+
+async def _execute_with_tracking(func: Callable, *args, **kwargs) -> Any:
+    """
+    Core function that handles automatic workflow tracking for reasoner calls.
+
+    Args:
+        func: The reasoner function to execute
+        *args: Positional arguments for the function
+        **kwargs: Keyword arguments for the function
+
+    Returns:
+        The result of the function execution
+    """
+    # Get current execution context
+    current_context = get_current_context()
+
+    # Get agent instance (from context or global registry)
+    agent_instance = get_current_agent_instance()
+
+    if not agent_instance:
+        # No agent context - execute without tracking
+        if asyncio.iscoroutinefunction(func):
+            return await func(*args, **kwargs)
+        else:
+            return func(*args, **kwargs)
+
+    workflow_handler = getattr(agent_instance, "workflow_handler", None)
+    reasoner_name = getattr(func, "__name__", "reasoner")
+
+    # Generate execution metadata
+    # Build a child context when executing under an existing workflow; otherwise create a root context
+    if current_context:
+        execution_context = current_context.create_child_context()
+        execution_context.reasoner_name = reasoner_name
+        parent_context = current_context
+    else:
+        workflow_name = reasoner_name
+        if hasattr(agent_instance, "node_id"):
+            workflow_name = f"{agent_instance.node_id}_{workflow_name}"
+        execution_context = ExecutionContext.new_root(
+            agent_node_id=getattr(agent_instance, "node_id", "agent"),
+            reasoner_name=workflow_name,
+        )
+        execution_context.reasoner_name = reasoner_name
+        execution_context.agent_instance = agent_instance
+        parent_context = None
+
+    # Align run/session metadata with the parent context so registration inherits the workflow run
+    if parent_context:
+        execution_context.run_id = parent_context.run_id
+        execution_context.session_id = parent_context.session_id
+        execution_context.caller_did = parent_context.caller_did
+        execution_context.target_did = parent_context.target_did
+        execution_context.agent_node_did = parent_context.agent_node_did
+    execution_context.agent_instance = agent_instance
+
+    if workflow_handler is not None:
+        execution_context = await workflow_handler._ensure_execution_registered(
+            execution_context, reasoner_name, parent_context
+        )
+
+    previous_agent_context = getattr(agent_instance, "_current_execution_context", None)
+    agent_instance._current_execution_context = execution_context
+
+    client = getattr(agent_instance, "client", None)
+    previous_client_context = None
+    if client is not None:
+        previous_client_context = getattr(client, "_current_workflow_context", None)
+        client._current_workflow_context = execution_context
+
+    token = None
+    start_time = time.time()
+    parent_execution_id = parent_context.execution_id if parent_context else None
+
+    sig = inspect.signature(func)
+    call_kwargs = dict(kwargs or {})
+    input_data: Dict[str, Any] = {}
+
+    # Prepare DID-aware execution context so VC generation works for decorator-driven calls
+    did_execution_context = None
+    agent_has_did = getattr(agent_instance, "did_enabled", False) and getattr(
+        agent_instance, "did_manager", None
+    )
+    if agent_has_did:
+        try:
+            session_id = execution_context.session_id or execution_context.workflow_id
+            did_execution_context = agent_instance.did_manager.create_execution_context(
+                execution_context.execution_id,
+                execution_context.workflow_id,
+                session_id,
+                "agent",
+                reasoner_name,
+            )
+            if did_execution_context and hasattr(
+                agent_instance, "_populate_execution_context_with_did"
+            ):
+                agent_instance._populate_execution_context_with_did(
+                    execution_context, did_execution_context
+                )
+        except Exception as exc:  # pragma: no cover - diagnostic only
+            if getattr(agent_instance, "dev_mode", False):
+                log_warn(f"Failed to build DID context for {reasoner_name}: {exc}")
+            did_execution_context = None
+
+    def _maybe_generate_vc(
+        status: str, result_payload: Any, duration_ms: int, error_message: Optional[str]
+    ) -> None:
+        """Fire-and-forget VC generation so decorator parity matches HTTP path."""
+        generate_vc = getattr(agent_instance, "_generate_vc_async", None)
+        vc_generator = getattr(agent_instance, "vc_generator", None)
+        if (
+            did_execution_context
+            and callable(generate_vc)
+            and hasattr(agent_instance, "_should_generate_vc")
+            and agent_instance._should_generate_vc(
+                reasoner_name, getattr(agent_instance, "_reasoner_vc_overrides", {})
+            )
+        ):
+            asyncio.create_task(
+                generate_vc(
+                    vc_generator,
+                    did_execution_context,
+                    reasoner_name,
+                    input_data,
+                    result_payload,
+                    status=status,
+                    error_message=error_message,
+                    duration_ms=duration_ms,
+                )
+            )
+
+    try:
+        # Execute function with new context
+        token = set_execution_context(execution_context)
+
+        # Inject execution_context if the function accepts it
+        if "execution_context" in sig.parameters:
+            call_kwargs.setdefault("execution_context", execution_context)
+
+        # 🔥 NEW: Automatic Pydantic model conversion (FastAPI-like behavior)
+        try:
+            if should_convert_args(func):
+                converted_args, converted_kwargs = convert_function_args(
+                    func, args, call_kwargs
+                )
+                args = converted_args
+                call_kwargs = converted_kwargs
+        except ValidationError as e:
+            # Re-raise validation errors with context
+            raise ValidationError(
+                f"Pydantic validation failed for reasoner '{func.__name__}': {e}",
+                model=getattr(e, "model", None),
+            ) from e
+        except Exception as e:
+            # Log conversion errors but continue with original args for backward compatibility
+            if hasattr(agent_instance, "dev_mode") and agent_instance.dev_mode:
+                log_warn(f"Failed to convert arguments for {func.__name__}: {e}")
+
+        input_data = _build_input_payload(sig, args, call_kwargs)
+
+        start_payload = {
+            "reasoner_name": reasoner_name,
+            "args": list(args),
+            "kwargs": dict(call_kwargs),
+            "input_data": input_data,
+            "parent_execution_id": parent_execution_id,
+        }
+        await asyncio.create_task(
+            _send_workflow_start(
+                agent_instance,
+                execution_context,
+                start_payload,
+            )
+        )
+
+        if asyncio.iscoroutinefunction(func):
+            result = await func(*args, **call_kwargs)
+        else:
+            result = func(*args, **call_kwargs)
+
+        duration_ms = int((time.time() - start_time) * 1000)
+        completion_payload = {
+            "input_data": input_data,
+            "parent_execution_id": parent_execution_id,
+        }
+        await asyncio.create_task(
+            _send_workflow_completion(
+                agent_instance,
+                execution_context,
+                result,
+                duration_ms,
+                completion_payload,
+            )
+        )
+        _maybe_generate_vc("success", result, duration_ms, None)
+        return result
+    except Exception as exc:
+        duration_ms = int((time.time() - start_time) * 1000)
+        error_payload = {
+            "input_data": input_data,
+            "parent_execution_id": parent_execution_id,
+        }
+        _maybe_generate_vc("error", None, duration_ms, str(exc))
+        await asyncio.create_task(
+            _send_workflow_error(
+                agent_instance,
+                execution_context,
+                str(exc),
+                duration_ms,
+                error_payload,
+            )
+        )
+        raise
+
+    finally:
+        if token is not None:
+            reset_execution_context(token)
+        agent_instance._current_execution_context = previous_agent_context
+        if client is not None:
+            client._current_workflow_context = previous_client_context
+
+
+def _build_input_payload(
+    signature: inspect.Signature, args: tuple, kwargs: Dict[str, Any]
+) -> Dict[str, Any]:
+    if not signature.parameters:
+        return dict(kwargs)
+
+    try:
+        bound = signature.bind_partial(*args, **kwargs)
+        bound.apply_defaults()
+    except Exception:
+        payload = {f"arg_{idx}": value for idx, value in enumerate(args)}
+        payload.update(kwargs)
+        return payload
+
+    payload = {}
+    for name, value in bound.arguments.items():
+        if name == "self":
+            continue
+        payload[name] = value
+    return payload
+
+
+def _compose_event_payload(
+    agent,
+    context: ExecutionContext,
+    reasoner_name: str,
+    status: str,
+    parent_execution_id: Optional[str],
+    input_data: Optional[Dict[str, Any]] = None,
+) -> Dict[str, Any]:
+    event: Dict[str, Any] = {
+        "execution_id": context.execution_id,
+        "workflow_id": context.workflow_id,
+        "run_id": context.run_id,
+        "reasoner_id": reasoner_name,
+        "agent_node_id": getattr(agent, "node_id", None),
+        "status": status,
+        "type": reasoner_name,
+        "parent_execution_id": parent_execution_id,
+        "parent_workflow_id": context.parent_workflow_id,
+    }
+    if input_data is not None:
+        event["input_data"] = input_data
+    return event
+
+
+def on_change(pattern: Union[str, List[str]]):
+    """
+    Decorator to mark a function as a memory event listener.
+
+    Args:
+        pattern: Memory pattern(s) to listen for changes
+
+    Returns:
+        Decorated function with memory event listener metadata
+    """
+
+    def decorator(func: Callable) -> Callable:
+        @functools.wraps(func)
+        async def wrapper(*args, **kwargs):
+            return await func(*args, **kwargs)
+
+        # Attach metadata to the function
+        wrapper._memory_event_listener = True
+        wrapper._memory_event_patterns = (
+            pattern if isinstance(pattern, list) else [pattern]
+        )
+        return wrapper
+
+    return decorator
+
+
+# Legacy support for old reasoner decorator signature
+async def _send_workflow_start(
+    agent, context: ExecutionContext, payload: Dict[str, Any]
+) -> None:
+    handler = getattr(agent, "workflow_handler", None)
+    if handler is None:
+        return
+    try:
+        reasoner_name = payload.get("reasoner_name", context.reasoner_name)
+        parent_execution_id = payload.get("parent_execution_id")
+        input_data = payload.get("input_data") or {}
+
+        if hasattr(handler, "notify_call_start"):
+            await handler.notify_call_start(
+                context.execution_id,
+                context,
+                reasoner_name,
+                input_data,
+                parent_execution_id=parent_execution_id,
+            )
+        elif hasattr(handler, "fire_and_forget_update"):
+            event_payload = _compose_event_payload(
+                agent,
+                context,
+                reasoner_name,
+                "running",
+                parent_execution_id,
+                input_data=input_data,
+            )
+            await handler.fire_and_forget_update(event_payload)
+    except Exception as exc:  # pragma: no cover - logging pathway
+        if getattr(agent, "dev_mode", False):
+            log_warn(f"Failed to emit workflow start: {exc}")
+
+
+async def _send_workflow_completion(
+    agent,
+    context: ExecutionContext,
+    result: Any,
+    duration_ms: int,
+    payload: Dict[str, Any],
+) -> None:
+    handler = getattr(agent, "workflow_handler", None)
+    if handler is None:
+        return
+    try:
+        parent_execution_id = payload.get("parent_execution_id")
+        input_data = payload.get("input_data")
+        reasoner_name = context.reasoner_name
+
+        if hasattr(handler, "notify_call_complete"):
+            await handler.notify_call_complete(
+                context.execution_id,
+                context.workflow_id,
+                result,
+                duration_ms,
+                context,
+                input_data=input_data,
+                parent_execution_id=parent_execution_id,
+            )
+        elif hasattr(handler, "fire_and_forget_update"):
+            event_payload = _compose_event_payload(
+                agent,
+                context,
+                reasoner_name,
+                "succeeded",
+                parent_execution_id,
+                input_data=input_data,
+            )
+            event_payload["result"] = result
+            event_payload["duration_ms"] = duration_ms
+            await handler.fire_and_forget_update(event_payload)
+    except Exception as exc:  # pragma: no cover - logging pathway
+        if getattr(agent, "dev_mode", False):
+            log_warn(f"Failed to emit workflow completion: {exc}")
+
+
+async def _send_workflow_error(
+    agent,
+    context: ExecutionContext,
+    message: str,
+    duration_ms: int,
+    payload: Dict[str, Any],
+) -> None:
+    handler = getattr(agent, "workflow_handler", None)
+    if handler is None:
+        return
+    try:
+        parent_execution_id = payload.get("parent_execution_id")
+        input_data = payload.get("input_data")
+        reasoner_name = context.reasoner_name
+
+        if hasattr(handler, "notify_call_error"):
+            await handler.notify_call_error(
+                context.execution_id,
+                context.workflow_id,
+                message,
+                duration_ms,
+                context,
+                input_data=input_data,
+                parent_execution_id=parent_execution_id,
+            )
+        elif hasattr(handler, "fire_and_forget_update"):
+            event_payload = _compose_event_payload(
+                agent,
+                context,
+                reasoner_name,
+                "failed",
+                parent_execution_id,
+                input_data=input_data,
+            )
+            event_payload["error"] = message
+            event_payload["duration_ms"] = duration_ms
+            await handler.fire_and_forget_update(event_payload)
+    except Exception as exc:  # pragma: no cover - logging pathway
+        if getattr(agent, "dev_mode", False):
+            log_warn(f"Failed to emit workflow error: {exc}")
+
+
+def legacy_reasoner(reasoner_id: str, input_schema: dict, output_schema: dict):
+    """
+    Legacy reasoner decorator for backward compatibility.
+
+    This is kept for compatibility with existing code that uses the old signature.
+    New code should use the enhanced @reasoner decorator.
+    """
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+
+        # Attach metadata to the function
+        wrapper._reasoner_def = ReasonerDefinition(
+            id=reasoner_id, input_schema=input_schema, output_schema=output_schema
+        )
+        return wrapper
+
+    return decorator