gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/__init__.py

@@ -0,0 +1,540 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""GatewaySDK's public package surface.
+
+The package root deliberately owns names, not feature initialization.  Python
+executes this file before every ``gatewaysdk.*`` submodule import, so eager
+re-exports here make a narrow Memory or tracing consumer initialize the runner,
+training, proxy, server, and GPU-monitoring graphs.  PEP 562 module attribute
+hooks keep the historical root-level imports while loading only the module that
+owns the requested name.
+"""
+
+from __future__ import annotations
+
+from importlib import import_module as _import_module
+from typing import TYPE_CHECKING as _TYPE_CHECKING
+
+if _TYPE_CHECKING:
+    # Static analyzers need concrete symbols; this branch is never evaluated by
+    # Python and therefore does not compromise the runtime import boundary.
+    from .adapter import *  # noqa: F401,F403
+    from .algorithm import *  # noqa: F401,F403
+    from .automations import AutomationClient, AutomationClientError  # noqa: F401
+    from .sessions import (  # noqa: F401
+        SessionLoadError,
+        SessionObservation,
+        SessionRecording,
+        SessionTrace,
+        load_session,
+    )
+    from .replay import (  # noqa: F401
+        ReplayCase,
+        ReplayCheck,
+        ReplaySuite,
+        ToolStubRegistry,
+        run_suite,
+    )
+    from .client import DevTaskLoader, GatewayClient  # noqa: F401
+    from .config import *  # noqa: F401,F403
+    from .context import (  # noqa: F401
+        GatewayContext,
+        GatewayContextConfig,
+        GatewayContextConflictError,
+        GatewayContextError,
+        render_agent_system,
+    )
+    from .emitter import *  # noqa: F401,F403
+    from .env_var import *  # noqa: F401,F403
+    from .execution import *  # noqa: F401,F403
+    from .instrumentation import (  # noqa: F401
+        Instrumentor,
+        InstrumentorRegistry,
+        MetricSchema,
+        instrument,
+        instrument_all,
+        uninstrument,
+        uninstrument_all,
+    )
+    from .integrations import (  # noqa: F401
+        GatewayToolAccess,
+        GatewayToolAccessConfig,
+        GatewayToolAccessError,
+        GatewayToolSubjectType,
+    )
+    from .integrations.gateway import GatewayIngestionClient, GatewayIngestionConfig  # noqa: F401
+    from .litagent import *  # noqa: F401,F403
+    from .llm_proxy import LLMProxy  # noqa: F401
+    from .logging import configure_logger  # noqa: F401
+    from .logging import setup as setup_logging  # noqa: F401
+    from .logging import setup_module as setup_module_logging  # noqa: F401
+    from .memory import (  # noqa: F401
+        GatewayMemory,
+        GatewayMemoryConfig,
+        GatewayMemoryConflictError,
+        GatewayMemoryError,
+    )
+    from .personas import Persona, PersonaError, PersonaReply, PersonasClient  # noqa: F401
+    from .run import (  # noqa: F401
+        GatewayRunError,
+        Metric,
+        Run,
+        episode,
+        experiment,
+        finish,
+        get_current_rollout,
+        get_current_run,
+        init_run,
+        log,
+        run,
+        score,
+        success,
+    )
+    from .runner import *  # noqa: F401,F403
+    from .security import (  # noqa: F401
+        GatewayAgentDescriptor,
+        GatewaySecurity,
+        GatewaySecurityConfig,
+        GatewaySecurityError,
+        GatewayToolApprovalRequiredError,
+        GatewayToolBlockedError,
+        GuardedMcpSession,
+    )
+    from .server import GatewayServer  # noqa: F401
+    from .store import *  # noqa: F401,F403
+    from .tracer import *  # noqa: F401,F403
+    from .tracing.identity import identify  # noqa: F401
+    from .trainer import *  # noqa: F401,F403
+    from .types import *  # noqa: F401,F403
+    from .users import UserEvent, UserProfile, UserProfileError, UsersClient  # noqa: F401
+    from .build import define_tool, register_artifact, register_prompt, register_skill  # noqa: F401
+
+__version__ = "0.3.2"
+
+
+# module: public names. Keep this explicit: wildcard ownership made additions
+# to an internal module silently become package-root API.
+_LAZY_EXPORT_GROUPS: dict[str, tuple[str, ...]] = {
+    ".adapter": (
+        "Adapter",
+        "LlmProxyTraceToTriplet",
+        "OtelTraceAdapter",
+        "TraceAdapter",
+        "TraceToMessages",
+        "TraceToTripletBase",
+        "TracerTraceToTriplet",
+    ),
+    ".algorithm": ("APO", "Algorithm", "Baseline", "FastAlgorithm", "GEPA", "GEPAAdapter", "VERL", "algo"),
+    ".automations": ("AutomationClient", "AutomationClientError"),
+    ".sessions": (
+        "SessionLoadError",
+        "SessionObservation",
+        "SessionRecording",
+        "SessionTrace",
+        "load_session",
+    ),
+    ".replay": (
+        "ReplayCase",
+        "ReplayCheck",
+        "ReplaySuite",
+        "ToolStubRegistry",
+        "run_suite",
+    ),
+    ".client": ("DevTaskLoader", "GatewayClient"),
+    ".config": ("lightning_cli",),
+    ".context": (
+        "GatewayContext",
+        "GatewayContextConfig",
+        "GatewayContextConflictError",
+        "GatewayContextError",
+        "render_agent_system",
+    ),
+    ".emitter": (
+        "emit_annotation",
+        "emit_exception",
+        "emit_message",
+        "emit_object",
+        "emit_reward",
+        "find_final_reward",
+        "find_reward_spans",
+        "get_message_value",
+        "get_object_value",
+        "get_reward_value",
+        "get_rewards_from_span",
+        "is_reward_span",
+        "operation",
+        "reward",
+    ),
+    ".env_var": ("LightningEnvVar", "resolve_bool_env_var", "resolve_int_env_var", "resolve_str_env_var"),
+    ".execution": (
+        "ClientServerExecutionStrategy",
+        "ExecutionEvent",
+        "ExecutionStrategy",
+        "MultiprocessingEvent",
+        "SharedMemoryExecutionStrategy",
+        "ThreadingEvent",
+    ),
+    ".instrumentation": (
+        "Instrumentor",
+        "InstrumentorRegistry",
+        "MetricSchema",
+        "instrument",
+        "instrument_all",
+        "uninstrument",
+        "uninstrument_all",
+    ),
+    ".integrations": (
+        "GatewayToolAccess",
+        "GatewayToolAccessConfig",
+        "GatewayToolAccessError",
+        "GatewayToolSubjectType",
+    ),
+    ".integrations.gateway": ("GatewayIngestionClient", "GatewayIngestionConfig"),
+    ".litagent": ("LitAgent", "llm_rollout", "prompt_rollout", "rollout"),
+    ".llm_proxy": ("LLMProxy",),
+    ".logging": ("configure_logger",),
+    ".memory": ("GatewayMemory", "GatewayMemoryConfig", "GatewayMemoryConflictError", "GatewayMemoryError"),
+    ".personas": ("Persona", "PersonaError", "PersonaReply", "PersonasClient"),
+    # Rollout intentionally remains owned by .types, matching the historical
+    # identity produced by the former final wildcard import.
+    ".run": (
+        "GatewayRunError",
+        "Metric",
+        "Run",
+        "episode",
+        "experiment",
+        "finish",
+        "get_current_rollout",
+        "get_current_run",
+        "init_run",
+        "log",
+        "run",
+        "score",
+        "success",
+    ),
+    ".runner": ("LegacyAgentRunner", "LitAgentRunner", "Runner"),
+    ".security": (
+        "GatewayAgentDescriptor",
+        "GatewaySecurity",
+        "GatewaySecurityConfig",
+        "GatewaySecurityError",
+        "GatewayToolApprovalRequiredError",
+        "GatewayToolBlockedError",
+        "GuardedMcpSession",
+    ),
+    ".server": ("GatewayServer",),
+    ".store": (
+        "CollectionBasedLightningStore",
+        "GatewayLightningStore",
+        "GatewayRunListener",
+        "GatewayTracker",
+        "InMemoryLightningStore",
+        "LightningStore",
+        "LightningStoreCapabilities",
+        "LightningStoreClient",
+        "LightningStoreServer",
+        "LightningStoreStatistics",
+        "LightningStoreThreaded",
+        "StorageListener",
+        "create_platform_store_client_from_env",
+        "create_platform_store_from_env",
+    ),
+    ".tracer": (
+        "DummyTracer",
+        "OtelTracer",
+        "Tracer",
+        "clear_active_tracer",
+        "get_active_tracer",
+        "set_active_tracer",
+    ),
+    ".tracing.identity": ("identify",),
+    ".trainer": ("Trainer", "build_component"),
+    ".types": (
+        "Attempt",
+        "AttemptStatus",
+        "AttemptedRollout",
+        "AttributeValue",
+        "Attributes",
+        "Dataset",
+        "EnqueueRolloutRequest",
+        "Event",
+        "FilterField",
+        "FilterOptions",
+        "GenericResponse",
+        "Hook",
+        "LLM",
+        "Link",
+        "NamedResources",
+        "OtelResource",
+        "PaginatedResult",
+        "ParallelWorkerBase",
+        "PromptTemplate",
+        "ProxyLLM",
+        "Resource",
+        "ResourceUnion",
+        "ResourcesUpdate",
+        "Rollout",
+        "RolloutConfig",
+        "RolloutLegacy",
+        "RolloutMode",
+        "RolloutRawResult",
+        "RolloutRawResultLegacy",
+        "RolloutStatus",
+        "SortOptions",
+        "Span",
+        "SpanAttributeNames",
+        "SpanContext",
+        "SpanCoreFields",
+        "SpanLike",
+        "SpanNames",
+        "SpanRecordingContext",
+        "StatusCode",
+        "Task",
+        "TaskIfAny",
+        "TaskInput",
+        "TraceState",
+        "TraceStatus",
+        "TracingConfig",
+        "Triplet",
+        "Worker",
+        "WorkerStatus",
+    ),
+    ".users": ("UserEvent", "UserProfile", "UserProfileError", "UsersClient"),
+    ".build": ("define_tool", "register_artifact", "register_prompt", "register_skill"),
+}
+
+_LAZY_EXPORTS = {
+    export_name: (module_name, export_name)
+    for module_name, export_names in _LAZY_EXPORT_GROUPS.items()
+    for export_name in export_names
+}
+_LAZY_EXPORTS.update(
+    {
+        "setup_logging": (".logging", "setup"),
+        "setup_module_logging": (".logging", "setup_module"),
+    }
+)
+
+
+_LAZY_MODULES: dict[str, str] = {
+    name: f".{name}"
+    for name in (
+        "adapter",
+        "algorithm",
+        "automations",
+        "sessions",
+        "replay",
+        "client",
+        "config",
+        "context",
+        "emitter",
+        "env_var",
+        "execution",
+        "instrumentation",
+        "integrations",
+        "litagent",
+        "logging",
+        "memory",
+        "personas",
+        "runner",
+        "security",
+        "semconv",
+        "server",
+        "store",
+        "tracer",
+        "tracing",
+        "types",
+        "users",
+        "utils",
+    )
+}
+
+_OPTIONAL_EXPORTS = {"LLMProxy", "Trainer", "build_component"}
+_OPTIONAL_EXPORT_EXTRAS = {
+    "LLMProxy": "proxy",
+    "Trainer": "trainer",
+    "build_component": "trainer",
+}
+
+# Preserve wildcard imports for the always-available surface.  Historically
+# optional proxy/trainer names appeared only when their extras happened to be
+# installed; explicit imports remain supported through ``__getattr__``.
+__all__ = sorted((set(_LAZY_EXPORTS) - _OPTIONAL_EXPORTS) | {"experiments", "init_experiments"})
+
+
+def _load_export(name: str) -> object:
+    module_name, attribute_name = _LAZY_EXPORTS[name]
+    try:
+        module = _import_module(module_name, __name__)
+    except ImportError as error:
+        if name in _OPTIONAL_EXPORTS:
+            # AttributeError keeps feature probes such as
+            # ``hasattr(gatewaysdk, "Trainer")`` compatible with the former
+            # optional import.  Direct attribute access still carries the
+            # installation instruction; ``from ... import`` converts it to
+            # the standard ImportError for a missing public name.
+            extra = _OPTIONAL_EXPORT_EXTRAS[name]
+            raise AttributeError(
+                f"{name} requires GatewaySDK's proxy/training dependencies. "
+                f"Install them with: pip install 'gatewaysdk[{extra}]'"
+            ) from error
+        raise
+
+    # Cache every export owned by this module.  Besides avoiding repeated
+    # __getattr__ calls, this restores the public ``run`` function after Python
+    # temporarily assigns the imported ``gatewaysdk.run`` module to that name.
+    for export_name, (owner_module, owner_attribute) in _LAZY_EXPORTS.items():
+        if owner_module == module_name and hasattr(module, owner_attribute):
+            globals()[export_name] = getattr(module, owner_attribute)
+    return globals()[name]
+
+
+def __getattr__(name: str) -> object:
+    if name in _LAZY_EXPORTS:
+        return _load_export(name)
+    if name in _LAZY_MODULES:
+        module = _import_module(_LAZY_MODULES[name], __name__)
+        globals()[name] = module
+        return module
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | set(_LAZY_EXPORTS) | set(_LAZY_MODULES))
+
+
+# =============================================================================
+# A/B Experiments — module-level proxy
+# =============================================================================
+
+
+class _ExperimentsNamespace:
+    """Proxy that delegates to the global ExperimentClient.
+
+    Allows `msk.experiments.register(...)` and `msk.experiments.get_variant(...)`
+    without requiring users to manage an ExperimentClient instance directly.
+
+    If no client has been initialized yet (via `init_experiments()`), calling
+    `register()` or `get_variant()` will lazily create a local-only client
+    (no exposure logging).
+    """
+
+    def _client(self):
+        from .experiments.client import ExperimentClient, _set_experiments_client, get_experiments_client
+
+        client = get_experiments_client()
+        if client is None:
+            client = ExperimentClient()
+            _set_experiments_client(client)
+        return client
+
+    def register(self, key, variants, traffic_percent=100, sync_to_backend=True):
+        """Register an A/B experiment. See ExperimentClient.register()."""
+        return self._client().register(key, variants, traffic_percent, sync_to_backend)
+
+    def get_variant(self, experiment_key, user_id):
+        """Get variant for a user. See ExperimentClient.get_variant()."""
+        return self._client().get_variant(experiment_key, user_id)
+
+    def is_registered(self, experiment_key):
+        """Check if an experiment is registered."""
+        return self._client().is_registered(experiment_key)
+
+    def list_experiments(self):
+        """List all registered experiment keys."""
+        return self._client().list_experiments()
+
+    def fetch(self, experiment_key, auto_register=True):
+        """Fetch an experiment from the backend. See ExperimentClient.fetch()."""
+        return self._client().fetch(experiment_key, auto_register)
+
+    def fetch_all(self, auto_register=True):
+        """Fetch all experiments from the backend. See ExperimentClient.fetch_all()."""
+        return self._client().fetch_all(auto_register)
+
+    def configure_backend(self, api_base_url, public_key, secret_key):
+        """Configure backend connectivity. See ExperimentClient.configure_backend()."""
+        return self._client().configure_backend(api_base_url, public_key, secret_key)
+
+
+# Store proxy under a private name so Python's import system can't clobber it
+# when 'gatewaysdk.experiments' (the subpackage) is later imported internally.
+_EXPERIMENTS_PROXY = _ExperimentsNamespace()
+experiments = _EXPERIMENTS_PROXY
+
+
+def init_experiments(
+    api_endpoint=None, public_key=None, secret_key=None, batch_size=100, flush_interval=1.0, api_base_url=None
+):
+    """Initialize the experiments system with backend connectivity.
+
+    Sets up the ExposureLogger (background thread that batches and sends
+    exposure events to the backend) and creates the global ExperimentClient.
+
+    Args:
+        api_endpoint: Full URL for the exposures batch API
+                     (e.g., "https://gateway.example.com/api/public/experiments/exposures/batch")
+                     If not provided, uses api_base_url to construct it.
+        public_key: Gateway public key (pk-lf-...)
+        secret_key: Gateway secret key (sk-lf-...)
+        batch_size: Max exposures per batch (default 100)
+        flush_interval: Seconds between flushes (default 1.0)
+        api_base_url: Base URL for the experiments API (e.g., "https://gateway.example.com").
+                     Used for fetching/syncing experiments. If not provided, derived from api_endpoint.
+
+    Returns:
+        The ExperimentClient instance.
+    """
+    # The imports above cause Python to register 'gatewaysdk.experiments' as the
+    # subpackage module in sys.modules, which overwrites 'gatewaysdk.experiments'
+    # attribute on this module. Restore the proxy immediately.
+    import sys
+
+    from .experiments.client import ExperimentClient, _set_experiments_client
+    from .experiments.exposure import ExposureLogger
+
+    _self = sys.modules.get(__name__)
+    if _self is not None:
+        _self.experiments = _EXPERIMENTS_PROXY
+
+    # Derive api_base_url from api_endpoint if not provided
+    if api_base_url is None and api_endpoint:
+        # Extract base URL from the full endpoint
+        # e.g., "https://gateway.example.com/api/public/experiments/exposures/batch"
+        # -> "https://gateway.example.com"
+        from urllib.parse import urlparse
+
+        parsed = urlparse(api_endpoint)
+        api_base_url = f"{parsed.scheme}://{parsed.netloc}"
+
+    # Derive api_endpoint from api_base_url if not provided
+    if api_endpoint is None and api_base_url:
+        api_endpoint = f"{api_base_url}/api/public/experiments/exposures/batch"
+
+    exposure_logger = None
+    if api_endpoint and public_key and secret_key:
+        exposure_logger = ExposureLogger(
+            api_endpoint=api_endpoint,
+            public_key=public_key,
+            secret_key=secret_key,
+            batch_size=batch_size,
+            flush_interval=flush_interval,
+        )
+        exposure_logger.start()
+
+    client = ExperimentClient(
+        exposure_logger=exposure_logger,
+        api_base_url=api_base_url,
+        public_key=public_key,
+        secret_key=secret_key,
+    )
+    _set_experiments_client(client)
+
+    # Transfer any experiments registered before init (use _EXPERIMENTS_PROXY
+    # directly — 'experiments' module-level name may be the subpackage at this point)
+    existing = _EXPERIMENTS_PROXY._client()
+    if existing is not client and hasattr(existing, "_experiments"):
+        for key, config in existing._experiments.items():
+            client._experiments[key] = config
+
+    return client

gatewaysdk/adapter/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .base import Adapter, OtelTraceAdapter, TraceAdapter
+from .messages import TraceToMessages
+from .triplet import LlmProxyTraceToTriplet, TracerTraceToTriplet, TraceToTripletBase
+
+__all__ = [
+    "TraceAdapter",
+    "OtelTraceAdapter",
+    "Adapter",
+    "TraceToTripletBase",
+    "TracerTraceToTriplet",
+    "LlmProxyTraceToTriplet",
+    "TraceToMessages",
+]

gatewaysdk/adapter/base.py

@@ -0,0 +1,94 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Generic, Sequence, TypeVar
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from gatewaysdk.types import Span
+
+T_from = TypeVar("T_from")
+T_to = TypeVar("T_to")
+
+
+class Adapter(Generic[T_from, T_to]):
+    """Base class for synchronous adapters that convert data from one format to another.
+
+    The class defines a minimal protocol so that adapters can be treated like callables while
+    still allowing subclasses to supply the concrete transformation logic.
+
+    !!! note
+        Subclasses must override [`adapt()`][gatewaysdk.Adapter.adapt] to provide
+        the actual conversion.
+
+    Type Variables:
+
+        T_from: Source data type supplied to the adapter.
+
+        T_to: Target data type produced by the adapter.
+
+    Examples:
+        >>> class IntToStrAdapter(Adapter[int, str]):
+        ...     def adapt(self, source: int) -> str:
+        ...         return str(source)
+        ...
+        >>> adapter = IntToStrAdapter()
+        >>> adapter(42)
+        '42'
+    """
+
+    def __call__(self, source: T_from, /) -> T_to:
+        """Convert the data to the target format.
+
+        This method delegates to [`adapt()`][gatewaysdk.Adapter.adapt] so that an
+        instance of [`Adapter`][gatewaysdk.Adapter] can be used like a standard
+        function.
+
+        Args:
+            source: Input data in the source format.
+
+        Returns:
+            Data converted to the target format.
+        """
+        return self.adapt(source)
+
+    def adapt(self, source: T_from, /) -> T_to:
+        """Convert the data to the target format.
+
+        Subclasses must override this method with the concrete transformation logic. The base
+        implementation raises `NotImplementedError` to make the requirement explicit.
+
+        Args:
+            source: Input data in the source format.
+
+        Returns:
+            Data converted to the target format.
+        """
+        raise NotImplementedError("Adapter.adapt() is not implemented")
+
+
+class OtelTraceAdapter(Adapter[Sequence[ReadableSpan], T_to], Generic[T_to]):
+    """Base class for adapters that convert OpenTelemetry trace spans into other formats.
+
+    This specialization of [`Adapter`][gatewaysdk.Adapter] expects a list of
+    `opentelemetry.sdk.trace.ReadableSpan` instances and produces any target format, such as
+    reinforcement learning trajectories, structured logs, or analytics-ready payloads.
+
+    Examples:
+        >>> class TraceToDictAdapter(OtelTraceAdapter[dict]):
+        ...     def adapt(self, spans: List[ReadableSpan]) -> dict:
+        ...         return {"count": len(spans)}
+        ...
+        >>> adapter = TraceToDictAdapter()
+        >>> adapter([span1, span2])
+        {'count': 2}
+    """
+
+
+class TraceAdapter(Adapter[Sequence[Span], T_to], Generic[T_to]):
+    """Base class for adapters that convert trace spans into other formats.
+
+    This class specializes [`Adapter`][gatewaysdk.Adapter] for working with
+    [`Span`][gatewaysdk.Span] instances emitted by GatewaySDK instrumentation.
+    Subclasses receive entire trace slices and return a format suited for the downstream consumer,
+    for example reinforcement learning training data or observability metrics.
+    """