fix(proxy_server.py): return 'allowed-model-region' in headers

This commit is contained in:
Krrish Dholakia 2024-05-13 08:46:44 -07:00
parent 514c5737f8
commit c3293474dd
3 changed files with 80 additions and 64 deletions

View file

@ -3762,6 +3762,7 @@ async def chat_completion(
"x-litellm-cache-key": cache_key,
"x-litellm-model-api-base": api_base,
"x-litellm-version": version,
"x-litellm-model-region": user_api_key_dict.allowed_model_region or "",
}
selected_data_generator = select_data_generator(
response=response,
@ -3778,6 +3779,9 @@ async def chat_completion(
fastapi_response.headers["x-litellm-cache-key"] = cache_key
fastapi_response.headers["x-litellm-model-api-base"] = api_base
fastapi_response.headers["x-litellm-version"] = version
fastapi_response.headers["x-litellm-model-region"] = (
user_api_key_dict.allowed_model_region or ""
)
### CALL HOOKS ### - modify outgoing data
response = await proxy_logging_obj.post_call_success_hook(
@ -4162,6 +4166,9 @@ async def embeddings(
fastapi_response.headers["x-litellm-cache-key"] = cache_key
fastapi_response.headers["x-litellm-model-api-base"] = api_base
fastapi_response.headers["x-litellm-version"] = version
fastapi_response.headers["x-litellm-model-region"] = (
user_api_key_dict.allowed_model_region or ""
)
return response
except Exception as e:
@ -4331,6 +4338,9 @@ async def image_generation(
fastapi_response.headers["x-litellm-cache-key"] = cache_key
fastapi_response.headers["x-litellm-model-api-base"] = api_base
fastapi_response.headers["x-litellm-version"] = version
fastapi_response.headers["x-litellm-model-region"] = (
user_api_key_dict.allowed_model_region or ""
)
return response
except Exception as e:
@ -4524,6 +4534,9 @@ async def audio_transcriptions(
fastapi_response.headers["x-litellm-cache-key"] = cache_key
fastapi_response.headers["x-litellm-model-api-base"] = api_base
fastapi_response.headers["x-litellm-version"] = version
fastapi_response.headers["x-litellm-model-region"] = (
user_api_key_dict.allowed_model_region or ""
)
return response
except Exception as e:
@ -4699,6 +4712,9 @@ async def moderations(
fastapi_response.headers["x-litellm-cache-key"] = cache_key
fastapi_response.headers["x-litellm-model-api-base"] = api_base
fastapi_response.headers["x-litellm-version"] = version
fastapi_response.headers["x-litellm-model-region"] = (
user_api_key_dict.allowed_model_region or ""
)
return response
except Exception as e: