From b15469b9fce459933b1408c84ea808d1f6448943 Mon Sep 17 00:00:00 2001 From: XiaoHuo888-hue Date: Wed, 19 Aug 2026 22:25:53 +0000 Subject: [PATCH] feat(sdk): add OrcaRouter gateway integration to wren-langchain and wren-pydantic Add a named OrcaRouter (https://www.orcarouter.ai) gateway wiring to both Wren AI SDKs. OrcaRouter is an OpenAI-compatible model gateway, so it plugs in through the existing OpenAI-compatible clients: - wren-langchain: new wren_langchain.orcarouter.create_orcarouter_chat_model() returns a ChatOpenAI pointed at https://api.orcarouter.ai/v1 with model orcarouter/auto by default; langchain_demo.py and langgraph_demo.py pick it up automatically when ORCAROUTER_API_KEY is set. - wren-pydantic: new wren_pydantic.orcarouter.create_orcarouter_model() returns a pydantic-ai OpenAIChatModel on the same endpoint; both example demos pick it up via ORCAROUTER_API_KEY. - Docs: 'Using OrcaRouter as the model gateway' sections in docs/core/sdk langchain.md and pydantic.md, plus the packaged SDK READMEs. - Tests: unit tests for both factories (missing key, defaults, env overrides). Verified: ruff check + format clean on both SDK trees; 88 + 103 unit tests pass; live round-trip to OrcaRouter returns 200 for both ChatOpenAI and OpenAIChatModel paths. Disclosure: I'm an engineer on the OrcaRouter team. Signed-off-by: XiaoHuo888-hue --- docs/core/sdk/langchain.md | 38 ++++++++++++++ docs/core/sdk/pydantic.md | 36 +++++++++++++ sdk/wren-langchain/README.md | 19 +++++++ sdk/wren-langchain/examples/langchain_demo.py | 26 ++++++++-- sdk/wren-langchain/examples/langgraph_demo.py | 40 ++++++++++++--- sdk/wren-langchain/pyproject.toml | 2 + .../src/wren_langchain/orcarouter.py | 51 +++++++++++++++++++ .../tests/unit/test_orcarouter.py | 34 +++++++++++++ sdk/wren-pydantic/README.md | 19 +++++++ .../examples/pydantic_ai_demo.py | 21 +++++++- .../examples/pydantic_ai_structured_demo.py | 19 ++++++- .../src/wren_pydantic/orcarouter.py | 47 +++++++++++++++++ .../tests/unit/test_orcarouter.py | 34 +++++++++++++ 13 files changed, 373 insertions(+), 13 deletions(-) create mode 100644 sdk/wren-langchain/src/wren_langchain/orcarouter.py create mode 100644 sdk/wren-langchain/tests/unit/test_orcarouter.py create mode 100644 sdk/wren-pydantic/src/wren_pydantic/orcarouter.py create mode 100644 sdk/wren-pydantic/tests/unit/test_orcarouter.py diff --git a/docs/core/sdk/langchain.md b/docs/core/sdk/langchain.md index 0c3609aade..b6a6c99dbd 100644 --- a/docs/core/sdk/langchain.md +++ b/docs/core/sdk/langchain.md @@ -62,6 +62,40 @@ print(result["messages"][-1].content) That's it. The toolkit reads your project's MDL, connection profile, and `instructions.md`; the system prompt teaches the agent the recommended workflow (fetch context → recall similar queries → write SQL → store the result). +## Using OrcaRouter as the model gateway + +The SDK's tools and system prompt are model-agnostic — any LangChain-compatible +chat model works. To route the agent through the [OrcaRouter](https://www.orcarouter.ai) +gateway instead of OpenAI, point a `ChatOpenAI` at OrcaRouter's OpenAI-compatible +endpoint: + +```python +import os + +from langchain_openai import ChatOpenAI + +model = ChatOpenAI( + model=os.environ.get("ORCAROUTER_MODEL", "orcarouter/auto"), + base_url=os.environ.get("ORCAROUTER_BASE_URL", "https://api.orcarouter.ai/v1"), + api_key=os.environ["ORCAROUTER_API_KEY"], + temperature=0, +) + +agent = create_agent( + model=model, + tools=toolkit.get_tools(), + system_prompt=toolkit.system_prompt(), +) +``` + +OrcaRouter is a unified model gateway: one key routes to 150+ models across +providers, with a single `orcarouter/auto` model id for smart default routing. +It also runs gateway-level, zero-trust security for AI agents on the same +endpoint — screening every prompt/response and governing every tool call on a +default-deny basis, with no application code changes. The runnable +[`examples/langchain_demo.py`](https://github.com/Canner/WrenAI/blob/main/sdk/wren-langchain/examples/langchain_demo.py) +picks OrcaRouter automatically when `ORCAROUTER_API_KEY` is set. + --- ## API Reference @@ -146,6 +180,10 @@ from langgraph.prebuilt import ToolNode, tools_condition from langchain.chat_models import init_chat_model tools = toolkit.get_tools() + +# Swap in an OrcaRouter-routed model the same way: +# export ORCAROUTER_API_KEY=sk-orca-... +# init_chat_model("openai:orcarouter/auto", base_url="https://api.orcarouter.ai/v1") llm = init_chat_model("openai:gpt-4o").bind_tools(tools) def chatbot(state: MessagesState) -> dict: diff --git a/docs/core/sdk/pydantic.md b/docs/core/sdk/pydantic.md index fab13b8983..41014ffabe 100644 --- a/docs/core/sdk/pydantic.md +++ b/docs/core/sdk/pydantic.md @@ -62,6 +62,42 @@ print(result.output) The toolkit reads your project's MDL, connection profile, and `instructions.md`; the instructions string teaches the agent the recommended workflow (recall → fetch context → write SQL → store the result). +## Using OrcaRouter as the model gateway + +The SDK's tools and instructions are model-agnostic — any Pydantic AI model works. +To route the agent through the [OrcaRouter](https://www.orcarouter.ai) gateway +instead of OpenAI, construct an `OpenAIChatModel` pointed at OrcaRouter's +OpenAI-compatible endpoint: + +```python +import os + +from pydantic_ai.models.openai import OpenAIChatModel +from pydantic_ai.providers.openai import OpenAIProvider + +model = OpenAIChatModel( + os.environ.get("ORCAROUTER_MODEL", "orcarouter/auto"), + provider=OpenAIProvider( + base_url=os.environ.get("ORCAROUTER_BASE_URL", "https://api.orcarouter.ai/v1"), + api_key=os.environ["ORCAROUTER_API_KEY"], + ), +) + +agent = Agent( + model, + instructions=toolkit.instructions(), + toolsets=[toolkit.toolset()], +) +``` + +OrcaRouter is a unified model gateway: one key routes to 150+ models across +providers, with a single `orcarouter/auto` model id for smart default routing. +It also runs gateway-level, zero-trust security for AI agents on the same +endpoint — screening every prompt/response and governing every tool call on a +default-deny basis, with no application code changes. The runnable +[`examples/pydantic_ai_demo.py`](https://github.com/Canner/WrenAI/blob/main/sdk/wren-pydantic/examples/pydantic_ai_demo.py) +picks OrcaRouter automatically when `ORCAROUTER_API_KEY` is set. + --- ## API Reference diff --git a/sdk/wren-langchain/README.md b/sdk/wren-langchain/README.md index 2dc463d58e..d69ce49076 100644 --- a/sdk/wren-langchain/README.md +++ b/sdk/wren-langchain/README.md @@ -26,6 +26,25 @@ Complete runnable demos: conditional edges). Use this when you need custom routing, state, or streaming. +### Routing through OrcaRouter + +[OrcaRouter](https://www.orcarouter.ai) is a unified model gateway with an +OpenAI-compatible endpoint. Point a `ChatOpenAI` at it to route the agent +through OrcaRouter — the examples above pick it up automatically when +`ORCAROUTER_API_KEY` is set: + +```bash +export ORCAROUTER_API_KEY=sk-orca-... # required +export ORCAROUTER_MODEL=orcarouter/auto # optional, default: orcarouter/auto +export ORCAROUTER_BASE_URL=https://api.orcarouter.ai/v1 # optional, default above +python examples/langchain_demo.py +``` + +OrcaRouter gives you one key for 150+ models across providers, and also runs +gateway-level, zero-trust security for AI agents on the same endpoint — +screening every prompt/response and governing every tool call on a +default-deny basis, with no application code changes. + ## Prerequisites This package assumes you have already used the Wren CLI to prepare a project: diff --git a/sdk/wren-langchain/examples/langchain_demo.py b/sdk/wren-langchain/examples/langchain_demo.py index 992a5a30f4..9cdad84419 100644 --- a/sdk/wren-langchain/examples/langchain_demo.py +++ b/sdk/wren-langchain/examples/langchain_demo.py @@ -18,7 +18,8 @@ for a one-shot DuckDB-backed demo project. - ``langchain-openai`` installed in the active venv: uv pip install langchain-openai - - ``OPENAI_API_KEY`` set in the environment. + - ``OPENAI_API_KEY`` (OpenAI) **or** ``ORCAROUTER_API_KEY`` + (OrcaRouter, https://www.orcarouter.ai) set in the environment. Usage ===== @@ -26,6 +27,10 @@ export PROJECT_PATH=/path/to/your-wren-project python examples/langchain_demo.py + # Route the agent through the OrcaRouter gateway instead of OpenAI: + export ORCAROUTER_API_KEY=sk-orca-... + python examples/langchain_demo.py + # Custom question: QUESTION="What's the gender distribution of users?" \\ python examples/langchain_demo.py @@ -63,6 +68,19 @@ ) from wren_langchain import WrenToolkit +from wren_langchain.orcarouter import create_orcarouter_chat_model + + +def build_chat_model() -> ChatOpenAI: + """Return a ChatOpenAI, routed through OrcaRouter when ``ORCAROUTER_API_KEY`` is set. + + OrcaRouter (https://www.orcarouter.ai) is an OpenAI-compatible gateway, so any + LangChain ``ChatOpenAI`` endpoint works. When no OrcaRouter key is present the + demo falls back to the default OpenAI model. + """ + if os.environ.get("ORCAROUTER_API_KEY"): + return create_orcarouter_chat_model() + return ChatOpenAI(model="gpt-4o", temperature=0) def main() -> None: @@ -72,8 +90,8 @@ def main() -> None: "PROJECT_PATH is required. Example:\n" " PROJECT_PATH=/Users/you/my-wren-project python examples/langchain_demo.py" ) - if not os.environ.get("OPENAI_API_KEY"): - sys.exit("OPENAI_API_KEY is required.") + if not (os.environ.get("OPENAI_API_KEY") or os.environ.get("ORCAROUTER_API_KEY")): + sys.exit("OPENAI_API_KEY (or ORCAROUTER_API_KEY) is required.") question = os.environ.get( "QUESTION", @@ -96,7 +114,7 @@ def main() -> None: # 2) Build the agent. Any LangChain-compatible chat model works here. agent = create_agent( - model=ChatOpenAI(model="gpt-4o", temperature=0), + model=build_chat_model(), tools=tools, system_prompt=prompt, ) diff --git a/sdk/wren-langchain/examples/langgraph_demo.py b/sdk/wren-langchain/examples/langgraph_demo.py index 4e7b34e1b8..57a614a2e6 100644 --- a/sdk/wren-langchain/examples/langgraph_demo.py +++ b/sdk/wren-langchain/examples/langgraph_demo.py @@ -26,8 +26,8 @@ │ (ToolNode) │ └──────────────┘ -Prereqs match ``langchain_demo.py``: a CLI-prepared Wren project, OPENAI_API_KEY, -and ``langchain-openai`` installed. +Prereqs match ``langchain_demo.py``: a CLI-prepared Wren project, OPENAI_API_KEY +(or ORCAROUTER_API_KEY), and ``langchain-openai`` installed. Usage ===== @@ -35,6 +35,10 @@ export PROJECT_PATH=/path/to/your-wren-project python examples/langgraph_demo.py + # Route the agent through the OrcaRouter gateway instead of OpenAI: + export ORCAROUTER_API_KEY=sk-orca-... + python examples/langgraph_demo.py + # Custom question + streaming view: QUESTION="..." STREAM=1 python examples/langgraph_demo.py """ @@ -63,6 +67,19 @@ from langgraph.prebuilt import ToolNode from wren_langchain import WrenToolkit +from wren_langchain.orcarouter import create_orcarouter_chat_model + + +def build_chat_model() -> ChatOpenAI: + """Return a ChatOpenAI, routed through OrcaRouter when ``ORCAROUTER_API_KEY`` is set. + + OrcaRouter (https://www.orcarouter.ai) is an OpenAI-compatible gateway, so any + LangChain ``ChatOpenAI`` endpoint works. When no OrcaRouter key is present the + demo falls back to the default OpenAI model. + """ + if os.environ.get("ORCAROUTER_API_KEY"): + return create_orcarouter_chat_model() + return ChatOpenAI(model="gpt-4o", temperature=0) class AgentState(TypedDict): @@ -76,11 +93,20 @@ class AgentState(TypedDict): messages: Annotated[list[BaseMessage], add_messages] -def build_app(toolkit: WrenToolkit, model_name: str = "gpt-4o"): - """Compile a ReAct graph that uses Wren tools.""" +def build_app(toolkit: WrenToolkit, model_name: str | None = None): + """Compile a ReAct graph that uses Wren tools. + + ``model_name`` defaults to ``gpt-4o`` unless ``ORCAROUTER_API_KEY`` is set, in + which case the graph routes through the OrcaRouter gateway. + """ tools = toolkit.get_tools() system_prompt = toolkit.system_prompt() - model_with_tools = ChatOpenAI(model=model_name, temperature=0).bind_tools(tools) + model = ( + build_chat_model() + if model_name is None + else ChatOpenAI(model=model_name, temperature=0) + ) + model_with_tools = model.bind_tools(tools) def agent_node(state: AgentState) -> dict: """Call the model. Inject the Wren system prompt only on the first turn.""" @@ -131,8 +157,8 @@ def main() -> None: "PROJECT_PATH is required. Example:\n" " PROJECT_PATH=/Users/you/my-wren-project python examples/langgraph_demo.py" ) - if not os.environ.get("OPENAI_API_KEY"): - sys.exit("OPENAI_API_KEY is required.") + if not (os.environ.get("OPENAI_API_KEY") or os.environ.get("ORCAROUTER_API_KEY")): + sys.exit("OPENAI_API_KEY (or ORCAROUTER_API_KEY) is required.") question = os.environ.get( "QUESTION", diff --git a/sdk/wren-langchain/pyproject.toml b/sdk/wren-langchain/pyproject.toml index 506a185035..324918fda5 100644 --- a/sdk/wren-langchain/pyproject.toml +++ b/sdk/wren-langchain/pyproject.toml @@ -84,6 +84,8 @@ dev = [ "ruff>=0.4", # Tests run real DuckDB + LanceDB integrations, so memory deps are required. "wrenai[memory]>=0.13.1", + # OrcaRouter gateway factory is tested directly; requires langchain-openai. + "langchain-openai>=0.1", ] [project.urls] diff --git a/sdk/wren-langchain/src/wren_langchain/orcarouter.py b/sdk/wren-langchain/src/wren_langchain/orcarouter.py new file mode 100644 index 0000000000..31e2baf4e5 --- /dev/null +++ b/sdk/wren-langchain/src/wren_langchain/orcarouter.py @@ -0,0 +1,51 @@ +"""OrcaRouter gateway integration for wren-langchain. + +[OrcaRouter](https://www.orcarouter.ai) is an OpenAI-compatible model gateway: one +key routes to 150+ models across providers, and the same endpoint runs +gateway-level, zero-trust security for AI agents. This module builds a +``langchain_openai.ChatOpenAI`` pointed at OrcaRouter's endpoint. +""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from langchain_openai import ChatOpenAI + +#: Default base URL for the OrcaRouter OpenAI-compatible endpoint. +DEFAULT_ORCAROUTER_BASE_URL = "https://api.orcarouter.ai/v1" +#: Default model id — OrcaRouter's smart auto-routing model. +DEFAULT_ORCAROUTER_MODEL = "orcarouter/auto" + + +def create_orcarouter_chat_model(*, temperature: float = 0) -> ChatOpenAI: + """Return a ``ChatOpenAI`` routed through OrcaRouter. + + Requires ``ORCAROUTER_API_KEY`` in the environment. ``ORCAROUTER_BASE_URL`` + and ``ORCAROUTER_MODEL`` override the defaults. + + Raises: + ImportError: if ``langchain-openai`` is not installed. + ValueError: if ``ORCAROUTER_API_KEY`` is not set. + """ + try: + from langchain_openai import ChatOpenAI # noqa: PLC0415 + except ImportError as exc: # pragma: no cover - exercised via dev extra in CI + raise ImportError( + "langchain-openai is required for OrcaRouter routing." + ) from exc + + api_key = os.environ.get("ORCAROUTER_API_KEY") + if not api_key: + raise ValueError( + "ORCAROUTER_API_KEY is required to use the OrcaRouter gateway." + ) + + return ChatOpenAI( + model=os.environ.get("ORCAROUTER_MODEL", DEFAULT_ORCAROUTER_MODEL), + base_url=os.environ.get("ORCAROUTER_BASE_URL", DEFAULT_ORCAROUTER_BASE_URL), + api_key=api_key, + temperature=temperature, + ) diff --git a/sdk/wren-langchain/tests/unit/test_orcarouter.py b/sdk/wren-langchain/tests/unit/test_orcarouter.py new file mode 100644 index 0000000000..b3ae8c9d07 --- /dev/null +++ b/sdk/wren-langchain/tests/unit/test_orcarouter.py @@ -0,0 +1,34 @@ +"""Tests for the OrcaRouter gateway factory.""" + +import pytest + +from wren_langchain.orcarouter import ( + DEFAULT_ORCAROUTER_BASE_URL, + DEFAULT_ORCAROUTER_MODEL, + create_orcarouter_chat_model, +) + +pytest.importorskip("langchain_openai") + + +def test_requires_api_key(monkeypatch): + monkeypatch.delenv("ORCAROUTER_API_KEY", raising=False) + with pytest.raises(ValueError, match="ORCAROUTER_API_KEY"): + create_orcarouter_chat_model() + + +def test_defaults(monkeypatch): + monkeypatch.setenv("ORCAROUTER_API_KEY", "sk-orca-test") + model = create_orcarouter_chat_model() + assert model.model_name == DEFAULT_ORCAROUTER_MODEL + assert model.openai_api_base == DEFAULT_ORCAROUTER_BASE_URL + assert model.openai_api_key.get_secret_value() == "sk-orca-test" + + +def test_env_overrides(monkeypatch): + monkeypatch.setenv("ORCAROUTER_API_KEY", "sk-orca-test") + monkeypatch.setenv("ORCAROUTER_MODEL", "anthropic/claude-sonnet-5") + monkeypatch.setenv("ORCAROUTER_BASE_URL", "https://proxy.example.com/v1") + model = create_orcarouter_chat_model() + assert model.model_name == "anthropic/claude-sonnet-5" + assert model.openai_api_base == "https://proxy.example.com/v1" diff --git a/sdk/wren-pydantic/README.md b/sdk/wren-pydantic/README.md index 57dfadc856..3984562977 100644 --- a/sdk/wren-pydantic/README.md +++ b/sdk/wren-pydantic/README.md @@ -30,6 +30,25 @@ Runnable demos: - [`examples/pydantic_ai_structured_demo.py`](./examples/pydantic_ai_structured_demo.py) — same shape with `output_type=` for structured / validated agent output. +### Routing through OrcaRouter + +[OrcaRouter](https://www.orcarouter.ai) is a unified model gateway with an +OpenAI-compatible endpoint. Point an `OpenAIChatModel` at it to route the +agent through OrcaRouter — the examples above pick it up automatically when +`ORCAROUTER_API_KEY` is set: + +```bash +export ORCAROUTER_API_KEY=sk-orca-... # required +export ORCAROUTER_MODEL=orcarouter/auto # optional, default: orcarouter/auto +export ORCAROUTER_BASE_URL=https://api.orcarouter.ai/v1 # optional, default above +python examples/pydantic_ai_demo.py +``` + +OrcaRouter gives you one key for 150+ models across providers, and also runs +gateway-level, zero-trust security for AI agents on the same endpoint — +screening every prompt/response and governing every tool call on a +default-deny basis, with no application code changes. + ## Prerequisites This package assumes you have already used the Wren CLI to prepare a project: diff --git a/sdk/wren-pydantic/examples/pydantic_ai_demo.py b/sdk/wren-pydantic/examples/pydantic_ai_demo.py index b152e4c862..b4c26345d3 100644 --- a/sdk/wren-pydantic/examples/pydantic_ai_demo.py +++ b/sdk/wren-pydantic/examples/pydantic_ai_demo.py @@ -8,6 +8,9 @@ Run: OPENAI_API_KEY=sk-... python examples/pydantic_ai_demo.py + + # Route the agent through the OrcaRouter gateway instead of OpenAI: + ORCAROUTER_API_KEY=sk-orca-... python examples/pydantic_ai_demo.py """ from __future__ import annotations @@ -17,8 +20,22 @@ from pathlib import Path from pydantic_ai import Agent +from pydantic_ai.models.openai import OpenAIChatModel from wren_pydantic import WrenToolkit +from wren_pydantic.orcarouter import create_orcarouter_model + + +def build_model() -> OpenAIChatModel: + """Return a model, routed through OrcaRouter when ``ORCAROUTER_API_KEY`` is set. + + OrcaRouter (https://www.orcarouter.ai) is an OpenAI-compatible gateway, so the + Pydantic AI ``OpenAIChatModel`` endpoint works. When no OrcaRouter key is present + the demo falls back to the default OpenAI model. + """ + if os.environ.get("ORCAROUTER_API_KEY"): + return create_orcarouter_model() + return OpenAIChatModel("gpt-4o") def main() -> None: @@ -31,10 +48,12 @@ def main() -> None: file=sys.stderr, ) sys.exit(1) + if not (os.environ.get("OPENAI_API_KEY") or os.environ.get("ORCAROUTER_API_KEY")): + sys.exit("OPENAI_API_KEY (or ORCAROUTER_API_KEY) is required.") toolkit = WrenToolkit.from_project(project_path) agent = Agent( - "openai:gpt-4o", + build_model(), instructions=toolkit.instructions(), toolsets=[toolkit.toolset()], ) diff --git a/sdk/wren-pydantic/examples/pydantic_ai_structured_demo.py b/sdk/wren-pydantic/examples/pydantic_ai_structured_demo.py index 4be8ce2e26..b6308460c2 100644 --- a/sdk/wren-pydantic/examples/pydantic_ai_structured_demo.py +++ b/sdk/wren-pydantic/examples/pydantic_ai_structured_demo.py @@ -15,8 +15,22 @@ from pydantic import BaseModel from pydantic_ai import Agent +from pydantic_ai.models.openai import OpenAIChatModel from wren_pydantic import WrenToolkit +from wren_pydantic.orcarouter import create_orcarouter_model + + +def build_model() -> OpenAIChatModel: + """Return an OpenAIChatModel, routed through OrcaRouter when ``ORCAROUTER_API_KEY`` is set. + + OrcaRouter (https://www.orcarouter.ai) is an OpenAI-compatible gateway, so the + Pydantic AI ``OpenAIChatModel`` endpoint works. When no OrcaRouter key is present + the demo falls back to the default OpenAI model. + """ + if os.environ.get("ORCAROUTER_API_KEY"): + return create_orcarouter_model() + return OpenAIChatModel("gpt-4o") class TopCustomers(BaseModel): @@ -33,9 +47,12 @@ def main() -> None: print(f"PROJECT_PATH={project_path} not a Wren project", file=sys.stderr) sys.exit(1) + if not (os.environ.get("OPENAI_API_KEY") or os.environ.get("ORCAROUTER_API_KEY")): + sys.exit("OPENAI_API_KEY (or ORCAROUTER_API_KEY) is required.") + toolkit = WrenToolkit.from_project(project_path) agent = Agent( - "openai:gpt-4o", + build_model(), instructions=toolkit.instructions(), toolsets=[toolkit.toolset()], output_type=TopCustomers, diff --git a/sdk/wren-pydantic/src/wren_pydantic/orcarouter.py b/sdk/wren-pydantic/src/wren_pydantic/orcarouter.py new file mode 100644 index 0000000000..fa493e7290 --- /dev/null +++ b/sdk/wren-pydantic/src/wren_pydantic/orcarouter.py @@ -0,0 +1,47 @@ +"""OrcaRouter gateway integration for wren-pydantic. + +[OrcaRouter](https://www.orcarouter.ai) is an OpenAI-compatible model gateway: one +key routes to 150+ models across providers, and the same endpoint runs +gateway-level, zero-trust security for AI agents. This module builds a +Pydantic AI ``OpenAIChatModel`` pointed at OrcaRouter's endpoint. +""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pydantic_ai.models.openai import OpenAIChatModel + +#: Default base URL for the OrcaRouter OpenAI-compatible endpoint. +DEFAULT_ORCAROUTER_BASE_URL = "https://api.orcarouter.ai/v1" +#: Default model id — OrcaRouter's smart auto-routing model. +DEFAULT_ORCAROUTER_MODEL = "orcarouter/auto" + + +def create_orcarouter_model() -> OpenAIChatModel: + """Return a Pydantic AI ``OpenAIChatModel`` routed through OrcaRouter. + + Requires ``ORCAROUTER_API_KEY`` in the environment. ``ORCAROUTER_BASE_URL`` + and ``ORCAROUTER_MODEL`` override the defaults. + + Raises: + ValueError: if ``ORCAROUTER_API_KEY`` is not set. + """ + from pydantic_ai.models.openai import OpenAIChatModel # noqa: PLC0415 + from pydantic_ai.providers.openai import OpenAIProvider # noqa: PLC0415 + + api_key = os.environ.get("ORCAROUTER_API_KEY") + if not api_key: + raise ValueError( + "ORCAROUTER_API_KEY is required to use the OrcaRouter gateway." + ) + + return OpenAIChatModel( + os.environ.get("ORCAROUTER_MODEL", DEFAULT_ORCAROUTER_MODEL), + provider=OpenAIProvider( + base_url=os.environ.get("ORCAROUTER_BASE_URL", DEFAULT_ORCAROUTER_BASE_URL), + api_key=api_key, + ), + ) diff --git a/sdk/wren-pydantic/tests/unit/test_orcarouter.py b/sdk/wren-pydantic/tests/unit/test_orcarouter.py new file mode 100644 index 0000000000..dfe0458889 --- /dev/null +++ b/sdk/wren-pydantic/tests/unit/test_orcarouter.py @@ -0,0 +1,34 @@ +"""Tests for the OrcaRouter gateway factory.""" + +import pytest + +from wren_pydantic.orcarouter import ( + DEFAULT_ORCAROUTER_BASE_URL, + DEFAULT_ORCAROUTER_MODEL, + create_orcarouter_model, +) + +pytest.importorskip("pydantic_ai") + + +def test_requires_api_key(monkeypatch): + monkeypatch.delenv("ORCAROUTER_API_KEY", raising=False) + with pytest.raises(ValueError, match="ORCAROUTER_API_KEY"): + create_orcarouter_model() + + +def test_defaults(monkeypatch): + monkeypatch.setenv("ORCAROUTER_API_KEY", "sk-orca-test") + model = create_orcarouter_model() + assert model.model_name == DEFAULT_ORCAROUTER_MODEL + assert model.provider.name == "openai" + assert model.provider.base_url == f"{DEFAULT_ORCAROUTER_BASE_URL}/" + + +def test_env_overrides(monkeypatch): + monkeypatch.setenv("ORCAROUTER_API_KEY", "sk-orca-test") + monkeypatch.setenv("ORCAROUTER_MODEL", "anthropic/claude-sonnet-5") + monkeypatch.setenv("ORCAROUTER_BASE_URL", "https://proxy.example.com/v1") + model = create_orcarouter_model() + assert model.model_name == "anthropic/claude-sonnet-5" + assert model.provider.base_url == "https://proxy.example.com/v1/"