phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 11:43:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(custom_llm.py): support async completion calls

Browse source
This commit is contained in:
Krrish Dholakia 2024-07-25 15:51:39 -07:00
parent 54e1ca29b7
commit fe503386ab
3 changed files with 50 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

25
litellm/tests/test_custom_llm.py
View file

@ -23,7 +23,7 @@ import httpx
from dotenv import load_dotenv
import litellm
from litellm import CustomLLM, completion, get_llm_provider
from litellm import CustomLLM, acompletion, completion, get_llm_provider
class MyCustomLLM(CustomLLM):
@ -35,6 +35,15 @@ class MyCustomLLM(CustomLLM):
) # type: ignore
class MyCustomAsyncLLM(CustomLLM):
async def acompletion(self, *args, **kwargs) -> litellm.ModelResponse:
return litellm.completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hello world"}],
mock_response="Hi!",
) # type: ignore
def test_get_llm_provider():
from litellm.utils import custom_llm_setup
@ -61,3 +70,17 @@ def test_simple_completion():
)
assert resp.choices[0].message.content == "Hi!"
@pytest.mark.asyncio
async def test_simple_acompletion():
my_custom_llm = MyCustomAsyncLLM()
litellm.custom_provider_map = [
{"provider": "custom_llm", "custom_handler": my_custom_llm}
]
resp = await acompletion(
model="custom_llm/my-fake-model",
messages=[{"role": "user", "content": "Hello world!"}],
)
assert resp.choices[0].message.content == "Hi!"