diff --git a/src/benchflow/providers/litellm_config.py b/src/benchflow/providers/litellm_config.py index acee4756..19600cfd 100644 --- a/src/benchflow/providers/litellm_config.py +++ b/src/benchflow/providers/litellm_config.py @@ -275,7 +275,11 @@ def _route_registered_provider( ) if provider_cfg.auth_type == "adc": - params = {"model": f"vertex_ai/{bare}"} + params = { + "model": f"vertex_ai/{bare}", + "vertex_project": env.get("GOOGLE_CLOUD_PROJECT", ""), + "vertex_location": env.get("GOOGLE_CLOUD_LOCATION", ""), + } return LiteLLMRoute( requested_model=model, model_alias=safe_model_alias(model), diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index fc53f118..e1b22c5a 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -1022,7 +1022,11 @@ async def _upload_runtime_files_to_sandbox( async def _ensure_sandbox_litellm( - sandbox: Any, *, venv_dir: str, install_timeout_sec: int = 600 + sandbox: Any, + *, + venv_dir: str, + extra_requirements: tuple[str, ...] = (), + install_timeout_sec: int = 600, ) -> str: vq = shlex.quote(venv_dir) # Prefer uv to bootstrap the venv: many sandbox base images ship a python3 @@ -1030,6 +1034,10 @@ async def _ensure_sandbox_litellm( # `python -m venv` and `pip install` fail. uv needs neither (it is the same # mechanism the openhands agent install already uses in-sandbox), with a # stdlib-venv fallback for images that have a working venv and lack uv. + requirements = " ".join( + shlex.quote(requirement) + for requirement in (LITELLM_VERSION_SPEC, "boto3>=1.40", *extra_requirements) + ) command = f""" set -eu export PATH="$HOME/.local/bin:$PATH" @@ -1041,7 +1049,7 @@ async def _ensure_sandbox_litellm( fi if [ -n "$UV" ]; then [ -x {vq}/bin/python ] || "$UV" venv {vq} >/dev/null 2>&1 - "$UV" pip install --python {vq}/bin/python -q '{LITELLM_VERSION_SPEC}' 'boto3>=1.40' + "$UV" pip install --python {vq}/bin/python -q {requirements} else PY="$(command -v python3 || command -v python)" if [ ! -x {vq}/bin/python ]; then @@ -1051,7 +1059,7 @@ async def _ensure_sandbox_litellm( ) fi {vq}/bin/python -m pip install -q --upgrade pip - {vq}/bin/python -m pip install -q '{LITELLM_VERSION_SPEC}' 'boto3>=1.40' + {vq}/bin/python -m pip install -q {requirements} fi {vq}/bin/python - <<'PY' import litellm @@ -1167,10 +1175,22 @@ async def _start_sandbox_litellm( runtime_dir=runtime_dir, config=config, ) + proxy_agent_env = dict(agent_env) + extra_requirements: tuple[str, ...] = () + if route.provider_name in {"google-vertex", "anthropic-vertex"}: + extra_requirements = ("google-cloud-aiplatform>=1.133.0,<2.0",) + adc_json = proxy_agent_env.pop("GOOGLE_APPLICATION_CREDENTIALS_JSON", None) + if adc_json: + adc_path = f"{runtime_dir}/application_default_credentials.json" + await _upload_text(sandbox, adc_json, adc_path, ".json") + proxy_agent_env["GOOGLE_APPLICATION_CREDENTIALS"] = adc_path python = await _ensure_sandbox_litellm( - sandbox, venv_dir=paths["venv"], install_timeout_sec=install_timeout_sec + sandbox, + venv_dir=paths["venv"], + extra_requirements=extra_requirements, + install_timeout_sec=install_timeout_sec, ) - env = dict(agent_env) + env = proxy_agent_env env.update( { "PYTHONPATH": f"{runtime_dir}:{env.get('PYTHONPATH', '')}", diff --git a/tests/test_litellm_config.py b/tests/test_litellm_config.py index 2611eecf..f4447954 100644 --- a/tests/test_litellm_config.py +++ b/tests/test_litellm_config.py @@ -37,6 +37,16 @@ def test_bedrock_model_honors_max_thinking_effort_env(): assert route.litellm_params["reasoning_effort"] == "high" +def test_vertex_route_passes_project_and_location_to_litellm(): + route = resolve_litellm_route( + "anthropic-vertex/claude-sonnet-4-6", + {"GOOGLE_CLOUD_PROJECT": "skillsbench", "GOOGLE_CLOUD_LOCATION": "global"}, + ) + + assert route.litellm_params["vertex_project"] == "skillsbench" + assert route.litellm_params["vertex_location"] == "global" + + def test_azure_openai_route_uses_resource_and_preview_version(): route = resolve_litellm_route( "azure-foundry-openai/gpt-5.5", diff --git a/tests/test_litellm_hardening.py b/tests/test_litellm_hardening.py index c39ee1e8..78c3b148 100644 --- a/tests/test_litellm_hardening.py +++ b/tests/test_litellm_hardening.py @@ -405,6 +405,49 @@ async def test_sandbox_litellm_install_uses_configured_setup_timeout(): assert sandbox.exec_timeouts[install_index] == 901 +@pytest.mark.asyncio +async def test_vertex_sandbox_litellm_has_google_runtime_and_local_adc(): + route = resolve_litellm_route( + "anthropic-vertex/claude-sonnet-4-6", + {"GOOGLE_CLOUD_PROJECT": "project", "GOOGLE_CLOUD_LOCATION": "region"}, + ) + sandbox = _FakeSandbox() + + await runtime_mod._start_sandbox_litellm( + sandbox=sandbox, + route=route, + master_key="sk-master", + agent_env={"GOOGLE_APPLICATION_CREDENTIALS_JSON": '{"type":"test"}'}, + session_id="s", + agent_name="claude-agent-acp", + ) + + install_command = next( + command + for command in sandbox.exec_calls + if "pip install" in command and "litellm" in command + ) + assert "google-cloud-aiplatform>=1.133.0,<2.0" in install_command + adc_path = next( + path + for path in sandbox.uploaded + if path.endswith("application_default_credentials.json") + ) + launch_path = next( + path for path in sandbox.uploaded if path.endswith("launch_config.json") + ) + assert ( + json.loads(sandbox.uploaded[launch_path])["env"][ + "GOOGLE_APPLICATION_CREDENTIALS" + ] + == adc_path + ) + assert ( + "GOOGLE_APPLICATION_CREDENTIALS_JSON" + not in json.loads(sandbox.uploaded[launch_path])["env"] + ) + + @pytest.mark.asyncio async def test_sandbox_litellm_stop_imports_usage_and_cleans_up(): route = resolve_litellm_route(