Source code for taimoe.platform.adapters.adk.handle

"""ADK-facing config handles."""

from __future__ import annotations

from typing import Any

from taimoe.platform.adapters.adk.callbacks import (
    create_prompt_callback,
    create_refresh_callback,
)
from taimoe.platform.cache import RuntimeConfigCache
from taimoe.platform.types import AgentRuntimeConfig


# SDK-internal bootstrap instruction. Returned only when the platform
# config has not arrived yet (e.g. taimoe.init() ran but couldn't reach
# the platform). Users never write this.
_BOOTSTRAP_INSTRUCTION = (
    "You are an AI agent managed by the Taimoe platform. Your operating "
    "instructions are being fetched from the platform. If you are seeing "
    "this text in a response, the platform sync has not completed — answer "
    "the user briefly and ask them to retry."
)


[docs] class TaimoeConfigUnavailableError(RuntimeError): """Raised when a handle is asked for a value the platform hasn't synced. We deliberately do *not* invent a fallback (e.g. a hardcoded model alias) — sending the wrong alias to the gateway would just produce a confusing 404 downstream. Failing loudly here points at the real cause: the agent isn't registered on the platform under this runtime_agent_id, or the platform was unreachable when taimoe.init() ran. """
class _InstructionProvider: """Callable that resolves the current instruction from the platform cache. Matches ADK's InstructionProvider = Callable[[ReadonlyContext], str]. Passed to LlmAgent(instruction=...). ADK invokes this on every LLM call, so prompt edits in the console go live as soon as the SDK syncs them. """ def __init__(self, handle: "TaimoeAgentHandle") -> None: self._handle = handle def __call__(self, _ctx: Any) -> str: config = self._handle.config if config and config.instruction: return config.instruction return _BOOTSTRAP_INSTRUCTION def __repr__(self) -> str: return f"<TaimoeInstruction {self._handle.runtime_agent_id}>"
[docs] class TaimoeAgentHandle: """Lazy handle that resolves the latest platform config for an ADK agent. Use directly in ``LlmAgent(...)``: .. code-block:: python agent = taimoe.agent("my_agent") root_agent = LlmAgent( name="my_agent", model=agent.model, instruction=agent.instruction, tools=agent.tools, ) Each attribute returns a value shaped for the ADK constructor field: ``instruction`` is a callable (InstructionProvider, ADK re-resolves per call), ``model`` and ``tools`` snapshot the current platform values at construction. The ``refresh_callback`` swaps ``agent.model`` / ``agent.tools`` in place on subsequent calls so platform edits take effect within one sync cycle. ``instruction`` and ``model`` are intentionally asymmetric in their "no config yet" behavior: - ``instruction`` falls back to a bootstrap string — it's user-visible and we want the agent to respond gracefully during cold start. - ``model`` raises :class:`TaimoeConfigUnavailableError` — it's a gateway-visible identifier and guessing would produce a confusing 404 downstream. """ def __init__( self, runtime_agent_id: str, cache: RuntimeConfigCache | None = None ) -> None: self.runtime_agent_id = runtime_agent_id self._cache = cache
[docs] def bind_cache(self, cache: RuntimeConfigCache) -> TaimoeAgentHandle: return TaimoeAgentHandle(self.runtime_agent_id, cache)
@property def config(self) -> AgentRuntimeConfig | None: if self._cache is None: return None return self._cache.get(self.runtime_agent_id) @property def instruction(self) -> _InstructionProvider: """ADK-compatible InstructionProvider. Resolved on every LLM call.""" return _InstructionProvider(self) @property def model(self) -> str: """Current model alias from the platform. Raises TaimoeConfigUnavailableError if the platform config hasn't been synced yet — see the class docstring for why this is louder than the instruction fallback. """ config = self.config if config and config.model_alias: return config.model_alias raise TaimoeConfigUnavailableError( f"No platform config for runtime_agent_id={self.runtime_agent_id!r}. " "Check that the agent exists in the console and is bound to this runtime." ) @property def tools(self) -> list[Any]: """Current tool list from the platform. Empty list when no config. Note: the platform currently sends tool *names/ids* (``list[str]``), not real ADK tool objects. ``LlmAgent.tools`` expects ``list[BaseTool]``, so passing ``agent.tools`` straight in will fail at ADK construction. Resolve the names to your tool instances first, e.g.:: adk_tools = [my_tool_registry[name] for name in agent.tools] Server-side tool-object resolution is a phase-2 feature. """ config = self.config if config and config.tools: return list(config.tools) return [] @property def generation_config(self) -> dict[str, Any]: """Current generation parameters (temperature, top_p, top_k, max_output_tokens, …) as defined in the console's Model Settings. Returns an empty dict when nothing has been configured — callers should treat absent keys as "use the model's default", not "set to 0". """ config = self.config if config and config.generation_config: return dict(config.generation_config) return {} def __repr__(self) -> str: return f"TaimoeAgentHandle(runtime_agent_id={self.runtime_agent_id!r})"
[docs] def refresh_callback(self, *args: Any, **kwargs: Any) -> None: """ADK ``before_agent_callback`` that swaps ``agent.model`` / ``agent.tools`` on every call. Wire this in addition to ``instruction=handle.instruction`` to get all three fields live-updating from the platform. """ return create_refresh_callback(self)(*args, **kwargs)
[docs] class TaimoePromptHandle: """Lazy handle that resolves just the prompt/instruction.""" def __init__( self, runtime_prompt_id: str, cache: RuntimeConfigCache | None = None ) -> None: self.runtime_prompt_id = runtime_prompt_id self._cache = cache @property def config(self) -> AgentRuntimeConfig | None: if self._cache is None: return None return self._cache.get(self.runtime_prompt_id) @property def prompt(self) -> str | None: config = self.config return config.instruction if config else None
[docs] def refresh_callback(self, *args: Any, **kwargs: Any) -> None: return create_prompt_callback(self)(*args, **kwargs)
[docs] class TaimoeHandle: """Global facade used by application code. Most callers should use :meth:`bind_to` to wire this handle to a :class:`taimoe.platform.registry.TaimoeRegistry` in one call. :meth:`bind_cache` / :meth:`bind_agent_syncer` are kept as lower-level knobs for tests and multi-tenant scenarios that need finer control. """ def __init__(self) -> None: self._cache: RuntimeConfigCache | None = None self._agent_syncer: Any | None = None
[docs] def bind_to(self, registry: Any) -> None: """Wire this global handle to a ``TaimoeRegistry``. Equivalent to calling ``bind_cache(registry.cache)`` and ``bind_agent_syncer(registry.agent_syncer)`` — kept as a single ergonomic call so application code doesn't need to know about the wiring contract. """ self._cache = registry.cache self._agent_syncer = registry.agent_syncer
[docs] def unbind(self) -> None: """Detach the cache and syncer references. Mainly useful in tests that construct multiple registries within one process; production code typically calls ``bind_to`` once at startup and never unbinds. """ self._cache = None self._agent_syncer = None
[docs] def bind_cache(self, cache: RuntimeConfigCache) -> None: self._cache = cache
[docs] def bind_agent_syncer(self, syncer: Any) -> None: """Wire in the runtime-less syncer so that every ``agent(name)`` call automatically subscribes that name to the background poller.""" self._agent_syncer = syncer
[docs] def agent(self, runtime_agent_id: str) -> TaimoeAgentHandle: # Auto-subscribe: the moment user code mentions an agent name we # start pulling its config in the background. Avoids the chicken-and- # egg where init() doesn't know which agents to poll yet. if self._agent_syncer is not None: self._agent_syncer.subscribe(runtime_agent_id) return TaimoeAgentHandle(runtime_agent_id, self._cache)
[docs] def prompt(self, runtime_prompt_id: str) -> TaimoePromptHandle: return TaimoePromptHandle(runtime_prompt_id, self._cache)
[docs] def init( self, *, runtime_id: str, base_url: str | None = None, platform_url: str | None = None, api_key: str | None = None, enabled: bool = True, raise_on_export_error: bool = False, timeout: float = 10.0, ) -> None: """Configure observability export for this process.""" from taimoe.platform.observability import init init( runtime_id=runtime_id, base_url=base_url, platform_url=platform_url, api_key=api_key, enabled=enabled, raise_on_export_error=raise_on_export_error, timeout=timeout, )
[docs] def trace(self, name: str, **kwargs: Any) -> Any: """Start an observability trace.""" from taimoe.platform.observability import trace return trace(name, **kwargs)
taimoe = TaimoeHandle()