import asyncio import json from typing import Any, Dict from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest from litellm import acompletion, completion from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler FAKE_API_BASE = "https://fake-cloudflare.example.com/client/v4/accounts/fake-acct/ai/v1" FAKE_API_KEY = "fake-cf-api-key" def _make_mock_response(json_data: Dict[str, Any]) -> MagicMock: mock = MagicMock(spec=httpx.Response) mock.status_code = 200 mock.headers = {"content-type": "application/json"} mock.json.return_value = json_data mock.text = json.dumps(json_data) return mock def _chat_response() -> Dict[str, Any]: return { "id": "chatcmpl-cf", "object": "chat.completion", "created": 1234567890, "model": "@cf/meta/llama-2-7b-chat-int8", "choices": [ { "index": 0, "message": { "role": "assistant", "content": "I am a large language model created to assist you.", }, "finish_reason": "stop", } ], "usage": {"prompt_tokens": 8, "completion_tokens": 11, "total_tokens": 19}, } def _tool_call_response() -> Dict[str, Any]: return { "id": "chatcmpl-cf-tools", "object": "chat.completion", "created": 1234567890, "model": "@cf/meta/llama-2-7b-chat-int8", "choices": [ { "index": 0, "message": { "role": "assistant", "content": None, "tool_calls": [ { "id": "call_1", "type": "function", "function": { "name": "get_weather", "arguments": '{"city": "New York"}', }, } ], }, "finish_reason": "tool_calls", } ], "usage": {"prompt_tokens": 20, "completion_tokens": 9, "total_tokens": 29}, } def _streaming_chunks() -> list[str]: base = { "id": "chatcmpl-cf", "object": "chat.completion.chunk", "created": 1234567890, "model": "@cf/meta/llama-2-7b-chat-int8", } return [ json.dumps({**base, "choices": [{"index": 0, "delta": {"content": "I am"}}]}), json.dumps( {**base, "choices": [{"index": 0, "delta": {"content": " a language"}}]} ), json.dumps( { **base, "choices": [ { "index": 0, "delta": {"content": " model."}, "finish_reason": "stop", } ], } ), ] @pytest.mark.parametrize("sync_mode", [True, False]) def test_completion_cloudflare(sync_mode): messages = [{"role": "user", "content": "what llm are you"}] mock_resp = _make_mock_response(_chat_response()) if sync_mode: with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post: response = completion( model="cloudflare/@cf/meta/llama-2-7b-chat-int8", 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/meta/llama-2-7b-chat-int8", 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() called_url = mock_post.call_args.kwargs.get("url") or mock_post.call_args.args[0] assert called_url.endswith("/ai/v1/chat/completions") assert "/ai/run/" not in called_url def test_completion_cloudflare_tool_calls_sent_to_openai_endpoint(): messages = [{"role": "user", "content": "weather in New York?"}] tools = [ { "type": "function", "function": { "name": "get_weather", "parameters": { "type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"], }, }, } ] mock_resp = _make_mock_response(_tool_call_response()) with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post: response = completion( model="cloudflare/@cf/meta/llama-2-7b-chat-int8", messages=messages, tools=tools, tool_choice="auto", api_base=FAKE_API_BASE, api_key=FAKE_API_KEY, ) mock_post.assert_called_once() sent_body = json.loads(mock_post.call_args.kwargs["data"]) assert sent_body["tools"] == tools assert sent_body["tool_choice"] == "auto" assert response.choices[0].finish_reason == "tool_calls" tool_calls = response.choices[0].message.tool_calls assert tool_calls is not None and len(tool_calls) == 1 assert tool_calls[0].function.name == "get_weather" @pytest.mark.parametrize("sync_mode", [True, False]) def test_completion_cloudflare_stream(sync_mode): messages = [{"role": "user", "content": "what llm are you"}] raw_chunks = _streaming_chunks() if sync_mode: def _iter_lines(): for chunk in raw_chunks: yield f"data: {chunk}" yield "data: [DONE]" mock_resp = MagicMock() mock_resp.iter_lines.return_value = _iter_lines() mock_resp.status_code = 200 mock_resp.headers = {"content-type": "text/event-stream"} with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post: response = completion( model="cloudflare/@cf/meta/llama-2-7b-chat-int8", messages=messages, max_tokens=15, stream=True, api_base=FAKE_API_BASE, api_key=FAKE_API_KEY, ) chunks_received = list(response) mock_post.assert_called_once() else: async def _aiter_lines(): for chunk in raw_chunks: yield f"data: {chunk}" yield "data: [DONE]" mock_resp = MagicMock() mock_resp.aiter_lines.return_value = _aiter_lines() mock_resp.status_code = 200 mock_resp.headers = {"content-type": "text/event-stream"} async def _run(): with patch.object( AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp ) as mock_post: resp = await acompletion( model="cloudflare/@cf/meta/llama-2-7b-chat-int8", messages=messages, max_tokens=15, stream=True, api_base=FAKE_API_BASE, api_key=FAKE_API_KEY, ) received = [] async for chunk in resp: received.append(chunk) mock_post.assert_called_once() return received chunks_received = asyncio.run(_run()) assert len(chunks_received) > 0 content = "".join( c.choices[0].delta.content for c in chunks_received if c.choices[0].delta.content ) assert "language" in content.lower()