Fix Cloudflare Workers AI chat response parsing

This commit is contained in:
velhud 2026-06-06 08:32:32 +02:00
parent 22186f457a
commit 04d7966e41
3 changed files with 96 additions and 3 deletions

View file

@ -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:

View file

@ -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"}]

View file

@ -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"