|
13 | 13 | # See the License for the specific language governing permissions and |
14 | 14 | # limitations under the License. |
15 | 15 |
|
| 16 | +import json |
| 17 | +import os |
16 | 18 | from contextlib import asynccontextmanager |
| 19 | +from pathlib import Path |
17 | 20 |
|
18 | 21 | import httpx |
19 | 22 |
|
| 23 | +from nemoguardrails.llm.clients.openai_chat_model import OpenAIChatModel |
20 | 24 | from nemoguardrails.llm.clients.openai_compatible import OpenAICompatibleClient |
21 | 25 |
|
| 26 | +_FIXTURES_DIR = Path(__file__).parent / "fixtures" |
| 27 | + |
| 28 | +LIVE_TEST_MODE = bool(os.environ.get("LIVE_TEST_MODE") or os.environ.get("TEST_LIVE_MODE")) |
| 29 | + |
| 30 | +OPENAI_BASE_URL = "https://api.openai.com/v1" |
| 31 | +OPENAI_DEFAULT_MODEL = "gpt-4o-mini" |
| 32 | +NIM_BASE_URL = "https://integrate.api.nvidia.com/v1" |
| 33 | +NIM_DEFAULT_MODEL = "nvidia/nemotron-3-nano-30b-a3b" |
| 34 | + |
22 | 35 |
|
23 | 36 | def make_client(**kwargs): |
24 | 37 | return OpenAICompatibleClient(base_url="https://api.openai.com/v1", api_key="sk-test", **kwargs) |
@@ -117,3 +130,151 @@ async def aiter_lines(self): |
117 | 130 | yield FakeResponse() |
118 | 131 |
|
119 | 132 | return mock, aread_calls |
| 133 | + |
| 134 | + |
| 135 | +def load_fixture(name): |
| 136 | + with open(_FIXTURES_DIR / name) as f: |
| 137 | + return json.load(f) |
| 138 | + |
| 139 | + |
| 140 | +def fixture_exists(name): |
| 141 | + return (_FIXTURES_DIR / name).exists() |
| 142 | + |
| 143 | + |
| 144 | +def _fixture_to_response(data, request): |
| 145 | + if isinstance(data, list): |
| 146 | + body = "".join(f"data: {json.dumps(c)}\n\n" for c in data) + "data: [DONE]\n\n" |
| 147 | + return httpx.Response( |
| 148 | + 200, |
| 149 | + content=body.encode("utf-8"), |
| 150 | + headers={"content-type": "text/event-stream"}, |
| 151 | + request=request, |
| 152 | + ) |
| 153 | + if isinstance(data, dict) and "status_code" in data and "body" in data: |
| 154 | + body = data.get("body") |
| 155 | + content = json.dumps(body).encode("utf-8") if body is not None else b"" |
| 156 | + return httpx.Response( |
| 157 | + data["status_code"], |
| 158 | + content=content, |
| 159 | + headers=data.get("response_headers") or {}, |
| 160 | + request=request, |
| 161 | + ) |
| 162 | + if isinstance(data, dict): |
| 163 | + return httpx.Response(200, json=data, request=request) |
| 164 | + raise ValueError(f"Unknown fixture shape: {type(data).__name__}") |
| 165 | + |
| 166 | + |
| 167 | +def _resolve_fixture(fixture_or_name): |
| 168 | + return load_fixture(fixture_or_name) if isinstance(fixture_or_name, str) else fixture_or_name |
| 169 | + |
| 170 | + |
| 171 | +def fixture_transport(fixture_or_name, on_request=None): |
| 172 | + """Build an httpx.MockTransport that serves a fixture on each request.""" |
| 173 | + data = _resolve_fixture(fixture_or_name) |
| 174 | + |
| 175 | + def handler(request): |
| 176 | + if on_request is not None: |
| 177 | + on_request(request) |
| 178 | + return _fixture_to_response(data, request) |
| 179 | + |
| 180 | + return httpx.MockTransport(handler) |
| 181 | + |
| 182 | + |
| 183 | +def sequenced_fixture_transport(fixtures_or_names, on_request=None): |
| 184 | + """Serve a sequence of fixtures, one per request. For multi-turn tests.""" |
| 185 | + data_list = [_resolve_fixture(f) for f in fixtures_or_names] |
| 186 | + idx = [0] |
| 187 | + |
| 188 | + def handler(request): |
| 189 | + if on_request is not None: |
| 190 | + on_request(request) |
| 191 | + if idx[0] >= len(data_list): |
| 192 | + raise AssertionError( |
| 193 | + f"sequenced_fixture_transport: request {idx[0] + 1} exceeds fixture count ({len(data_list)})" |
| 194 | + ) |
| 195 | + response = _fixture_to_response(data_list[idx[0]], request) |
| 196 | + idx[0] += 1 |
| 197 | + return response |
| 198 | + |
| 199 | + return httpx.MockTransport(handler) |
| 200 | + |
| 201 | + |
| 202 | +@asynccontextmanager |
| 203 | +async def simulated_model( |
| 204 | + fixture_or_name, |
| 205 | + *, |
| 206 | + model_name="gpt-4o-mini", |
| 207 | + base_url="https://api.openai.com/v1", |
| 208 | + api_key="sk-test-simulated", |
| 209 | + on_request=None, |
| 210 | + transport=None, |
| 211 | + **client_kwargs, |
| 212 | +): |
| 213 | + """Context manager yielding an OpenAIChatModel whose HTTP transport is |
| 214 | + backed by a fixture. Exercises the full client stack (serialization, |
| 215 | + httpx, SSE parsing, response parsing) without a real API call. |
| 216 | + """ |
| 217 | + if transport is None: |
| 218 | + transport = fixture_transport(fixture_or_name, on_request=on_request) |
| 219 | + async with httpx.AsyncClient(transport=transport) as http: |
| 220 | + client = OpenAICompatibleClient( |
| 221 | + base_url=base_url, |
| 222 | + api_key=api_key, |
| 223 | + http_client=http, |
| 224 | + max_retries=0, |
| 225 | + **client_kwargs, |
| 226 | + ) |
| 227 | + yield OpenAIChatModel(client=client, model=model_name) |
| 228 | + |
| 229 | + |
| 230 | +@asynccontextmanager |
| 231 | +async def simulated_model_sequenced( |
| 232 | + fixtures, |
| 233 | + *, |
| 234 | + model_name="gpt-4o-mini", |
| 235 | + base_url="https://api.openai.com/v1", |
| 236 | + on_request=None, |
| 237 | + **client_kwargs, |
| 238 | +): |
| 239 | + """Like simulated_model but serves a sequence of fixtures (for multi-turn).""" |
| 240 | + transport = sequenced_fixture_transport(fixtures, on_request=on_request) |
| 241 | + async with simulated_model( |
| 242 | + None, |
| 243 | + model_name=model_name, |
| 244 | + base_url=base_url, |
| 245 | + on_request=on_request, |
| 246 | + transport=transport, |
| 247 | + **client_kwargs, |
| 248 | + ) as model: |
| 249 | + yield model |
| 250 | + |
| 251 | + |
| 252 | +def live_mode_enabled(provider: str) -> bool: |
| 253 | + if not LIVE_TEST_MODE: |
| 254 | + return False |
| 255 | + env_var = {"openai": "OPENAI_API_KEY", "nim": "NVIDIA_API_KEY"}[provider] |
| 256 | + return bool(os.environ.get(env_var)) |
| 257 | + |
| 258 | + |
| 259 | +@asynccontextmanager |
| 260 | +async def live_openai_model(*, model_name=OPENAI_DEFAULT_MODEL): |
| 261 | + """Yield an OpenAIChatModel backed by the real OpenAI API.""" |
| 262 | + async with httpx.AsyncClient() as http: |
| 263 | + client = OpenAICompatibleClient( |
| 264 | + base_url=OPENAI_BASE_URL, |
| 265 | + api_key=os.environ.get("OPENAI_API_KEY"), |
| 266 | + http_client=http, |
| 267 | + ) |
| 268 | + yield OpenAIChatModel(client=client, model=model_name) |
| 269 | + |
| 270 | + |
| 271 | +@asynccontextmanager |
| 272 | +async def live_nim_model(*, model_name=NIM_DEFAULT_MODEL): |
| 273 | + """Yield a NIM OpenAIChatModel backed by the real NIM API.""" |
| 274 | + async with httpx.AsyncClient() as http: |
| 275 | + client = OpenAICompatibleClient( |
| 276 | + base_url=NIM_BASE_URL, |
| 277 | + api_key=os.environ.get("NVIDIA_API_KEY"), |
| 278 | + http_client=http, |
| 279 | + ) |
| 280 | + yield OpenAIChatModel(client=client, model=model_name) |
0 commit comments