d8026347e5
Renames: - platform/ → workspace-server/ (Go module path stays as "platform" for external dep compat — will update after plugin module republish) - workspace-template/ → workspace/ Removed (moved to separate repos or deleted): - PLAN.md — internal roadmap (move to private project board) - HANDOFF.md, AGENTS.md — one-time internal session docs - .claude/ — gitignored entirely (local agent config) - infra/cloudflare-worker/ → Molecule-AI/molecule-tenant-proxy - org-templates/molecule-dev/ → standalone template repo - .mcp-eval/ → molecule-mcp-server repo - test-results/ — ephemeral, gitignored Security scrubbing: - Cloudflare account/zone/KV IDs → placeholders - Real EC2 IPs → <EC2_IP> in all docs - CF token prefix, Neon project ID, Fly app names → redacted - Langfuse dev credentials → parameterized - Personal runner username/machine name → generic Community files: - CONTRIBUTING.md — build, test, branch conventions - CODE_OF_CONDUCT.md — Contributor Covenant 2.1 All Dockerfiles, CI workflows, docker-compose, railway.toml, render.yaml, README, CLAUDE.md updated for new directory names. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
393 lines
15 KiB
Python
393 lines
15 KiB
Python
"""Google ADK adapter for Molecule AI workspace runtime.
|
|
|
|
Wraps Google's Agent Development Kit (google-adk v1.x) as a Molecule AI
|
|
WorkspaceAdapter, bridging the A2A protocol to Google ADK's runner/session
|
|
model.
|
|
|
|
Google ADK concepts used
|
|
------------------------
|
|
- ``google.adk.agents.LlmAgent`` — An LLM-backed agent with instructions and
|
|
optional tools. Declared with ``model``, ``name``, and ``instruction``.
|
|
- ``google.adk.runners.Runner`` — Drives one or more agents inside a session;
|
|
``run_async()`` streams ``Event`` objects, including the final response text.
|
|
- ``google.adk.sessions.InMemorySessionService`` — Manages session state in
|
|
memory. Each ``Runner`` owns a single ``InMemorySessionService`` instance.
|
|
|
|
Runtime-config keys (all optional)
|
|
------------------------------------
|
|
``max_output_tokens`` — int, default 8192. Forwarded to the ADK ``GenerateContentConfig``.
|
|
``temperature`` — float, default 1.0.
|
|
``agent_name`` — str, default ``"molecule-adk-agent"``.
|
|
|
|
Environment variables
|
|
---------------------
|
|
``GOOGLE_API_KEY`` — Google AI Studio key (required for ``gemini-*`` models).
|
|
``GOOGLE_GENAI_USE_VERTEXAI`` — set to ``"1"`` to use Vertex AI instead of AI
|
|
Studio. In that case supply
|
|
``GOOGLE_CLOUD_PROJECT`` and
|
|
``GOOGLE_CLOUD_LOCATION`` as well.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from a2a.server.agent_execution import AgentExecutor, RequestContext
|
|
from a2a.server.events import EventQueue
|
|
from a2a.utils import new_agent_text_message
|
|
|
|
from adapter_base import AdapterConfig, BaseAdapter
|
|
|
|
if TYPE_CHECKING:
|
|
pass
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Constants
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_DEFAULT_AGENT_NAME = "molecule-adk-agent"
|
|
_DEFAULT_MAX_OUTPUT_TOKENS = 8192
|
|
_DEFAULT_TEMPERATURE = 1.0
|
|
_NO_TEXT_MSG = "Error: message contained no text content."
|
|
_NO_RESPONSE_MSG = "(no response generated)"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GoogleADKA2AExecutor
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class GoogleADKA2AExecutor(AgentExecutor):
|
|
"""A2A executor backed by a Google ADK ``Runner``.
|
|
|
|
Each executor instance owns a single ``Runner`` and ``InMemorySessionService``.
|
|
Sessions are created on first use and reused across subsequent turns
|
|
(the session_id is derived from the A2A context_id so each task gets a
|
|
stable, isolated session).
|
|
|
|
Parameters
|
|
----------
|
|
model:
|
|
ADK model identifier, e.g. ``"gemini-2.0-flash"`` or
|
|
``"gemini-1.5-pro"``.
|
|
system_prompt:
|
|
Optional instruction prepended to every conversation. Passed to
|
|
``LlmAgent(instruction=...)``.
|
|
agent_name:
|
|
Internal ADK agent name. Defaults to ``_DEFAULT_AGENT_NAME``.
|
|
max_output_tokens:
|
|
Token cap forwarded to ``GenerateContentConfig``.
|
|
temperature:
|
|
Sampling temperature forwarded to ``GenerateContentConfig``.
|
|
heartbeat:
|
|
Optional ``HeartbeatLoop`` instance (unused directly but stored for
|
|
future heartbeat integration).
|
|
_runner:
|
|
Inject a pre-built ``Runner`` — for testing only. When provided,
|
|
the real ADK ``Runner`` is never constructed.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model: str,
|
|
system_prompt: str | None = None,
|
|
agent_name: str = _DEFAULT_AGENT_NAME,
|
|
max_output_tokens: int = _DEFAULT_MAX_OUTPUT_TOKENS,
|
|
temperature: float = _DEFAULT_TEMPERATURE,
|
|
heartbeat: Any = None,
|
|
_runner: Any = None,
|
|
) -> None:
|
|
self.model = model
|
|
self.system_prompt = system_prompt
|
|
self.agent_name = agent_name
|
|
self.max_output_tokens = max_output_tokens
|
|
self.temperature = temperature
|
|
self._heartbeat = heartbeat
|
|
self._sessions_created: set[str] = set()
|
|
|
|
if _runner is not None:
|
|
# Test injection — skip building the real ADK objects.
|
|
self._runner = _runner
|
|
else:
|
|
self._runner = self._build_runner()
|
|
|
|
# ------------------------------------------------------------------
|
|
# Internal helpers
|
|
# ------------------------------------------------------------------
|
|
|
|
def _build_runner(self) -> Any: # pragma: no cover — requires real ADK
|
|
"""Construct a Google ADK ``Runner`` with an ``LlmAgent``.
|
|
|
|
Lazy-imports ``google.adk`` so the rest of the workspace runtime
|
|
doesn't pull in google-adk on startup (it's only needed when this
|
|
executor is actually instantiated by ``GoogleADKAdapter.create_executor``).
|
|
"""
|
|
from google.adk.agents import LlmAgent
|
|
from google.adk.runners import Runner
|
|
from google.adk.sessions import InMemorySessionService
|
|
|
|
agent = LlmAgent(
|
|
name=self.agent_name,
|
|
model=self.model,
|
|
instruction=self.system_prompt or "",
|
|
)
|
|
|
|
session_service = InMemorySessionService()
|
|
runner = Runner(
|
|
agent=agent,
|
|
app_name=self.agent_name,
|
|
session_service=session_service,
|
|
)
|
|
return runner
|
|
|
|
async def _ensure_session(self, session_id: str, user_id: str) -> None:
|
|
"""Create a session in the service if it doesn't exist yet."""
|
|
if session_id in self._sessions_created:
|
|
return
|
|
session_service = self._runner.session_service
|
|
existing = await session_service.get_session(
|
|
app_name=self.agent_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
)
|
|
if existing is None:
|
|
await session_service.create_session(
|
|
app_name=self.agent_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
)
|
|
self._sessions_created.add(session_id)
|
|
|
|
def _extract_text(self, context: RequestContext) -> str:
|
|
"""Pull plain text out of the A2A message parts."""
|
|
from shared_runtime import extract_message_text
|
|
return extract_message_text(context)
|
|
|
|
def _build_content(self, user_text: str) -> Any:
|
|
"""Wrap user text in an ADK-compatible ``Content`` object."""
|
|
from google.genai.types import Content, Part
|
|
return Content(role="user", parts=[Part(text=user_text)])
|
|
|
|
# ------------------------------------------------------------------
|
|
# AgentExecutor interface
|
|
# ------------------------------------------------------------------
|
|
|
|
async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
|
|
"""Run a single ADK turn and enqueue the reply as an A2A Message.
|
|
|
|
Sequence:
|
|
1. Extract user text from A2A message parts.
|
|
2. Ensure an ADK session exists for this context_id.
|
|
3. Call ``runner.run_async()`` and collect all response events.
|
|
4. Concatenate final-response text; fall back to ``_NO_RESPONSE_MSG``
|
|
when the model produces no output.
|
|
5. Enqueue the reply via ``event_queue``.
|
|
"""
|
|
user_text = self._extract_text(context)
|
|
if not user_text:
|
|
parts = getattr(getattr(context, "message", None), "parts", None)
|
|
logger.warning("GoogleADKA2AExecutor: no text in message parts: %s", parts)
|
|
await event_queue.enqueue_event(new_agent_text_message(_NO_TEXT_MSG))
|
|
return
|
|
|
|
session_id = getattr(context, "context_id", None) or "default-session"
|
|
user_id = "molecule-user"
|
|
|
|
try:
|
|
await self._ensure_session(session_id, user_id)
|
|
|
|
content = self._build_content(user_text)
|
|
response_parts: list[str] = []
|
|
|
|
async for event in self._runner.run_async(
|
|
session_id=session_id,
|
|
user_id=user_id,
|
|
new_message=content,
|
|
):
|
|
# Collect text from final-response events
|
|
if not getattr(event, "is_final_response", lambda: False)():
|
|
continue
|
|
candidate_response = getattr(event, "response", None)
|
|
if candidate_response is None:
|
|
continue
|
|
for part in getattr(
|
|
getattr(candidate_response, "content", None) or MissingContent(),
|
|
"parts", []
|
|
):
|
|
text = getattr(part, "text", None)
|
|
if text:
|
|
response_parts.append(text)
|
|
|
|
final_text = "".join(response_parts).strip() or _NO_RESPONSE_MSG
|
|
await event_queue.enqueue_event(new_agent_text_message(final_text))
|
|
|
|
except Exception as exc:
|
|
logger.error(
|
|
"GoogleADKA2AExecutor: execution error [model=%s]: %s",
|
|
self.model,
|
|
type(exc).__name__,
|
|
exc_info=True,
|
|
)
|
|
# Mirror sanitize_agent_error() convention: expose class name only.
|
|
await event_queue.enqueue_event(
|
|
new_agent_text_message(f"Agent error: {type(exc).__name__}")
|
|
)
|
|
|
|
async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
|
|
"""Cancel a running task — emits canceled state per A2A protocol."""
|
|
from a2a.types import TaskState, TaskStatus, TaskStatusUpdateEvent
|
|
|
|
await event_queue.enqueue_event(
|
|
TaskStatusUpdateEvent(
|
|
status=TaskStatus(state=TaskState.canceled),
|
|
final=True,
|
|
)
|
|
)
|
|
|
|
|
|
class MissingContent:
|
|
"""Sentinel to avoid AttributeError when response.content is None."""
|
|
parts: list = []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GoogleADKAdapter
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class GoogleADKAdapter(BaseAdapter):
|
|
"""Molecule AI workspace adapter for Google ADK (google-adk v1.x).
|
|
|
|
Implements the full ``BaseAdapter`` lifecycle:
|
|
- ``setup()`` — validates config and runs ``_common_setup()``.
|
|
- ``create_executor()`` — returns a ``GoogleADKA2AExecutor`` configured
|
|
from ``AdapterConfig``.
|
|
"""
|
|
|
|
# Stored by setup(); consumed by create_executor()
|
|
_setup_result: Any = None
|
|
|
|
# ------------------------------------------------------------------
|
|
# Identity
|
|
# ------------------------------------------------------------------
|
|
|
|
@staticmethod
|
|
def name() -> str:
|
|
"""Runtime identifier — matches the ``runtime`` field in config.yaml."""
|
|
return "google-adk"
|
|
|
|
@staticmethod
|
|
def display_name() -> str:
|
|
"""Human-readable name shown in the Molecule AI UI."""
|
|
return "Google ADK"
|
|
|
|
@staticmethod
|
|
def description() -> str:
|
|
"""Short description of this adapter's capabilities."""
|
|
return (
|
|
"Google Agent Development Kit (ADK) adapter. "
|
|
"Runs LLM agents via Google Gemini models using the official "
|
|
"google-adk Python SDK (Apache-2.0)."
|
|
)
|
|
|
|
@staticmethod
|
|
def get_config_schema() -> dict:
|
|
"""JSON Schema for runtime_config fields rendered in the Config tab."""
|
|
return {
|
|
"type": "object",
|
|
"properties": {
|
|
"agent_name": {
|
|
"type": "string",
|
|
"default": _DEFAULT_AGENT_NAME,
|
|
"description": "Internal ADK agent name",
|
|
},
|
|
"max_output_tokens": {
|
|
"type": "integer",
|
|
"default": _DEFAULT_MAX_OUTPUT_TOKENS,
|
|
"description": "Maximum output tokens for the Gemini model",
|
|
},
|
|
"temperature": {
|
|
"type": "number",
|
|
"default": _DEFAULT_TEMPERATURE,
|
|
"minimum": 0.0,
|
|
"maximum": 2.0,
|
|
"description": "Sampling temperature",
|
|
},
|
|
},
|
|
"additionalProperties": False,
|
|
}
|
|
|
|
# ------------------------------------------------------------------
|
|
# Lifecycle
|
|
# ------------------------------------------------------------------
|
|
|
|
async def setup(self, config: AdapterConfig) -> None:
|
|
"""Validate config and run the shared platform setup pipeline.
|
|
|
|
Raises ``RuntimeError`` if the required API key is not set and
|
|
Vertex AI mode is not active.
|
|
|
|
Args:
|
|
config: ``AdapterConfig`` populated by the workspace runtime.
|
|
"""
|
|
use_vertex = os.environ.get("GOOGLE_GENAI_USE_VERTEXAI", "").strip() in ("1", "true", "True")
|
|
api_key = os.environ.get("GOOGLE_API_KEY", "").strip()
|
|
|
|
if not use_vertex and not api_key:
|
|
raise RuntimeError(
|
|
"GoogleADKAdapter requires GOOGLE_API_KEY (for AI Studio) or "
|
|
"GOOGLE_GENAI_USE_VERTEXAI=1 with GOOGLE_CLOUD_PROJECT set."
|
|
)
|
|
|
|
logger.info(
|
|
"GoogleADKAdapter.setup: model=%s vertex=%s", config.model, use_vertex
|
|
)
|
|
|
|
self._setup_result = await self._common_setup(config)
|
|
|
|
async def create_executor(self, config: AdapterConfig) -> GoogleADKA2AExecutor:
|
|
"""Build and return a ``GoogleADKA2AExecutor`` for A2A integration.
|
|
|
|
Uses the system prompt assembled by ``_common_setup()`` in ``setup()``.
|
|
Runtime-config keys ``agent_name``, ``max_output_tokens``, and
|
|
``temperature`` are respected when present.
|
|
|
|
Args:
|
|
config: ``AdapterConfig`` populated by the workspace runtime.
|
|
|
|
Returns:
|
|
A ready-to-use ``GoogleADKA2AExecutor`` instance.
|
|
"""
|
|
rc = config.runtime_config or {}
|
|
|
|
# Strip provider prefix from model, e.g. "google:gemini-2.0-flash" → "gemini-2.0-flash"
|
|
model = config.model
|
|
if ":" in model:
|
|
model = model.split(":", 1)[1]
|
|
|
|
system_prompt = (
|
|
self._setup_result.system_prompt
|
|
if self._setup_result is not None
|
|
else config.system_prompt or ""
|
|
)
|
|
|
|
return GoogleADKA2AExecutor(
|
|
model=model,
|
|
system_prompt=system_prompt,
|
|
agent_name=rc.get("agent_name", _DEFAULT_AGENT_NAME),
|
|
max_output_tokens=int(rc.get("max_output_tokens", _DEFAULT_MAX_OUTPUT_TOKENS)),
|
|
temperature=float(rc.get("temperature", _DEFAULT_TEMPERATURE)),
|
|
heartbeat=config.heartbeat,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Module-level alias required by the adapter autodiscovery loader
|
|
# ---------------------------------------------------------------------------
|
|
|
|
Adapter = GoogleADKAdapter
|