fix gcs bucket auth

This commit is contained in:
Ishaan Jaff 2024-11-06 15:00:06 -08:00
parent 66c1ee09cf
commit 087905656e
3 changed files with 12 additions and 16 deletions

View file

@ -29,7 +29,7 @@ else:
class GCSLoggingConfig(TypedDict):
bucket_name: str
vertex_instance: VertexBase
path_service_account: str
path_service_account: Optional[str]
class GCSBucketLogger(GCSBucketBase):
@ -173,7 +173,7 @@ class GCSBucketLogger(GCSBucketBase):
)
bucket_name: str
path_service_account: str
path_service_account: Optional[str]
if standard_callback_dynamic_params is not None:
verbose_logger.debug("Using dynamic GCS logging")
verbose_logger.debug(
@ -193,10 +193,6 @@ class GCSBucketLogger(GCSBucketBase):
raise ValueError(
"GCS_BUCKET_NAME is not set in the environment, but GCS Bucket is being used as a logging callback. Please set 'GCS_BUCKET_NAME' in the environment."
)
if _path_service_account is None:
raise ValueError(
"GCS_PATH_SERVICE_ACCOUNT is not set in the environment, but GCS Bucket is being used as a logging callback. Please set 'GCS_PATH_SERVICE_ACCOUNT' in the environment."
)
bucket_name = _bucket_name
path_service_account = _path_service_account
vertex_instance = await self.get_or_create_vertex_instance(
@ -208,10 +204,6 @@ class GCSBucketLogger(GCSBucketBase):
raise ValueError(
"GCS_BUCKET_NAME is not set in the environment, but GCS Bucket is being used as a logging callback. Please set 'GCS_BUCKET_NAME' in the environment."
)
if self.path_service_account_json is None:
raise ValueError(
"GCS_PATH_SERVICE_ACCOUNT is not set in the environment, but GCS Bucket is being used as a logging callback. Please set 'GCS_PATH_SERVICE_ACCOUNT' in the environment."
)
bucket_name = self.BUCKET_NAME
path_service_account = self.path_service_account_json
vertex_instance = await self.get_or_create_vertex_instance(
@ -224,7 +216,9 @@ class GCSBucketLogger(GCSBucketBase):
path_service_account=path_service_account,
)
async def get_or_create_vertex_instance(self, credentials: str) -> VertexBase:
async def get_or_create_vertex_instance(
self, credentials: Optional[str]
) -> VertexBase:
"""
This function is used to get the Vertex instance for the GCS Bucket Logger.
It checks if the Vertex instance is already created and cached, if not it creates a new instance and caches it.
@ -233,6 +227,9 @@ class GCSBucketLogger(GCSBucketBase):
VertexBase,
)
if credentials is None:
credentials = "IAM_AUTH"
if credentials not in self.vertex_instances:
vertex_instance = VertexBase()
await vertex_instance._ensure_access_token_async(

View file

@ -33,7 +33,7 @@ class GCSBucketBase(CustomLogger):
async def construct_request_headers(
self,
service_account_json: str,
service_account_json: Optional[str],
vertex_instance: Optional[VertexBase] = None,
) -> Dict[str, str]:
from litellm import vertex_chat_completion

View file

@ -5,7 +5,6 @@ model_list:
api_key: os.environ/OPENAI_API_KEY
api_base: https://exampleopenaiendpoint-production.up.railway.app/
general_settings:
alerting: ["slack"]
alerting_threshold: 0.001
general_settings:
store_model_in_db: true
public_routes: ["LiteLLMRoutes.public_routes", "/spend/calculate", "/model/info"]