From 2ad77d9bf608fc4b0a1fe74ec9240055a398486e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sun, 31 Aug 2025 21:17:10 -0700 Subject: [PATCH] feat(ollama/completion): output parse thinking content on streaming + non-streaming for ollama completion calls Completes 'thinking' param support for ollama --- litellm/llms/ollama/chat/transformation.py | 29 +- .../llms/ollama/completion/transformation.py | 102 +++++-- litellm/proxy/_new_secret_config.yaml | 2 +- .../test_ollama_completion_transformation.py | 264 +++++++++++++++++- 4 files changed, 353 insertions(+), 44 deletions(-) diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 2ee7d06ae5..c70fb97af7 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -504,34 +504,23 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): reasoning_content = chunk["message"].get("thinking") self.finished_reasoning_content = True elif chunk["message"].get("content") is not None: - if "" in chunk["message"].get("content"): - reasoning_content = ( - chunk["message"].get("content").replace("", "") - ) + message_content = chunk["message"].get("content") + if "" in message_content: + message_content = message_content.replace("", "") self.started_reasoning_content = True - if ( - "" in chunk["message"].get("content") - and self.started_reasoning_content - ): - reasoning_content = chunk["message"].get("content") - remaining_content = ( - chunk["message"].get("content").split("") - ) - if len(remaining_content) > 1: - content = remaining_content[1] + if "" in message_content and self.started_reasoning_content: + message_content = message_content.replace("", "") self.finished_reasoning_content = True if ( - self.started_reasoning_content is True - and self.finished_reasoning_content is False + self.started_reasoning_content + and not self.finished_reasoning_content ): - reasoning_content = ( - chunk["message"].get("content").replace("", "") - ) + reasoning_content = message_content else: - content = chunk["message"].get("content") + content = message_content delta = Delta( content=content, diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index 4f7be507cc..2654d9461e 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -19,13 +19,13 @@ from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMExcepti from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues, ChatCompletionUsageBlock from litellm.types.utils import ( + Delta, GenericStreamingChunk, ModelInfoBase, ModelResponse, ModelResponseStream, ProviderField, StreamingChoices, - Delta, ) from ..common_utils import OllamaError, _convert_image @@ -92,9 +92,9 @@ class OllamaConfig(BaseConfig): repeat_penalty: Optional[float] = None temperature: Optional[float] = None seed: Optional[int] = None - stop: Optional[ - list - ] = None # stop is a list based on this - https://github.com/ollama/ollama/pull/442 + stop: Optional[list] = ( + None # stop is a list based on this - https://github.com/ollama/ollama/pull/442 + ) tfs_z: Optional[float] = None num_predict: Optional[int] = None top_k: Optional[int] = None @@ -154,6 +154,7 @@ class OllamaConfig(BaseConfig): "stop", "response_format", "max_completion_tokens", + "reasoning_effort", ] def map_openai_params( @@ -166,19 +167,21 @@ class OllamaConfig(BaseConfig): for param, value in non_default_params.items(): if param == "max_tokens" or param == "max_completion_tokens": optional_params["num_predict"] = value - if param == "stream": + elif param == "stream": optional_params["stream"] = value - if param == "temperature": + elif param == "temperature": optional_params["temperature"] = value - if param == "seed": + elif param == "seed": optional_params["seed"] = value - if param == "top_p": + elif param == "top_p": optional_params["top_p"] = value - if param == "frequency_penalty": + elif param == "frequency_penalty": optional_params["frequency_penalty"] = value - if param == "stop": + elif param == "stop": optional_params["stop"] = value - if param == "response_format" and isinstance(value, dict): + elif param == "reasoning_effort" and value is not None: + optional_params["think"] = True + elif param == "response_format" and isinstance(value, dict): if value["type"] == "json_object": optional_params["format"] = "json" elif value["type"] == "json_schema": @@ -258,12 +261,17 @@ class OllamaConfig(BaseConfig): api_key: Optional[str] = None, json_mode: Optional[bool] = None, ) -> ModelResponse: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + _parse_content_for_reasoning, + ) + response_json = raw_response.json() ## RESPONSE OBJECT model_response.choices[0].finish_reason = "stop" if request_data.get("format", "") == "json": # Check if response field exists and is not empty before parsing JSON response_text = response_json.get("response", "") + if not response_text or not response_text.strip(): # Handle empty response gracefully - set empty content message = litellm.Message(content="") @@ -288,7 +296,9 @@ class OllamaConfig(BaseConfig): "id": f"call_{str(uuid.uuid4())}", "function": { "name": function_call["name"], - "arguments": json.dumps(function_call["arguments"]), + "arguments": json.dumps( + function_call["arguments"] + ), }, "type": "function", } @@ -305,11 +315,26 @@ class OllamaConfig(BaseConfig): model_response.choices[0].finish_reason = "stop" except json.JSONDecodeError: # If JSON parsing fails, treat as regular text response - message = litellm.Message(content=response_text) + ## output parse reasoning content from response_text + reasoning_content: Optional[str] = None + content: Optional[str] = None + if response_text is not None: + reasoning_content, content = _parse_content_for_reasoning( + response_text + ) + message = litellm.Message( + content=content, reasoning_content=reasoning_content + ) model_response.choices[0].message = message # type: ignore model_response.choices[0].finish_reason = "stop" else: - model_response.choices[0].message.content = response_json["response"] # type: ignore + response_text = response_json.get("response", "") + content: Optional[str] = None + reasoning_content: Optional[str] = None + if response_text is not None: + reasoning_content, content = _parse_content_for_reasoning(response_text) + model_response.choices[0].message.content = content # type: ignore + model_response.choices[0].message.reasoning_content = reasoning_content # type: ignore model_response.created = int(time.time()) model_response.model = "ollama/" + model _prompt = request_data.get("prompt", "") @@ -434,12 +459,21 @@ class OllamaConfig(BaseConfig): class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): + def __init__( + self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False + ): + super().__init__(streaming_response, sync_stream, json_mode) + self.started_reasoning_content: bool = False + self.finished_reasoning_content: bool = False + def _handle_string_chunk( self, str_line: str ) -> Union[GenericStreamingChunk, ModelResponseStream]: return self.chunk_parser(json.loads(str_line)) - def chunk_parser(self, chunk: dict) -> Union[GenericStreamingChunk, ModelResponseStream]: + def chunk_parser( + self, chunk: dict + ) -> Union[GenericStreamingChunk, ModelResponseStream]: try: if "error" in chunk: raise Exception(f"Ollama Error - {chunk}") @@ -469,12 +503,42 @@ class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): ) elif chunk["response"]: text = chunk["response"] - return GenericStreamingChunk( - text=text, - is_finished=is_finished, - finish_reason="stop", + reasoning_content: Optional[str] = None + content: Optional[str] = None + if text is not None: + if "" in text: + text = text.replace("", "") + self.started_reasoning_content = True + elif "" in text: + text = text.replace("", "") + self.finished_reasoning_content = True + + if ( + self.started_reasoning_content + and not self.finished_reasoning_content + ): + reasoning_content = text + else: + content = text + + return ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta( + reasoning_content=reasoning_content, content=content + ), + ) + ], + finish_reason=finish_reason, usage=None, ) + # return GenericStreamingChunk( + # text=text, + # is_finished=is_finished, + # finish_reason="stop", + # usage=None, + # ) elif "thinking" in chunk and not chunk["response"]: # Return reasoning content as ModelResponseStream so UIs can render it thinking_content = chunk.get("thinking") or "" diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index adfee33eba..324b486630 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -15,7 +15,7 @@ model_list: mode: chat - model_name: ollama-deepseek-r1 litellm_params: - model: ollama_chat/deepseek-r1:1.5b + model: ollama/deepseek-r1:1.5b model_info: mode: chat diff --git a/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py b/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py index 985d51f99d..452f5a9402 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py +++ b/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py @@ -159,6 +159,261 @@ class TestOllamaConfig: assert result.choices[0]["finish_reason"] == "stop" # No usage assertions here as we don't need to test them in every case + def test_transform_response_with_thinking_tags(self): + """Test that responses with ... tags parse reasoning content correctly.""" + # Initialize config + config = OllamaConfig() + + # Create mock response with thinking tags + raw_response = MagicMock() + raw_response.json.return_value = { + "response": "I need to think about this problem step by stepHere is my answer", + "prompt_eval_count": 15, + "eval_count": 8, + } + + # Create properly structured model response object + model_response = ModelResponse( + id="test_id", + choices=[{"message": Message(content="")}], + ) + + # Create mock encoding + mock_encoding = MagicMock() + mock_encoding.encode.return_value = [1, 2, 3] + + # Transform response + result = config.transform_response( + model="llama2", + raw_response=raw_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=mock_encoding, + ) + + # Verify reasoning content is extracted + assert ( + result.choices[0]["message"].reasoning_content + == "I need to think about this problem step by step" + ) + assert result.choices[0]["message"].content == "Here is my answer" + assert result.choices[0]["finish_reason"] == "stop" + + def test_transform_response_with_thinking_tags_alternative(self): + """Test that responses with ... tags parse reasoning content correctly.""" + # Initialize config + config = OllamaConfig() + + # Create mock response with thinking tags (alternative format) + raw_response = MagicMock() + raw_response.json.return_value = { + "response": "Let me analyze this carefullyThe solution is X", + } + + # Create properly structured model response object + model_response = ModelResponse( + id="test_id", + choices=[{"message": Message(content="")}], + ) + + # Create mock encoding + mock_encoding = MagicMock() + mock_encoding.encode.return_value = [1, 2, 3] + + # Transform response + result = config.transform_response( + model="llama2", + raw_response=raw_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=mock_encoding, + ) + + # Verify reasoning content is extracted + assert ( + result.choices[0]["message"].reasoning_content + == "Let me analyze this carefully" + ) + assert result.choices[0]["message"].content == "The solution is X" + assert result.choices[0]["finish_reason"] == "stop" + + def test_transform_response_with_multiline_thinking_tags(self): + """Test that responses with multiline thinking content work correctly.""" + # Initialize config + config = OllamaConfig() + + # Create mock response with multiline thinking content + raw_response = MagicMock() + raw_response.json.return_value = { + "response": "\nThis is a complex problem.\nI need to break it down:\n1. First step\n2. Second step\nBased on my analysis, the answer is Y", + } + + # Create properly structured model response object + model_response = ModelResponse( + id="test_id", + choices=[{"message": Message(content="")}], + ) + + # Create mock encoding + mock_encoding = MagicMock() + mock_encoding.encode.return_value = [1, 2, 3] + + # Transform response + result = config.transform_response( + model="llama2", + raw_response=raw_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=mock_encoding, + ) + + # Verify multiline reasoning content is extracted + expected_reasoning = "\nThis is a complex problem.\nI need to break it down:\n1. First step\n2. Second step\n" + assert result.choices[0]["message"].reasoning_content == expected_reasoning + assert ( + result.choices[0]["message"].content + == "Based on my analysis, the answer is Y" + ) + assert result.choices[0]["finish_reason"] == "stop" + + def test_transform_response_thinking_only(self): + """Test response with only thinking content and no additional content.""" + # Initialize config + config = OllamaConfig() + + # Create mock response with only thinking content + raw_response = MagicMock() + raw_response.json.return_value = { + "response": "Just internal thoughts, no response", + } + + # Create properly structured model response object + model_response = ModelResponse( + id="test_id", + choices=[{"message": Message(content="")}], + ) + + # Create mock encoding + mock_encoding = MagicMock() + mock_encoding.encode.return_value = [1, 2, 3] + + # Transform response + result = config.transform_response( + model="llama2", + raw_response=raw_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=mock_encoding, + ) + + # Verify reasoning content is extracted and content is empty + assert ( + result.choices[0]["message"].reasoning_content + == "Just internal thoughts, no response" + ) + assert result.choices[0]["message"].content == "" + assert result.choices[0]["finish_reason"] == "stop" + + def test_transform_response_json_mode_with_thinking_tags(self): + """Test JSON mode with thinking tags - should handle as text when JSON parsing fails.""" + # Initialize config + config = OllamaConfig() + + # Create mock response with thinking tags in JSON mode + raw_response = MagicMock() + raw_response.json.return_value = { + "response": "Planning my JSON responseThis is not valid JSON", + } + + # Create properly structured model response object + model_response = ModelResponse( + id="test_id", + choices=[{"message": Message(content="")}], + ) + + # Create mock encoding + mock_encoding = MagicMock() + mock_encoding.encode.return_value = [1, 2, 3] + + # Transform response + result = config.transform_response( + model="llama2", + raw_response=raw_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={"format": "json"}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=mock_encoding, + ) + + # Verify reasoning content is extracted even in JSON mode when JSON parsing fails + assert ( + result.choices[0]["message"].reasoning_content + == "Planning my JSON response" + ) + assert result.choices[0]["message"].content == "This is not valid JSON" + assert result.choices[0]["finish_reason"] == "stop" + + def test_transform_response_no_thinking_tags(self): + """Test that responses without thinking tags work normally.""" + # Initialize config + config = OllamaConfig() + + # Create mock response without thinking tags + raw_response = MagicMock() + raw_response.json.return_value = { + "response": "Regular response without any thinking tags", + } + + # Create properly structured model response object + model_response = ModelResponse( + id="test_id", + choices=[{"message": Message(content="")}], + ) + + # Create mock encoding + mock_encoding = MagicMock() + mock_encoding.encode.return_value = [1, 2, 3] + + # Transform response + result = config.transform_response( + model="llama2", + raw_response=raw_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=mock_encoding, + ) + + # Verify no reasoning content is extracted + assert result.choices[0]["message"].reasoning_content is None + assert ( + result.choices[0]["message"].content + == "Regular response without any thinking tags" + ) + assert result.choices[0]["finish_reason"] == "stop" + class TestOllamaTextCompletionResponseIterator: def test_chunk_parser_with_thinking_field(self): @@ -199,10 +454,11 @@ class TestOllamaTextCompletionResponseIterator: result = iterator.chunk_parser(normal_chunk) - assert result["text"] == "Hello world" - assert result["is_finished"] is False - assert result["finish_reason"] == "stop" - assert result["usage"] is None + # Updated to handle ModelResponseStream return type + assert isinstance(result, ModelResponseStream) + assert result.choices and result.choices[0].delta is not None + assert result.choices[0].delta.content == "Hello world" + assert getattr(result.choices[0].delta, "reasoning_content", None) is None def test_chunk_parser_done_chunk(self): """Test that done chunks work correctly."""