Source code for taimoe.platform.init

"""Top-level ``taimoe.init()`` entry point.

Environment variables consumed when the matching argument is omitted:

================================  ===========================================
``TAIMOE_VIRTUAL_KEY``            API key issued by the Taimoe Gateway.
``TAIMOE_RUNTIME_TOKEN``          **Deprecated** alias for
                                  ``TAIMOE_VIRTUAL_KEY`` — emits a
                                  ``DeprecationWarning`` and will be removed
                                  in a future minor release.
``TAIMOE_PLATFORM_URL``           Base URL for the Taimoe Gateway. Defaults
                                  to ``http://localhost:8000`` for local
                                  development; **production callers must
                                  pass this explicitly** — the default is a
                                  footgun on hosts where no such service
                                  exists locally and the failure modes are
                                  silent (degraded sync, dropped audit).
``TAIMOE_RUNTIME_NAME``           Runtime identifier for the runtime-bound
                                  sync path. Omit for runtime-less mode.
================================  ===========================================
"""

from __future__ import annotations

import logging
import warnings
from typing import Optional

from taimoe.platform._sync_types import SyncErrorHandler
from taimoe.platform.adapters.adk import taimoe
from taimoe.platform.instrumentation import is_adk_available, patch_all
from taimoe.platform.registry import TaimoeRegistry
from taimoe.platform.utils.env import get_env_or_raise

logger = logging.getLogger(__name__)

# Global registry instance to hold the initialized state
_global_registry: Optional[TaimoeRegistry] = None

#: Fallback runtime_id used for observability when the caller doesn't supply
#: one. Kept as a plain string — we no longer sniff ``sys.modules`` for
#: ``google.adk`` because the result depends on import order and silently
#: misattributes spans. Callers that care should pass ``runtime_name``.
_DEFAULT_RUNTIME_ID = "runtime-less"

#: Default ``TAIMOE_PLATFORM_URL`` for local development. See the module
#: docstring for why production callers must pass ``platform_url`` instead.
_DEFAULT_PLATFORM_URL = "http://localhost:8000"


