Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/google_endpoints/agents_endpoints.py: 74%
94 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""
2Google AI Studio Managed Agents API Proxy Endpoints.
4Exposes Gemini's /v1beta/agents surface through the LiteLLM proxy so that
5user curl commands transfer 1-to-1 by swapping the host + auth header.
7Routes:
8 POST /v1beta/agents -> acreate_agent
9 GET /v1beta/agents -> alist_agents
10 GET /v1beta/agents/{name} -> aget_agent
11 DELETE /v1beta/agents/{name} -> adelete_agent
12 GET /v1beta/agents/{name}/versions -> alist_agent_versions
14These are distinct from the A2A agent registry at /v1/agents.
15"""
17import json
18from typing import Final
20from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
21from fastapi.responses import ORJSONResponse
23from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
24from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
25from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
26from litellm.proxy.common_utils.http_parsing_utils import (
27 _read_request_body,
28 _safe_get_request_query_params,
29)
31router: Final = APIRouter(tags=["gemini managed agents"])
34def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool:
35 return (
36 user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
37 or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
38 )
41def _enforce_caller_supplied_provider_key(
42 data: dict,
43 user_api_key_dict: UserAPIKeyAuth,
44) -> None:
45 """
46 SECURITY: refuse to use the proxy's shared GOOGLE_API_KEY / GEMINI_API_KEY
47 env fallback for non-admin callers on Gemini managed-agent CRUD endpoints.
49 These endpoints are part of ``llm_api_routes`` so any authenticated LLM key
50 can reach them, but unlike ``/v1beta/models/...:generateContent`` they are
51 *not* routed through ``model_list`` — the only credential source is either
52 the per-request ``litellm_params_template`` or the env var fallback. Without
53 this guard, any ordinary proxy user could list, create, or delete managed
54 agents inside the operator's Gemini project using the operator's key.
56 Proxy admins (master key) keep the env-fallback convenience for ops use.
57 """
58 if _is_proxy_admin(user_api_key_dict): 58 ↛ 60line 58 didn't jump to line 60 because the condition on line 58 was always true
59 return
60 if data.get("api_key"):
61 return
62 raise HTTPException(
63 status_code=status.HTTP_401_UNAUTHORIZED,
64 detail=(
65 "Gemini managed-agent endpoints require a caller-supplied "
66 "Gemini api_key (via 'litellm_params_template'). Falling back to "
67 "the proxy's GOOGLE_API_KEY / GEMINI_API_KEY env vars is only "
68 "permitted for proxy admins."
69 ),
70 )
73def _merge_query_params_into_data(data: dict, request: Request) -> dict:
74 """
75 For GET/DELETE endpoints that cannot carry a JSON body, read a
76 JSON-encoded ``litellm_params_template`` query parameter and merge its
77 contents into *data*, without overwriting keys that are already present
78 (e.g. path params like ``name`` or the fixed ``custom_llm_provider``).
80 This mirrors the ``litellm_params_template`` handling in
81 ``create_gemini_agent`` and is the supported way for multi-tenant
82 callers to supply per-request credentials on non-POST endpoints:
84 .. code-block:: bash
86 curl "http://localhost:4000/v1beta/agents?litellm_params_template=%7B%22api_key%22%3A%22AIza...%22%7D" \\
87 -H "Authorization: Bearer sk-..."
89 Credentials MUST NOT be passed as plain flat query parameters (e.g.
90 ``?api_key=AIza...``) because URL query strings appear verbatim in
91 web-server access logs, CDN edge logs, browser history, and Referer
92 headers. Use the ``litellm_params_template`` JSON body field on POST
93 requests, or the JSON-encoded query parameter above for GET/DELETE.
94 """
95 query_params: Final = _safe_get_request_query_params(request)
96 if not query_params: 96 ↛ 99line 96 didn't jump to line 99 because the condition on line 96 was always true
97 return data
99 raw_template: Final = query_params.get("litellm_params_template")
100 if raw_template:
101 try:
102 template = json.loads(raw_template) if isinstance(raw_template, str) else raw_template
103 except (json.JSONDecodeError, ValueError):
104 template = {}
105 if isinstance(template, dict):
106 for key, value in template.items():
107 data.setdefault(key, value)
109 return data
112def _proxy_server_imports():
113 from litellm.proxy.proxy_server import ( # noqa: PLC0415
114 general_settings,
115 llm_router,
116 proxy_config,
117 proxy_logging_obj,
118 select_data_generator,
119 user_api_base,
120 user_max_tokens,
121 user_model,
122 user_request_timeout,
123 user_temperature,
124 version,
125 )
127 return dict(
128 general_settings=general_settings,
129 llm_router=llm_router,
130 proxy_config=proxy_config,
131 proxy_logging_obj=proxy_logging_obj,
132 select_data_generator=select_data_generator,
133 user_api_base=user_api_base,
134 user_max_tokens=user_max_tokens,
135 user_model=user_model,
136 user_request_timeout=user_request_timeout,
137 user_temperature=user_temperature,
138 version=version,
139 )
142@router.post(
143 "/v1beta/agents",
144 dependencies=[Depends(user_api_key_auth)],
145 response_class=ORJSONResponse,
146)
147async def create_gemini_agent(
148 request: Request,
149 fastapi_response: Response,
150 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
151):
152 """
153 Create a named custom agent on the Gemini side.
155 Example:
156 ```bash
157 curl -X POST "http://localhost:4000/v1beta/agents" \\
158 -H "Authorization: Bearer sk-..." \\
159 -H "Content-Type: application/json" \\
160 -d '{
161 "name": "my-custom-slides-agent",
162 "base_agent": "waverunner",
163 "instructions": "You are a helpful assistant that creates slides.",
164 "base_environment": {
165 "type": "remote",
166 "sources": [
167 {"type": "gcs", "source": "gs://eap-templates/slides-skill",
168 "target": "/.agents/skills/slides-skill"}
169 ]
170 }
171 }'
172 ```
173 """
174 srv: Final = _proxy_server_imports()
175 data: Final = await _read_request_body(request=request)
176 # Merge litellm_params_template (e.g. custom_llm_provider, api_key) into the request
177 litellm_params_template: Final = data.pop("litellm_params_template", None) or {}
178 if isinstance(litellm_params_template, dict): 178 ↛ 182line 178 didn't jump to line 182 because the condition on line 178 was always true
179 for key, value in litellm_params_template.items(): 179 ↛ 180line 179 didn't jump to line 180 because the loop on line 179 never started
180 if key not in data:
181 data[key] = value
182 data.setdefault("custom_llm_provider", "gemini")
183 _enforce_caller_supplied_provider_key(data, user_api_key_dict)
185 processor: Final = ProxyBaseLLMRequestProcessing(data=data)
186 try:
187 return await processor.base_process_llm_request(
188 request=request,
189 fastapi_response=fastapi_response,
190 user_api_key_dict=user_api_key_dict,
191 route_type="acreate_agent",
192 proxy_logging_obj=srv["proxy_logging_obj"],
193 llm_router=srv["llm_router"],
194 general_settings=srv["general_settings"],
195 proxy_config=srv["proxy_config"],
196 select_data_generator=srv["select_data_generator"],
197 model=None,
198 user_model=srv["user_model"],
199 user_temperature=srv["user_temperature"],
200 user_request_timeout=srv["user_request_timeout"],
201 user_max_tokens=srv["user_max_tokens"],
202 user_api_base=srv["user_api_base"],
203 version=srv["version"],
204 )
205 except Exception as e:
206 raise await processor._handle_llm_api_exception(
207 e=e,
208 user_api_key_dict=user_api_key_dict,
209 proxy_logging_obj=srv["proxy_logging_obj"],
210 version=srv["version"],
211 )
214@router.get(
215 "/v1beta/agents",
216 dependencies=[Depends(user_api_key_auth)],
217 response_class=ORJSONResponse,
218)
219async def list_gemini_agents(
220 request: Request,
221 fastapi_response: Response,
222 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
223):
224 """
225 List all custom agents on the Gemini side.
227 Pass per-request Gemini credentials via the JSON-encoded
228 ``litellm_params_template`` query parameter. Flat query parameters
229 (e.g. ``?api_key=AIza...``) are intentionally ignored — see
230 ``_merge_query_params_into_data`` for the rationale.
232 ```bash
233 curl "http://localhost:4000/v1beta/agents?litellm_params_template=%7B%22api_key%22%3A%22AIza...%22%7D" \\
234 -H "Authorization: Bearer sk-..."
235 ```
236 """
237 srv: Final = _proxy_server_imports()
238 data: Final[dict] = {"custom_llm_provider": "gemini"}
239 _merge_query_params_into_data(data, request)
240 _enforce_caller_supplied_provider_key(data, user_api_key_dict)
242 processor: Final = ProxyBaseLLMRequestProcessing(data=data)
243 try:
244 return await processor.base_process_llm_request(
245 request=request,
246 fastapi_response=fastapi_response,
247 user_api_key_dict=user_api_key_dict,
248 route_type="alist_agents",
249 proxy_logging_obj=srv["proxy_logging_obj"],
250 llm_router=srv["llm_router"],
251 general_settings=srv["general_settings"],
252 proxy_config=srv["proxy_config"],
253 select_data_generator=srv["select_data_generator"],
254 model=None,
255 user_model=srv["user_model"],
256 user_temperature=srv["user_temperature"],
257 user_request_timeout=srv["user_request_timeout"],
258 user_max_tokens=srv["user_max_tokens"],
259 user_api_base=srv["user_api_base"],
260 version=srv["version"],
261 )
262 except Exception as e:
263 raise await processor._handle_llm_api_exception(
264 e=e,
265 user_api_key_dict=user_api_key_dict,
266 proxy_logging_obj=srv["proxy_logging_obj"],
267 version=srv["version"],
268 )
271@router.get(
272 "/v1beta/agents/{name}",
273 dependencies=[Depends(user_api_key_auth)],
274 response_class=ORJSONResponse,
275)
276async def get_gemini_agent(
277 request: Request,
278 name: str,
279 fastapi_response: Response,
280 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
281):
282 """
283 Get a specific custom agent by name.
285 Pass per-request Gemini credentials via the JSON-encoded
286 ``litellm_params_template`` query parameter. Flat query parameters
287 (e.g. ``?api_key=AIza...``) are intentionally ignored — see
288 ``_merge_query_params_into_data`` for the rationale.
290 ```bash
291 curl "http://localhost:4000/v1beta/agents/my-custom-slides-agent?litellm_params_template=%7B%22api_key%22%3A%22AIza...%22%7D" \\
292 -H "Authorization: Bearer sk-..."
293 ```
294 """
295 srv: Final = _proxy_server_imports()
296 data: Final = {"name": name, "custom_llm_provider": "gemini"}
297 _merge_query_params_into_data(data, request)
298 _enforce_caller_supplied_provider_key(data, user_api_key_dict)
300 processor: Final = ProxyBaseLLMRequestProcessing(data=data)
301 try:
302 return await processor.base_process_llm_request(
303 request=request,
304 fastapi_response=fastapi_response,
305 user_api_key_dict=user_api_key_dict,
306 route_type="aget_agent",
307 proxy_logging_obj=srv["proxy_logging_obj"],
308 llm_router=srv["llm_router"],
309 general_settings=srv["general_settings"],
310 proxy_config=srv["proxy_config"],
311 select_data_generator=srv["select_data_generator"],
312 model=None,
313 user_model=srv["user_model"],
314 user_temperature=srv["user_temperature"],
315 user_request_timeout=srv["user_request_timeout"],
316 user_max_tokens=srv["user_max_tokens"],
317 user_api_base=srv["user_api_base"],
318 version=srv["version"],
319 )
320 except Exception as e:
321 raise await processor._handle_llm_api_exception(
322 e=e,
323 user_api_key_dict=user_api_key_dict,
324 proxy_logging_obj=srv["proxy_logging_obj"],
325 version=srv["version"],
326 )
329@router.delete(
330 "/v1beta/agents/{name}",
331 dependencies=[Depends(user_api_key_auth)],
332 response_class=ORJSONResponse,
333)
334async def delete_gemini_agent(
335 request: Request,
336 name: str,
337 fastapi_response: Response,
338 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
339):
340 """
341 Delete a custom agent by name.
343 Pass per-request Gemini credentials via the JSON-encoded
344 ``litellm_params_template`` query parameter. Flat query parameters
345 (e.g. ``?api_key=AIza...``) are intentionally ignored — see
346 ``_merge_query_params_into_data`` for the rationale.
348 ```bash
349 curl -X DELETE "http://localhost:4000/v1beta/agents/my-custom-slides-agent?litellm_params_template=%7B%22api_key%22%3A%22AIza...%22%7D" \\
350 -H "Authorization: Bearer sk-..."
351 ```
352 """
353 srv: Final = _proxy_server_imports()
354 data: Final = {"name": name, "custom_llm_provider": "gemini"}
355 _merge_query_params_into_data(data, request)
356 _enforce_caller_supplied_provider_key(data, user_api_key_dict)
358 processor: Final = ProxyBaseLLMRequestProcessing(data=data)
359 try:
360 return await processor.base_process_llm_request(
361 request=request,
362 fastapi_response=fastapi_response,
363 user_api_key_dict=user_api_key_dict,
364 route_type="adelete_agent",
365 proxy_logging_obj=srv["proxy_logging_obj"],
366 llm_router=srv["llm_router"],
367 general_settings=srv["general_settings"],
368 proxy_config=srv["proxy_config"],
369 select_data_generator=srv["select_data_generator"],
370 model=None,
371 user_model=srv["user_model"],
372 user_temperature=srv["user_temperature"],
373 user_request_timeout=srv["user_request_timeout"],
374 user_max_tokens=srv["user_max_tokens"],
375 user_api_base=srv["user_api_base"],
376 version=srv["version"],
377 )
378 except Exception as e:
379 raise await processor._handle_llm_api_exception(
380 e=e,
381 user_api_key_dict=user_api_key_dict,
382 proxy_logging_obj=srv["proxy_logging_obj"],
383 version=srv["version"],
384 )
387@router.get(
388 "/v1beta/agents/{name}/versions",
389 dependencies=[Depends(user_api_key_auth)],
390 response_class=ORJSONResponse,
391)
392async def list_gemini_agent_versions(
393 request: Request,
394 name: str,
395 fastapi_response: Response,
396 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
397):
398 """
399 List versions of a custom agent.
401 Pass per-request Gemini credentials via the JSON-encoded
402 ``litellm_params_template`` query parameter. Flat query parameters
403 (e.g. ``?api_key=AIza...``) are intentionally ignored — see
404 ``_merge_query_params_into_data`` for the rationale.
406 ```bash
407 curl "http://localhost:4000/v1beta/agents/my-custom-slides-agent/versions?litellm_params_template=%7B%22api_key%22%3A%22AIza...%22%7D" \\
408 -H "Authorization: Bearer sk-..."
409 ```
410 """
411 srv: Final = _proxy_server_imports()
412 data: Final = {"name": name, "custom_llm_provider": "gemini"}
413 _merge_query_params_into_data(data, request)
414 _enforce_caller_supplied_provider_key(data, user_api_key_dict)
416 processor: Final = ProxyBaseLLMRequestProcessing(data=data)
417 try:
418 return await processor.base_process_llm_request(
419 request=request,
420 fastapi_response=fastapi_response,
421 user_api_key_dict=user_api_key_dict,
422 route_type="alist_agent_versions",
423 proxy_logging_obj=srv["proxy_logging_obj"],
424 llm_router=srv["llm_router"],
425 general_settings=srv["general_settings"],
426 proxy_config=srv["proxy_config"],
427 select_data_generator=srv["select_data_generator"],
428 model=None,
429 user_model=srv["user_model"],
430 user_temperature=srv["user_temperature"],
431 user_request_timeout=srv["user_request_timeout"],
432 user_max_tokens=srv["user_max_tokens"],
433 user_api_base=srv["user_api_base"],
434 version=srv["version"],
435 )
436 except Exception as e:
437 raise await processor._handle_llm_api_exception(
438 e=e,
439 user_api_key_dict=user_api_key_dict,
440 proxy_logging_obj=srv["proxy_logging_obj"],
441 version=srv["version"],
442 )