phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(main.py): pass extra headers through for async calls

Browse source
This commit is contained in:
Krrish Dholakia 2024-05-27 19:11:40 -07:00
parent 24eb79da91
commit 6b50e656b8
3 changed files with 18 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

1
litellm/main.py
View file

@ -294,6 +294,7 @@ async def acompletion(
"api_version": api_version, "api_version": api_version,
"api_key": api_key, "api_key": api_key,
"model_list": model_list, "model_list": model_list,
"extra_headers": extra_headers,
"acompletion": True, # assuming this is a required parameter "acompletion": True, # assuming this is a required parameter
} }
if custom_llm_provider is None: if custom_llm_provider is None:

13
litellm/proxy/_super_secret_config.yaml
View file

@ -24,8 +24,8 @@ general_settings:
database_connection_timeout: 60 database_connection_timeout: 60
health_check_interval: 300 health_check_interval: 300
ui_access_mode: all ui_access_mode: all
litellm_settings: # litellm_settings:
json_logs: true # json_logs: true
model_list: model_list:
- litellm_params: - litellm_params:
api_base: http://0.0.0.0:8080 api_base: http://0.0.0.0:8080
@ -43,5 +43,14 @@ model_list:
api_version: '2023-05-15' api_version: '2023-05-15'
model: azure/chatgpt-v-2 model: azure/chatgpt-v-2
model_name: gpt-3.5-turbo model_name: gpt-3.5-turbo
- litellm_params:
model: anthropic.claude-3-sonnet-20240229-v1:0
model_name: bedrock-anthropic-claude-3
- litellm_params:
api_base: https://openai-gpt-4-test-v-1.openai.azure.com/
api_key: os.environ/AZURE_API_KEY
api_version: '2023-05-15'
model: azure/chatgpt-v-2
model_name: gpt-3.5-turbo
router_settings: router_settings:
enable_pre_call_checks: true enable_pre_call_checks: true

11
litellm/tests/test_bedrock_completion.py
View file

@ -13,7 +13,7 @@ import pytest
import litellm import litellm
from litellm import embedding, completion, completion_cost, Timeout, ModelResponse from litellm import embedding, completion, completion_cost, Timeout, ModelResponse
from litellm import RateLimitError from litellm import RateLimitError
from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler
from unittest.mock import patch, AsyncMock, Mock from unittest.mock import patch, AsyncMock, Mock
# litellm.num_retries = 3 # litellm.num_retries = 3
@ -519,20 +519,21 @@ def test_bedrock_ptu():
mock_client_post.assert_called_once() mock_client_post.assert_called_once()
def test_bedrock_extra_headers(): @pytest.mark.asyncio
async def test_bedrock_extra_headers():
""" """
Check if a url with 'modelId' passed in, is created correctly Check if a url with 'modelId' passed in, is created correctly
Reference: https://github.com/BerriAI/litellm/issues/3805 Reference: https://github.com/BerriAI/litellm/issues/3805
""" """
client = HTTPHandler() client = AsyncHTTPHandler()
with patch.object(client, "post", new=Mock()) as mock_client_post: with patch.object(client, "post", new=AsyncMock()) as mock_client_post:
litellm.set_verbose = True litellm.set_verbose = True
from openai.types.chat import ChatCompletion from openai.types.chat import ChatCompletion
try: try:
response = litellm.completion( response = await litellm.acompletion(
model="anthropic.claude-3-sonnet-20240229-v1:0", model="anthropic.claude-3-sonnet-20240229-v1:0",
messages=[{"role": "user", "content": "What's AWS?"}], messages=[{"role": "user", "content": "What's AWS?"}],
client=client, client=client,