Merge pull request #40268 from BerriAI/litellm_fireworks_responses_reasoning_instructions

fix(fireworks_ai): fold instructions and developer items into one leading system message on the Responses path
This commit is contained in:
Mateo Wang 2026-09-08 17:16:52 -07:00 committed by GitHub
commit 402351d980
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 266 additions and 10 deletions

View file

@ -1,10 +1,10 @@
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from urllib.parse import unquote
import httpx
from openai.types.responses import EasyInputMessageParam, ResponseInputItemParam
from openai.types.responses import EasyInputMessageParam, ResponseInputContentParam, ResponseInputItemParam
from litellm.llms.fireworks_ai.common_utils import (
resolve_fireworks_api_key,
@ -31,6 +31,17 @@ def _session_params(litellm_params: GenericLiteLLMParams) -> Mapping[str, object
)
_INSTRUCTION_ROLES: Final = frozenset({"system", "developer"})
def _role(item: ResponseInputItemParam) -> str | None:
match item:
case {"role": str(role)}:
return role
case _:
return None
def _developer_item_as_system(item: ResponseInputItemParam) -> ResponseInputItemParam:
if "role" not in item or item["role"] != "developer":
return item
@ -43,6 +54,69 @@ def _developer_items_as_system(input: str | ResponseInputParam) -> str | Respons
return [_developer_item_as_system(item) for item in input]
def _text_part(part: ResponseInputContentParam) -> str | None:
match part:
case {"type": "input_text", "text": str(text)}:
return text
case _:
return None
def _text_only_content(item: ResponseInputItemParam) -> str | None:
match item:
case {"role": "system" | "developer", "content": str(text)}:
return text
case {"role": "system" | "developer", "content": [*parts]}:
texts: Final = tuple(map(_text_part, parts))
return None if any(text is None for text in texts) else "\n\n".join(text for text in texts if text)
case _:
return None
def _leading_instruction_block_length(roles: Sequence[str | None]) -> int:
return next((index for index, role in enumerate(roles) if role not in _INSTRUCTION_ROLES), len(roles))
def _closing_instruction_block_start(roles: Sequence[str | None], leading_length: int) -> int:
last_conversation_index: Final = next(
(index for index in range(len(roles) - 1, leading_length - 1, -1) if roles[index] not in _INSTRUCTION_ROLES),
None,
)
if last_conversation_index is None or roles[last_conversation_index] != "assistant":
return len(roles)
return last_conversation_index + 1
def _hoisted_indices(roles: Sequence[str | None]) -> tuple[int, ...]:
leading_length: Final = _leading_instruction_block_length(roles)
closing_start: Final = _closing_instruction_block_start(roles, leading_length)
return tuple(
index for index, role in enumerate(roles[:closing_start]) if index < leading_length or role == "developer"
)
def _with_instruction_items_folded(
input: str | ResponseInputParam, instructions: str | None
) -> tuple[str | None, str | ResponseInputParam]:
if isinstance(input, str):
return instructions, input
items: Final = tuple(input)
folded: Final = MappingProxyType(
{
index: text
for index in _hoisted_indices(tuple(map(_role, items)))
if (text := _text_only_content(items[index])) is not None
}
)
joined: Final = "\n\n".join(chunk for chunk in (instructions, *folded.values()) if chunk)
return (
instructions if not folded else joined or None,
[ # mutable-ok: the base class takes the input items as a list
_developer_item_as_system(item) for index, item in enumerate(items) if index not in folded
],
)
class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
@property
def custom_llm_provider(self) -> LlmProviders:
@ -68,9 +142,6 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
base: Final = (api_base or get_secret_str("FIREWORKS_API_BASE") or FIREWORKS_AI_DEFAULT_API_BASE).rstrip("/")
return f"{base}/responses"
def _validate_input_param(self, input: str | ResponseInputParam) -> str | ResponseInputParam:
return _developer_items_as_system(super()._validate_input_param(input))
def transform_responses_api_request(
self,
model: str,
@ -79,10 +150,25 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
litellm_params: GenericLiteLLMParams,
headers: dict, # mutable-ok: overrides the base class signature
) -> dict: # mutable-ok: overrides the base class signature
instructions_param: Final[object] = response_api_optional_request_params.get("instructions")
validated_input: Final = self._validate_input_param(input)
instructions, folded_input = (
_with_instruction_items_folded(validated_input, instructions_param)
if isinstance(instructions_param, str | None)
else (instructions_param, _developer_items_as_system(validated_input))
)
instruction_entries: Final = () if instructions is None else (("instructions", instructions),)
folded_params: Final = { # mutable-ok: the base class takes the optional params as a dict
key: value
for key, value in (
*((key, value) for key, value in response_api_optional_request_params.items() if key != "instructions"),
*instruction_entries,
)
}
return super().transform_responses_api_request(
model=resolve_fireworks_resource_name(model),
input=input,
response_api_optional_request_params=response_api_optional_request_params,
input=folded_input,
response_api_optional_request_params=folded_params,
litellm_params=litellm_params,
headers=headers,
)

View file

@ -162,7 +162,7 @@ def test_responses_call_forwards_previous_response_id_and_store() -> None:
assert body["input"][0]["call_id"] == "call_abc123"
def test_responses_call_sends_developer_items_as_system_messages() -> None:
def test_responses_call_folds_developer_items_into_instructions() -> None:
client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3"))
with patch(HTTPX_CLIENT_FACTORY, return_value=client):
litellm.responses(
@ -175,13 +175,183 @@ def test_responses_call_sends_developer_items_as_system_messages() -> None:
api_key="fw-test-key",
)
_, _, body = _sent_request(client)
assert body["instructions"] == "Answer with exactly one word."
assert tuple(body["input"]) == (
{"role": "user", "content": "Hi there"},
{"role": "system", "content": "Answer with exactly one word.", "type": "message"},
{"role": "user", "content": [{"type": "input_text", "text": "What is the capital of France?"}]},
)
def test_responses_call_folds_instructions_and_developer_item_into_instructions_with_reasoning_replayed() -> None:
client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/qwen3p8-2p4t-a95b"))
with patch(HTTPX_CLIENT_FACTORY, return_value=client):
litellm.responses(
model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b",
instructions="You are a coding agent running in the Codex CLI.",
input=[ # mutable-ok: the Responses API takes input as a JSON list
{
"role": "developer",
"content": [{"type": "input_text", "text": "<permissions instructions>read-only</permissions instructions>"}],
},
{"role": "user", "content": [{"type": "input_text", "text": "What is the capital of France?"}]},
{"id": "rs_1", "type": "reasoning", "summary": [{"type": "summary_text", "text": "A trivial question."}]},
{
"id": "msg_1",
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "Paris is the capital of France.", "annotations": []}],
},
{"role": "user", "content": [{"type": "input_text", "text": "And of Spain?"}]},
],
store=False,
api_key="fw-test-key",
)
_, _, body = _sent_request(client)
assert body["instructions"] == (
"You are a coding agent running in the Codex CLI.\n\n<permissions instructions>read-only</permissions instructions>"
)
assert tuple(body["input"]) == (
{"role": "user", "content": [{"type": "input_text", "text": "What is the capital of France?"}]},
{"id": "rs_1", "type": "reasoning", "summary": [{"type": "summary_text", "text": "A trivial question."}]},
{
"id": "msg_1",
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "Paris is the capital of France.", "annotations": []}],
},
{"role": "user", "content": [{"type": "input_text", "text": "And of Spain?"}]},
)
def test_responses_call_folds_instructions_and_developer_item_with_previous_response_id() -> None:
client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/qwen3p8-2p4t-a95b"))
with patch(HTTPX_CLIENT_FACTORY, return_value=client):
litellm.responses(
model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b",
instructions="You are a terse assistant.",
input=[ # mutable-ok: the Responses API takes input as a JSON list
{"role": "developer", "content": "Answer with exactly one word."},
{"role": "user", "content": "And of Spain?"},
],
previous_response_id="resp_0e946f2d46bf4b49bf8b29ff78083583",
store=True,
api_key="fw-test-key",
)
_, _, body = _sent_request(client)
assert body["instructions"] == "You are a terse assistant.\n\nAnswer with exactly one word."
assert body["previous_response_id"] == "resp_0e946f2d46bf4b49bf8b29ff78083583"
assert tuple(body["input"]) == ({"role": "user", "content": "And of Spain?"},)
def test_responses_call_keeps_a_closing_developer_item_after_an_assistant_turn_in_place() -> None:
client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/qwen3p8-2p4t-a95b"))
assistant_turn: Final = {
"id": "msg_1",
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "Paris.", "annotations": []}],
}
with patch(HTTPX_CLIENT_FACTORY, return_value=client):
litellm.responses(
model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b",
instructions="Be terse.",
input=[ # mutable-ok: the Responses API takes input as a JSON list
{"role": "developer", "content": "Answer with exactly one word."},
{"role": "user", "content": "What is the capital of France?"},
assistant_turn,
{"role": "developer", "content": "Now restate it in French."},
],
api_key="fw-test-key",
)
_, _, body = _sent_request(client)
assert body["instructions"] == "Be terse.\n\nAnswer with exactly one word."
assert tuple(body["input"]) == (
{"role": "user", "content": "What is the capital of France?"},
assistant_turn,
{"role": "system", "content": "Now restate it in French.", "type": "message"},
)
def test_responses_call_keeps_a_mid_conversation_system_item_in_place() -> None:
client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3"))
with patch(HTTPX_CLIENT_FACTORY, return_value=client):
litellm.responses(
model="fireworks_ai/accounts/fireworks/models/kimi-k3",
input=[ # mutable-ok: the Responses API takes input as a JSON list
{"role": "user", "content": "Hi there"},
{"role": "system", "content": "Switch to French."},
{"role": "user", "content": "What is the capital of France?"},
],
api_key="fw-test-key",
)
_, _, body = _sent_request(client)
assert "instructions" not in body
assert tuple(body["input"]) == (
{"role": "user", "content": "Hi there"},
{"role": "system", "content": "Switch to French."},
{"role": "user", "content": "What is the capital of France?"},
)
def test_responses_call_keeps_a_developer_item_with_non_text_parts_in_place_as_a_system_item() -> None:
client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/qwen3p8-2p4t-a95b"))
developer_item: Final = {
"role": "developer",
"content": [
{"type": "input_text", "text": "Match the style of this reference image."},
{"type": "input_image", "image_url": "data:image/png;base64,iVBORw0KGgo=", "detail": "auto"},
],
}
with patch(HTTPX_CLIENT_FACTORY, return_value=client):
litellm.responses(
model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b",
instructions="Answer with one word.",
input=[developer_item, {"role": "user", "content": "What is the capital of France?"}], # mutable-ok: JSON list
store=False,
api_key="fw-test-key",
)
_, _, body = _sent_request(client)
assert body["instructions"] == "Answer with one word."
assert tuple(body["input"]) == (
{"role": "system", "content": developer_item["content"], "type": "message"},
{"role": "user", "content": "What is the capital of France?"},
)
def test_responses_call_forwards_string_input_and_instructions_unchanged() -> None:
client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3"))
with patch(HTTPX_CLIENT_FACTORY, return_value=client):
litellm.responses(
model="fireworks_ai/accounts/fireworks/models/kimi-k3",
instructions="Answer with exactly one word.",
input="What is the capital of France?",
api_key="fw-test-key",
)
_, _, body = _sent_request(client)
assert body["instructions"] == "Answer with exactly one word."
assert body["input"] == "What is the capital of France?"
def test_transform_request_forwards_non_string_instructions_and_input_untouched() -> None:
developer_item: Final = {"role": "developer", "content": "Answer with exactly one word."}
user_item: Final = {"role": "user", "content": "What is the capital of France?"}
request: Final = FireworksAIResponsesAPIConfig().transform_responses_api_request(
model="accounts/fireworks/models/kimi-k3",
input=cast(ResponseInputParam, [developer_item, user_item]), # mutable-ok: JSON list
response_api_optional_request_params={"instructions": ["not", "a", "string"]}, # mutable-ok: base takes a dict
litellm_params=GenericLiteLLMParams(),
headers={}, # mutable-ok: base takes a dict
)
assert request["instructions"] == ["not", "a", "string"]
assert tuple(request["input"]) == (
{"role": "system", "content": "Answer with exactly one word.", "type": "message"},
user_item,
)
def test_responses_call_maps_pydantic_developer_items_and_replays_pydantic_output_items() -> None:
client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3"))
pydantic_input: Final = cast(
@ -205,8 +375,8 @@ def test_responses_call_maps_pydantic_developer_items_and_replays_pydantic_outpu
model="fireworks_ai/accounts/fireworks/models/kimi-k3", input=pydantic_input, api_key="fw-test-key"
)
_, _, body = _sent_request(client)
assert body["instructions"] == "Answer with exactly one word."
assert tuple(body["input"]) == (
{"role": "system", "content": "Answer with exactly one word.", "type": "message"},
{"id": "rs_1", "summary": [], "type": "reasoning"},
{
"id": "fc_1",