scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 145

fix(vertex_ai): keep legacy bucket_name in credential resolution and add GCS_BATCH_BUCKET_NAME env var (#42803)

Browse source 
* fix(vertex_ai): map legacy bucket_name to gcs_bucket_name and add GCS_BATCH_BUCKET_NAME env var

* refactor(router): keep legacy bucket_name as a credential field instead of a validator

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(vertex_ai): pass the RAG corpus bucket to the file upload instead of hopping through GCS_BUCKET_NAME

* fix(vertex_ai): accept existing_file_id in the RAG Engine store step so ingest() runs end to end

---------

Co-authored-by: Mubashir Osmani <mubashir@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
mubashir1osmani 2026-09-25 18:32:21 -04:00 • committed by GitHub
parent 191305e6d4
commit 8b68c3cd09
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 203 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm/llms/vertex_ai/files/handler.py
View file

@ -53,14 +53,18 @@ class VertexAIFilesHandler(GCSBucketBase):
Sources them from the deployment's ``litellm_params`` (``gcs_bucket_name`` /
``bucket_name`` and ``vertex_credentials``), mirroring the write path in
``VertexAIFilesConfig._get_configured_bucket_name``, and falls back to the global
``GCS_BUCKET_NAME`` / ``GCS_PATH_SERVICE_ACCOUNT`` env vars. This lets Vertex batch
run entirely at the model-group level, so output written to a per-model bucket is
readable without setting the global env vars.
``VertexAIFilesConfig._get_configured_bucket_name``, and falls back to the
``GCS_BATCH_BUCKET_NAME`` then ``GCS_BUCKET_NAME`` / ``GCS_PATH_SERVICE_ACCOUNT``
env vars. This lets Vertex batch run entirely at the model-group level, so output
written to a per-model bucket is readable without setting the global env vars.
"""
params: Final[Mapping[str, object]] = litellm_params or {}
bucket_candidate: Final = params.get("gcs_bucket_name") or params.get("bucket_name")
configured_bucket_name = bucket_candidate if isinstance(bucket_candidate, str) else os.getenv("GCS_BUCKET_NAME")
configured_bucket_name = (
bucket_candidate
if isinstance(bucket_candidate, str)
else os.getenv("GCS_BATCH_BUCKET_NAME") or os.getenv("GCS_BUCKET_NAME")
)
credentials: Final = params.get("vertex_credentials") or vertex_credentials
if isinstance(credentials, dict):

5
litellm/llms/vertex_ai/files/transformation.py
View file

@ -961,7 +961,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
def _get_configured_bucket_name(self, litellm_params: dict) -> str:
bucket_name: Final = (
litellm_params.get("gcs_bucket_name") or litellm_params.get("bucket_name") or os.getenv("GCS_BUCKET_NAME")
litellm_params.get("gcs_bucket_name")
or litellm_params.get("bucket_name")
or os.getenv("GCS_BATCH_BUCKET_NAME")
or os.getenv("GCS_BUCKET_NAME")
)
if not bucket_name:
raise ValueError("GCS bucket_name is required")

49
litellm/llms/vertex_ai/rag_engine/ingestion.py
View file

@ -122,41 +122,26 @@ class VertexAIRAGIngestion(BaseRAGIngestion):
"""
import litellm
# Set GCS_BUCKET_NAME env var for litellm.files.create_file
# The handler uses this to determine where to upload
original_bucket: Final = os.environ.get("GCS_BUCKET_NAME")
if self.gcs_bucket:
os.environ["GCS_BUCKET_NAME"] = self.gcs_bucket
file_tuple: Final = (filename, file_content, content_type)
try:
# Create file tuple for litellm.files.acreate_file
file_tuple: Final = (filename, file_content, content_type)
verbose_logger.debug(
"Uploading file to GCS via litellm.files.acreate_file: %s (bucket: %s)", filename, self.gcs_bucket
)
verbose_logger.debug(
"Uploading file to GCS via litellm.files.acreate_file: %s (bucket: %s)", filename, self.gcs_bucket
)
response: Final = await litellm.acreate_file(
file=file_tuple,
purpose="assistants",
custom_llm_provider="vertex_ai",
gcs_bucket_name=self.gcs_bucket,
vertex_project=self.vertex_project,
vertex_location=self.vertex_location,
vertex_credentials=self.vertex_credentials,
)
# Upload to GCS using LiteLLM's file upload
response: Final = await litellm.acreate_file(
file=file_tuple,
purpose="assistants", # Purpose for file storage
custom_llm_provider="vertex_ai",
vertex_project=self.vertex_project,
vertex_location=self.vertex_location,
vertex_credentials=self.vertex_credentials,
)
gcs_uri: Final = response.id
verbose_logger.info("Uploaded file to GCS: %s", gcs_uri)
# The response.id should be the GCS URI
gcs_uri: Final = response.id
verbose_logger.info("Uploaded file to GCS: %s", gcs_uri)
return gcs_uri
finally:
# Restore original env var
if original_bucket is not None:
os.environ["GCS_BUCKET_NAME"] = original_bucket
elif "GCS_BUCKET_NAME" in os.environ:
del os.environ["GCS_BUCKET_NAME"]
return gcs_uri
async def _import_file_to_corpus_via_sdk(
self,
@ -259,6 +244,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion):
content_type: str | None,
chunks: list[str],
embeddings: list[list[float]] | None,
existing_file_id: str | None = None,
) -> tuple[str | None, str | None]:
"""
Store content in Vertex AI RAG corpus.
@ -274,6 +260,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion):
content_type: MIME type
chunks: Ignored - Vertex AI handles chunking
embeddings: Ignored - Vertex AI handles embedding
existing_file_id: Existing provider file ID, unsupported for Vertex AI RAG Engine
Returns:
Tuple of (corpus_id, gcs_uri)

1
litellm/types/router.py
View file

@ -345,6 +345,7 @@ class CredentialLiteLLMParams(BaseModel):
## OBJECT STORAGE (files / batches) ##
gcs_bucket_name: str | None = None
bucket_name: str | None = None
## AWS BEDROCK / SAGEMAKER ##
aws_access_key_id: str | None = None

35
tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py
View file

@ -208,6 +208,7 @@ class TestVertexAIFilesHandler:
assert service_account == "/model/sa.json"
def test_resolve_read_gcs_config_falls_back_to_env(self, monkeypatch):
monkeypatch.delenv("GCS_BATCH_BUCKET_NAME", raising=False)
monkeypatch.setenv("GCS_BUCKET_NAME", "env-default-bucket")
monkeypatch.setenv("GCS_PATH_SERVICE_ACCOUNT", "/env/sa.json")
@ -216,6 +217,40 @@ class TestVertexAIFilesHandler:
assert bucket == "env-default-bucket"
assert service_account == "/env/sa.json"
def test_resolve_read_gcs_config_prefers_batch_env_over_logging_env(self, monkeypatch):
monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket")
monkeypatch.setenv("GCS_BUCKET_NAME", "logging-bucket")
bucket, _ = self.handler._resolve_read_gcs_config(litellm_params={}, vertex_credentials=None)
assert bucket == "batch-bucket"
def test_resolve_read_gcs_config_prefers_per_model_bucket_over_batch_env(self, monkeypatch):
monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket")
bucket, _ = self.handler._resolve_read_gcs_config(
litellm_params={"gcs_bucket_name": "my-model-bucket"},
vertex_credentials=None,
)
assert bucket == "my-model-bucket"
def test_resolve_read_gcs_config_prefers_gcs_bucket_name_over_legacy(self):
bucket, _ = self.handler._resolve_read_gcs_config(
litellm_params={"gcs_bucket_name": "my-model-bucket", "bucket_name": "legacy-bucket"},
vertex_credentials=None,
)
assert bucket == "my-model-bucket"
def test_resolve_read_gcs_config_accepts_legacy_bucket_name_alone(self):
bucket, _ = self.handler._resolve_read_gcs_config(
litellm_params={"bucket_name": "legacy-bucket"},
vertex_credentials=None,
)
assert bucket == "legacy-bucket"
def test_resolve_read_gcs_config_serializes_dict_credentials(self, monkeypatch):
monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False)

11
tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py
View file

@ -1186,10 +1186,21 @@ class TestConfiguredBucketNameResolution:
assert config._get_configured_bucket_name({"gcs_bucket_name": "new", "bucket_name": "legacy"}) == "new"
def test_should_fall_back_to_env(self, config, monkeypatch):
monkeypatch.delenv("GCS_BATCH_BUCKET_NAME", raising=False)
monkeypatch.setenv("GCS_BUCKET_NAME", "env-bucket")
assert config._get_configured_bucket_name({}) == "env-bucket"
def test_should_prefer_batch_env_over_logging_env(self, config, monkeypatch):
monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket")
monkeypatch.setenv("GCS_BUCKET_NAME", "logging-bucket")
assert config._get_configured_bucket_name({}) == "batch-bucket"
def test_should_prefer_litellm_params_over_batch_env(self, config, monkeypatch):
monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket")
assert config._get_configured_bucket_name({"gcs_bucket_name": "per-model-bucket"}) == "per-model-bucket"
def test_should_raise_when_no_bucket_anywhere(self, config, monkeypatch):
monkeypatch.delenv("GCS_BATCH_BUCKET_NAME", raising=False)
monkeypatch.delenv("GCS_BUCKET_NAME", raising=False)
with pytest.raises(ValueError, match="GCS bucket_name is required"):
config._get_configured_bucket_name({})

0
tests/unit/llms/vertex_ai/rag_engine/__init__.py Normal file
View file

76
tests/unit/llms/vertex_ai/rag_engine/test_ingestion.py Normal file
View file

@ -0,0 +1,76 @@
import asyncio
import sys
from types import ModuleType, SimpleNamespace
import litellm
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.llms.vertex_ai.files.transformation import VertexAIFilesConfig
from litellm.llms.vertex_ai.rag_engine.ingestion import VertexAIRAGIngestion
def _ingestion_for_bucket(bucket: str) -> VertexAIRAGIngestion:
return VertexAIRAGIngestion(
{
"vector_store": {
"custom_llm_provider": "vertex_ai",
"vector_store_id": "corpus-123",
"vertex_project": "test-project",
"vertex_location": "us-central1",
"gcs_bucket": bucket,
}
}
)
def test_upload_lands_in_the_corpus_bucket_when_batch_bucket_env_is_set(monkeypatch):
monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket")
monkeypatch.setenv("GCS_BUCKET_NAME", "logging-bucket")
resolver = VertexAIFilesConfig()
async def acreate_file_through_real_bucket_resolver(**kwargs):
bucket = resolver._get_configured_bucket_name(get_litellm_params(**kwargs))
return SimpleNamespace(id=f"gs://{bucket}/{kwargs['file'][0]}")
monkeypatch.setattr(litellm, "acreate_file", acreate_file_through_real_bucket_resolver)
uri = asyncio.run(_ingestion_for_bucket("rag-bucket")._upload_file_to_gcs(b"doc", "doc.txt", "text/plain"))
assert uri == "gs://rag-bucket/doc.txt"
def _vertexai_sdk_stub(import_calls: list[dict[str, object]]) -> ModuleType:
rag = ModuleType("vertexai.rag")
rag.TransformationConfig = lambda chunking_config: chunking_config
rag.ChunkingConfig = lambda chunk_size, chunk_overlap: (chunk_size, chunk_overlap)
def import_files(**kwargs):
import_calls.append(kwargs)
return SimpleNamespace(imported_rag_files_count=1)
rag.import_files = import_files
vertexai = ModuleType("vertexai")
vertexai.init = lambda project, location: None
vertexai.rag = rag
return vertexai
def test_ingest_runs_end_to_end_through_the_base_pipeline(monkeypatch):
monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket")
resolver = VertexAIFilesConfig()
import_calls: list[dict[str, object]] = []
stub = _vertexai_sdk_stub(import_calls)
monkeypatch.setitem(sys.modules, "vertexai", stub)
monkeypatch.setitem(sys.modules, "vertexai.rag", stub.rag)
async def acreate_file_through_real_bucket_resolver(**kwargs):
bucket = resolver._get_configured_bucket_name(get_litellm_params(**kwargs))
return SimpleNamespace(id=f"gs://{bucket}/{kwargs['file'][0]}")
monkeypatch.setattr(litellm, "acreate_file", acreate_file_through_real_bucket_resolver)
result = asyncio.run(_ingestion_for_bucket("rag-bucket").ingest(file_data=("doc.txt", b"doc", "text/plain")))
assert (result["status"], result["vector_store_id"], result["file_id"]) == ("completed", "corpus-123", "gs://rag-bucket/doc.txt")
assert [(c["corpus_name"], c["paths"]) for c in import_calls] == [
("projects/test-project/locations/us-central1/ragCorpora/corpus-123", ["gs://rag-bucket/doc.txt"])
]

45
tests/unit/test_router/test_router.py
View file

@ -6470,6 +6470,51 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name():
assert credentials["custom_llm_provider"] == "vertex_ai"
def test_get_deployment_credentials_with_provider_keeps_legacy_bucket_name():
router = litellm.Router(
model_list=[
{
"model_name": "vertex-gemini",
"litellm_params": {
"model": "vertex_ai/gemini-3.5-flash",
"vertex_project": "my-project",
"vertex_location": "global",
"bucket_name": "my-legacy-bucket",
},
}
],
)
credentials = router.get_deployment_credentials_with_provider(model_id="vertex-gemini")
assert credentials is not None
assert credentials["bucket_name"] == "my-legacy-bucket"
assert "gcs_bucket_name" not in credentials
def test_get_deployment_credentials_with_provider_keeps_both_bucket_keys():
router = litellm.Router(
model_list=[
{
"model_name": "vertex-gemini",
"litellm_params": {
"model": "vertex_ai/gemini-3.5-flash",
"vertex_project": "my-project",
"vertex_location": "global",
"gcs_bucket_name": "new-bucket",
"bucket_name": "legacy-bucket",
},
}
],
)
credentials = router.get_deployment_credentials_with_provider(model_id="vertex-gemini")
assert credentials is not None
assert credentials["gcs_bucket_name"] == "new-bucket"
assert credentials["bucket_name"] == "legacy-bucket"
def test_get_deployment_credentials_with_provider_resolves_credential_name():
"""
Test that get_deployment_credentials_with_provider correctly resolves

4
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -32945,6 +32945,8 @@ export interface components {
azure_username?: string | null;
/** Bedrock Tags */
bedrock_tags?: unknown[] | null;
/** Bucket Name */
bucket_name?: string | null;
/** Budget Duration */
budget_duration?: string | null;
/** Cache Creation Input Audio Token Cost */
@ -46730,6 +46732,8 @@ export interface components {
azure_username?: string | null;
/** Bedrock Tags */
bedrock_tags?: unknown[] | null;
/** Bucket Name */
bucket_name?: string | null;
/** Budget Duration */
budget_duration?: string | null;
/** Cache Creation Input Audio Token Cost */