mirror of https://github.com/aliasrobotics/cai.git
87 lines
2.6 KiB
Python
87 lines
2.6 KiB
Python
import pytest
|
|
from openai import NOT_GIVEN
|
|
|
|
from cai.sdk.agents.model_settings import ModelSettings
|
|
from cai.sdk.agents.models.chatcompletions.litellm_adapter import (
|
|
fetch_response_litellm_openai,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_litellm_streaming_fetch_opens_one_completion(monkeypatch):
|
|
calls = []
|
|
sentinel_stream = object()
|
|
|
|
async def fake_acompletion(**kwargs):
|
|
calls.append(kwargs.copy())
|
|
return sentinel_stream
|
|
|
|
monkeypatch.setattr(
|
|
"cai.sdk.agents.models.chatcompletions.litellm_adapter.litellm.acompletion",
|
|
fake_acompletion,
|
|
)
|
|
|
|
response, stream = await fetch_response_litellm_openai(
|
|
kwargs={"model": "deepseek/deepseek-v4-pro", "messages": [], "stream": True},
|
|
model_name="deepseek/deepseek-v4-pro",
|
|
model_settings=ModelSettings(),
|
|
tool_choice=NOT_GIVEN,
|
|
stream=True,
|
|
parallel_tool_calls=False,
|
|
)
|
|
|
|
assert stream is sentinel_stream
|
|
assert response.model == "deepseek/deepseek-v4-pro"
|
|
assert len(calls) == 1
|
|
assert calls[0]["stream"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_litellm_streaming_tool_call_id_retry_opens_one_retry_stream(monkeypatch):
|
|
calls = []
|
|
sentinel_stream = object()
|
|
|
|
async def fake_acompletion(**kwargs):
|
|
calls.append(kwargs.copy())
|
|
if len(calls) == 1:
|
|
raise Exception("Invalid 'messages': tool_call_id string too long maximum length")
|
|
return sentinel_stream
|
|
|
|
monkeypatch.setattr(
|
|
"cai.sdk.agents.models.chatcompletions.litellm_adapter.litellm.acompletion",
|
|
fake_acompletion,
|
|
)
|
|
|
|
long_id = "call_" + "x" * 80
|
|
kwargs = {
|
|
"model": "deepseek/deepseek-v4-pro",
|
|
"messages": [
|
|
{"role": "tool", "tool_call_id": long_id, "content": "ok"},
|
|
{
|
|
"role": "assistant",
|
|
"tool_calls": [
|
|
{
|
|
"id": long_id,
|
|
"type": "function",
|
|
"function": {"name": "probe", "arguments": "{}"},
|
|
}
|
|
],
|
|
},
|
|
],
|
|
"stream": True,
|
|
}
|
|
|
|
_response, stream = await fetch_response_litellm_openai(
|
|
kwargs=kwargs,
|
|
model_name="deepseek/deepseek-v4-pro",
|
|
model_settings=ModelSettings(),
|
|
tool_choice=NOT_GIVEN,
|
|
stream=True,
|
|
parallel_tool_calls=False,
|
|
)
|
|
|
|
assert stream is sentinel_stream
|
|
assert len(calls) == 2
|
|
assert kwargs["messages"][0]["tool_call_id"] == long_id[:40]
|
|
assert kwargs["messages"][1]["tool_calls"][0]["id"] == long_id[:40]
|