Updated examples/adk_agent.py

Author: KellySit

examples/google_adk/adk_agent.py

@@ -6,10 +6,32 @@
 """
 
 from __future__ import annotations
+
+import asyncio
+import logging
+from collections.abc import Mapping
 from typing import Any, Dict, TypedDict, cast
+
 from agentlightning import LLM, LitAgent, NamedResources
 from agentlightning.types import Rollout, RolloutRawResult
 
+logger = logging.getLogger(__name__)
+
+try:  # pragma: no cover - import guarded for optional dependency.
+    from google.genai import types as genai_types
+
+    from google.adk.agents.llm_agent import LlmAgent
+    from google.adk.apps.app import App
+    from google.adk.models.lite_llm import LiteLlm
+    from google.adk.runners import InMemoryRunner
+    from google.adk.utils.context_utils import Aclosing
+
+    _HAS_GOOGLE_ADK = True
+except ImportError:  # pragma: no cover
+    genai_types = None  # type: ignore[assignment]
+    LlmAgent = App = LiteLlm = InMemoryRunner = Aclosing = None  # type: ignore[assignment]
+    _HAS_GOOGLE_ADK = False
+
 
 class AdkTask(TypedDict):
     """
@@ -29,29 +51,172 @@ class AdkTask(TypedDict):
 
 
 class LitAdkAgent(LitAgent[AdkTask]):