[docs] def init( api_key: str | None = None, runtime_name: str | None = None, platform_url: str | None = None, runtime_id: str | None = None, framework: str | None = None, framework_version: str | None = None, poll_interval_seconds: float = 300, on_sync_error: SyncErrorHandler | None = None, instrument: bool | None = None, ) -> TaimoeRegistry: """Initialize the Taimoe Platform SDK. Two operating modes (see :class:`TaimoeRegistry` for details): 1. **Runtime-bound** — pass ``runtime_name``. Pulls every agent under that runtime via one batched HTTP call. 2. **Runtime-less** — omit ``runtime_name``. Pulls individual agents on demand whenever user code does ``taimoe.agent("foo")``. Both modes can run together: with ``runtime_name`` set, on-demand ``taimoe.agent(name)`` calls still get individually polled. Args: api_key: Gateway API key. Falls back to ``TAIMOE_VIRTUAL_KEY`` then (with a deprecation warning) ``TAIMOE_RUNTIME_TOKEN``. runtime_name: Runtime identifier; enables batched sync. Falls back to ``TAIMOE_RUNTIME_NAME``. platform_url: Gateway base URL. Falls back to ``TAIMOE_PLATFORM_URL`` and finally to ``http://localhost:8000`` (development only). runtime_id: Override for the runtime identifier used in observability spans. Defaults to ``runtime_name`` when omitted. framework: Framework name advertised in the discovery manifest (e.g. ``"google-adk"``). framework_version: Framework version advertised in the manifest. poll_interval_seconds: How often the background syncers refresh. on_sync_error: Callback invoked when a background sync iteration fails after exhausting its in-loop retries. instrument: Whether to apply auto-instrumentation (currently Google ADK). When ``None`` (default), the SDK auto-detects: instrument iff ``google.adk`` is already imported in this process. Pass ``True`` / ``False`` to override. Idempotent: calling ``init()`` twice in the same process returns the existing registry and logs a warning. Tests should call :func:`_reset_for_tests` to clear state between runs. """ global _global_registry if _global_registry is not None: logger.warning( "taimoe.init() has already been called. Returning existing registry." ) return _global_registry api_key = _resolve_api_key(api_key) platform_url = _resolve_platform_url(platform_url) runtime_name = runtime_name or get_env_or_raise("TAIMOE_RUNTIME_NAME", default="") or None if not api_key: logger.warning( "No API key provided to taimoe.init(). " "Governance and Sync features may fail." ) registry = TaimoeRegistry( runtime_name=runtime_name, platform_url=platform_url, platform_token=api_key, runtime_id=runtime_id, poll_interval_seconds=poll_interval_seconds, framework=framework, framework_version=framework_version, on_sync_error=on_sync_error, ) _global_registry = registry # Explicit binding (was a side effect of TaimoeRegistry.__init__ before). taimoe.bind_to(registry) # Observability shares the registry's resolved runtime_id so audit logs # and trace spans always agree. Empty string falls back to the SDK-wide # constant rather than something heuristic. obs_runtime_id = registry.runtime_id or _DEFAULT_RUNTIME_ID from taimoe.platform.observability import init as obs_init obs_init( runtime_id=obs_runtime_id, base_url=platform_url, api_key=api_key, ) # Run one synchronous sync before returning so that handles created # immediately after init() (e.g. as LlmAgent constructor args) see real # platform config rather than SDK bootstrap fallbacks. If the platform # is unreachable we still proceed — the background syncer will retry. registry.sync_once(ignore_errors=True) # Start background sync automatically. The first background tick fires # after one poll-interval (RuntimeSyncer/AgentSyncer wait-then-sync), # so there's no race with the sync_once above. registry.start_sync() _apply_instrumentation(instrument) logger.info( "Taimoe Platform SDK initialized for runtime %r", runtime_name or "(none)" ) return registry
[docs] def get_registry() -> TaimoeRegistry: """Return the global registry instance, or raise if not initialized.""" if _global_registry is None: raise RuntimeError("Taimoe SDK is not initialized. Call taimoe.init() first.") return _global_registry
def _reset_for_tests() -> None: """Clear the global registry so a test can re-init from scratch. Internal — production code should treat ``init()`` as one-shot. Stops the background syncer cleanly, unbinds the global ``taimoe`` handle, then drops the registry reference. """ global _global_registry if _global_registry is not None: try: _global_registry.stop_sync(timeout=1.0) except Exception as exc: # noqa: BLE001 logger.debug("Failed to stop sync during reset: %s", exc) taimoe.unbind() _global_registry = None # --- Resolution helpers ---------------------------------------------------- def _resolve_api_key(supplied: str | None) -> str | None: """Pick the API key from arg → ``TAIMOE_VIRTUAL_KEY`` → deprecated alias. Logs a ``DeprecationWarning`` when only the legacy ``TAIMOE_RUNTIME_TOKEN`` env var is present. """ if supplied: return supplied primary = get_env_or_raise("TAIMOE_VIRTUAL_KEY", default="") if primary: return primary legacy = get_env_or_raise("TAIMOE_RUNTIME_TOKEN", default="") if legacy: warnings.warn( "TAIMOE_RUNTIME_TOKEN is deprecated; rename it to " "TAIMOE_VIRTUAL_KEY. The fallback will be removed in a future " "minor release.", DeprecationWarning, stacklevel=3, ) return legacy return None def _resolve_platform_url(supplied: str | None) -> str: """Pick the platform URL from arg → env → development default.""" if supplied: return supplied return get_env_or_raise("TAIMOE_PLATFORM_URL", default=_DEFAULT_PLATFORM_URL) def _apply_instrumentation(instrument: bool | None) -> None: """Apply auto-instrumentation per the caller's policy. ``instrument=None`` → auto-detect (instrument iff ``google.adk`` is already imported). Explicit True / False overrides the detection. """ if instrument is False: logger.debug("instrument=False: skipping auto-instrumentation.") return if instrument is None and not is_adk_available(): logger.debug( "google.adk not imported; skipping auto-instrumentation. " "Pass instrument=True to force." ) return results = patch_all() for library, ok in results.items(): logger.info("Auto-instrumentation for %s: %s", library, "applied" if ok else "skipped")