fix(openrouter): pass custom fields through SDK extra_body

This commit is contained in:
shashank-100
2026-09-06 13:13:27 +05:30
parent e25ab65e69
commit d05053ed60
2 changed files with 34 additions and 2 deletions
+2 -2
View File
@@ -167,7 +167,7 @@ class ChatOpenRouter(BaseChatModel):
top_p=self.top_p,
seed=self.seed,
extra_headers=extra_headers,
**(self.extra_body or {}),
extra_body=self.extra_body,
)
choice = self._get_first_choice(response)
@@ -199,7 +199,7 @@ class ChatOpenRouter(BaseChatModel):
type='json_schema',
),
extra_headers=extra_headers,
**(self.extra_body or {}),
extra_body=self.extra_body,
)
choice = self._get_first_choice(response)
+32
View File
@@ -1,7 +1,9 @@
"""Regression tests for OpenRouter client setup and response handling."""
import json
from unittest.mock import AsyncMock, patch
import httpx
import pytest
from openai.types.chat import ChatCompletion, ChatCompletionMessage
from openai.types.chat.chat_completion import Choice
@@ -16,6 +18,36 @@ class Answer(BaseModel):
answer: str
@pytest.mark.parametrize('structured', [False, True])
@pytest.mark.parametrize('extra_body', [None, {}, {'provider': {'order': ['test-provider']}, 'transforms': ['middle-out']}])
async def test_extra_body_reaches_http_request(structured: bool, extra_body: dict | None):
"""Send provider-specific fields through the real SDK into the JSON request body."""
requests: list[httpx.Request] = []
def handle_request(request: httpx.Request) -> httpx.Response:
requests.append(request)
return httpx.Response(200, json=_completion(content='{"answer":"ok"}' if structured else 'ok').model_dump())
async with httpx.AsyncClient(transport=httpx.MockTransport(handle_request)) as client:
llm = ChatOpenRouter(model='openai/gpt-4o', api_key='test-key', http_client=client, extra_body=extra_body)
result = await llm.ainvoke([UserMessage(content='question')], output_format=Answer if structured else None)
assert result.completion == (Answer(answer='ok') if structured else 'ok')
assert len(requests) == 1
body = json.loads(requests[0].content)
assert body['model'] == 'openai/gpt-4o'
assert body['messages'] == [{'role': 'user', 'content': 'question'}]
assert 'extra_body' not in body
if extra_body:
assert body['provider'] == {'order': ['test-provider']}
assert body['transforms'] == ['middle-out']
else:
assert 'provider' not in body
assert 'transforms' not in body
if structured:
assert body['response_format']['type'] == 'json_schema'
def _completion(*, content: str | None = 'ok', choices: bool = True) -> ChatCompletion:
return ChatCompletion(
id='chatcmpl-test',