Created
March 27, 2026 16:55
-
-
Save clutchski/3ce63c17cc36cac4eda65594732033ca to your computer and use it in GitHub Desktop.
ADK Python + Braintrust Gateway: LlmAgent with LiteLLM, parallel tool use across providers
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| """Test ADK LlmAgent with LiteLLM through the Braintrust gateway. | |
| Uses LlmAgent (not raw model calls) with tools, routing all providers | |
| through the Braintrust gateway via LiteLLM's api_base parameter. | |
| """ | |
| import asyncio | |
| import datetime | |
| import os | |
| from zoneinfo import ZoneInfo | |
| from google.adk.agents import LlmAgent | |
| from google.adk.models.lite_llm import LiteLlm | |
| from google.adk.runners import Runner | |
| from google.adk.sessions import InMemorySessionService | |
| from google.genai import types | |
| GATEWAY_URL = "https://gateway.braintrust.dev" | |
| BT_API_KEY = os.environ["BRAINTRUST_API_KEY"] | |
| MODELS = [ | |
| "openai/gpt-4o-mini", | |
| "anthropic/claude-haiku-4-5", | |
| ] | |
| # -- Tools ------------------------------------------------------------------ | |
| def get_weather(city: str) -> dict: | |
| """Retrieves the current weather for a given city.""" | |
| weather_data = { | |
| "new york": "Sunny, 24C (75F), light breeze from the west.", | |
| "london": "Overcast, 14C (57F), chance of rain.", | |
| "tokyo": "Partly cloudy, 28C (82F), humid.", | |
| "san francisco": "Foggy, 16C (61F), clearing by afternoon.", | |
| } | |
| report = weather_data.get(city.lower()) | |
| if report: | |
| return {"result": f"Weather in {city}: {report}"} | |
| return {"result": f"No weather data for '{city}'."} | |
| def get_current_time(city: str) -> dict: | |
| """Returns the current time in a given city.""" | |
| tz_map = { | |
| "new york": "America/New_York", | |
| "london": "Europe/London", | |
| "tokyo": "Asia/Tokyo", | |
| "san francisco": "America/Los_Angeles", | |
| } | |
| tz_name = tz_map.get(city.lower()) | |
| if tz_name: | |
| now = datetime.datetime.now(ZoneInfo(tz_name)) | |
| return {"result": f"Current time in {city}: {now.strftime('%Y-%m-%d %H:%M:%S %Z')}"} | |
| return {"result": f"Unknown timezone for '{city}'."} | |
| # -- Run -------------------------------------------------------------------- | |
| async def run_agent(model_name: str): | |
| print(f"\n{'='*60}") | |
| print(f"Model: {model_name}") | |
| print(f"{'='*60}") | |
| agent = LlmAgent( | |
| name="bt_gateway_agent", | |
| model=LiteLlm( | |
| model=model_name, | |
| api_base=GATEWAY_URL, | |
| api_key=BT_API_KEY, | |
| ), | |
| instruction="You are a helpful assistant. Be concise.", | |
| tools=[get_weather, get_current_time], | |
| ) | |
| session_service = InMemorySessionService() | |
| runner = Runner(agent=agent, app_name="test", session_service=session_service) | |
| session = await session_service.create_session(app_name="test", user_id="test-user") | |
| query = "What's the weather and time in San Francisco?" | |
| print(f"\n--- User: {query}") | |
| content = types.Content( | |
| role="user", parts=[types.Part.from_text(text=query)] | |
| ) | |
| try: | |
| async for event in runner.run_async( | |
| user_id="test-user", session_id=session.id, new_message=content | |
| ): | |
| if event.content and event.content.parts: | |
| for part in event.content.parts: | |
| if part.text: | |
| print(f" [text] {part.text}") | |
| elif part.function_call: | |
| print(f" [tool call] {part.function_call.name}({dict(part.function_call.args)})") | |
| elif part.function_response: | |
| print(f" [tool result] {part.function_response.name} -> {dict(part.function_response.response)}") | |
| if event.is_final_response(): | |
| final = "" | |
| if event.content and event.content.parts: | |
| final = " ".join(p.text for p in event.content.parts if p.text) | |
| print(f"--- Agent: {final}") | |
| except Exception as e: | |
| print(f"--- ERROR: {e}") | |
| async def main(): | |
| for model in MODELS: | |
| await run_agent(model) | |
| if __name__ == "__main__": | |
| asyncio.run(main()) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment