mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
Fix Cloudflare Workers AI chat response parsing
This commit is contained in:
parent
22186f457a
commit
04d7966e41
3 changed files with 96 additions and 3 deletions
|
|
@ -71,7 +71,7 @@ class CloudflareChatConfig(BaseConfig):
|
|||
)
|
||||
headers = {
|
||||
"accept": "application/json",
|
||||
"content-type": "apbplication/json",
|
||||
"content-type": "application/json",
|
||||
"Authorization": "Bearer " + api_key,
|
||||
}
|
||||
return headers
|
||||
|
|
@ -149,9 +149,12 @@ class CloudflareChatConfig(BaseConfig):
|
|||
) -> ModelResponse:
|
||||
completion_response = raw_response.json()
|
||||
|
||||
# Support both "response" and "response_text" keys (newer models like Nemotron use "response_text")
|
||||
# Support Workers AI response variants:
|
||||
# - result.response
|
||||
# - result.response_text
|
||||
# - result.choices[0].message.content for OpenAI-compatible chat models
|
||||
result = completion_response["result"]
|
||||
model_response.choices[0].message.content = result.get("response") if result.get("response") is not None else result.get("response_text", "") # type: ignore
|
||||
model_response.choices[0].message.content = self._get_response_text(result) # type: ignore
|
||||
|
||||
prompt_tokens = litellm.utils.get_token_count(messages=messages, model=model)
|
||||
completion_tokens = len(
|
||||
|
|
@ -168,6 +171,28 @@ class CloudflareChatConfig(BaseConfig):
|
|||
setattr(model_response, "usage", usage)
|
||||
return model_response
|
||||
|
||||
@staticmethod
|
||||
def _get_response_text(result: dict) -> str:
|
||||
response = result.get("response")
|
||||
if response is not None:
|
||||
return response
|
||||
response_text = result.get("response_text")
|
||||
if response_text is not None:
|
||||
return response_text
|
||||
choices = result.get("choices")
|
||||
if isinstance(choices, list) and len(choices) > 0:
|
||||
first_choice = choices[0]
|
||||
if isinstance(first_choice, dict):
|
||||
message = first_choice.get("message")
|
||||
if isinstance(message, dict):
|
||||
content = message.get("content")
|
||||
if content is not None:
|
||||
return content
|
||||
text = first_choice.get("text")
|
||||
if text is not None:
|
||||
return text
|
||||
return ""
|
||||
|
||||
def get_error_class(
|
||||
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
|
||||
) -> BaseLLMException:
|
||||
|
|
|
|||
|
|
@ -35,6 +35,24 @@ def _chat_response() -> Dict[str, Any]:
|
|||
}
|
||||
|
||||
|
||||
def _chat_choices_response() -> Dict[str, Any]:
|
||||
return {
|
||||
"result": {
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "I am a large language model created to assist you.",
|
||||
}
|
||||
}
|
||||
],
|
||||
},
|
||||
"success": True,
|
||||
"errors": [],
|
||||
"messages": [],
|
||||
}
|
||||
|
||||
|
||||
def _streaming_chunks() -> list[str]:
|
||||
return [
|
||||
json.dumps({"response": "I am"}),
|
||||
|
|
@ -86,6 +104,41 @@ def test_completion_cloudflare(sync_mode):
|
|||
assert "language model" in response.choices[0].message.content.lower()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
def test_completion_cloudflare_choices_response(sync_mode):
|
||||
messages = [{"role": "user", "content": "what llm are you"}]
|
||||
mock_resp = _make_mock_response(_chat_choices_response())
|
||||
|
||||
if sync_mode:
|
||||
with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post:
|
||||
response = completion(
|
||||
model="cloudflare/@cf/openai/gpt-oss-20b",
|
||||
messages=messages,
|
||||
max_tokens=15,
|
||||
api_base=FAKE_API_BASE,
|
||||
api_key=FAKE_API_KEY,
|
||||
)
|
||||
mock_post.assert_called_once()
|
||||
else:
|
||||
with patch.object(
|
||||
AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp
|
||||
) as mock_post:
|
||||
response = asyncio.run(
|
||||
acompletion(
|
||||
model="cloudflare/@cf/openai/gpt-oss-20b",
|
||||
messages=messages,
|
||||
max_tokens=15,
|
||||
api_base=FAKE_API_BASE,
|
||||
api_key=FAKE_API_KEY,
|
||||
)
|
||||
)
|
||||
mock_post.assert_called_once()
|
||||
|
||||
assert response is not None
|
||||
assert response.choices[0].message.content is not None
|
||||
assert "language model" in response.choices[0].message.content.lower()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
def test_completion_cloudflare_stream(sync_mode):
|
||||
messages = [{"role": "user", "content": "what llm are you"}]
|
||||
|
|
|
|||
|
|
@ -25,3 +25,18 @@ def test_get_complete_url_encodes_model_path_segment():
|
|||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
|
||||
|
||||
def test_validate_environment_sets_json_content_type():
|
||||
config = CloudflareChatConfig()
|
||||
|
||||
headers = config.validate_environment(
|
||||
headers={},
|
||||
model="@cf/meta/llama-2-7b-chat-int8",
|
||||
messages=[{"role": "user", "content": "ping"}],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
api_key="cf-key",
|
||||
)
|
||||
|
||||
assert headers["content-type"] == "application/json"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue