8000 ADK Web Async Agent Compatibility · google/adk-python@23a4eaa · GitHub
[go: up one dir, main page]

Skip to content

Commit 23a4eaa

Browse files
author
fbiebl
committed
ADK Web Async Agent Compatibility
1 parent a4d432a commit 23a4eaa

File tree

7 files changed

+433
-7
lines changed

7 files changed

+433
-7
lines changed
Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
# ADK Web Async Agent Compatibility
2+
3+
This minimal example demonstrates the technical foundation for making async agents compatible with ADK Web interface.
4+
5+
## Problem Solved
6+
7+
ADK Web had two main compatibility issues with async agents:
8+
9+
1. **MCP Tools Event Loop Conflicts**: "Event loop is closed" errors with uvloop
10+
2. **Session State Customization**: No way to load custom data before template processing
11+
12+
## Solution
13+
14+
The ADK now provides:
15+
16+
1. **Automatic MCP uvloop compatibility** in the MCP tool implementation
17+
2. **Session state preprocessor callback** for custom session data
18+
19+
## Usage
20+
21+
Create an agent with optional `session_state_preprocessor` function:
22+
23+
```python
24+
# agent.py
25+
async def session_state_preprocessor(state):
26+
"""Called before template processing - load custom data here"""
27+
# Load user data, project context, etc.
28+
return state
29+
30+
def create_adk_web_agent():
31+
"""Standard ADK Web agent factory function"""
32+
return your_async_agent
33+
```
34+
35+
## Key Features
36+
37+
- **MCP Tools**: Work automatically in ADK Web (uvloop compatibility handled by ADK)
38+
- **Session Preprocessor**: Load database data, set defaults, add custom variables
39+
- **Template Variables**: Use standard ADK `{variable}` format with your custom data
40+
- **ADK Web Detection**: ADK Web calls `create_adk_web_agent()` and provides `user_id="1"` default
41+
42+
## Files
43+
44+
- `README.md` - This documentation
45+
- `agent.py` - Minimal async agent with preprocessor example
46+
- `main.py` - Simple test script
47+
48+
## Test
49+
50+
```bash
51+
# Test the agent
52+
python main.py
53+
54+
# Use with ADK Web
55+
adk web --port 8088
56+
```
57+
58+
That's it! The ADK handles all the complexity automatically.
Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
"""
2+
ADK Web Async Agent Compatibility Example.
3+
4+
Minimal example demonstrating the technical foundation for
5+
making async agents compatible with ADK Web interface.
6+
"""
Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
"""
2+
Minimal ADK Web Async Agent Example.
3+
4+
Demonstrates the technical foundation for ADK Web compatibility:
5+
1. MCP tools work automatically (uvloop compatibility handled by ADK)
6+
2. Session state preprocessor for custom data before template processing
7+
"""
8+
9+
import asyncio
10+
import logging
11+
from typing import Dict, Any, Optional
12+
13+
from google.adk.agents.llm_agent import LlmAgent
14+
from google.adk.models.lite_llm import LiteLlm
15+
16+
logger = logging.getLogger(__name__)
17+
18+
19+
async def session_state_preprocessor(state: Dict[str, Any]) -> Dict[str, Any]:
20+
"""
21+
Session state preprocessor - called by ADK Web before template processing.
22+
23+
This is where you load custom data that becomes available in templates.
24+
25+
Args:
26+
state: Current session state (includes user_id="1" from ADK Web)
27+
28+
Returns:
29+
Enhanced session state with your custom data
30+
"""
31+
logger.info("🔧 Session preprocessor called")
32+
33+
# Example: Load user data based on user_id
34+
user_id = state.get("user_id")
35+
if user_id:
36+
# In real app: user_data = await load_from_database(user_id)
37+
state["user_name"] = f"Demo User {user_id}"
38+
state["user_role"] = "developer"
39+
40+
# Example: Add application context
41+
state["app_name"] = "ADK Web Demo"
42+
state["features"] = "MCP tools, async operations"
43+
44+
logger.info(f"✅ Session enhanced: {list(state.keys())}")
45+
return state
46+
47+
48+
async def create_async_agent() -> LlmAgent:
49+
"""Create the async agent with MCP tools."""
50+
51+
# Simple instruction using variables from session_state_preprocessor
52+
instruction = """
53+
You are an ADK Web compatible async agent.
54+
55+
User: {user_name} (ID: {user_id}, Role: {user_role})
56+
App: {app_name}
57+
Features: {features}
58+
59+
You demonstrate:
60+
- MCP tools working in ADK Web (automatic uvloop compatibility)
61+
- Session state preprocessor for custom template data
62+
- Standard ADK template variables {variable}
63+
"""
64+
65+
agent = LlmAgent(
66+
name="adk_web_async_demo",
67+
model=LiteLlm(model="openai/gpt-4o", stream=True),
68+
instruction=instruction,
69+
description="Minimal example of ADK Web async compatibility",
70+
tools=[] # MCP tools would be added here automatically
71+
)
72+
73+
logger.info("✅ Async agent created for ADK Web")
74+
return agent
75+
76+
77+
def create_adk_web_agent() -> Optional[LlmAgent]:
78+
"""
79+
ADK Web entry point - called by ADK Web interface.
80+
81+
This function is called by ADK Web when it loads the agent.
82+
The mere fact that this function is called indicates ADK Web mode.
83+
84+
Returns:
85+
Agent instance for ADK Web usage
86+
"""
87+
try:
88+
logger.info("🌐 ADK Web agent creation requested...")
89+
90+
# Handle async creation in sync context
91+
try:
92+
loop = asyncio.get_running_loop()
93+
# If event loop exists, use thread executor
94+
import concurrent.futures
95+
96+
def run_creation():
97+
new_loop = asyncio.new_event_loop()
98+
asyncio.set_event_loop(new_loop)
99+
try:
100+
return new_loop.run_until_complete(create_async_agent())
101+
finally:
102+
new_loop.close()
103+
104+
with concurrent.futures.ThreadPoolExecutor() as executor:
105+
agent = executor.submit(run_creation).result()
106+
107+
except RuntimeError:
108+
# No event loop, use asyncio.run
109+
agent = asyncio.run(create_async_agent())
110+
111+
logger.info("✅ ADK Web 10BC0 agent created successfully")
112+
return agent
113+
114+
except Exception as e:
115+
logger.error(f"❌ ADK Web agent creation failed: {e}")
116+
return None
117+
118+
119+
# Export for ADK Web
120+
__all__ = ["create_adk_web_agent", "session_state_preprocessor"]
Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
"""
2+
Test script for ADK Web async agent compatibility.
3+
4+
Simple test to verify the agent works correctly.
5+
"""
6+
7+
import asyncio
8+
import logging
9+
import os
10+
11+
from agent import create_adk_web_agent, session_state_preprocessor
12+
13+
# Setup logging
14+
logging.basicConfig(level=logging.INFO)
15+
logger = logging.getLogger(__name__)
16+
17+
18+
async def test_session_preprocessor():
19+
"""Test the session state preprocessor."""
20+
logger.info("🧪 Testing session state preprocessor...")
21+
22+
# Simulate ADK Web session state
23+
test_state = {"user_id": "1"} # ADK Web default
24+
25+
# Call preprocessor
26+
enhanced_state = await session_state_preprocessor(test_state)
27+
28+
# Verify enhancement
29+
expected_keys = ["user_id", "user_name", "user_role", "app_name", "features"]
30+
for key in expected_keys:
31+
assert key in enhanced_state, f"Missing key: {key}"
32+
33+
logger.info(f"✅ Session state enhanced: {enhanced_state}")
34+
35+
36+
def test_agent_creation():
37+
"""Test agent creation for ADK Web."""
38+
logger.info("🧪 Testing agent creation...")
39+
40+
# Create agent (simulates ADK Web calling create_adk_web_agent)
41+
agent = create_adk_web_agent()
42+
43+
assert agent is not None, "Agent creation failed"
44+
assert agent.name == "adk_web_async_demo"
45+
46+
logger.info(f"✅ Agent created: {agent.name}")
47+
return agent
48+
49+
50+
async def main():
51+
"""Run all tests."""
52+
logger.info("🚀 Starting ADK Web async compatibility tests...")
53+
54+
try:
55+
# Test session preprocessor
56+
await test_session_preprocessor()
57+
58+
# Test agent creation
59+
agent = test_agent_creation()
60+
61+
logger.info("🎉 All tests passed!")
62+
logger.info("💡 Ready for ADK Web: adk web --port 8088")
63+
64+
except Exception as e:
65+
logger.error(f"❌ Test failed: {e}")
66+
raise
67+
68+
69+
if __name__ == "__main__":
70+
asyncio.run(main())

