diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index ed6d167a118..7c81a43e7ab 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -1,5 +1,5 @@ import json -from typing import Any, List, Literal, Optional, Tuple, Union, cast +from typing import Any, AsyncIterator, Dict, Iterator, List, Literal, Optional, Tuple, Union, cast import httpx @@ -10,6 +10,9 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, ) +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _extract_reasoning_content, +) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( AllMessageValues, @@ -23,6 +26,7 @@ from litellm.types.utils import ( Function, Message, ModelResponse, + ModelResponseStream, ProviderSpecificModelInfo, ) from litellm.utils import ( @@ -31,7 +35,7 @@ from litellm.utils import ( supports_tool_choice, ) -from ...openai.chat.gpt_transformation import OpenAIGPTConfig +from ...openai.chat.gpt_transformation import OpenAIGPTConfig, OpenAIChatCompletionStreamingHandler from ..common_utils import FireworksAIException @@ -399,10 +403,33 @@ class FireworksAIConfig(OpenAIGPTConfig): ) ) + ## Extract ... reasoning from content into reasoning_content. + ## Applied to all Fireworks models — only activates when tags are present. + for choice in response.choices: + _msg = cast(Choices, choice).message + if _msg.content is not None and getattr(_msg, "reasoning_content", None) is None: + _msg_dict = {"content": _msg.content} + reasoning_content, content = _extract_reasoning_content(_msg_dict) + if reasoning_content is not None: + _msg.reasoning_content = reasoning_content + _msg.content = content + response._hidden_params = {"additional_headers": additional_headers} return response + def get_model_response_iterator( + self, + streaming_response: Union[Iterator[str], AsyncIterator[str], ModelResponse], + sync_stream: bool, + json_mode: Optional[bool] = False, + ) -> Any: + return FireworksAIChatCompletionStreamingHandler( + streaming_response=streaming_response, + sync_stream=sync_stream, + json_mode=json_mode, + ) + def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: @@ -459,3 +486,62 @@ class FireworksAIConfig(OpenAIGPTConfig): or get_secret_str("FIREWORKSAI_API_KEY") or get_secret_str("FIREWORKS_AI_TOKEN") ) + + +class FireworksAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): + """ + Streaming handler for Fireworks AI that extracts ... tags + from delta content into reasoning_content, mirroring DeepSeek / Ollama behavior. + + Applied to all Fireworks models — only activates when tags are + actually present in the stream, so models that don't emit them are unaffected. + """ + + started_reasoning_content: bool = False + finished_reasoning_content: bool = False + + def chunk_parser(self, chunk: dict) -> ModelResponseStream: + try: + choices = chunk.get("choices", []) + choices = self._map_reasoning_to_reasoning_content(choices) + + for choice in choices: + delta: Dict[str, Any] = choice.get("delta", {}) + content: Optional[str] = delta.get("content") + + # Extract tags into reasoning_content when present. + if content is not None and delta.get("reasoning_content") is None: + if "" in content: + content = content.replace("", "") + self.started_reasoning_content = True + + if "" in content and self.started_reasoning_content: + # Split on : part before → reasoning, part after → content + parts = content.split("", 1) + reasoning_chunk = parts[0] + content_after = parts[1] if len(parts) > 1 else "" + self.finished_reasoning_content = True + + delta["reasoning_content"] = reasoning_chunk + delta["content"] = content_after if content_after else None + elif self.started_reasoning_content and not self.finished_reasoning_content: + # Mid-think chunk — move content to reasoning_content + delta["reasoning_content"] = content + delta["content"] = None + else: + delta["content"] = content + + choice["delta"] = delta + + kwargs: Dict[str, Any] = { + "id": chunk.get("id"), + "object": "chat.completion.chunk", + "created": chunk.get("created"), + "model": chunk.get("model"), + "choices": choices, + } + if "usage" in chunk and chunk["usage"] is not None: + kwargs["usage"] = chunk["usage"] + return ModelResponseStream(**kwargs) + except Exception as e: + raise e diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 323443b2e15..2985d3cd16d 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -11,7 +11,10 @@ sys.path.insert( ) # Adds the parent directory to the system path from litellm import supports_reasoning -from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig +from litellm.llms.fireworks_ai.chat.transformation import ( + FireworksAIConfig, + FireworksAIChatCompletionStreamingHandler, +) from litellm.types.llms.openai import ChatCompletionToolCallFunctionChunk from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message @@ -232,3 +235,281 @@ def test_transform_messages_helper_removes_provider_specific_fields(): ) for msg in out: assert "provider_specific_fields" not in msg + + +# --------------------------------------------------------------------------- +# tag extraction — non-streaming transform_response +# --------------------------------------------------------------------------- + + +def _make_raw_response(body: dict) -> MagicMock: + """Build a minimal httpx.Response-like mock from a dict body.""" + mock_resp = MagicMock(spec=httpx.Response) + mock_resp.json.return_value = body + mock_resp.text = json.dumps(body) + mock_resp.status_code = 200 + mock_resp.headers = {} + return mock_resp + + +def _make_completion_body(content: str) -> dict: + return { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1700000000, + "model": "accounts/fireworks/models/kimi-k2.5", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": content}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + } + + +def test_transform_response_extracts_think_tags(): + """Non-streaming: ... in content → reasoning_content + clean content.""" + config = FireworksAIConfig() + body = _make_completion_body("step one\nstep twoThe answer is 42.") + raw = _make_raw_response(body) + + from litellm.types.utils import ModelResponse + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + model_response = ModelResponse() + logging_obj = MagicMock(spec=LiteLLMLoggingObj) + logging_obj.post_call = MagicMock() + + response = config.transform_response( + model="fireworks_ai/accounts/fireworks/models/kimi-k2.5", + raw_response=raw, + model_response=model_response, + logging_obj=logging_obj, + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + api_key=None, + ) + + msg = response.choices[0].message + assert msg.reasoning_content == "step one\nstep two" + assert msg.content == "The answer is 42." + + +def test_transform_response_no_think_tags_unchanged(): + """Non-streaming: content without tags is not modified.""" + config = FireworksAIConfig() + body = _make_completion_body("Just a plain response.") + raw = _make_raw_response(body) + + from litellm.types.utils import ModelResponse + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + model_response = ModelResponse() + logging_obj = MagicMock(spec=LiteLLMLoggingObj) + logging_obj.post_call = MagicMock() + + response = config.transform_response( + model="fireworks_ai/accounts/fireworks/models/kimi-k2.5", + raw_response=raw, + model_response=model_response, + logging_obj=logging_obj, + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + api_key=None, + ) + + msg = response.choices[0].message + assert msg.content == "Just a plain response." + assert getattr(msg, "reasoning_content", None) is None + + +def test_transform_response_existing_reasoning_content_not_overwritten(): + """Non-streaming: explicit reasoning_content field is preserved as-is.""" + config = FireworksAIConfig() + body = _make_completion_body("The answer.") + # Inject reasoning_content directly in the raw response message + body["choices"][0]["message"]["reasoning_content"] = "pre-existing reasoning" + raw = _make_raw_response(body) + + from litellm.types.utils import ModelResponse + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + model_response = ModelResponse() + logging_obj = MagicMock(spec=LiteLLMLoggingObj) + logging_obj.post_call = MagicMock() + + response = config.transform_response( + model="fireworks_ai/accounts/fireworks/models/kimi-k2.5", + raw_response=raw, + model_response=model_response, + logging_obj=logging_obj, + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + api_key=None, + ) + + msg = response.choices[0].message + assert msg.reasoning_content == "pre-existing reasoning" + assert msg.content == "The answer." + + +# --------------------------------------------------------------------------- +# tag extraction — streaming chunk_parser +# --------------------------------------------------------------------------- + + +def test_streaming_chunk_parser_no_think_tags(): + """Streaming: plain content chunks pass through unchanged.""" + handler = FireworksAIChatCompletionStreamingHandler( + streaming_response=iter([]), sync_stream=True + ) + chunk = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "accounts/fireworks/models/kimi-k2.5", + "choices": [{"index": 0, "delta": {"content": "Hello world"}, "finish_reason": None}], + } + result = handler.chunk_parser(chunk) + assert result.choices[0].delta.content == "Hello world" + assert getattr(result.choices[0].delta, "reasoning_content", None) is None + + +def test_streaming_chunk_parser_open_think_tag(): + """Streaming: chunk containing starts reasoning accumulation.""" + handler = FireworksAIChatCompletionStreamingHandler( + streaming_response=iter([]), sync_stream=True + ) + chunk = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "accounts/fireworks/models/kimi-k2.5", + "choices": [{"index": 0, "delta": {"content": "start of reasoning"}, "finish_reason": None}], + } + result = handler.chunk_parser(chunk) + assert handler.started_reasoning_content is True + assert handler.finished_reasoning_content is False + assert result.choices[0].delta.reasoning_content == "start of reasoning" + assert result.choices[0].delta.content is None + + +def test_streaming_chunk_parser_mid_think_chunk(): + """Streaming: mid-think chunk (no open/close tag) routed to reasoning_content.""" + handler = FireworksAIChatCompletionStreamingHandler( + streaming_response=iter([]), sync_stream=True + ) + handler.started_reasoning_content = True + handler.finished_reasoning_content = False + + chunk = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "accounts/fireworks/models/kimi-k2.5", + "choices": [{"index": 0, "delta": {"content": "middle of reasoning"}, "finish_reason": None}], + } + result = handler.chunk_parser(chunk) + assert result.choices[0].delta.reasoning_content == "middle of reasoning" + assert result.choices[0].delta.content is None + + +def test_streaming_chunk_parser_close_think_tag(): + """Streaming: chunk with splits reasoning from content correctly.""" + handler = FireworksAIChatCompletionStreamingHandler( + streaming_response=iter([]), sync_stream=True + ) + handler.started_reasoning_content = True + handler.finished_reasoning_content = False + + chunk = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "accounts/fireworks/models/kimi-k2.5", + "choices": [ + { + "index": 0, + "delta": {"content": "final reasoning bitActual answer here"}, + "finish_reason": None, + } + ], + } + result = handler.chunk_parser(chunk) + assert handler.finished_reasoning_content is True + assert result.choices[0].delta.reasoning_content == "final reasoning bit" + assert result.choices[0].delta.content == "Actual answer here" + + +def test_streaming_chunk_parser_content_after_think_closed(): + """Streaming: chunks after are plain content.""" + handler = FireworksAIChatCompletionStreamingHandler( + streaming_response=iter([]), sync_stream=True + ) + handler.started_reasoning_content = True + handler.finished_reasoning_content = True + + chunk = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "accounts/fireworks/models/kimi-k2.5", + "choices": [{"index": 0, "delta": {"content": "More answer text"}, "finish_reason": None}], + } + result = handler.chunk_parser(chunk) + assert result.choices[0].delta.content == "More answer text" + assert getattr(result.choices[0].delta, "reasoning_content", None) is None + + +def test_streaming_chunk_parser_think_tags_in_single_chunk(): + """Streaming: single chunk with full ... is split correctly.""" + handler = FireworksAIChatCompletionStreamingHandler( + streaming_response=iter([]), sync_stream=True + ) + chunk = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "accounts/fireworks/models/kimi-k2.5", + "choices": [ + { + "index": 0, + "delta": {"content": "I thinkHere is the answer"}, + "finish_reason": None, + } + ], + } + result = handler.chunk_parser(chunk) + assert handler.started_reasoning_content is True + assert handler.finished_reasoning_content is True + assert result.choices[0].delta.reasoning_content == "I think" + assert result.choices[0].delta.content == "Here is the answer" + + +def test_streaming_chunk_parser_no_think_tags_any_model(): + """Streaming: content without tags passes through unchanged regardless of model.""" + handler = FireworksAIChatCompletionStreamingHandler( + streaming_response=iter([]), sync_stream=True + ) + chunk = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "accounts/fireworks/models/llama-v3-70b-instruct", + "choices": [{"index": 0, "delta": {"content": "Plain response"}, "finish_reason": None}], + } + result = handler.chunk_parser(chunk) + assert result.choices[0].delta.content == "Plain response" + assert getattr(result.choices[0].delta, "reasoning_content", None) is None + assert handler.started_reasoning_content is False