diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index ea1191a5331..59cd4c821e4 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -388,10 +388,14 @@ class OllamaChatConfig(BaseConfig): "eval_count", litellm.token_counter(text=response_json["message"]["content"]), ) - model_response.usage = litellm.Usage( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=prompt_tokens + completion_tokens, + setattr( + model_response, + "usage", + litellm.Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ), ) return model_response @@ -477,7 +481,7 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): function_args = tool_call.get("function").get("arguments") if function_args is not None and len(function_args) > 0: is_function_call_complete = self._is_function_call_complete(function_args) - if is_function_call_complete and not tool_call.get("id"): + if is_function_call_complete and "id" not in tool_call: tool_call["id"] = str(uuid.uuid4()) # PROCESS REASONING CONTENT diff --git a/tests/unit/llms/ollama/test_ollama_chat_transformation.py b/tests/unit/llms/ollama/test_ollama_chat_transformation.py index 8cb365b61ba..3f80c84bfd6 100644 --- a/tests/unit/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_chat_transformation.py @@ -7,7 +7,9 @@ from typing import Final, cast import pytest from pydantic import BaseModel -sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))) +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) +) from litellm.llms.ollama.chat.transformation import ( OllamaChatConfig, @@ -45,9 +47,9 @@ class TestOllamaChatConfigResponseFormat: expected_schema_structure = TestEvent.model_json_schema() transformed_format.pop("additionalProperties") - assert transformed_format == expected_schema_structure, ( - f"Transformed schema does not match expected. Got: {transformed_format}, Expected: {expected_schema_structure}" - ) + assert ( + transformed_format == expected_schema_structure + ), f"Transformed schema does not match expected. Got: {transformed_format}, Expected: {expected_schema_structure}" def test_map_openai_params_with_dict_json_schema(self): config = OllamaChatConfig() @@ -67,9 +69,9 @@ class TestOllamaChatConfigResponseFormat: ) assert "format" in optional_params - assert optional_params["format"] == direct_schema, ( - f"Schema from dict did not pass through correctly. Got: {optional_params['format']}, Expected: {direct_schema}" - ) + assert ( + optional_params["format"] == direct_schema + ), f"Schema from dict did not pass through correctly. Got: {optional_params['format']}, Expected: {direct_schema}" def test_map_openai_params_with_json_object(self): optional_params = get_optional_params( @@ -79,9 +81,9 @@ class TestOllamaChatConfigResponseFormat: ) assert "format" in optional_params - assert optional_params["format"] == "json", ( - f"Expected 'json' for type 'json_object', got: {optional_params['format']}" - ) + assert ( + optional_params["format"] == "json" + ), f"Expected 'json' for type 'json_object', got: {optional_params['format']}" def test_transform_request_loads_config_parameters(self): """Test that transform_request loads config parameters without overriding existing optional_params""" @@ -149,7 +151,9 @@ class TestOllamaChatConfigResponseFormat: config = OllamaChatConfig() # Test message with content as string - messages = cast(list[AllMessageValues], [{"role": "user", "content": "Hello world!"}]) + messages = cast( + list[AllMessageValues], [{"role": "user", "content": "Hello world!"}] + ) result = config.transform_request( model="llama2", @@ -198,7 +202,9 @@ class TestOllamaChatConfigResponseFormat: {"type": "text", "text": "What's in this image?"}, { "type": "image_url", - "image_url": {"url": "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ..."}, + "image_url": { + "url": "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ..." + }, }, ], } @@ -238,7 +244,9 @@ class TestOllamaChatConfigResponseFormat: {"type": "text", "text": "Compare these images:"}, { "type": "image_url", - "image_url": {"url": "data:image/jpeg;base64,image1data..."}, + "image_url": { + "url": "data:image/jpeg;base64,image1data..." + }, }, {"type": "text", "text": " and "}, { @@ -363,7 +371,9 @@ class TestOllamaToolCalling: assert optional_params["tools"] == tools # Should NOT trigger the broken fallback assert "functions_unsupported_model" not in optional_params - assert "format" not in optional_params or optional_params.get("format") != "json" + assert ( + "format" not in optional_params or optional_params.get("format") != "json" + ) def test_finish_reason_tool_calls_non_streaming(self): """Test that finish_reason is set to 'tool_calls' when tool_calls present. @@ -519,9 +529,9 @@ class TestOllamaFinishReasonLength: json_mode=False, ) - assert result.choices[0].finish_reason == "length", ( - f"Expected 'length' when done_reason='length', got '{result.choices[0].finish_reason}'" - ) + assert ( + result.choices[0].finish_reason == "length" + ), f"Expected 'length' when done_reason='length', got '{result.choices[0].finish_reason}'" def test_finish_reason_stop_non_streaming(self): """Non-streaming: done_reason='stop' (natural finish) must stay 'stop'.""" @@ -560,9 +570,9 @@ class TestOllamaFinishReasonLength: json_mode=False, ) - assert result.choices[0].finish_reason == "stop", ( - f"Expected 'stop' for natural finish, got '{result.choices[0].finish_reason}'" - ) + assert ( + result.choices[0].finish_reason == "stop" + ), f"Expected 'stop' for natural finish, got '{result.choices[0].finish_reason}'" def test_finish_reason_length_streaming(self): """Streaming: done_reason='length' in final chunk must produce finish_reason='length'.""" @@ -583,9 +593,9 @@ class TestOllamaFinishReasonLength: result = iterator.chunk_parser(done_chunk) - assert result.choices[0].finish_reason == "length", ( - f"Expected 'length' when done_reason='length', got '{result.choices[0].finish_reason}'" - ) + assert ( + result.choices[0].finish_reason == "length" + ), f"Expected 'length' when done_reason='length', got '{result.choices[0].finish_reason}'" def test_finish_reason_stop_streaming(self): """Streaming: done_reason='stop' in final chunk must produce finish_reason='stop'.""" @@ -603,9 +613,9 @@ class TestOllamaFinishReasonLength: result = iterator.chunk_parser(done_chunk) - assert result.choices[0].finish_reason == "stop", ( - f"Expected 'stop' for natural finish, got '{result.choices[0].finish_reason}'" - ) + assert ( + result.choices[0].finish_reason == "stop" + ), f"Expected 'stop' for natural finish, got '{result.choices[0].finish_reason}'" def test_finish_reason_tool_calls_streamed_before_done_chunk(self): """Streaming: tool_calls arriving mid-stream (not on the done chunk) must @@ -626,7 +636,9 @@ class TestOllamaFinishReasonLength: "message": { "role": "assistant", "content": "", - "tool_calls": [{"function": {"name": "get_weather", "arguments": {"city": "San Francisco"}}}], + "tool_calls": [ + {"function": {"name": "get_weather", "arguments": {"city": "San Francisco"}}} + ], }, "done": False, } @@ -641,9 +653,9 @@ class TestOllamaFinishReasonLength: assert tool_result.choices[0].delta.tool_calls is not None done_result = iterator.chunk_parser(done_chunk) - assert done_result.choices[0].finish_reason == "tool_calls", ( - f"Expected 'tool_calls' when tool_calls were streamed earlier, got '{done_result.choices[0].finish_reason}'" - ) + assert ( + done_result.choices[0].finish_reason == "tool_calls" + ), f"Expected 'tool_calls' when tool_calls were streamed earlier, got '{done_result.choices[0].finish_reason}'" class TestOllamaReasoningContentStreaming: @@ -688,7 +700,7 @@ class TestOllamaReasoningContentStreaming: # Verify none of them are None for i, rc in enumerate(reasoning_contents): - assert rc is not None, f"Chunk {i + 1} reasoning_content should not be None" + assert rc is not None, f"Chunk {i+1} reasoning_content should not be None" def test_thinking_to_content_transition(self): """ @@ -706,7 +718,9 @@ class TestOllamaReasoningContentStreaming: "done": False, } result1 = iterator.chunk_parser(thinking_chunk) - assert result1.choices[0].delta.reasoning_content == "Let me think about this..." + assert ( + result1.choices[0].delta.reasoning_content == "Let me think about this..." + ) assert result1.choices[0].delta.content is None # Then: regular content chunk @@ -1057,7 +1071,9 @@ def test_get_ollama_params(): } print("Converted params", converted_params) for key in expected_params.keys(): - assert expected_params[key] == converted_params[key], f"{converted_params} != {expected_params}" + assert ( + expected_params[key] == converted_params[key] + ), f"{converted_params} != {expected_params}" except Exception as e: pytest.fail(f"Error occurred: {e}") @@ -1098,7 +1114,9 @@ mock_ollama_embedding_response = EmbeddingResponse(model="ollama/nomic-embed-tex def test_ollama_embeddings(mock_embeddings): # assert that ollama_embeddings is called with the right parameters try: - embeddings = litellm.embedding(model="ollama/nomic-embed-text", input=["hello world"]) + embeddings = litellm.embedding( + model="ollama/nomic-embed-text", input=["hello world"] + ) print(embeddings) mock_embeddings.assert_called_once_with( api_base="http://localhost:11434", @@ -1120,7 +1138,9 @@ def test_ollama_embeddings(mock_embeddings): def test_ollama_aembeddings(mock_aembeddings): # assert that ollama_aembeddings is called with the right parameters try: - embeddings = asyncio.run(litellm.aembedding(model="ollama/nomic-embed-text", input=["hello world"])) + embeddings = asyncio.run( + litellm.aembedding(model="ollama/nomic-embed-text", input=["hello world"]) + ) print(embeddings) mock_aembeddings.assert_called_once_with( api_base="http://localhost:11434",