gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/emitter/reward.py

@@ -0,0 +1,320 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Helpers for emitting reward spans and integrating with AgentOps telemetry."""
+
+import asyncio
+import inspect
+import json
+import logging
+import warnings
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Sequence,
+    TypedDict,
+    TypeVar,
+    cast,
+)
+
+from pydantic import TypeAdapter
+
+from gatewaysdk.semconv import AGL_ANNOTATION, LightningSpanAttributes, RewardPydanticModel
+from gatewaysdk.types import SpanCoreFields, SpanLike
+from gatewaysdk.utils.otel import filter_and_unflatten_attributes
+
+# NOTE: emit_annotation is imported lazily in emit_reward() to avoid circular import
+# The cycle is: emitter/__init__ → annotation → tracer → store → gateway → reward → annotation
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "reward",
+    "emit_reward",
+    "get_reward_value",
+    "get_rewards_from_span",
+    "is_reward_span",
+    "find_reward_spans",
+    "find_final_reward",
+]
+
+
+class RewardDimension(TypedDict):
+    """Type representing a single dimension in a multi-dimensional reward."""
+
+    name: str
+    value: float
+
+
+class _RewardSpanData(TypedDict):
+    type: Literal["reward"]
+    value: Optional[float]
+
+
+_FnType = TypeVar("_FnType", bound=Callable[..., Any])
+
+
+def _agentops_initialized() -> bool:
+    """Return `True` when the AgentOps client has been configured."""
+    import agentops
+
+    return agentops.get_client().initialized
+
+
+def reward(fn: _FnType) -> _FnType:
+    """Decorate a reward function so its outputs are tracked as spans.
+
+    The decorator integrates with AgentOps when it is available and falls back to
+    the built-in telemetry otherwise. Both synchronous and asynchronous functions
+    are supported transparently.
+
+    Deprecated:
+        This decorator is deprecated. Use [`emit_reward`][gatewaysdk.emit_reward] instead.
+
+    Args:
+        fn: Callable that produces a numeric reward.
+
+    Returns:
+        Wrapped callable that preserves the original signature.
+    """
+
+    from agentops.sdk.decorators import operation
+
+    def wrap_result(result: Optional[float]) -> _RewardSpanData:
+        """Normalize the reward value into the span payload format."""
+        if result is None:
+            return {"type": "reward", "value": None}
+        if not isinstance(result, (float, int)):  # type: ignore
+            warnings.warn(f"Reward is ignored because it is not a number: {result}")
+            return {"type": "reward", "value": None}
+        return {"type": "reward", "value": float(result)}
+
+    # Check if the function is async
+    is_async = asyncio.iscoroutinefunction(fn) or inspect.iscoroutinefunction(fn)
+
+    if is_async:
+
+        async def wrapper_async(*args: Any, **kwargs: Any) -> Any:
+            if not _agentops_initialized():
+                # Track the reward without AgentOps
+                result = await fn(*args, **kwargs)
+                emit_reward(cast(float, result))
+                return result
+
+            result: Optional[float] = None
+
+            @operation
+            async def agentops_reward_operation() -> _RewardSpanData:
+                # The reward function we are interested in tracing
+                # It takes zero inputs and return a formatted dict
+                nonlocal result
+                result = await fn(*args, **kwargs)
+                return wrap_result(result)
+
+            await agentops_reward_operation()
+            return result
+
+        return wrapper_async  # type: ignore
+
+    else:
+
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            if not _agentops_initialized():
+                # Track the reward without AgentOps
+                result = fn(*args, **kwargs)
+                emit_reward(cast(float, result))
+                return result
+
+            result: Optional[float] = None
+
+            @operation
+            def agentops_reward_operation() -> _RewardSpanData:
+                nonlocal result
+                result = fn(*args, **kwargs)
+                return wrap_result(result)
+
+            agentops_reward_operation()
+            return result
+
+        return wrapper  # type: ignore
+
+
+def emit_reward(
+    reward: float | Dict[str, Any],
+    *,
+    primary_key: str | None = None,
+    attributes: Dict[str, Any] | None = None,
+    propagate: bool = True,
+) -> SpanCoreFields:
+    """Emit a reward value as an OpenTelemetry span.
+
+    Examples:
+        Emit a single-dimensional reward:
+        >>> emit_reward(1.0)
+
+        Emit multi-dimensional rewards:
+        >>> emit_reward({"task_completion": 1.0, "efficiency": 0.8}, primary_key="task_completion")
+
+        Emit a reward with additional attributes (for example linking to another response span):
+        >>> from gatewaysdk.utils.otel import make_link_attributes
+        >>> emit_reward(0.5, attributes=make_link_attributes({"gen_ai.response.id": "response-123"}))
+
+        Or adding tags onto the reward span:
+        >>> from gatewaysdk.utils.otel import make_tag_attributes
+        >>> emit_reward(0.7, attributes=make_tag_attributes(["fast", "reliable"]))
+
+    Args:
+        reward: Numeric reward to record. Integers and booleans are converted to
+            floating point numbers for consistency.
+            Use a dictionary to represent a multi-dimensional reward.
+        attributes: Other optional span attributes.
+        propagate: Whether to propagate the span to exporters automatically.
+
+    Returns:
+        Span core fields capturing the recorded reward.
+    """
+    # Lazy import to avoid circular dependency
+    from .annotation import emit_annotation
+
+    logger.debug(f"Emitting reward: {reward}")
+    reward_dimensions: List[RewardDimension] = []
+    if isinstance(reward, dict):
+        reward_dict: Dict[str, float] = {}
+        for k, v in reward.items():
+            if isinstance(v, (int, bool)):
+                reward_dict[k] = float(v)
+            elif isinstance(v, float):
+                reward_dict[k] = v
+            else:
+                raise ValueError(f"Reward value must be a number, got: {type(v)} for key {k}")
+        if primary_key is None:
+            raise ValueError("When emitting a multi-dimensional reward as a dict, primary_key must be provided.")
+        if primary_key not in reward_dict:
+            raise ValueError(f"Primary key '{primary_key}' not found in reward dict keys: {list(reward_dict.keys())}")
+        reward_dimensions.append(RewardDimension(name=primary_key, value=reward_dict[primary_key]))
+        for k, v in reward_dict.items():
+            if k != primary_key:
+                reward_dimensions.append(RewardDimension(name=k, value=v))
+    else:
+        if isinstance(reward, (int, bool)):
+            reward = float(reward)
+        elif not isinstance(reward, float):  # pyright: ignore[reportUnnecessaryIsInstance]
+            raise TypeError(f"Reward must be a number, got: {type(reward)}")
+        reward_dimensions.append(RewardDimension(name="primary", value=reward))
+
+    return emit_annotation(
+        {LightningSpanAttributes.REWARD.value: reward_dimensions, **(attributes or {})}, propagate=propagate
+    )
+
+
+def get_reward_value(span: SpanLike) -> Optional[float]:
+    """Extract the reward value from a span, if available.
+
+    Args:
+        span: Span object produced by AgentOps or GatewaySDK emitters.
+
+    Returns:
+        The primary reward encoded in the span or `None` when the span does not represent a reward.
+    """
+    # v0.3+ emit reward format
+    reward_list = get_rewards_from_span(span)
+    if reward_list:
+        # Reward list is ordered and the first element is the primary reward
+        return reward_list[0].value
+
+    for key in [
+        "agentops.task.output",  # newer versions of agentops
+        "agentops.entity.output",
+    ]:
+        reward_dict: Dict[str, Any] | None = None
+        if span.attributes:
+            output = span.attributes.get(key)
+            if output:
+                if isinstance(output, dict):
+                    reward_dict = cast(Dict[str, Any], output)
+                elif isinstance(output, str):
+                    try:
+                        reward_dict = cast(Dict[str, Any], json.loads(output))
+                    except json.JSONDecodeError:
+                        reward_dict = None
+
+        if reward_dict and reward_dict.get("type") == "reward":
+            reward_value = reward_dict.get("value", None)
+            if reward_value is None:
+                return None
+            if not isinstance(reward_value, float):
+                logger.error(f"Reward is not a number, got: {type(reward_value)}. This may cause undefined behaviors.")
+            logger.warning(
+                f"Extracted reward {reward_value} from AgentOps. This format is deprecated, please migrate to using `emit_reward`."
+            )
+            return cast(float, reward_value)
+
+    # v0.2 emit reward format
+    if span.name == AGL_ANNOTATION and span.attributes:
+        reward_value = span.attributes.get("reward", None)
+        if reward_value is None:
+            return None
+        if not isinstance(reward_value, float):
+            logger.error(f"Reward is not a number, got: {type(reward_value)}. This may cause undefined behaviors.")
+        logger.warning(
+            f"Extracted reward {reward_value} from a legacy version of reward span. You might have inconsistent gatewaysdk versions."
+        )
+        return cast(float, reward_value)
+
+    return None
+
+
+def get_rewards_from_span(span: SpanLike) -> List[RewardPydanticModel]:
+    """Extract the reward as a list from a span, if available.
+
+    Args:
+        span: Span object produced by AgentOps or GatewaySDK emitters.
+
+    Returns:
+        A list of reward dimensions encoded in the span or an empty list when the span does not represent a reward.
+    """
+    if span.attributes and any(key.startswith(LightningSpanAttributes.REWARD.value) for key in span.attributes):
+        reward_attr = filter_and_unflatten_attributes(
+            cast(Any, span.attributes or {}), LightningSpanAttributes.REWARD.value
+        )
+        recovered_rewards = TypeAdapter(List[RewardPydanticModel]).validate_python(reward_attr)
+        return recovered_rewards
+    else:
+        return []
+
+
+def is_reward_span(span: SpanLike) -> bool:
+    """Return ``True`` when the provided span encodes a reward value."""
+    maybe_reward = get_reward_value(span)
+    return maybe_reward is not None
+
+
+def find_reward_spans(spans: Sequence[SpanLike]) -> List[SpanLike]:
+    """Return all reward spans in the provided sequence.
+
+    Args:
+        spans: Sequence containing [`ReadableSpan`](https://opentelemetry.io/docs/concepts/signals/traces/) objects or mocked span-like values.
+
+    Returns:
+        List of spans that could be parsed as rewards.
+    """
+    return [span for span in spans if is_reward_span(span)]
+
+
+def find_final_reward(spans: Sequence[SpanLike]) -> Optional[float]:
+    """Return the last reward value present in the provided spans.
+
+    Args:
+        spans: Sequence containing [`ReadableSpan`](https://opentelemetry.io/docs/concepts/signals/traces/) objects or mocked span-like values.
+
+    Returns:
+        Reward value from the latest reward span, or `None` when none are found.
+    """
+    for span in reversed(spans):
+        reward = get_reward_value(span)
+        if reward is not None:
+            return reward
+    return None

gatewaysdk/env_var.py

@@ -0,0 +1,156 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Environment variable managements."""
+
+from __future__ import annotations
+
+import os
+from enum import Enum
+from typing import overload
+
+__all__ = [
+    "LightningEnvVar",
+    "resolve_bool_env_var",
+    "resolve_int_env_var",
+    "resolve_str_env_var",
+]
+
+
+class LightningEnvVar(Enum):
+    """Environment variables for GatewaySDK."""
+
+    AGL_EMITTER_DEBUG = "AGL_EMITTER_DEBUG"
+    """Enable debug logging for the emitter."""
+
+    AGL_MANAGED_STORE = "AGL_MANAGED_STORE"
+    """If yes, the [`ExecutionStrategy`][gatewaysdk.ExecutionStrategy]
+    constructs LightningStore wrappers automatically. When `False` the provided
+    `store` is passed directly to the bundles, allowing callers to manage
+    store wrappers manually."""
+
+    AGL_CURRENT_ROLE = "AGL_CURRENT_ROLE"
+    """Which side(s) to run in this process. Used in
+    [`ClientServerExecutionStrategy`][gatewaysdk.ClientServerExecutionStrategy]."""
+
+    AGL_SERVER_HOST = "AGL_SERVER_HOST"
+    """Interface the [`LightningStoreServer`][gatewaysdk.LightningStoreServer]
+    binds to when running the algorithm bundle locally."""
+
+    AGL_SERVER_PORT = "AGL_SERVER_PORT"
+    """Port the [`LightningStoreServer`][gatewaysdk.LightningStoreServer] listens to."""
+
+
+_TRUTHY_VALUES = {"1", "true", "yes", "on"}
+_FALSY_VALUES = {"0", "false", "no", "off"}
+
+
+@overload
+def resolve_bool_env_var(env_var: LightningEnvVar, override: bool, fallback: bool) -> bool: ...
+
+
+@overload
+def resolve_bool_env_var(env_var: LightningEnvVar, *, fallback: bool) -> bool: ...
+
+
+@overload
+def resolve_bool_env_var(
+    env_var: LightningEnvVar, override: bool | None = None, fallback: bool | None = None
+) -> bool | None: ...
+
+
+def resolve_bool_env_var(
+    env_var: LightningEnvVar, override: bool | None = None, fallback: bool | None = None
+) -> bool | None:
+    """Resolve a boolean environment variable.
+
+    Args:
+        env_var: The environment variable to resolve.
+        override: Optional override supplied by the caller.
+        fallback: Default value if the environment variable is not set.
+    """
+
+    if override is not None:
+        return override
+
+    env_value = os.getenv(env_var.value)
+    if env_value is None:
+        return fallback
+
+    normalized = env_value.strip().lower()
+    if normalized in _TRUTHY_VALUES:
+        return True
+    if normalized in _FALSY_VALUES:
+        return False
+
+    raise ValueError(f"{env_var.value} must be one of {_TRUTHY_VALUES} or {_FALSY_VALUES}")
+
+
+@overload
+def resolve_int_env_var(env_var: LightningEnvVar, override: int, fallback: int) -> int: ...
+
+
+@overload
+def resolve_int_env_var(env_var: LightningEnvVar, *, fallback: int) -> int: ...
+
+
+@overload
+def resolve_int_env_var(
+    env_var: LightningEnvVar, override: int | None = None, fallback: int | None = None
+) -> int | None: ...
+
+
+def resolve_int_env_var(
+    env_var: LightningEnvVar, override: int | None = None, fallback: int | None = None
+) -> int | None:
+    """Resolve an integer environment variable.
+
+    Args:
+        env_var: The environment variable to resolve.
+        override: Optional override supplied by the caller.
+        fallback: Default value if the environment variable is not set.
+    """
+    if override is not None:
+        return override
+
+    env_value = os.getenv(env_var.value)
+    if env_value is None:
+        return fallback
+
+    try:
+        return int(env_value)
+    except ValueError:
+        raise ValueError(f"{env_var.value} must be an integer")
+
+
+@overload
+def resolve_str_env_var(env_var: LightningEnvVar, override: str, fallback: str) -> str: ...
+
+
+@overload
+def resolve_str_env_var(env_var: LightningEnvVar, *, fallback: str) -> str: ...
+
+
+@overload
+def resolve_str_env_var(
+    env_var: LightningEnvVar, override: str | None = None, fallback: str | None = None
+) -> str | None: ...
+
+
+def resolve_str_env_var(
+    env_var: LightningEnvVar, override: str | None = None, fallback: str | None = None
+) -> str | None:
+    """Resolve a string environment variable.
+
+    Args:
+        env_var: The environment variable to resolve.
+        override: Optional override supplied by the caller.
+        fallback: Default value if the environment variable is not set.
+    """
+    if override is not None:
+        return override
+
+    env_value = os.getenv(env_var.value)
+    if env_value is None:
+        return fallback
+
+    return env_value

gatewaysdk/environment/__init__.py

@@ -0,0 +1,108 @@
+"""Mock environments as an execution substrate — ``gatewaysdk.environment``.
+
+    import gatewaysdk.environment as gwenv
+
+    env = gwenv.init("refund-flows-env")          # slug/id, READY head
+    with env:                                     # fresh world + ambient interception
+        my_agent("handle ticket #4415")           # define_tool calls hit the mock world
+    assert env.verify("refund_row_created").passed
+
+Seams (spec §4.4): ``world.as_registry()`` (replay's registry contract),
+``world.wrap_tools()``, ``world.execute()``, and the ambient ``define_tool``
+wrapper inside ``with env:`` / ``env.branch()``. Harbor-kind environments
+(task.toml trees) expose ``env.tasks`` / ``env.task(name)`` docker sandboxes
+instead. ``gwenv.init_from_path(dir)`` boots a local bundle offline.
+
+Note: :mod:`gatewaysdk.environments` (plural) is the platform HUB client
+(read/inspect/trigger); this module is the local execution runtime built on
+top of it.
+"""
+
+import importlib
+from typing import TYPE_CHECKING
+
+from . import runtime
+from ._bundle import EnvManifest, HarborTask
+from .errors import (
+    EnvironmentKindError,
+    EnvironmentNotConfigured,
+    EnvironmentNotReady,
+    GatewayEnvError,
+    HermeticViolation,
+    ToolRouteError,
+    VerifierError,
+    WorldStartError,
+)
+
+# The heavy seams live in ``core`` (platform hub + tracing clients), ``_harbor``
+# (docker CLI) and ``_world`` (the replay registry contract). Bind them LAZILY
+# via PEP 562 so a lean consumer — e.g. the benchmark runner's mock_tools engine,
+# which vendors only this ``environment`` subpackage and imports ``_bundle`` /
+# ``_world`` (worlds + verifiers) directly — can ``import gatewaysdk.environment``
+# without dragging ``gatewaysdk.environments`` / ``gatewaysdk.build`` / docker /
+# replay into the image. First access to any of these names imports its module
+# on demand, so full-SDK behaviour is unchanged (just deferred).
+_LAZY_ATTRS = {
+    "ENV_ENVIRONMENT_PREFIX": "core",
+    "Environment": "core",
+    "Report": "core",
+    "RouteRecord": "core",
+    "get_active_world": "core",
+    "init": "core",
+    "init_from_path": "core",
+    "CliDockerRuntime": "_harbor",
+    "DockerRuntime": "_harbor",
+    "ExecResult": "_harbor",
+    "HarborWorld": "_harbor",
+    "VerifierResult": "_world",
+    "World": "_world",
+    "WorldRegistry": "_world",
+}
+
+if TYPE_CHECKING:  # pragma: no cover — static analysers still see the symbols
+    from ._harbor import CliDockerRuntime, DockerRuntime, ExecResult, HarborWorld
+    from ._world import VerifierResult, World, WorldRegistry
+    from .core import (
+        ENV_ENVIRONMENT_PREFIX,
+        Environment,
+        Report,
+        RouteRecord,
+        get_active_world,
+        init,
+        init_from_path,
+    )
+
+
+def __getattr__(name: str):
+    module = _LAZY_ATTRS.get(name)
+    if module is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    return getattr(importlib.import_module(f".{module}", __name__), name)
+
+__all__ = [
+    "ENV_ENVIRONMENT_PREFIX",
+    "CliDockerRuntime",
+    "DockerRuntime",
+    "EnvManifest",
+    "Environment",
+    "EnvironmentKindError",
+    "EnvironmentNotConfigured",
+    "EnvironmentNotReady",
+    "ExecResult",
+    "GatewayEnvError",
+    "HarborTask",
+    "HarborWorld",
+    "HermeticViolation",
+    "Report",
+    "RouteRecord",
+    "ToolRouteError",
+    "VerifierError",
+    "VerifierResult",
+    "World",
+    "WorldRegistry",
+    "WorldStartError",
+    "get_active_world",
+    "init",
+    "init_from_path",
+    "runtime",
+]