From a4eeb50b0e5c9d192e91ed2c3576defc0c35c23c Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 6 Nov 2024 15:06:46 -0800
Subject: [PATCH] allow iam auth for gcs logging

---
 litellm/integrations/gcs_bucket/gcs_bucket.py | 24 ++++++++++++++-----
 litellm/proxy/proxy_config.yaml               |  5 ++--
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py
index 172185fff25..111730d1f0d 100644
--- a/litellm/integrations/gcs_bucket/gcs_bucket.py
+++ b/litellm/integrations/gcs_bucket/gcs_bucket.py
@@ -26,6 +26,9 @@ else:
     VertexBase = Any
 
 
+IAM_AUTH_KEY = "IAM_AUTH"
+
+
 class GCSLoggingConfig(TypedDict):
     bucket_name: str
     vertex_instance: VertexBase
@@ -227,18 +230,27 @@ class GCSBucketLogger(GCSBucketBase):
             VertexBase,
         )
 
-        if credentials is None:
-            credentials = "IAM_AUTH"
-
-        if credentials not in self.vertex_instances:
+        _in_memory_key = self._get_in_memory_key_for_vertex_instance(credentials)
+        if _in_memory_key not in self.vertex_instances:
             vertex_instance = VertexBase()
             await vertex_instance._ensure_access_token_async(
                 credentials=credentials,
                 project_id=None,
                 custom_llm_provider="vertex_ai",
             )
-            self.vertex_instances[credentials] = vertex_instance
-        return self.vertex_instances[credentials]
+            self.vertex_instances[_in_memory_key] = vertex_instance
+        return self.vertex_instances[_in_memory_key]
+
+    def _get_in_memory_key_for_vertex_instance(self, credentials: Optional[str]) -> str:
+        """
+        Returns key to use for caching the Vertex instance in-memory.
+
+        When using Vertex with Key based logging, we need to cache the Vertex instance in-memory.
+
+        - If a credentials string is provided, it is used as the key.
+        - If no credentials string is provided, "IAM_AUTH" is used as the key.
+        """
+        return credentials or IAM_AUTH_KEY
 
     async def download_gcs_object(self, object_name: str, **kwargs):
         """
diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml
index a54e6a3b6a8..9767677cf32 100644
--- a/litellm/proxy/proxy_config.yaml
+++ b/litellm/proxy/proxy_config.yaml
@@ -5,6 +5,5 @@ model_list:
       api_key: os.environ/OPENAI_API_KEY
       api_base: https://exampleopenaiendpoint-production.up.railway.app/
 
-general_settings:
-  store_model_in_db: true
-  public_routes: ["LiteLLMRoutes.public_routes", "/spend/calculate", "/model/info"]
+litellm_settings:
+  callbacks: ["gcs_bucket"]
\ No newline at end of file