diff --git a/assemblyai/__version__.py b/assemblyai/__version__.py index e4adfb83..f49459c7 100644 --- a/assemblyai/__version__.py +++ b/assemblyai/__version__.py @@ -1 +1 @@ -__version__ = "1.6.0" +__version__ = "1.6.1" diff --git a/assemblyai/llm_gateway/v1/models.py b/assemblyai/llm_gateway/v1/models.py index 8f1a3e60..761096ab 100644 --- a/assemblyai/llm_gateway/v1/models.py +++ b/assemblyai/llm_gateway/v1/models.py @@ -49,7 +49,8 @@ class Config: class LLMGatewayDefaultParameters(_LLMGatewayResponseModel): temperature: Optional[float] = None top_p: Optional[float] = None - frequency_penalty: Optional[int] = None + frequency_penalty: Optional[float] = None + presence_penalty: Optional[float] = None class LLMGatewayTopProvider(_LLMGatewayResponseModel): @@ -118,22 +119,22 @@ def set_collection_default(cls, v): class LLMGatewayCacheCreation(_LLMGatewayResponseModel): - ephemeral_5m_input_tokens: int - ephemeral_1h_input_tokens: int + ephemeral_5m_input_tokens: Optional[int] = None + ephemeral_1h_input_tokens: Optional[int] = None class LLMGatewayPromptTokensDetails(_LLMGatewayResponseModel): - cached_tokens: int - audio_tokens: int + cached_tokens: Optional[int] = None + audio_tokens: Optional[int] = None cache_creation: Optional[LLMGatewayCacheCreation] = None cache_write_tokens: Optional[int] = None class LLMGatewayCompletionTokensDetails(_LLMGatewayResponseModel): - reasoning_tokens: int - audio_tokens: int - accepted_prediction_tokens: int - rejected_prediction_tokens: int + reasoning_tokens: Optional[int] = None + audio_tokens: Optional[int] = None + accepted_prediction_tokens: Optional[int] = None + rejected_prediction_tokens: Optional[int] = None class LLMGatewayUsage(_LLMGatewayResponseModel): @@ -232,6 +233,7 @@ class LLMGatewayChunkToolCall(_LLMGatewayResponseModel): class LLMGatewayChunkDelta(_LLMGatewayResponseModel): content: Optional[str] = None role: Optional[str] = None + thinking: Optional[str] = None tool_calls: Optional[List[LLMGatewayChunkToolCall]] = None diff --git a/tests/unit/test_llm_gateway.py b/tests/unit/test_llm_gateway.py index 93961923..e1f7b292 100644 --- a/tests/unit/test_llm_gateway.py +++ b/tests/unit/test_llm_gateway.py @@ -757,3 +757,148 @@ def test_stream_chunk_exposes_tool_call_delta(httpx_mock: HTTPXMock): assert tool_call.function.name == "get_weather" # arguments stream in fragments; the SDK doesn't reassemble them assert tool_call.function.arguments == '{"ci' + + +def test_chat_completions_tolerates_partial_token_details(httpx_mock: HTTPXMock): + # Given a completion whose usage has partial details (e.g. reasoning models + # returning only reasoning_tokens without audio/speculative prediction tokens, + # and prompt caching returning only cached_tokens without audio tokens) + response = dict( + _COMPLETION_RESPONSE, + usage={ + "prompt_tokens": 100, + "completion_tokens": 50, + "total_tokens": 150, + "prompt_tokens_details": { + "cached_tokens": 80, + "cache_creation": {"ephemeral_5m_input_tokens": 40}, + }, + "completion_tokens_details": { + "reasoning_tokens": 30, + }, + }, + ) + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=response, + ) + + # When creating a completion + result = aai.LLMGateway().chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Explain quantum computing."}], + ) + + # Then partial details parse without error and missing fields are None + assert result.usage.prompt_tokens_details is not None + assert result.usage.prompt_tokens_details.cached_tokens == 80 + assert result.usage.prompt_tokens_details.audio_tokens is None + assert result.usage.prompt_tokens_details.cache_creation is not None + assert ( + result.usage.prompt_tokens_details.cache_creation.ephemeral_5m_input_tokens + == 40 + ) + assert ( + result.usage.prompt_tokens_details.cache_creation.ephemeral_1h_input_tokens + is None + ) + + assert result.usage.completion_tokens_details is not None + assert result.usage.completion_tokens_details.reasoning_tokens == 30 + assert result.usage.completion_tokens_details.audio_tokens is None + assert result.usage.completion_tokens_details.accepted_prediction_tokens is None + assert result.usage.completion_tokens_details.rejected_prediction_tokens is None + + +def test_stream_chunk_exposes_thinking_delta(httpx_mock: HTTPXMock): + # Given a streaming chunk carrying a thinking delta + chunk = { + "id": "chatcmpl-thinking-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "claude-sonnet-5", + "choices": [ + { + "index": 0, + "delta": { + "role": "assistant", + "thinking": "Let's break down this problem step by step.", + }, + "finish_reason": None, + } + ], + } + import json as _json + + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + stream=IteratorStream( + [f"data: {_json.dumps(chunk)}\n\n".encode(), b"data: [DONE]\n\n"] + ), + headers={"content-type": "text/event-stream"}, + ) + + # When consuming the stream + chunks = list( + aai.LLMGateway().chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hello"}], + stream=True, + ) + ) + + # Then thinking delta is parsed and reachable + delta = chunks[0].choices[0].delta + assert delta.thinking == "Let's break down this problem step by step." + + +def test_list_models_tolerates_float_penalties_and_presence_penalty( + httpx_mock: HTTPXMock, +): + # Given a models endpoint returning float frequency_penalty and presence_penalty + response = { + "data": [ + { + "id": "gpt-4o", + "name": "GPT-4o", + "description": "OpenAI's GPT-4o", + "default_parameters": { + "temperature": 0.7, + "top_p": 0.95, + "frequency_penalty": 0.5, + "presence_penalty": 0.2, + }, + "supported_parameters": ["temperature", "frequency_penalty"], + "top_provider": { + "is_moderated": True, + "context_length": 128000, + "max_completion_tokens": 4096, + }, + "context_length": 128000, + "pricing": { + "global": {"completions": 10.0, "prompt": 2.5}, + }, + "creator": "openai", + "retirement_date": 0, + "default_provider": {"id": "openai", "name": "Open AI"}, + } + ] + } + httpx_mock.add_response( + url=MODELS_URL, + method="GET", + status_code=httpx.codes.OK, + json=response, + ) + + # When listing models + result = aai.LLMGateway().models.list() + + # Then float penalties are preserved + model = result.data[0] + assert model.default_parameters.frequency_penalty == 0.5 + assert model.default_parameters.presence_penalty == 0.2