fix(platform): Also handle client_name inside _acompletion (#808)

Author: daavoo

src/any_llm/providers/platform/platform.py

@@ -103,6 +103,10 @@ async def _acompletion(
         params: CompletionParams,
         **kwargs: Any,
     ) -> ChatCompletion | AsyncIterator[ChatCompletionChunk]:
+        client_name = kwargs.pop("client_name", None)
+        if self.client_name is None:
+            self.client_name = client_name
+
         start_time = time.perf_counter()
 
         # List of providers that don't support stream_options and automatically return token usage.

tests/unit/providers/test_platform_provider.py

@@ -1158,3 +1158,80 @@ async def test_usage_event_includes_version_header(
 
     assert "User-Agent" in headers
     assert headers["User-Agent"] == f"python-any-llm/{__version__}"
+
+
+@pytest.mark.asyncio
+@patch("any_llm_platform_client.AnyLLMPlatformClient.get_decrypted_provider_key")
+@patch("any_llm.providers.platform.platform.post_completion_usage_event")
+async def test_acompletion_handles_client_name_in_kwargs(
+    mock_post_usage: AsyncMock,
+    mock_get_decrypted_provider_key: Mock,
+    any_llm_key: str,
+    mock_decrypted_provider_key: DecryptedProviderKey,
+    mock_completion: ChatCompletion,
+) -> None:
+    """Test that _acompletion correctly handles client_name passed in kwargs."""
+    mock_get_decrypted_provider_key.return_value = mock_decrypted_provider_key
+
+    provider_instance = PlatformProvider(api_key=any_llm_key)
+    provider_instance.provider = OpenaiProvider
+    provider_instance.provider._acompletion = AsyncMock(return_value=mock_completion)  # type: ignore[method-assign]
+
+    # Create completion params
+    params = CompletionParams(
+        model_id="gpt-4",
+        messages=[{"role": "user", "content": "Hello"}],
+        stream=False,
+    )
+
+    client_name = "test-client-from-kwargs"
+
+    # Call _acompletion with client_name in kwargs
+    await provider_instance._acompletion(params, client_name=client_name)
+
+    # Verify self.client_name was updated
+    assert provider_instance.client_name == client_name
+
+    # Verify post_completion_usage_event was called with the correct client_name
+    mock_post_usage.assert_called_once()
+    call_args = mock_post_usage.call_args
+    assert call_args.kwargs["client_name"] == client_name
+
+
+@pytest.mark.asyncio
+@patch("any_llm_platform_client.AnyLLMPlatformClient.get_decrypted_provider_key")
+@patch("any_llm.providers.platform.platform.post_completion_usage_event")
+async def test_acompletion_does_not_overwrite_existing_client_name(
+    mock_post_usage: AsyncMock,
+    mock_get_decrypted_provider_key: Mock,
+    any_llm_key: str,
+    mock_decrypted_provider_key: DecryptedProviderKey,
+    mock_completion: ChatCompletion,
+) -> None:
+    """Test that _acompletion does not overwrite an existing client_name if one is already set."""
+    mock_get_decrypted_provider_key.return_value = mock_decrypted_provider_key
+
+    initial_client_name = "initial-client"
+    provider_instance = PlatformProvider(api_key=any_llm_key, client_name=initial_client_name)
+    provider_instance.provider = OpenaiProvider
+    provider_instance.provider._acompletion = AsyncMock(return_value=mock_completion)  # type: ignore[method-assign]
+
+    # Create completion params
+    params = CompletionParams(
+        model_id="gpt-4",
+        messages=[{"role": "user", "content": "Hello"}],
+        stream=False,
+    )
+
+    new_client_name = "new-client-from-kwargs"
+
+    # Call _acompletion with a new client_name in kwargs
+    await provider_instance._acompletion(params, client_name=new_client_name)
+
+    # Verify self.client_name was NOT updated
+    assert provider_instance.client_name == initial_client_name
+
+    # Verify post_completion_usage_event was called with the INITIAL client_name
+    mock_post_usage.assert_called_once()
+    call_args = mock_post_usage.call_args
+    assert call_args.kwargs["client_name"] == initial_client_name