From 1fbfc09b443337b531be863b8a631b8c3bd8e136 Mon Sep 17 00:00:00 2001 From: Joe Cheng Date: Fri, 2 Aug 2024 20:51:03 -0700 Subject: [PATCH] Add unit test --- litellm/tests/stream_chunk_testdata.py | 543 +++++++++++++++++++++ litellm/tests/test_stream_chunk_builder.py | 23 + 2 files changed, 566 insertions(+) create mode 100644 litellm/tests/stream_chunk_testdata.py diff --git a/litellm/tests/stream_chunk_testdata.py b/litellm/tests/stream_chunk_testdata.py new file mode 100644 index 00000000000..6be9d1ebdff --- /dev/null +++ b/litellm/tests/stream_chunk_testdata.py @@ -0,0 +1,543 @@ +from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponse, + StreamingChoices, +) + +chunks = [ + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="To answer", + role="assistant", + function_call=None, + tool_calls=None, + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" your", role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" question about", + role=None, + function_call=None, + tool_calls=None, + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" how", role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" many rows are in the ", + role=None, + function_call=None, + tool_calls=None, + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="'users' table, I", + role=None, + function_call=None, + tool_calls=None, + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="'ll", role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" need to", role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" run", role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" a SQL query.", + role=None, + function_call=None, + tool_calls=None, + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" Let", role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" me", role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" ", role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="do that for", + role=None, + function_call=None, + tool_calls=None, + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" you.", role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="", + role=None, + function_call=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id="toolu_01H3AjkLpRtGQrof13CBnWfK", + function=Function(arguments="", name="sql_query"), + type="function", + index=1, + ) + ], + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="", + role=None, + function_call=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id=None, + function=Function(arguments="", name=None), + type="function", + index=1, + ) + ], + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="", + role=None, + function_call=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id=None, + function=Function(arguments='{"', name=None), + type="function", + index=1, + ) + ], + ), + logprobs=None, + ) + ], + created=1722656357, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="", + role=None, + function_call=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id=None, + function=Function(arguments='query": ', name=None), + type="function", + index=1, + ) + ], + ), + logprobs=None, + ) + ], + created=1722656357, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="", + role=None, + function_call=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id=None, + function=Function(arguments='"SELECT C', name=None), + type="function", + index=1, + ) + ], + ), + logprobs=None, + ) + ], + created=1722656357, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="", + role=None, + function_call=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id=None, + function=Function(arguments="OUNT(*", name=None), + type="function", + index=1, + ) + ], + ), + logprobs=None, + ) + ], + created=1722656357, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="", + role=None, + function_call=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id=None, + function=Function(arguments=") ", name=None), + type="function", + index=1, + ) + ], + ), + logprobs=None, + ) + ], + created=1722656357, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="", + role=None, + function_call=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id=None, + function=Function(arguments="FROM use", name=None), + type="function", + index=1, + ) + ], + ), + logprobs=None, + ) + ], + created=1722656357, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="", + role=None, + function_call=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id=None, + function=Function(arguments='rs;"}', name=None), + type="function", + index=1, + ) + ], + ), + logprobs=None, + ) + ], + created=1722656357, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason="tool_calls", + index=0, + delta=Delta( + content=None, role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656357, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), +] diff --git a/litellm/tests/test_stream_chunk_builder.py b/litellm/tests/test_stream_chunk_builder.py index 342b070ae70..78d2617f1e7 100644 --- a/litellm/tests/test_stream_chunk_builder.py +++ b/litellm/tests/test_stream_chunk_builder.py @@ -18,6 +18,8 @@ from openai import OpenAI import litellm from litellm import completion, stream_chunk_builder +import litellm.tests.stream_chunk_testdata + dotenv.load_dotenv() user_message = "What is the current weather in Boston?" @@ -196,3 +198,24 @@ def test_stream_chunk_builder_litellm_usage_chunks(): # assert prompt tokens are the same assert gemini_pt == stream_rebuilt_pt + + +def test_stream_chunk_builder_litellm_mixed_calls(): + response = stream_chunk_builder(litellm.tests.stream_chunk_testdata.chunks) + assert ( + response.choices[0].message.content + == "To answer your question about how many rows are in the 'users' table, I'll need to run a SQL query. Let me do that for you." + ) + + print(response.choices[0].message.tool_calls[0].to_dict()) + + assert len(response.choices[0].message.tool_calls) == 1 + assert response.choices[0].message.tool_calls[0].to_dict() == { + "index": 1, + "function": { + "arguments": '{"query": "SELECT COUNT(*) FROM users;"}', + "name": "sql_query", + }, + "id": "toolu_01H3AjkLpRtGQrof13CBnWfK", + "type": "function", + }