Skip to content

Commit b18692d

Browse files
committed
fix: defer load_model() to first invocation for Container build compatibility
In Container builds, the Python module loads at container startup before any request context exists. API-key-based providers use @requires_api_key which needs a workload access token only available within request context, causing ValueError at import time. Defer load_model() using lazy initialization in all affected templates: - Strands: get_or_create_agent() singleton - LangChain/LangGraph: get_or_create_model() singleton - Google ADK / OpenAI Agents: ensure_credentials_loaded() guard
1 parent 90c14d4 commit b18692d

File tree

5 files changed

+86
-36
lines changed

5 files changed

+86
-36
lines changed

src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap

Lines changed: 43 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1554,12 +1554,18 @@ def add_numbers(a: int, b: int) -> int:
15541554
return a + b
15551555
15561556
1557-
# Set environment variables for model authentication
1558-
load_model()
1559-
15601557
# Get MCP Toolset
15611558
mcp_toolset = [get_streamable_http_mcp_client()]
15621559
1560+
_credentials_loaded = False
1561+
1562+
def ensure_credentials_loaded():
1563+
global _credentials_loaded
1564+
if not _credentials_loaded:
1565+
load_model()
1566+
_credentials_loaded = True
1567+
1568+
15631569
# Agent Definition
15641570
agent = Agent(
15651571
model=MODEL_ID,
@@ -1572,6 +1578,7 @@ agent = Agent(
15721578
15731579
# Session and Runner
15741580
async def setup_session_and_runner(user_id, session_id):
1581+
ensure_credentials_loaded()
15751582
session_service = InMemorySessionService()
15761583
session = await session_service.create_session(
15771584
app_name=APP_NAME, user_id=user_id, session_id=session_id
@@ -1807,8 +1814,13 @@ from mcp_client.client import get_streamable_http_mcp_client
18071814
app = BedrockAgentCoreApp()
18081815
log = app.logger
18091816
1810-
# Instantiate model
1811-
llm = load_model()
1817+
_llm = None
1818+
1819+
def get_or_create_model():
1820+
global _llm
1821+
if _llm is None:
1822+
_llm = load_model()
1823+
return _llm
18121824
18131825
18141826
# Define a simple function tool
@@ -1833,7 +1845,7 @@ async def invoke(payload, context):
18331845
mcp_tools = await mcp_client.get_tools()
18341846
18351847
# Define the agent using create_react_agent
1836-
graph = create_react_agent(llm, tools=mcp_tools + tools)
1848+
graph = create_react_agent(get_or_create_model(), tools=mcp_tools + tools)
18371849
18381850
# Process the user prompt
18391851
prompt = payload.get("prompt", "What can you help me with?")
@@ -2138,12 +2150,17 @@ from mcp_client.client import get_streamable_http_mcp_client
21382150
app = BedrockAgentCoreApp()
21392151
log = app.logger
21402152
2141-
# Set environment variables for model authentication
2142-
load_model()
2143-
21442153
# Get MCP Server
21452154
mcp_server = get_streamable_http_mcp_client()
21462155
2156+
_credentials_loaded = False
2157+
2158+
def ensure_credentials_loaded():
2159+
global _credentials_loaded
2160+
if not _credentials_loaded:
2161+
load_model()
2162+
_credentials_loaded = True
2163+
21472164
21482165
# Define a simple function tool
21492166
@function_tool
@@ -2154,6 +2171,7 @@ def add_numbers(a: int, b: int) -> int:
21542171
21552172
# Define the agent execution
21562173
async def main(query):
2174+
ensure_credentials_loaded()
21572175
try:
21582176
async with mcp_server as server:
21592177
active_servers = [server] if server else []
@@ -2405,14 +2423,19 @@ def agent_factory():
24052423
return get_or_create_agent
24062424
get_or_create_agent = agent_factory()
24072425
{{else}}
2408-
# Create agent
2409-
agent = Agent(
2410-
model=load_model(),
2411-
system_prompt="""
2412-
You are a helpful assistant. Use tools when appropriate.
2413-
""",
2414-
tools=tools+[mcp_client]
2415-
)
2426+
_agent = None
2427+
2428+
def get_or_create_agent():
2429+
global _agent
2430+
if _agent is None:
2431+
_agent = Agent(
2432+
model=load_model(),
2433+
system_prompt="""
2434+
You are a helpful assistant. Use tools when appropriate.
2435+
""",
2436+
tools=tools+[mcp_client]
2437+
)
2438+
return _agent
24162439
{{/if}}
24172440
24182441
@@ -2424,8 +2447,10 @@ async def invoke(payload, context):
24242447
session_id = getattr(context, 'session_id', 'default-session')
24252448
user_id = getattr(context, 'user_id', 'default-user')
24262449
agent = get_or_create_agent(session_id, user_id)
2427-
2450+
{{else}}
2451+
agent = get_or_create_agent()
24282452
{{/if}}
2453+
24292454
# Execute and format response
24302455
stream = agent.stream_async(payload.get("prompt"))
24312456

src/assets/python/googleadk/base/main.py

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,12 +22,18 @@ def add_numbers(a: int, b: int) -> int:
2222
return a + b
2323

2424

25-
# Set environment variables for model authentication
26-
load_model()
27-
2825
# Get MCP Toolset
2926
mcp_toolset = [get_streamable_http_mcp_client()]
3027

28+
_credentials_loaded = False
29+
30+
def ensure_credentials_loaded():
31+
global _credentials_loaded
32+
if not _credentials_loaded:
33+
load_model()
34+
_credentials_loaded = True
35+
36+
3137
# Agent Definition
3238
agent = Agent(
3339
model=MODEL_ID,
@@ -40,6 +46,7 @@ def add_numbers(a: int, b: int) -> int:
4046

4147
# Session and Runner
4248
async def setup_session_and_runner(user_id, session_id):
49+
ensure_credentials_loaded()
4350
session_service = InMemorySessionService()
4451
session = await session_service.create_session(
4552
app_name=APP_NAME, user_id=user_id, session_id=session_id

src/assets/python/langchain_langgraph/base/main.py

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,13 @@
99
app = BedrockAgentCoreApp()
1010
log = app.logger
1111

12-
# Instantiate model
13-
llm = load_model()
12+
_llm = None
13+
14+
def get_or_create_model():
15+
global _llm
16+
if _llm is None:
17+
_llm = load_model()
18+
return _llm
1419

1520

1621
# Define a simple function tool
@@ -35,7 +40,7 @@ async def invoke(payload, context):
3540
mcp_tools = await mcp_client.get_tools()
3641

3742
# Define the agent using create_react_agent
38-
graph = create_react_agent(llm, tools=mcp_tools + tools)
43+
graph = create_react_agent(get_or_create_model(), tools=mcp_tools + tools)
3944

4045
# Process the user prompt
4146
prompt = payload.get("prompt", "What can you help me with?")

src/assets/python/openaiagents/base/main.py

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,17 @@
77
app = BedrockAgentCoreApp()
88
log = app.logger
99

10-
# Set environment variables for model authentication
11-
load_model()
12-
1310
# Get MCP Server
1411
mcp_server = get_streamable_http_mcp_client()
1512

13+
_credentials_loaded = False
14+
15+
def ensure_credentials_loaded():
16+
global _credentials_loaded
17+
if not _credentials_loaded:
18+
load_model()
19+
_credentials_loaded = True
20+
1621

1722
# Define a simple function tool
1823
@function_tool
@@ -23,6 +28,7 @@ def add_numbers(a: int, b: int) -> int:
2328

2429
# Define the agent execution
2530
async def main(query):
31+
ensure_credentials_loaded()
2632
try:
2733
async with mcp_server as server:
2834
active_servers = [server] if server else []

src/assets/python/strands/base/main.py

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -42,14 +42,19 @@ def get_or_create_agent(session_id, user_id):
4242
return get_or_create_agent
4343
get_or_create_agent = agent_factory()
4444
{{else}}
45-
# Create agent
46-
agent = Agent(
47-
model=load_model(),
48-
system_prompt="""
49-
You are a helpful assistant. Use tools when appropriate.
50-
""",
51-
tools=tools+[mcp_client]
52-
)
45+
_agent = None
46+
47+
def get_or_create_agent():
48+
global _agent
49+
if _agent is None:
50+
_agent = Agent(
51+
model=load_model(),
52+
system_prompt="""
53+
You are a helpful assistant. Use tools when appropriate.
54+
""",
55+
tools=tools+[mcp_client]
56+
)
57+
return _agent
5358
{{/if}}
5459

5560

@@ -61,8 +66,10 @@ async def invoke(payload, context):
6166
session_id = getattr(context, 'session_id', 'default-session')
6267
user_id = getattr(context, 'user_id', 'default-user')
6368
agent = get_or_create_agent(session_id, user_id)
64-
69+
{{else}}
70+
agent = get_or_create_agent()
6571
{{/if}}
72+
6673
# Execute and format response
6774
stream = agent.stream_async(payload.get("prompt"))
6875

0 commit comments

Comments
 (0)