src/google/adk/cli/fast_api.py

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,53 @@
8787

8888
_EVAL_SET_FILE_EXTENSION = ".evalset.json"
8989

90+
# Global agents directory for session state preprocessor access
91+
_agents_dir = None
92+
93+
94+
async def _apply_session_state_preprocessor(
95+
app_name: str,
96+
state: dict[str, Any]
97+
) -> dict[str, Any]:
98+
"""
99+
Apply session state preprocessor callback if agent provides one.
100+
101+
This allows agents to customize session state before template processing.
102+
Useful for loading user data, setting project defaults, etc.
103+
104+
Args:
105+
app_name: Name of the agent/app
106+
state: Current session state
107+
108+
Returns:
109+
Updated session state with agent-specific modifications
110+
"""
111+
try:
112+
# Try to import and call agent's session preprocessor
113+
# Use the global agents_dir that was set during app initialization
114+
if _agents_dir:
115+
agent_module = AgentLoader(_agents_dir).load_agent(app_name)
116+
else:
117+
# Fallback: try to load without agents_dir (may fail)
118+
agent_module = AgentLoader().load_agent(app_name)
119+
120+
# Look for session_state_preprocessor function
121+
if hasattr(agent_module, 'session_state_preprocessor'):
122+
preprocessor = getattr(agent_module, 'session_state_preprocessor')
123+
124+
if callable(preprocessor):
125+
# Call preprocessor (support both sync and async)
126+
if asyncio.iscoroutinefunction(preprocessor):
127+
state = await preprocessor(state)
128+
else:
129+
state = preprocessor(state)
130+
131+
except Exception as e:
132+
logger.warning(f"Session state preprocessor not available or failed: {e}")
133+
# Continue without preprocessor - this is optional
134+
135+
return state
136+
90137

