scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

feat: adds reasoning support to fireworks AI models

Browse source
This commit is contained in:
anuragrao04 2026-04-23 16:13:34 +05:30
parent 63ba912b47
commit 89aa8bb9bf
No known key found for this signature in database
GPG key ID: 9FA0B6F4AB81880A
2 changed files with 370 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

90
litellm/llms/fireworks_ai/chat/transformation.py
View file

@ -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 <think>...</think> 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 <think>...</think> tags
from delta content into reasoning_content, mirroring DeepSeek / Ollama behavior.
Applied to all Fireworks models — only activates when <think> 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 <think> tags into reasoning_content when present.
if content is not None and delta.get("reasoning_content") is None:
if "<think>" in content:
content = content.replace("<think>", "")
self.started_reasoning_content = True
if "</think>" in content and self.started_reasoning_content:
# Split on </think>: part before → reasoning, part after → content
parts = content.split("</think>", 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

283
tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py
View file

@ -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
# ---------------------------------------------------------------------------
# <think> 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: <think>...</think> in content → reasoning_content + clean content."""
config = FireworksAIConfig()
body = _make_completion_body("<think>step one\nstep two</think>The 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 <think> 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."
# ---------------------------------------------------------------------------
# <think> 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 <think> 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": "<think>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 </think> 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 bit</think>Actual 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 </think> 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 <think>...</think> 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": "<think>I think</think>Here 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 <think> 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