mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
Add unit test
This commit is contained in:
Joe Cheng
parent
33f4411f17
commit
1fbfc09b44
2 changed files with 566 additions and 0 deletions
543
litellm/tests/stream_chunk_testdata.py
Normal file
543
litellm/tests/stream_chunk_testdata.py
Normal file
|
|
@ -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,
|
||||
),
|
||||
]
|
||||
|
|
@ -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",
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue