litellm/litellm/proxy/google_endpoints/agents_endpoints.py

"""
Google AI Studio Managed Agents API Proxy Endpoints.

Exposes Gemini's /v1beta/agents surface through the LiteLLM proxy so that
user curl commands transfer 1-to-1 by swapping the host + auth header.

Routes:
  POST   /v1beta/agents                     -> acreate_agent
  GET    /v1beta/agents                     -> alist_agents
  GET    /v1beta/agents/{name}              -> aget_agent
  DELETE /v1beta/agents/{name}              -> adelete_agent
  GET    /v1beta/agents/{name}/versions     -> alist_agent_versions

These are distinct from the A2A agent registry at /v1/agents.
"""

import json
from typing import Final

from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
from fastapi.responses import ORJSONResponse

from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_utils.http_parsing_utils import (
    _read_request_body,
    _safe_get_request_query_params,
)

router: Final = APIRouter(tags=["gemini managed agents"])


def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool:
    return (
        user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
        or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
    )


def _enforce_caller_supplied_provider_key(
    data: dict,
    user_api_key_dict: UserAPIKeyAuth,
) -> None:
    """
    SECURITY: refuse to use the proxy's shared GOOGLE_API_KEY / GEMINI_API_KEY
    env fallback for non-admin callers on Gemini managed-agent CRUD endpoints.

    These endpoints are part of ``llm_api_routes`` so any authenticated LLM key
    can reach them, but unlike ``/v1beta/models/...:generateContent`` they are
    *not* routed through ``model_list`` — the only credential source is either
    the per-request ``litellm_params_template`` or the env var fallback. Without
    this guard, any ordinary proxy user could list, create, or delete managed
    agents inside the operator's Gemini project using the operator's key.

    Proxy admins (master key) keep the env-fallback convenience for ops use.
    """
    if _is_proxy_admin(user_api_key_dict):
        return
    if data.get("api_key"):
        return
    raise HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail=(
            "Gemini managed-agent endpoints require a caller-supplied "
            "Gemini api_key (via 'litellm_params_template'). Falling back to "
            "the proxy's GOOGLE_API_KEY / GEMINI_API_KEY env vars is only "
            "permitted for proxy admins."
        ),
    )


def _merge_query_params_into_data(data: dict, request: Request) -> dict:
    """
    For GET/DELETE endpoints that cannot carry a JSON body, read a
    JSON-encoded ``litellm_params_template`` query parameter and merge its
    contents into *data*, without overwriting keys that are already present
    (e.g. path params like ``name`` or the fixed ``custom_llm_provider``).

    This mirrors the ``litellm_params_template`` handling in
    ``create_gemini_agent`` and is the supported way for multi-tenant
    callers to supply per-request credentials on non-POST endpoints:

    .. code-block:: bash

        curl "http://localhost:4000/v1beta/agents?litellm_params_template=%7B%22api_key%22%3A%22AIza...%22%7D" \\
            -H "Authorization: Bearer sk-..."

    Credentials MUST NOT be passed as plain flat query parameters (e.g.
    ``?api_key=AIza...``) because URL query strings appear verbatim in
    web-server access logs, CDN edge logs, browser history, and Referer
    headers. Use the ``litellm_params_template`` JSON body field on POST
    requests, or the JSON-encoded query parameter above for GET/DELETE.
    """
    query_params: Final = _safe_get_request_query_params(request)
    if not query_params:
        return data

    raw_template: Final = query_params.get("litellm_params_template")
    if raw_template:
        try:
            template = json.loads(raw_template) if isinstance(raw_template, str) else raw_template
        except (json.JSONDecodeError, ValueError):
            template = {}
        if isinstance(template, dict):
            for key, value in template.items():
                data.setdefault(key, value)

    return data


def _proxy_server_imports():
    from litellm.proxy.proxy_server import (  # noqa: PLC0415
        general_settings,
        llm_router,
        proxy_config,
        proxy_logging_obj,
        select_data_generator,
        user_api_base,
        user_max_tokens,
        user_model,
        user_request_timeout,
        user_temperature,
        version,
    )

    return dict(
        general_settings=general_settings,
        llm_router=llm_router,
        proxy_config=proxy_config,
        proxy_logging_obj=proxy_logging_obj,
        select_data_generator=select_data_generator,
        user_api_base=user_api_base,
        user_max_tokens=user_max_tokens,
        user_model=user_model,
        user_request_timeout=user_request_timeout,
        user_temperature=user_temperature,
        version=version,
    )


