diff --git a/docs/my-website/docs/completion/stream.md b/docs/my-website/docs/completion/stream.md index e5fa50ab5a..b954e60ad1 100644 --- a/docs/my-website/docs/completion/stream.md +++ b/docs/my-website/docs/completion/stream.md @@ -30,4 +30,42 @@ async def test_get_response(): response = asyncio.run(test_get_response()) print(response) +``` + +## Async Streaming +We've implemented an `__anext__()` function in the streaming object returned. This +### Usage +``` +from litellm import acompletion +import asyncio + +def logger_fn(model_call_object: dict): + print(f"LOGGER FUNCTION: {model_call_object}") + + +user_message = "Hello, how are you?" +messages = [{"content": user_message, "role": "user"}] + +# # test on ai21 completion call +async def ai21_async_completion_call(): + try: + response = completion( + model="j2-ultra", messages=messages, stream=True, logger_fn=logger_fn + ) + print(f"response: {response}") + complete_response = "" + start_time = time.time() + # Change for loop to async for loop + async for chunk in response: + chunk_time = time.time() + print(f"time since initial request: {chunk_time - start_time:.5f}") + print(chunk["choices"][0]["delta"]) + complete_response += chunk["choices"][0]["delta"]["content"] + if complete_response == "": + raise Exception("Empty response received") + except: + print(f"error occurred: {traceback.format_exc()}") + pass + +asyncio.run(ai21_async_completion_call()) ``` \ No newline at end of file diff --git a/litellm/__pycache__/__init__.cpython-311.pyc b/litellm/__pycache__/__init__.cpython-311.pyc index 1417494660..a95ce65c85 100644 Binary files a/litellm/__pycache__/__init__.cpython-311.pyc and b/litellm/__pycache__/__init__.cpython-311.pyc differ diff --git a/litellm/__pycache__/main.cpython-311.pyc b/litellm/__pycache__/main.cpython-311.pyc index 3b9ce290a9..d7f053f664 100644 Binary files a/litellm/__pycache__/main.cpython-311.pyc and b/litellm/__pycache__/main.cpython-311.pyc differ diff --git a/litellm/__pycache__/utils.cpython-311.pyc b/litellm/__pycache__/utils.cpython-311.pyc index 56b13fe11a..538c697581 100644 Binary files a/litellm/__pycache__/utils.cpython-311.pyc and b/litellm/__pycache__/utils.cpython-311.pyc differ diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py index 90828d20cb..e5b9c3993e 100644 --- a/litellm/tests/test_streaming.py +++ b/litellm/tests/test_streaming.py @@ -9,7 +9,7 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path import litellm -from litellm import completion +from litellm import completion, acompletion litellm.logging = False litellm.set_verbose = False @@ -217,3 +217,25 @@ def test_together_ai_completion_call_starcoder(): # except: # print(f"error occurred: {traceback.format_exc()}") # pass +#### Test Async streaming + +# # test on ai21 completion call +async def ai21_async_completion_call(): + try: + response = completion( + model="j2-ultra", messages=messages, stream=True, logger_fn=logger_fn + ) + print(f"response: {response}") + complete_response = "" + start_time = time.time() + # Change for loop to async for loop + async for chunk in response: + chunk_time = time.time() + print(f"time since initial request: {chunk_time - start_time:.5f}") + print(chunk["choices"][0]["delta"]) + complete_response += chunk["choices"][0]["delta"]["content"] + if complete_response == "": + raise Exception("Empty response received") + except: + print(f"error occurred: {traceback.format_exc()}") + pass \ No newline at end of file