fix(custom): support extraHeaders for OpenAI-compatible endpoints
This commit is contained in:
@@ -364,6 +364,7 @@ def _make_provider(config: Config):
|
|||||||
api_key=p.api_key if p else "no-key",
|
api_key=p.api_key if p else "no-key",
|
||||||
api_base=config.get_api_base(model) or "http://localhost:8000/v1",
|
api_base=config.get_api_base(model) or "http://localhost:8000/v1",
|
||||||
default_model=model,
|
default_model=model,
|
||||||
|
extra_headers=p.extra_headers if p else None,
|
||||||
)
|
)
|
||||||
# Azure OpenAI: direct Azure OpenAI endpoint with deployment name
|
# Azure OpenAI: direct Azure OpenAI endpoint with deployment name
|
||||||
elif provider_name == "azure_openai":
|
elif provider_name == "azure_openai":
|
||||||
|
|||||||
@@ -13,14 +13,25 @@ from nanobot.providers.base import LLMProvider, LLMResponse, ToolCallRequest
|
|||||||
|
|
||||||
class CustomProvider(LLMProvider):
|
class CustomProvider(LLMProvider):
|
||||||
|
|
||||||
def __init__(self, api_key: str = "no-key", api_base: str = "http://localhost:8000/v1", default_model: str = "default"):
|
def __init__(
|
||||||
|
self,
|
||||||
|
api_key: str = "no-key",
|
||||||
|
api_base: str = "http://localhost:8000/v1",
|
||||||
|
default_model: str = "default",
|
||||||
|
extra_headers: dict[str, str] | None = None,
|
||||||
|
):
|
||||||
super().__init__(api_key, api_base)
|
super().__init__(api_key, api_base)
|
||||||
self.default_model = default_model
|
self.default_model = default_model
|
||||||
# Keep affinity stable for this provider instance to improve backend cache locality.
|
# Keep affinity stable for this provider instance to improve backend cache locality,
|
||||||
|
# while still letting users attach provider-specific headers for custom gateways.
|
||||||
|
default_headers = {
|
||||||
|
"x-session-affinity": uuid.uuid4().hex,
|
||||||
|
**(extra_headers or {}),
|
||||||
|
}
|
||||||
self._client = AsyncOpenAI(
|
self._client = AsyncOpenAI(
|
||||||
api_key=api_key,
|
api_key=api_key,
|
||||||
base_url=api_base,
|
base_url=api_base,
|
||||||
default_headers={"x-session-affinity": uuid.uuid4().hex},
|
default_headers=default_headers,
|
||||||
)
|
)
|
||||||
|
|
||||||
async def chat(self, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None = None,
|
async def chat(self, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None = None,
|
||||||
|
|||||||
@@ -6,7 +6,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||||||
import pytest
|
import pytest
|
||||||
from typer.testing import CliRunner
|
from typer.testing import CliRunner
|
||||||
|
|
||||||
from nanobot.cli.commands import app
|
from nanobot.cli.commands import _make_provider, app
|
||||||
from nanobot.config.schema import Config
|
from nanobot.config.schema import Config
|
||||||
from nanobot.providers.litellm_provider import LiteLLMProvider
|
from nanobot.providers.litellm_provider import LiteLLMProvider
|
||||||
from nanobot.providers.openai_codex_provider import _strip_model_prefix
|
from nanobot.providers.openai_codex_provider import _strip_model_prefix
|
||||||
@@ -199,6 +199,33 @@ def test_openai_codex_strip_prefix_supports_hyphen_and_underscore():
|
|||||||
assert _strip_model_prefix("openai_codex/gpt-5.1-codex") == "gpt-5.1-codex"
|
assert _strip_model_prefix("openai_codex/gpt-5.1-codex") == "gpt-5.1-codex"
|
||||||
|
|
||||||
|
|
||||||
|
def test_make_provider_passes_extra_headers_to_custom_provider():
|
||||||
|
config = Config.model_validate(
|
||||||
|
{
|
||||||
|
"agents": {"defaults": {"provider": "custom", "model": "gpt-4o-mini"}},
|
||||||
|
"providers": {
|
||||||
|
"custom": {
|
||||||
|
"apiKey": "test-key",
|
||||||
|
"apiBase": "https://example.com/v1",
|
||||||
|
"extraHeaders": {
|
||||||
|
"APP-Code": "demo-app",
|
||||||
|
"x-session-affinity": "sticky-session",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
with patch("nanobot.providers.custom_provider.AsyncOpenAI") as mock_async_openai:
|
||||||
|
_make_provider(config)
|
||||||
|
|
||||||
|
kwargs = mock_async_openai.call_args.kwargs
|
||||||
|
assert kwargs["api_key"] == "test-key"
|
||||||
|
assert kwargs["base_url"] == "https://example.com/v1"
|
||||||
|
assert kwargs["default_headers"]["APP-Code"] == "demo-app"
|
||||||
|
assert kwargs["default_headers"]["x-session-affinity"] == "sticky-session"
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def mock_agent_runtime(tmp_path):
|
def mock_agent_runtime(tmp_path):
|
||||||
"""Mock agent command dependencies for focused CLI tests."""
|
"""Mock agent command dependencies for focused CLI tests."""
|
||||||
|
|||||||
Reference in New Issue
Block a user