agenta 0.50.6__py3-none-any.whl → 0.51.0__py3-none-any.whl

agenta/__init__.py

@@ -21,13 +21,15 @@ from .sdk.types import (
     PromptTemplate,
 )
 
+from .sdk.agenta_init import Config, AgentaSingleton, init as _init
 from .sdk.utils.logging import get_module_logger
+from .sdk.utils.costs import calculate_token_usage
 from .sdk.tracing import Tracing, get_tracer
-from .sdk.decorators.tracing import instrument
 from .sdk.tracing.conventions import Reference
-from .sdk.decorators.routing import entrypoint, app, route
-from .sdk.agenta_init import Config, AgentaSingleton, init as _init
-from .sdk.utils.costs import calculate_token_usage
+from .sdk.decorators.tracing import instrument
+from .sdk.decorators.running import workflow, workflows
+from .sdk.decorators.serving import entrypoint, app, route
+from .sdk.context.running import workflow_mode_enabled
 from .sdk.litellm import litellm as callbacks
 from .sdk.managers.apps import AppManager
 from .sdk.managers.vault import VaultManager
@@ -52,6 +54,7 @@ tracer = get_tracer(tracing)
 
 def init(
     host: Optional[str] = None,
+    api_url: Optional[str] = None,
     api_key: Optional[str] = None,
     config_fname: Optional[str] = None,
     redact: Optional[Callable[..., Any]] = None,
@@ -62,7 +65,7 @@ def init(
     global api, async_api, tracing, tracer  # pylint: disable=global-statement
 
     _init(
-        host=host,
+        host=host or api_url,
         api_key=api_key,
         config_fname=config_fname,
         redact=redact,

agenta/sdk/__init__.py

@@ -22,8 +22,9 @@ from .types import (
 
 from .tracing import Tracing, get_tracer
 from .decorators.tracing import instrument
+from .decorators.running import workflow, workflows
+from .decorators.serving import entrypoint, app, route
 from .tracing.conventions import Reference
-from .decorators.routing import entrypoint, app, route
 from .agenta_init import Config, AgentaSingleton, init as _init
 from .utils.costs import calculate_token_usage
 from .managers.apps import AppManager

agenta/sdk/agenta_init.py

@@ -9,7 +9,7 @@ from agenta.sdk.utils.logging import get_module_logger
 from agenta.client.client import AgentaApi, AsyncAgentaApi
 
 from agenta.sdk.tracing import Tracing
-from agenta.sdk.context.routing import routing_context
+from agenta.sdk.context.serving import serving_context
 
 
 log = get_module_logger(__name__)
@@ -149,7 +149,7 @@ class Config:
         return self.default_parameters
 
     def __getattr__(self, key):
-        context = routing_context.get()
+        context = serving_context.get()
 
         parameters = context.parameters
 

agenta/sdk/context/running.py

@@ -0,0 +1,39 @@
+from typing import Dict, Any, Optional
+from contextvars import ContextVar
+from contextlib import contextmanager
+
+from pydantic import BaseModel
+
+workflow_mode_enabled_context: ContextVar[bool] = ContextVar(
+    "ag.workflow_context_enabled",
+    default=False,
+)
+
+
+@contextmanager
+def workflow_mode_enabled():
+    token = workflow_mode_enabled_context.set(True)
+    try:
+        yield
+    finally:
+        workflow_mode_enabled_context.reset(token)
+
+
+@contextmanager
+async def async_workflow_mode_enabled():
+    token = workflow_mode_enabled_context.set(True)
+    try:
+        yield
+    finally:
+        workflow_mode_enabled_context.reset(token)
+
+
+class WorkflowRegistryContext(BaseModel):
+    version: Optional[str] = None
+    handlers: Dict[str, Any] = {}
+
+
+workflow_registry_context: ContextVar[WorkflowRegistryContext] = ContextVar(
+    "ag.workflow_registry",
+    default=WorkflowRegistryContext(),
+)

agenta/sdk/context/{routing.py → serving.py}

@@ -12,16 +12,16 @@ class RoutingContext(BaseModel):
     mock: Optional[str] = None
 
 
-routing_context = ContextVar("routing_context", default=RoutingContext())
+serving_context = ContextVar("serving_context", default=RoutingContext())
 
 
 @contextmanager
-def routing_context_manager(
+def serving_context_manager(
     *,
     context: Optional[RoutingContext] = None,
 ):
-    token = routing_context.set(context)
+    token = serving_context.set(context)
     try:
         yield
     finally:
-        routing_context.reset(token)
+        serving_context.reset(token)

agenta/sdk/context/tracing.py

@@ -13,9 +13,13 @@ class TracingContext(BaseModel):
     parameters: Optional[Dict[str, Any]] = None
     references: Optional[Dict[str, Any]] = None
     link: Optional[Dict[str, Any]] = None
+    type: Optional[str] = None
 
 
-tracing_context = ContextVar("tracing_context", default=TracingContext())
+tracing_context = ContextVar(
+    "ag.tracing_context",
+    default=TracingContext(),
+)
 
 
 @contextmanager
@@ -28,3 +32,25 @@ def tracing_context_manager(
         yield
     finally:
         tracing_context.reset(token)
+
+
+class TracingExporterContext(BaseModel):
+    credentials: Optional[str] = None
+
+
+tracing_exporter_context = ContextVar(
+    "ag.tracing_exporter_context",
+    default=TracingExporterContext(),
+)
+
+
+@contextmanager
+def tracing_exporter_context_manager(
+    *,
+    context: Optional[TracingExporterContext] = None,
+):
+    token = tracing_exporter_context.set(context)
+    try:
+        yield
+    finally:
+        tracing_exporter_context.reset(token)

agenta/sdk/decorators/running.py

@@ -0,0 +1,134 @@
+from typing import Any, Callable, List, Awaitable, Dict, Optional
+import asyncio
+from functools import wraps
+from inspect import iscoroutinefunction
+from copy import deepcopy
+
+from decorator import decorator
+
+from agenta.sdk.context.running import (
+    workflow_mode_enabled_context,
+    workflow_registry_context,
+    WorkflowRegistryContext,
+)
+from agenta.sdk.workflows.types import (
+    WorkflowServiceRequest,
+    WorkflowServiceResponse,
+    WorkflowServiceInterface,
+    WorkflowRevision,
+    Schema,
+)
+from agenta.sdk.middleware.base import WorkflowMiddlewareDecorator
+
+from agenta.sdk.middleware.auth import AuthMiddleware
+from agenta.sdk.middleware.flags import FlagsMiddleware
+from agenta.sdk.middleware.adapt import AdaptMiddleware
+
+
+LATEST_VERSION = "2025.07.14"
+DEFAULT_SCHEMAS = lambda: {}  # pylint: disable=unnecessary-lambda-assignment
+
+
+class workflows:
+    @classmethod
+    def get_registry(cls) -> WorkflowRegistryContext:
+        return deepcopy(workflow_registry_context.get())
+
+
+class workflow:  # pylint: disable=invalid-name
+    def __init__(
+        self,
+        version: Optional[str] = None,
+        schemas: Optional[Dict[str, Schema]] = None,
+    ):
+        self.middleware: List[WorkflowMiddlewareDecorator] = [
+            AuthMiddleware,
+            FlagsMiddleware,
+            AdaptMiddleware,
+        ]
+
+        self.version = version or LATEST_VERSION
+        self.schemas = schemas or DEFAULT_SCHEMAS()
+
+    def __call__(self, func: Callable[..., Any]) -> Callable[..., Any]:
+        is_async = iscoroutinefunction(func)
+
+        workflow_registry = workflow_registry_context.get()
+
+        workflow_registry.version = workflow_registry.version or self.version
+
+        if is_async:
+
+            @wraps(func)
+            async def async_wrapper(*args, **kwargs):
+                if workflow_mode_enabled_context.get():
+                    return await self._wrapped_async(func)(*args, **kwargs)
+                return await func(*args, **kwargs)
+
+            workflow_registry.handlers = {
+                "invoke": async_wrapper,
+                "inspect": self.make_interface_wrapper(self.version, self.schemas),
+            }
+
+            return async_wrapper
+
+        else:
+
+            @wraps(func)
+            def sync_wrapper(*args, **kwargs):
+                if workflow_mode_enabled_context.get():
+                    return self._wrapped_async(func)(*args, **kwargs)
+                return func(*args, **kwargs)
+
+            return sync_wrapper
+
+    def _wrapped_async(
+        self,
+        func: Callable[..., Any],
+    ) -> Callable[..., Awaitable[WorkflowServiceResponse]]:
+        @decorator
+        async def async_wrapper(func, *args, **kwargs):
+            result = (
+                await func(*args, **kwargs)
+                if iscoroutinefunction(func)
+                else await asyncio.to_thread(func, *args, **kwargs)
+            )
+
+            return result
+
+        @wraps(func)
+        async def wrapper(*args, **kwargs):
+            handler = async_wrapper(func, *args, **kwargs)
+            request: WorkflowServiceRequest = (
+                kwargs.pop("request")
+                if "request" in kwargs
+                else args[0]
+                if len(args) > 0
+                else None
+            )
+            revision: WorkflowRevision = (
+                kwargs.pop("revision")
+                if "revision" in kwargs
+                else args[1]
+                if len(args) > 1
+                else None
+            )
+
+            _handler = handler
+            for middleware in reversed(self.middleware):
+                _handler = middleware(_handler)
+
+            return await _handler(
+                request=request,
+                revision=revision,
+            )
+
+        return wrapper
+
+    def make_interface_wrapper(self, path, schemas):
+        async def interface_wrapper() -> WorkflowServiceResponse:
+            return WorkflowServiceInterface(
+                schemas={path: schemas},
+            )
+
+        return interface_wrapper

agenta/sdk/decorators/{routing.py → serving.py}

@@ -14,11 +14,11 @@ from agenta.sdk.middleware.inline import InlineMiddleware
 from agenta.sdk.middleware.vault import VaultMiddleware
 from agenta.sdk.middleware.config import ConfigMiddleware
 from agenta.sdk.middleware.otel import OTelMiddleware
-from agenta.sdk.middleware.auth import AuthMiddleware
+from agenta.sdk.middleware.auth import AuthHTTPMiddleware
 from agenta.sdk.middleware.cors import CORSMiddleware
 
-from agenta.sdk.context.routing import (
-    routing_context_manager,
+from agenta.sdk.context.serving import (
+    serving_context_manager,
     RoutingContext,
 )
 from agenta.sdk.context.tracing import (
@@ -145,7 +145,7 @@ class entrypoint:
             app.add_middleware(InlineMiddleware)
             app.add_middleware(VaultMiddleware)
             app.add_middleware(ConfigMiddleware)
-            app.add_middleware(AuthMiddleware)
+            app.add_middleware(AuthHTTPMiddleware)
             app.add_middleware(OTelMiddleware)
             app.add_middleware(CORSMiddleware)
         ### ------------------ #
@@ -326,7 +326,7 @@ class entrypoint:
         inline = state.inline
         mock = state.mock
 
-        with routing_context_manager(
+        with serving_context_manager(
             context=RoutingContext(
                 parameters=parameters,
                 secrets=secrets,

agenta/sdk/decorators/tracing.py

@@ -169,8 +169,14 @@ class instrument:  # pylint: disable=invalid-name
 
             ag.tracing.credentials.put(trace_id, context.credentials)
 
+            trace_type = context.type or "invocation"
+            span_type = self.type or "task"
+
             span.set_attributes(
-                attributes={"node": self.type},
+                attributes={
+                    "node": span_type,
+                    "tree": trace_type,
+                },
                 namespace="type",
             )
 

agenta/sdk/litellm/mockllm.py

@@ -1,16 +1,69 @@
 from typing import Optional, Protocol, Any
 from os import environ
+from contextlib import contextmanager
 
 from agenta.sdk.utils.logging import get_module_logger
 
 from agenta.sdk.litellm.mocks import MOCKS
-from agenta.sdk.context.routing import routing_context
+from agenta.sdk.context.serving import serving_context
 
 AGENTA_LITELLM_MOCK = environ.get("AGENTA_LITELLM_MOCK") or None
 
 log = get_module_logger(__name__)
 
 
+ENV_KEYS_TO_CLEAR = [
+    # anything that could inject Lambda/ECS role creds
+    "AWS_SESSION_TOKEN",
+    "AWS_SECURITY_TOKEN",
+    "AWS_WEB_IDENTITY_TOKEN_FILE",
+    "AWS_ROLE_ARN",
+    "AWS_CONTAINER_CREDENTIALS_FULL_URI",
+    "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI",
+    "AWS_ECS_CONTAINER_AUTHORIZATION_TOKEN",
+    "AWS_PROFILE",
+    "AWS_SHARED_CREDENTIALS_FILE",
+    "AWS_CONFIG_FILE",
+]
+
+ENV_KEYS_FROM_USER = [
+    "AWS_ACCESS_KEY_ID",
+    "AWS_SECRET_ACCESS_KEY",
+    "AWS_SESSION_TOKEN",
+    "AWS_REGION",
+    "AWS_DEFAULT_REGION",
+]
+
+
+@contextmanager
+def user_aws_credentials_from(ps: dict):
+    old = {}
+    try:
+        # Save original state of ALL keys we'll modify
+        for k in ENV_KEYS_TO_CLEAR + ENV_KEYS_FROM_USER:
+            if k in environ:
+                old[k] = environ[k]
+
+        # Clear AWS role credentials
+        for k in ENV_KEYS_TO_CLEAR:
+            environ.pop(k, None)
+
+        # Set user credentials
+        for k in ENV_KEYS_FROM_USER:
+            if ps.get(k.upper()) is None:
+                environ.pop(k, None)
+            else:
+                environ[k] = ps[k.upper()]
+        yield
+    finally:
+        # Restore all environment variables to original state
+        for k in ENV_KEYS_TO_CLEAR + ENV_KEYS_FROM_USER:
+            if k in old:
+                environ[k] = old[k]
+            else:
+                environ.pop(k, None)
+
+
 class LitellmProtocol(Protocol):
     async def acompletion(self, *args: Any, **kwargs: Any) -> Any:
         ...
@@ -20,7 +73,7 @@ litellm: Optional[LitellmProtocol] = None  # pylint: disable=invalid-name
 
 
 async def acompletion(*args, **kwargs):
-    mock = AGENTA_LITELLM_MOCK or routing_context.get().mock
+    mock = AGENTA_LITELLM_MOCK or serving_context.get().mock
 
     if mock:
         log.debug("Mocking litellm: %s.", mock)

agenta/sdk/managers/config.py

@@ -7,7 +7,7 @@ from pydantic import BaseModel
 
 from agenta.sdk.utils.logging import get_module_logger
 from agenta.sdk.managers.shared import SharedManager
-from agenta.sdk.context.routing import routing_context
+from agenta.sdk.context.serving import serving_context
 
 T = TypeVar("T", bound=BaseModel)
 
@@ -45,7 +45,7 @@ class ConfigManager:
             Only one of these should be provided.
         """
 
-        context = routing_context.get()
+        context = serving_context.get()
 
         parameters = context.parameters
 

agenta/sdk/managers/secrets.py

@@ -1,14 +1,14 @@
 import re
 from typing import Optional, Dict, Any, List
 
-from agenta.sdk.context.routing import routing_context
+from agenta.sdk.context.serving import serving_context
 from agenta.sdk.assets import model_to_provider_mapping as _standard_providers
 
 
 class SecretsManager:
     @staticmethod
     def get_from_route() -> Optional[List[Dict[str, Any]]]:
-        context = routing_context.get()
+        context = serving_context.get()
 
         secrets = context.secrets
 

agenta/sdk/managers/vault.py

@@ -1,12 +1,12 @@
 from typing import Optional, Dict, Any
 
-from agenta.sdk.context.routing import routing_context
+from agenta.sdk.context.serving import serving_context
 
 
 class VaultManager:
     @staticmethod
     def get_from_route() -> Optional[Dict[str, Any]]:
-        context = routing_context.get()
+        context = serving_context.get()
 
         secrets = context.secrets