91138
class ApiServerSpanExporter(export.SpanExporter):
92139

@@ -202,6 +249,10 @@ def get_fast_api_app(
202249
trace_to_cloud: bool = False,
203250
lifespan: Optional[Lifespan[FastAPI]] = None,
204251
) -> FastAPI:
252+
# Store agents_dir globally for session state preprocessor access
253+
global _agents_dir
254+
_agents_dir = agents_dir
255+
205256
# InMemory tracing dict.
206257
trace_dict: dict[str, Any] = {}
207258
session_trace_dict: dict[str, Any] = {}
@@ -400,6 +451,19 @@ async def create_session_with_id(
400451
status_code=400, detail=f"Session already exists: {session_id}"
401452
)
402453
logger.info("New session created: %s", session_id)
454+
455+
# Initialize state with ADK Web development defaults
456+
if state is None:
457+
state = {}
458+
459+
# Call session state preprocessor callback if agent provides one
460+
state = await _apply_session_state_preprocessor(app_name, state)
461+
462+
# Add basic ADK Web development defaults
463+
# These defaults only apply when values aren't already set
464+
if "user_id" not in state:
465+
state["user_id"] = "1" # Default ADK Web user for development
466+
403467
return await session_service.create_session(
404468
app_name=app_name, user_id=user_id, state=state, session_id=session_id
405469
)
@@ -415,6 +479,14 @@ async def create_session(
415479
events: Optional[list[Event]] = None,
416480
) -> Session:
417481
logger.info("New session created")
482+
483+
# Initialize state with ADK Web development defaults
484+
if state is None:
485+
state = {}
486+
487+
# Call session state preprocessor callback if agent provides one
488+
state = await _apply_session_state_preprocessor(app_name, state)
489+
418490
session = await session_service.create_session(
419491
app_name=app_name, user_id=user_id, state=state
420492
)

0 commit comments

Comments
 (0)
0