From 548b2b686118306dd751cd6a1815982ad675d70f Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 6 Apr 2024 17:55:26 -0700 Subject: [PATCH] test - async claude streaming --- litellm/tests/test_streaming.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py index ed62987528..a8e3f76807 100644 --- a/litellm/tests/test_streaming.py +++ b/litellm/tests/test_streaming.py @@ -831,22 +831,25 @@ def test_bedrock_claude_3_streaming(): pytest.fail(f"Error occurred: {e}") -def test_claude_3_streaming_finish_reason(): +@pytest.mark.asyncio +async def test_claude_3_streaming_finish_reason(): try: litellm.set_verbose = True messages = [ {"role": "system", "content": "Be helpful"}, {"role": "user", "content": "What do you know?"}, ] - response: ModelResponse = completion( # type: ignore + response: ModelResponse = await litellm.acompletion( # type: ignore model="claude-3-opus-20240229", messages=messages, stream=True, + max_tokens=10, ) complete_response = "" # Add any assertions here to check the response num_finish_reason = 0 - for idx, chunk in enumerate(response): + async for chunk in response: + print(f"chunk: {chunk}") if isinstance(chunk, ModelResponse): if chunk.choices[0].finish_reason is not None: num_finish_reason += 1