-    """Basic ADK + LitAgent example."""
-
-    def rollout(self, task: AdkTask, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
-        # get the llm resource
-        llm: LLM = cast(LLM, resources.get("main_llm"))
-
+    """ADK-backed agent that produces observability-friendly rollouts."""
+
+    def rollout(self, task: AdkTask, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:  # type: ignore[override]
+        """Synchronous entry point – forwards to ``rollout_async``."""
+
+        if not _HAS_GOOGLE_ADK:  # pragma: no cover
+            raise RuntimeError(
+                "google-adk>=0.3.0 is required to run this example. "
+                "Install it with `pip install google-adk` or enable the "
+                "`adk` optional dependency group."
+            )
+
+        try:
+            return asyncio.run(self.rollout_async(task, resources, rollout))
+        except RuntimeError as exc:  # pragma: no cover - defensive path
+            if "asyncio.run()" in str(exc):
+                raise RuntimeError(
+                    "LitAdkAgent.rollout cannot be executed while an event loop "
+                    "is already running. Call `rollout_async` instead."
+                ) from exc
+            raise
+
+    async def rollout_async(  # type: ignore[override]
+        self,
+        task: AdkTask,
+        resources: NamedResources,
+        rollout: Rollout,
+    ) -> RolloutRawResult:
+        """Runs a single rollout by delegating to ADK's orchestration runtime."""
+
+        if not _HAS_GOOGLE_ADK:  # pragma: no cover
+            raise RuntimeError(
+                "google-adk>=0.3.0 is required to run this example. "
+                "Install it with `pip install google-adk` or enable the "
+                "`adk` optional dependency group."
+            )
+
+        llm = cast(LLM, resources.get("main_llm"))
         question = task["question"]
-        app_id = task["app_id"]
         truth = task["ground_truth"]
+        app_id = task["app_id"] or "adk_agent_app"
+
+        adk_model = self._build_adk_model(llm)
+        adk_agent = self._build_adk_agent(adk_model, task)
+        app = App(name=app_id, root_agent=adk_agent)
+
+        runner = InMemoryRunner(app=app)
+        try:
+            session_state = self._build_session_state(task)
+            session = await runner.session_service.create_session(
+                app_name=app.name,
+                user_id="agent_lightning_user",
+                state=session_state,
+            )
+
+            message = genai_types.Content(  # type: ignore[union-attr]
+                role="user",
+                parts=[genai_types.Part.from_text(text=question)],  # type: ignore[union-attr]
+            )
+
+            last_response = ""
+            async with Aclosing(  # type: ignore[union-attr]
+                runner.run_async(
+                    user_id=session.user_id,
+                    session_id=session.id,
+                    new_message=message,
+                )
+            ) as agen:
+                async for event in agen:
+                    if event.content and event.content.parts:
+                        text = "".join(part.text or "" for part in event.content.parts).strip()
+                        if text:
+                            last_response = text
+
+            reward = self._compute_reward(last_response, truth)
+            return reward
+        except Exception as exc:  # pragma: no cover - surfaced via training logs
+            logger.exception("ADK rollout failed: %s", exc)
+            return 0.0
+        finally:
+            await runner.close()
+
+    @staticmethod
+    def _build_adk_model(llm: LLM) -> LiteLlm:  # type: ignore[valid-type]
+        """Create a LiteLLM-backed ADK model from the Agent-Lightning LLM resource."""
+
+        sampling_params = llm.sampling_parameters or {}
+        temperature: float | None = None
+        top_p: float | None = None
+        if isinstance(sampling_params, Mapping):
+            temperature = sampling_params.get("temperature")
+            top_p = sampling_params.get("top_p")
+
+        llm_kwargs: Dict[str, Any] = {}
+        if llm.endpoint:
+            llm_kwargs["api_base"] = llm.endpoint
+        if llm.api_key:
+            llm_kwargs["api_key"] = llm.api_key
+        if temperature is not None:
+            llm_kwargs["temperature"] = float(temperature)
+        if top_p is not None:
+            llm_kwargs["top_p"] = float(top_p)
+
+        return LiteLlm(model=llm.model, **llm_kwargs)  # type: ignore[arg-type]
+
+    def _build_adk_agent(self, adk_model: LiteLlm, task: AdkTask) -> LlmAgent:  # type: ignore[valid-type]
+        """Construct the ADK agent definition for this rollout."""
+
+        instruction = self._compose_instruction(task)
+        agent_name = task["app_id"] or "adk_agent"
+
+        return LlmAgent(
+            name=agent_name,
+            model=adk_model,
+            instruction=instruction,
+        )
+
+    @staticmethod
+    def _build_session_state(task: AdkTask) -> Dict[str, Any]:
+        """Extract session state from the task metadata."""
+
+        session_state: Dict[str, Any] = {}
+        meta = task.get("meta")
+        if isinstance(meta, Mapping):
+            session_state.update(meta)
+        return session_state
+
+    def _compose_instruction(self, task: AdkTask) -> str:
+        """Compose the instruction prompt for the ADK agent."""
+
+        base_instruction = (
+            "You are an enterprise assistant integrated with Google ADK. "
+            "Use available tools to reason carefully and produce concise, factual answers. "
+            "Explain limitations when information is missing."
+        )
+
+        meta = task.get("meta")
+        if isinstance(meta, Mapping):
+            # Recognise common fields that may contain additional guidance.
+            for key in ("instruction", "goal", "context", "description"):
+                if meta.get(key):
+                    return f"{base_instruction}\n\nAdditional context:\n{meta[key]}"
+
+        return base_instruction
+
+    @staticmethod
+    def _compute_reward(answer: str, truth: str) -> float:
+        """Simple lexical reward comparing the ADK answer with ground truth."""
+
+        if not answer or not truth:
+            return 0.0
+
+        answer_lower = answer.lower()
+        truth_lower = truth.lower()
+        return 1.0 if truth_lower in answer_lower else 0.0
 
-        # TODO: hook up real ADK orchestration later
-        # for now just simulate an action string
-        action = f"adk://{app_id}?plan={question}"
-
-        # quick check for correctness
-        reward = 1.0 if truth and truth.lower() in action.lower() else 0.0
 
-        return reward
+# Minimal smoke-test entry point (manual run)
+if __name__ == "__main__":  # pragma: no cover - manual verification helper
+    import sys
 
+    if not _HAS_GOOGLE_ADK:
+        sys.exit(
+            "google-adk is not installed. Run `pip install google-adk` before executing this script."
+        )
 
-# Minimal smoke-test entry point (manual run)
-if __name__ == "__main__":
-    # very minimal test run
     sample_task: AdkTask = {
         "question": "Create a calendar event for Monday 10am titled 'Standup'",
         "app_id": "sample_calendar_app",
@@ -62,15 +227,17 @@ def rollout(self, task: AdkTask, resources: NamedResources, rollout: Rollout) ->
     resources: NamedResources = {
         "main_llm": LLM(
             endpoint="http://localhost:8000/v1",
-            model="meta-llama/Meta-Llama-3-8B-Instruct"
+            model="gpt-4.1-mini",
+            sampling_parameters={"temperature": 0.0},
+            api_key="dummy-key",
         ),
     }
 
     class DummyRollout:
         pass
 
     agent = LitAdkAgent()
-    result = agent.rollout(sample_task, resources, cast(Rollout, DummyRollout()))
+    result = asyncio.run(agent.rollout_async(sample_task, resources, cast(Rollout, DummyRollout())))
     print("Reward:", result)
 
 

vendor/google_adk/_yaml/__init__.py

@@ -0,0 +1,33 @@
+# This is a stub package designed to roughly emulate the _yaml
+# extension module, which previously existed as a standalone module
+# and has been moved into the `yaml` package namespace.
+# It does not perfectly mimic its old counterpart, but should get
+# close enough for anyone who's relying on it even when they shouldn't.
+import yaml
+
+# in some circumstances, the yaml module we imoprted may be from a different version, so we need
+# to tread carefully when poking at it here (it may not have the attributes we expect)
+if not getattr(yaml, '__with_libyaml__', False):
+    from sys import version_info
+
+    exc = ModuleNotFoundError if version_info >= (3, 6) else ImportError
+    raise exc("No module named '_yaml'")
+else:
+    from yaml._yaml import *
+    import warnings
+    warnings.warn(
+        'The _yaml extension module is now located at yaml._yaml'
+        ' and its location is subject to change.  To use the'
+        ' LibYAML-based parser and emitter, import from `yaml`:'
+        ' `from yaml import CLoader as Loader, CDumper as Dumper`.',
+        DeprecationWarning
+    )
+    del warnings
+    # Don't `del yaml` here because yaml is actually an existing
+    # namespace member of _yaml.
+
+__name__ = '_yaml'
+# If the module is top-level (i.e. not a part of any specific package)
+# then the attribute should be set to ''.
+# https://docs.python.org/3.8/library/types.html
+__package__ = ''

vendor/google_adk/alembic-1.16.5.dist-info/INSTALLER

@@ -0,0 +1 @@
+pip