@router.post(
    "/v1beta/agents",
    dependencies=[Depends(user_api_key_auth)],
    response_class=ORJSONResponse,
)
async def create_gemini_agent(
    request: Request,
    fastapi_response: Response,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Create a named custom agent on the Gemini side.

    Example:
    ```bash
    curl -X POST "http://localhost:4000/v1beta/agents" \\
        -H "Authorization: Bearer sk-..." \\
        -H "Content-Type: application/json" \\
        -d '{
            "name": "my-custom-slides-agent",
            "base_agent": "waverunner",
            "instructions": "You are a helpful assistant that creates slides.",
            "base_environment": {
                "type": "remote",
                "sources": [
                    {"type": "gcs", "source": "gs://eap-templates/slides-skill",
                     "target": "/.agents/skills/slides-skill"}
                ]
            }
        }'
    ```
    """
    srv: Final = _proxy_server_imports()
    data: Final = await _read_request_body(request=request)
    # Merge litellm_params_template (e.g. custom_llm_provider, api_key) into the request
    litellm_params_template: Final = data.pop("litellm_params_template", None) or {}
    if isinstance(litellm_params_template, dict):
        for key, value in litellm_params_template.items():
            if key not in data:
                data[key] = value
    data.setdefault("custom_llm_provider", "gemini")
    _enforce_caller_supplied_provider_key(data, user_api_key_dict)

    processor: Final = ProxyBaseLLMRequestProcessing(data=data)
    try:
        return await processor.base_process_llm_request(
            request=request,
            fastapi_response=fastapi_response,
            user_api_key_dict=user_api_key_dict,
            route_type="acreate_agent",
            proxy_logging_obj=srv["proxy_logging_obj"],
            llm_router=srv["llm_router"],
            general_settings=srv["general_settings"],
            proxy_config=srv["proxy_config"],
            select_data_generator=srv["select_data_generator"],
            model=None,
            user_model=srv["user_model"],
            user_temperature=srv["user_temperature"],
            user_request_timeout=srv["user_request_timeout"],
            user_max_tokens=srv["user_max_tokens"],
            user_api_base=srv["user_api_base"],
            version=srv["version"],
        )
    except Exception as e:
        raise await processor._handle_llm_api_exception(
            e=e,
            user_api_key_dict=user_api_key_dict,
            proxy_logging_obj=srv["proxy_logging_obj"],
            version=srv["version"],
        )


@router.get(
    "/v1beta/agents",
    dependencies=[Depends(user_api_key_auth)],
    response_class=ORJSONResponse,
)
async def list_gemini_agents(
    request: Request,
    fastapi_response: Response,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    List all custom agents on the Gemini side.

    Pass per-request Gemini credentials via the JSON-encoded
    ``litellm_params_template`` query parameter. Flat query parameters
    (e.g. ``?api_key=AIza...``) are intentionally ignored — see
    ``_merge_query_params_into_data`` for the rationale.

    ```bash
    curl "http://localhost:4000/v1beta/agents?litellm_params_template=%7B%22api_key%22%3A%22AIza...%22%7D" \\
        -H "Authorization: Bearer sk-..."
    ```
    """
    srv: Final = _proxy_server_imports()
    data: Final[dict] = {"custom_llm_provider": "gemini"}
    _merge_query_params_into_data(data, request)
    _enforce_caller_supplied_provider_key(data, user_api_key_dict)

    processor: Final = ProxyBaseLLMRequestProcessing(data=data)
    try:
        return await processor.base_process_llm_request(
            request=request,
            fastapi_response=fastapi_response,
            user_api_key_dict=user_api_key_dict,
            route_type="alist_agents",
            proxy_logging_obj=srv["proxy_logging_obj"],
            llm_router=srv["llm_router"],
            general_settings=srv["general_settings"],
            proxy_config=srv["proxy_config"],
            select_data_generator=srv["select_data_generator"],
            model=None,
            user_model=srv["user_model"],
            user_temperature=srv["user_temperature"],
            user_request_timeout=srv["user_request_timeout"],
            user_max_tokens=srv["user_max_tokens"],
            user_api_base=srv["user_api_base"],
            version=srv["version"],
        )
    except Exception as e:
        raise await processor._handle_llm_api_exception(
            e=e,
            user_api_key_dict=user_api_key_dict,
            proxy_logging_obj=srv["proxy_logging_obj"],
            version=srv["version"],
        )


@router.get(
    "/v1beta/agents/{name}",
    dependencies=[Depends(user_api_key_auth)],
    response_class=ORJSONResponse,
)
async def get_gemini_agent(
    request: Request,
    name: str,
    fastapi_response: Response,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Get a specific custom agent by name.

    Pass per-request Gemini credentials via the JSON-encoded
    ``litellm_params_template`` query parameter. Flat query parameters
    (e.g. ``?api_key=AIza...``) are intentionally ignored — see
    ``_merge_query_params_into_data`` for the rationale.

    ```bash
    curl "http://localhost:4000/v1beta/agents/my-custom-slides-agent?litellm_params_template=%7B%22api_key%22%3A%22AIza...%22%7D" \\
        -H "Authorization: Bearer sk-..."
    ```
    """
    srv: Final = _proxy_server_imports()
    data: Final = {"name": name, "custom_llm_provider": "gemini"}
    _merge_query_params_into_data(data, request)
    _enforce_caller_supplied_provider_key(data, user_api_key_dict)

    processor: Final = ProxyBaseLLMRequestProcessing(data=data)
    try:
        return await processor.base_process_llm_request(
            request=request,
            fastapi_response=fastapi_response,
            user_api_key_dict=user_api_key_dict,
            route_type="aget_agent",
            proxy_logging_obj=srv["proxy_logging_obj"],
            llm_router=srv["llm_router"],
            general_settings=srv["general_settings"],
            proxy_config=srv["proxy_config"],
            select_data_generator=srv["select_data_generator"],
            model=None,
            user_model=srv["user_model"],
            user_temperature=srv["user_temperature"],
            user_request_timeout=srv["user_request_timeout"],
            user_max_tokens=srv["user_max_tokens"],
            user_api_base=srv["user_api_base"],
            version=srv["version"],
        )
    except Exception as e:
        raise await processor._handle_llm_api_exception(
            e=e,
            user_api_key_dict=user_api_key_dict,
            proxy_logging_obj=srv["proxy_logging_obj"],
            version=srv["version"],
        )


@router.delete(
    "/v1beta/agents/{name}",
    dependencies=[Depends(user_api_key_auth)],
    response_class=ORJSONResponse,
)
async def delete_gemini_agent(
    request: Request,
    name: str,
    fastapi_response: Response,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Delete a custom agent by name.

    Pass per-request Gemini credentials via the JSON-encoded
    ``litellm_params_template`` query parameter. Flat query parameters
    (e.g. ``?api_key=AIza...``) are intentionally ignored — see
    ``_merge_query_params_into_data`` for the rationale.

    ```bash
    curl -X DELETE "http://localhost:4000/v1beta/agents/my-custom-slides-agent?litellm_params_template=%7B%22api_key%22%3A%22AIza...%22%7D" \\
        -H "Authorization: Bearer sk-..."
    ```
    """
    srv: Final = _proxy_server_imports()
    data: Final = {"name": name, "custom_llm_provider": "gemini"}
    _merge_query_params_into_data(data, request)
    _enforce_caller_supplied_provider_key(data, user_api_key_dict)

    processor: Final = ProxyBaseLLMRequestProcessing(data=data)
    try:
        return await processor.base_process_llm_request(
            request=request,
            fastapi_response=fastapi_response,
            user_api_key_dict=user_api_key_dict,
            route_type="adelete_agent",
            proxy_logging_obj=srv["proxy_logging_obj"],
            llm_router=srv["llm_router"],
            general_settings=srv["general_settings"],
            proxy_config=srv["proxy_config"],
            select_data_generator=srv["select_data_generator"],
            model=None,
            user_model=srv["user_model"],
            user_temperature=srv["user_temperature"],
            user_request_timeout=srv["user_request_timeout"],
            user_max_tokens=srv["user_max_tokens"],
            user_api_base=srv["user_api_base"],
            version=srv["version"],
        )
    except Exception as e:
        raise await processor._handle_llm_api_exception(
            e=e,
            user_api_key_dict=user_api_key_dict,
            proxy_logging_obj=srv["proxy_logging_obj"],
            version=srv["version"],
        )


@router.get(
    "/v1beta/agents/{name}/versions",
    dependencies=[Depends(user_api_key_auth)],
    response_class=ORJSONResponse,
)
async def list_gemini_agent_versions(
    request: Request,
    name: str,
    fastapi_response: Response,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    List versions of a custom agent.

    Pass per-request Gemini credentials via the JSON-encoded
    ``litellm_params_template`` query parameter. Flat query parameters
    (e.g. ``?api_key=AIza...``) are intentionally ignored — see
    ``_merge_query_params_into_data`` for the rationale.

    ```bash
    curl "http://localhost:4000/v1beta/agents/my-custom-slides-agent/versions?litellm_params_template=%7B%22api_key%22%3A%22AIza...%22%7D" \\
        -H "Authorization: Bearer sk-..."
    ```
    """
    srv: Final = _proxy_server_imports()
    data: Final = {"name": name, "custom_llm_provider": "gemini"}
    _merge_query_params_into_data(data, request)
    _enforce_caller_supplied_provider_key(data, user_api_key_dict)

    processor: Final = ProxyBaseLLMRequestProcessing(data=data)
    try:
        return await processor.base_process_llm_request(
            request=request,
            fastapi_response=fastapi_response,
            user_api_key_dict=user_api_key_dict,
            route_type="alist_agent_versions",
            proxy_logging_obj=srv["proxy_logging_obj"],
            llm_router=srv["llm_router"],
            general_settings=srv["general_settings"],
            proxy_config=srv["proxy_config"],
            select_data_generator=srv["select_data_generator"],
            model=None,
            user_model=srv["user_model"],
            user_temperature=srv["user_temperature"],
            user_request_timeout=srv["user_request_timeout"],
            user_max_tokens=srv["user_max_tokens"],
            user_api_base=srv["user_api_base"],
            version=srv["version"],
        )
    except Exception as e:
        raise await processor._handle_llm_api_exception(
            e=e,
            user_api_key_dict=user_api_key_dict,
            proxy_logging_obj=srv["proxy_logging_obj"],
            version=srv["version"],
        )