diff --git a/tests/llm_translation/test_cometapi_chat_transformation.py b/tests/llm_translation/test_cometapi_chat_transformation.py deleted file mode 100644 index eb81e513f40..00000000000 --- a/tests/llm_translation/test_cometapi_chat_transformation.py +++ /dev/null @@ -1,113 +0,0 @@ -""" -Unit tests for CometAPI Chat Configuration - -Tests the CometAPIChatConfig class methods using mocks -""" - -import os - -import pytest - -# Integration test example (requires real API key) - - -def test_cometapi_streaming_integration(): - """ - Integration test for streaming - requires real API key - Run with: pytest -k test_cometapi_streaming_integration -s - """ - from litellm import completion - - # Try to get API key from multiple environment variables - api_key = ( - os.getenv("COMETAPI_API_KEY") - or os.getenv("COMETAPI_KEY") - or os.getenv("COMET_API_KEY") - ) - - if not api_key: - pytest.skip("COMETAPI_API_KEY not set - skipping streaming integration test") - - try: - print( - f"🔍 Testing streaming with API key: {api_key[:6]}...{api_key[-4:]} (length: {len(api_key)})" - ) - print(f"🔍 API base URL: {os.getenv('COMETAPI_API_BASE', 'default')}") - - # test streaming API call - response = completion( - model="cometapi/gpt-3.5-turbo", - messages=[{"role": "user", "content": "Count from 1 to 5"}], - api_key=api_key, - max_tokens=50, - stream=True, - ) - - # collect streaming response - chunks = [] - content_parts = [] - - for chunk in response: - chunks.append(chunk) - if chunk.choices[0].delta.content: - content_parts.append(chunk.choices[0].delta.content) - - # Verify we received at least one chunk and content - assert len(chunks) > 0, "Should receive at least one chunk" - assert len(content_parts) > 0, "Should receive content in chunks" - - full_content = "".join(content_parts) - assert len(full_content.strip()) > 0, "Should have non-empty content" - - print(f"✅ Received {len(chunks)} chunks") - print(f"✅ Full content: {full_content}") - - except Exception as e: - print(f"❌ Streaming integration test error details:") - print(f" Error type: {type(e).__name__}") - print(f" Error message: {str(e)}") - if hasattr(e, "status_code"): - print(f" Status code: {e.status_code}") - if hasattr(e, "response"): - print(f" Response: {e.response}") - - # Re-raise with more context for pytest - pytest.fail(f"Streaming integration test failed: {type(e).__name__}: {str(e)}") - - -def test_cometapi_with_custom_base_url(): - """ - Test CometAPI with custom base URL - """ - from litellm import completion - - api_key = ( - os.getenv("COMETAPI_API_KEY") - or os.getenv("COMETAPI_KEY") - or os.getenv("COMET_API_KEY") - ) - - custom_base_url = os.getenv("COMETAPI_API_BASE", "https://api.cometapi.com/v1") - - if not api_key: - pytest.skip("COMETAPI_API_KEY not set - skipping custom base URL test") - - try: - response = completion( - model="cometapi/gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello"}], - api_key=api_key, - api_base=custom_base_url, - max_tokens=5, - ) - - assert response.choices[0].message.content - print(f"✅ Custom base URL test passed: {response.choices[0].message.content}") - - except Exception as e: - pytest.fail(f"Custom base URL test failed: {str(e)}") - - -if __name__ == "__main__": - # Quick test runner - pytest.main([__file__, "-v"]) diff --git a/tests/llm_translation/test_xai.py b/tests/llm_translation/test_xai.py deleted file mode 100644 index 74ddf6b7536..00000000000 --- a/tests/llm_translation/test_xai.py +++ /dev/null @@ -1,118 +0,0 @@ -import json -import os -from datetime import datetime -from unittest.mock import AsyncMock, patch - -import httpx -import pytest -from base_llm_unit_tests import BaseLLMChatTest, BaseReasoningLLMTests - -from litellm import Choices, EmbeddingResponse, Message, ModelResponse, Usage, completion -from litellm.llms.xai.chat.transformation import XAI_API_BASE, XAIChatConfig - - - - - - - - - - - - -def test_xai_message_name_filtering(): - messages = [ - { - "role": "system", - "content": "*I press the green button*", - "name": "example_user", - }, - {"role": "user", "content": "Hello", "name": "John"}, - {"role": "assistant", "content": "Hello", "name": "Jane"}, - ] - response = completion( - model="xai/grok-3-mini-beta", - messages=messages, - ) - assert response is not None - assert response.choices[0].message.content is not None - - -class TestXAIReasoningEffort(BaseReasoningLLMTests): - def get_base_completion_call_args(self): - return { - "model": "xai/grok-4.6", - "messages": [ - {"role": "user", "content": "What is 17 times 23? Work it out step by step, then give the number."} - ], - } - - -class TestXAIChat(BaseLLMChatTest): - def get_base_completion_call_args(self): - return { - "model": "xai/grok-3-mini-beta", - } - - test_web_search = None - - -def test_xai_streaming_with_include_usage(): - """ - Test that xAI streaming correctly handles usage in the last chunk - when stream_options={"include_usage": True} is set. - - xAI sends usage in a chunk with empty choices array, which should be - handled by XAIChatCompletionStreamingHandler. - """ - try: - response = completion( - model="xai/grok-4-1-fast-non-reasoning", - messages=[ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "Say hello in one word"}, - ], - stream=True, - stream_options={"include_usage": True}, - max_tokens=10, - ) - - chunks = [] - usage_chunk = None - - for chunk in response: - chunks.append(chunk) - if hasattr(chunk, "usage") and chunk.usage is not None: - usage_chunk = chunk - - # Verify we got chunks - assert len(chunks) > 0, "Should receive streaming chunks" - - # Verify usage was included in one of the chunks - assert usage_chunk is not None, "Should receive usage in streaming chunks" - - # Verify usage has expected fields - assert hasattr( - usage_chunk.usage, "prompt_tokens" - ), "Usage should have prompt_tokens" - assert hasattr( - usage_chunk.usage, "completion_tokens" - ), "Usage should have completion_tokens" - assert hasattr( - usage_chunk.usage, "total_tokens" - ), "Usage should have total_tokens" - - # Verify usage values are positive - assert usage_chunk.usage.prompt_tokens > 0, "prompt_tokens should be positive" - assert ( - usage_chunk.usage.completion_tokens > 0 - ), "completion_tokens should be positive" - assert usage_chunk.usage.total_tokens > 0, "total_tokens should be positive" - - print(f"✓ Successfully received usage in streaming chunk: {usage_chunk.usage}") - - except Exception as e: - if "API key" in str(e) or "authentication" in str(e).lower(): - pytest.skip(f"Skipping test due to API key issue: {str(e)}") - raise diff --git a/tests/unit/llms/cometapi/chat/test_cometapi_chat_transformation.py b/tests/unit/llms/cometapi/chat/test_cometapi_chat_transformation.py index 607648dd6c9..ae84ac65f51 100644 --- a/tests/unit/llms/cometapi/chat/test_cometapi_chat_transformation.py +++ b/tests/unit/llms/cometapi/chat/test_cometapi_chat_transformation.py @@ -4,10 +4,14 @@ Unit tests for CometAPI Chat Configuration Tests the CometAPIChatConfig class methods using mocks """ +import json +from typing import Final +import httpx import pytest +import respx - +import litellm from litellm.llms.cometapi.chat.transformation import ( CometAPIChatCompletionStreamingHandler, CometAPIConfig, @@ -44,6 +48,7 @@ class TestCometAPIChatCompletionStreamingHandler: assert result.usage.completion_tokens == chunk["usage"]["completion_tokens"] assert result.usage.total_tokens == chunk["usage"]["total_tokens"] assert len(result.choices) == 1 + assert result.choices[0]["delta"]["content"] == "test content" assert result.choices[0]["delta"]["reasoning_content"] == "test reasoning" def test_chunk_parser_error_response(self): @@ -175,6 +180,42 @@ class TestCometAPIConfig: assert error.status_code == 400 +@pytest.mark.respx(assert_all_called=True) +def test_completion_uses_custom_api_base(respx_mock: respx.MockRouter) -> None: + api_base: Final = "https://cometapi.example/custom/v1" + route: Final = respx_mock.post(f"{api_base}/chat/completions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "chatcmpl-cometapi", + "object": "chat.completion", + "created": 1234567890, + "model": "gpt-5.6", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "OK"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3}, + }, + ) + ) + + response: Final = litellm.completion( + model="cometapi/gpt-5.6", + messages=[{"role": "user", "content": "Say OK"}], + api_key="cometapi-test-key", + api_base=api_base, + ) + + request_body: Final = json.loads(route.calls.last.request.content) + assert str(route.calls.last.request.url) == f"{api_base}/chat/completions" + assert request_body["messages"] == [{"role": "user", "content": "Say OK"}] + assert response.choices[0].message.content == "OK" + + # Integration test example (requires real API key) diff --git a/tests/unit/llms/databricks/test_databricks.py b/tests/unit/llms/databricks/test_databricks.py index 3301d48c3e9..526f0f72e93 100644 --- a/tests/unit/llms/databricks/test_databricks.py +++ b/tests/unit/llms/databricks/test_databricks.py @@ -1,11 +1,13 @@ import asyncio import json import sys -from typing import Any, Dict, List +from types import ModuleType, SimpleNamespace +from typing import Any, Dict, Final, List from unittest.mock import MagicMock, Mock, patch import httpx import pytest +import respx import litellm from litellm._version import version @@ -21,6 +23,21 @@ except ImportError: databricks_sdk_installed = False +def _install_databricks_sdk(monkeypatch: pytest.MonkeyPatch, host: str, token: str) -> None: + config: Final = SimpleNamespace( + host=host, + authenticate=lambda: {"Authorization": f"Bearer {token}"}, + ) + sdk_module: Final = ModuleType("databricks.sdk") + package_module: Final = ModuleType("databricks") + package_module.__path__ = [] + setattr(sdk_module, "WorkspaceClient", lambda: SimpleNamespace(config=config)) + setattr(sdk_module, "useragent", SimpleNamespace(with_partner=lambda _: None)) + setattr(package_module, "sdk", sdk_module) + monkeypatch.setitem(sys.modules, "databricks", package_module) + monkeypatch.setitem(sys.modules, "databricks.sdk", sdk_module) + + def mock_chat_response() -> Dict[str, Any]: return { "id": "chatcmpl_3f78f09a-489c-4b8d-a587-f162c7497891", @@ -536,6 +553,62 @@ def test_embeddings_uses_databricks_sdk_if_api_key_and_base_not_specified(monkey ) +@pytest.mark.respx(assert_all_called=True) +def test_chat_completion_uses_databricks_sdk_credentials( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.delenv("DATABRICKS_API_BASE", raising=False) + monkeypatch.delenv("DATABRICKS_API_KEY", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + host: Final = "https://my.workspace.cloud.databricks.com" + token: Final = "sdk-test-token" + _install_databricks_sdk(monkeypatch, host, token) + messages: Final = [{"role": "user", "content": "How are you?"}] + route: Final = respx_mock.post(f"{host}/serving-endpoints/chat/completions").mock( + return_value=httpx.Response(200, json=mock_chat_response()) + ) + + response: Final = litellm.completion( + model="databricks/databricks-gpt-5-6-sol", + messages=messages, + temperature=0.5, + ) + + request_body: Final = json.loads(route.calls.last.request.content) + assert str(route.calls.last.request.url) == f"{host}/serving-endpoints/chat/completions" + assert route.calls.last.request.headers["Authorization"] == f"Bearer {token}" + assert request_body["messages"] == messages + assert request_body["temperature"] == 0.5 + assert response.choices[0].message.content == mock_chat_response()["choices"][0]["message"]["content"] + + +@pytest.mark.respx(assert_all_called=True) +def test_embedding_uses_databricks_sdk_credentials( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.delenv("DATABRICKS_API_BASE", raising=False) + monkeypatch.delenv("DATABRICKS_API_KEY", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + host: Final = "https://my.workspace.cloud.databricks.com" + token: Final = "sdk-test-token" + _install_databricks_sdk(monkeypatch, host, token) + inputs: Final = ["hello"] + route: Final = respx_mock.post(f"{host}/serving-endpoints/embeddings").mock( + return_value=httpx.Response(200, json=mock_embedding_response()) + ) + + response: Final = litellm.embedding( + model="databricks/databricks-qwen3-embedding-0-6b", + input=inputs, + ) + + request_body: Final = json.loads(route.calls.last.request.content) + assert str(route.calls.last.request.url) == f"{host}/serving-endpoints/embeddings" + assert route.calls.last.request.headers["Authorization"] == f"Bearer {token}" + assert request_body["input"] == inputs + assert response.to_dict()["data"] == mock_embedding_response()["data"] def test_completion_with_prompt_caching_anthropic_model(monkeypatch): diff --git a/tests/unit/llms/xai/test_xai_chat_transformation.py b/tests/unit/llms/xai/test_xai_chat_transformation.py index 3065f24ab7b..570bf9beeba 100644 --- a/tests/unit/llms/xai/test_xai_chat_transformation.py +++ b/tests/unit/llms/xai/test_xai_chat_transformation.py @@ -1,9 +1,11 @@ +import json import os from typing import Final from unittest.mock import Mock, patch import httpx import pytest +import respx import litellm from litellm.llms.xai.chat.transformation import ( @@ -167,6 +169,102 @@ class TestXAIUsageNormalization: assert usage["total_tokens"] == 200 +@pytest.mark.respx(assert_all_called=True) +def test_completion_strips_message_names_from_request(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(f"{XAI_API_BASE}/chat/completions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "chatcmpl-xai", + "object": "chat.completion", + "created": 1234567890, + "model": "grok-4.20-beta-latest", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "OK"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3}, + }, + ) + ) + + response: Final = litellm.completion( + model="xai/grok-4.20-beta-latest", + messages=[ + {"role": "system", "content": "Be concise", "name": "system_prompt"}, + {"role": "user", "content": "Say OK", "name": "caller"}, + ], + api_key="xai-test-key", + api_base=XAI_API_BASE, + ) + + request_body: Final = json.loads(route.calls.last.request.content) + assert request_body["messages"] == [ + {"role": "system", "content": "Be concise"}, + {"role": "user", "content": "Say OK", "name": "caller"}, + ] + assert response.choices[0].message.content == "OK" + + +@pytest.mark.respx(assert_all_called=True) +def test_streaming_include_usage_surfaces_xai_final_usage_chunk(respx_mock: respx.MockRouter) -> None: + frames: Final = ( + { + "id": "chatcmpl-xai-stream", + "object": "chat.completion.chunk", + "created": 1234567890, + "model": "grok-4.20-beta-latest", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "OK"}, "finish_reason": None}], + }, + { + "id": "chatcmpl-xai-stream", + "object": "chat.completion.chunk", + "created": 1234567890, + "model": "grok-4.20-beta-latest", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + }, + { + "id": "chatcmpl-xai-stream", + "object": "chat.completion.chunk", + "created": 1234567890, + "model": "grok-4.20-beta-latest", + "choices": [], + "usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3}, + }, + ) + sse_body: Final = "".join(f"data: {json.dumps(frame)}\n\n" for frame in frames) + "data: [DONE]\n\n" + respx_mock.post(f"{XAI_API_BASE}/chat/completions").mock( + return_value=httpx.Response( + 200, + content=sse_body.encode(), + headers={"content-type": "text/event-stream"}, + ) + ) + + chunks: Final = list( + litellm.completion( + model="xai/grok-4.20-beta-latest", + messages=[{"role": "user", "content": "Say OK"}], + stream=True, + stream_options={"include_usage": True}, + api_key="xai-test-key", + api_base=XAI_API_BASE, + ) + ) + + usage_chunks: Final = [chunk for chunk in chunks if getattr(chunk, "usage", None) is not None] + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "OK" + assert len(usage_chunks) == 1 + assert ( + usage_chunks[0].usage.prompt_tokens, + usage_chunks[0].usage.completion_tokens, + usage_chunks[0].usage.total_tokens, + ) == (2, 1, 3) + + class TestXAIChatWebSearchBilling: _TOOL_DETAILS = { "web_search_calls": 3,