From 8d819daafab3d456b1156671122f058bf0677bee Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 19 Aug 2023 21:17:17 -0700 Subject: [PATCH] fixes to embedding logging --- litellm/main.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index a01bea4ad..ea2dd9f25 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -676,9 +676,10 @@ def batch_completion(*args, **kwargs): @timeout( # type: ignore 60 ) ## set timeouts, in case calls hang (e.g. Azure) - default is 60s, override with `force_timeout` -def embedding(model, input=[], azure=False, force_timeout=60, logger_fn=None): +def embedding(model, input=[], azure=False, force_timeout=60, litellm_call_id=None, logger_fn=None): try: response = None + logging = Logging(model=model, messages=input, optional_params={}, litellm_params={"azure": azure, "force_timeout": force_timeout, "logger_fn": logger_fn, "litellm_call_id": litellm_call_id}) if azure == True: # azure configs openai.api_type = "azure" @@ -686,7 +687,7 @@ def embedding(model, input=[], azure=False, force_timeout=60, logger_fn=None): openai.api_version = get_secret("AZURE_API_VERSION") openai.api_key = get_secret("AZURE_API_KEY") ## LOGGING - logging.pre_call(model=model, input=input, azure=azure, logger_fn=logger_fn) + logging.pre_call(input=input, api_key=openai.api_key, additional_args={"api_type": openai.api_type, "api_base": openai.api_base, "api_version": openai.api_version}) ## EMBEDDING CALL response = openai.Embedding.create(input=input, engine=model) print_verbose(f"response_value: {str(response)[:50]}") @@ -696,19 +697,16 @@ def embedding(model, input=[], azure=False, force_timeout=60, logger_fn=None): openai.api_version = None openai.api_key = get_secret("OPENAI_API_KEY") ## LOGGING - logging(model=model, input=input, azure=azure, logger_fn=logger_fn) + logging.pre_call(input=input, api_key=openai.api_key, additional_args={"api_type": openai.api_type, "api_base": openai.api_base, "api_version": openai.api_version}) ## EMBEDDING CALL response = openai.Embedding.create(input=input, model=model) print_verbose(f"response_value: {str(response)[:50]}") else: - logging(model=model, input=input, azure=azure, logger_fn=logger_fn) args = locals() raise ValueError(f"No valid embedding model args passed in - {args}") return response except Exception as e: - # log the original exception - logging(model=model, input=input, azure=azure, logger_fn=logger_fn, exception=e) ## Map to OpenAI Exception raise exception_type(model=model, original_exception=e, custom_llm_provider="azure" if azure==True else None) raise e