ai-parrot-client-nvidia 0.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,26 @@
1
+ Metadata-Version: 2.4
2
+ Name: ai-parrot-client-nvidia
3
+ Version: 0.2.0
4
+ Summary: NVIDIA NIM LLM client for AI-Parrot (parrot.clients.nvidia)
5
+ Author-email: Jesus Lara <jesuslara@phenobarbital.info>
6
+ License-Expression: MIT
7
+ Requires-Python: >=3.11
8
+ Description-Content-Type: text/markdown
9
+ Requires-Dist: ai-parrot>=0.29.0
10
+
11
+ # ai-parrot-client-nvidia
12
+
13
+ NVIDIA NIM LLM client satellite for
14
+ [AI-Parrot](https://github.com/phenobarbital/ai-parrot).
15
+
16
+ Provides `parrot.clients.nvidia.NvidiaClient` (built on the
17
+ OpenAI-compatible `OpenAIBaseClient`). Registers itself with `LLMFactory`
18
+ via the `parrot.clients` entry point group — no import of this package is
19
+ required for core AI-Parrot to know it exists once installed.
20
+
21
+ ```bash
22
+ uv pip install ai-parrot-client-nvidia
23
+ ```
24
+
25
+ See `sdd/specs/pep-420-llm-clients.spec.md` (FEAT-523) for the extraction
26
+ this package was split from.
@@ -0,0 +1,16 @@
1
+ # ai-parrot-client-nvidia
2
+
3
+ NVIDIA NIM LLM client satellite for
4
+ [AI-Parrot](https://github.com/phenobarbital/ai-parrot).
5
+
6
+ Provides `parrot.clients.nvidia.NvidiaClient` (built on the
7
+ OpenAI-compatible `OpenAIBaseClient`). Registers itself with `LLMFactory`
8
+ via the `parrot.clients` entry point group — no import of this package is
9
+ required for core AI-Parrot to know it exists once installed.
10
+
11
+ ```bash
12
+ uv pip install ai-parrot-client-nvidia
13
+ ```
14
+
15
+ See `sdd/specs/pep-420-llm-clients.spec.md` (FEAT-523) for the extraction
16
+ this package was split from.
@@ -0,0 +1,30 @@
1
+ [build-system]
2
+ requires = ["setuptools>=77.0.0", "wheel>=0.44.0"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "ai-parrot-client-nvidia"
7
+ version = "0.2.0"
8
+ description = "NVIDIA NIM LLM client for AI-Parrot (parrot.clients.nvidia)"
9
+ readme = "README.md"
10
+ requires-python = ">=3.11"
11
+ license = "MIT"
12
+ authors = [
13
+ {name = "Jesus Lara", email = "jesuslara@phenobarbital.info"}
14
+ ]
15
+ dependencies = [
16
+ "ai-parrot>=0.29.0",
17
+ # No extra SDK: NvidiaClient extends OpenAIBaseClient (core) and uses
18
+ # only tenacity (already a core ai-parrot base dependency).
19
+ ]
20
+
21
+ [project.entry-points."parrot.clients"]
22
+ nvidia = "parrot.clients.nvidia:NvidiaClient"
23
+
24
+ [tool.setuptools.packages.find]
25
+ where = ["src"]
26
+ include = ["parrot*"]
27
+ namespaces = true
28
+
29
+ [tool.uv.sources]
30
+ ai-parrot = { workspace = true }
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,26 @@
1
+ Metadata-Version: 2.4
2
+ Name: ai-parrot-client-nvidia
3
+ Version: 0.2.0
4
+ Summary: NVIDIA NIM LLM client for AI-Parrot (parrot.clients.nvidia)
5
+ Author-email: Jesus Lara <jesuslara@phenobarbital.info>
6
+ License-Expression: MIT
7
+ Requires-Python: >=3.11
8
+ Description-Content-Type: text/markdown
9
+ Requires-Dist: ai-parrot>=0.29.0
10
+
11
+ # ai-parrot-client-nvidia
12
+
13
+ NVIDIA NIM LLM client satellite for
14
+ [AI-Parrot](https://github.com/phenobarbital/ai-parrot).
15
+
16
+ Provides `parrot.clients.nvidia.NvidiaClient` (built on the
17
+ OpenAI-compatible `OpenAIBaseClient`). Registers itself with `LLMFactory`
18
+ via the `parrot.clients` entry point group — no import of this package is
19
+ required for core AI-Parrot to know it exists once installed.
20
+
21
+ ```bash
22
+ uv pip install ai-parrot-client-nvidia
23
+ ```
24
+
25
+ See `sdd/specs/pep-420-llm-clients.spec.md` (FEAT-523) for the extraction
26
+ this package was split from.
@@ -0,0 +1,11 @@
1
+ README.md
2
+ pyproject.toml
3
+ src/ai_parrot_client_nvidia.egg-info/PKG-INFO
4
+ src/ai_parrot_client_nvidia.egg-info/SOURCES.txt
5
+ src/ai_parrot_client_nvidia.egg-info/dependency_links.txt
6
+ src/ai_parrot_client_nvidia.egg-info/entry_points.txt
7
+ src/ai_parrot_client_nvidia.egg-info/requires.txt
8
+ src/ai_parrot_client_nvidia.egg-info/top_level.txt
9
+ src/parrot/clients/nvidia/__init__.py
10
+ src/parrot/clients/nvidia/client.py
11
+ src/parrot/clients/nvidia/models.py
@@ -0,0 +1,2 @@
1
+ [parrot.clients]
2
+ nvidia = parrot.clients.nvidia:NvidiaClient
@@ -0,0 +1,4 @@
1
+ from .client import NvidiaClient
2
+ from .models import NvidiaModel
3
+
4
+ __all__ = ["NvidiaClient", "NvidiaModel"]
@@ -0,0 +1,674 @@
1
+ """Nvidia NIM client for AI-Parrot.
2
+
3
+ Extends OpenAIBaseClient to route requests through Nvidia's OpenAI-compatible
4
+ NIM gateway at https://integrate.api.nvidia.com/v1.
5
+
6
+ All completion, streaming, tool-calling, retry, and invoke logic is inherited
7
+ from OpenAIBaseClient unchanged. Two Nvidia-specific affordances are added:
8
+
9
+ 1. The ``enable_thinking`` keyword on ``ask`` / ``ask_stream`` that injects
10
+ ``chat_template_kwargs`` into ``extra_body`` for reasoning-capable models
11
+ such as ``nvidia/nemotron-3-nano-omni-30b-a3b-reasoning``.
12
+ 2. A ``free_tier`` flag (default ``True``) that throttles outbound requests to
13
+ Nvidia's free-endpoint quota of 40 requests per minute. Set
14
+ ``free_tier=False`` — or ``NVIDIA_FREE_TIER=false`` in the environment — for
15
+ paid/self-hosted NIM endpoints that carry no such cap.
16
+ """
17
+
18
+ import asyncio
19
+ import contextvars
20
+ from collections import deque
21
+ from enum import Enum
22
+ from typing import Any, AsyncIterator, Deque, Dict, Optional
23
+
24
+ from navconfig import config
25
+ from tenacity import (
26
+ AsyncRetrying,
27
+ retry_if_exception_type,
28
+ retry_if_not_exception_type,
29
+ stop_after_attempt,
30
+ wait_exponential,
31
+ )
32
+
33
+ from ...exceptions import ParrotError
34
+ from ...models import AIMessage
35
+ from ..openai_base import OpenAIBaseClient
36
+ from .models import NvidiaModel
37
+
38
+ #: Requests-per-minute quota enforced on Nvidia's free NIM endpoints.
39
+ FREE_TIER_RPM: int = 40
40
+
41
+ #: Length of the rate-limit window, in seconds.
42
+ RATE_LIMIT_WINDOW: float = 60.0
43
+
44
+ # Context variable that carries enable_thinking / clear_thinking flags from
45
+ # ask / ask_stream down to _chat_completion without altering the parent's
46
+ # call signatures. Using a ContextVar is safe for concurrent async calls
47
+ # because each asyncio Task inherits an isolated copy of the context.
48
+ _thinking_ctx: contextvars.ContextVar[Dict[str, Any]] = contextvars.ContextVar("_nvidia_thinking_ctx", default={})
49
+
50
+ # Sampling parameters that ``OpenAIBaseClient.ask`` cannot accept. Its signature is
51
+ # fixed and carries no ``**kwargs``, so passing ``top_p``/``seed`` to it raises
52
+ # TypeError. They ride the same ContextVar channel as the thinking flags and
53
+ # are merged into the request inside ``_chat_completion``.
54
+ _sampling_ctx: contextvars.ContextVar[Dict[str, Any]] = contextvars.ContextVar("_nvidia_sampling_ctx", default={})
55
+
56
+ #: Sampling parameters this client can inject that the parent cannot forward.
57
+ #: ``top_p`` is deliberately NOT defaulted from ``self.top_p``: ``AbstractClient``
58
+ #: assigns it ``0.2`` when the caller says nothing (base.py:319), and silently
59
+ #: sending that would change results for every existing Nvidia call.
60
+ #:
61
+ #: These are all *native* OpenAI parameters, so they are passed straight to
62
+ #: ``client.chat.completions.create()`` as keyword arguments.
63
+ INJECTABLE_SAMPLING_PARAMS: tuple[str, ...] = ("top_p", "seed")
64
+
65
+ #: NIM-only request parameters, which must travel inside ``extra_body`` rather
66
+ #: than as keyword arguments: the OpenAI SDK's ``create()`` has an explicit
67
+ #: signature and raises ``TypeError`` on any keyword it does not define, so
68
+ #: passing these directly would break the call before it left the process.
69
+ #:
70
+ #: ``reasoning_budget`` caps how many tokens a reasoning model may spend on
71
+ #: ``reasoning_content`` before it must start answering. It is only sent when
72
+ #: explicitly requested — NIM applies its own per-model default otherwise, and
73
+ #: inventing a number here would silently cap every reasoning call.
74
+ INJECTABLE_EXTRA_BODY_PARAMS: tuple[str, ...] = ("reasoning_budget",)
75
+
76
+
77
+ class NvidiaRateLimitError(ParrotError):
78
+ """Raised when a free-tier slot could not be acquired within ``max_wait``.
79
+
80
+ Only raised when :class:`NvidiaClient` was constructed with an explicit
81
+ ``rate_limit_max_wait``. With the default (``None``) the client waits as
82
+ long as necessary instead of raising.
83
+
84
+ Args:
85
+ message: Human-readable error description.
86
+ *args: Forwarded to :class:`~parrot.exceptions.ParrotError`.
87
+ retry_after: Seconds the caller should wait before retrying, when known.
88
+ **kwargs: Forwarded to :class:`~parrot.exceptions.ParrotError`.
89
+
90
+ Attributes:
91
+ retry_after: Seconds until a slot is expected to free up, or ``None``.
92
+ """
93
+
94
+ def __init__(
95
+ self,
96
+ message: str,
97
+ *args,
98
+ retry_after: Optional[float] = None,
99
+ **kwargs,
100
+ ) -> None:
101
+ super().__init__(message, *args, **kwargs)
102
+ self.retry_after: Optional[float] = retry_after
103
+
104
+
105
+ class SlidingWindowRateLimiter:
106
+ """Async sliding-window rate limiter.
107
+
108
+ Admits at most ``limit`` acquisitions in any trailing ``window`` seconds.
109
+ Unlike a fixed-window counter, a sliding window cannot be defeated by
110
+ bursting across a window boundary, which matters because Nvidia measures
111
+ its 40 rpm free-tier quota continuously.
112
+
113
+ Timestamps come from the running event loop's monotonic clock
114
+ (``loop.time()``), so the limiter is unaffected by wall-clock adjustments.
115
+
116
+ The limiter is *not* shared between instances: each :class:`NvidiaClient`
117
+ owns its own window. Nvidia enforces the quota per account, so N clients
118
+ sharing one API key can collectively exceed 40 rpm.
119
+
120
+ Args:
121
+ limit: Maximum number of acquisitions per window. Must be >= 1.
122
+ window: Window length in seconds. Must be > 0.
123
+
124
+ Raises:
125
+ ValueError: If ``limit`` < 1 or ``window`` <= 0.
126
+
127
+ Example::
128
+
129
+ limiter = SlidingWindowRateLimiter(limit=40, window=60.0)
130
+ waited = await limiter.acquire()
131
+ """
132
+
133
+ #: Added to every computed sleep so float imprecision can never cause a
134
+ #: spin where the woken waiter finds the oldest hit still inside the window.
135
+ _EPSILON: float = 0.001
136
+
137
+ def __init__(self, limit: int, window: float = RATE_LIMIT_WINDOW) -> None:
138
+ if limit < 1:
139
+ raise ValueError(f"limit must be >= 1, got {limit!r}")
140
+ if window <= 0:
141
+ raise ValueError(f"window must be > 0, got {window!r}")
142
+ self._limit: int = int(limit)
143
+ self._window: float = float(window)
144
+ self._hits: Deque[float] = deque()
145
+ self._lock: asyncio.Lock = asyncio.Lock()
146
+
147
+ @property
148
+ def limit(self) -> int:
149
+ """Maximum acquisitions allowed per window."""
150
+ return self._limit
151
+
152
+ @property
153
+ def window(self) -> float:
154
+ """Window length in seconds."""
155
+ return self._window
156
+
157
+ def _prune(self, now: float) -> None:
158
+ """Drop recorded hits that have aged out of the trailing window.
159
+
160
+ Args:
161
+ now: Current monotonic time, from ``loop.time()``.
162
+ """
163
+ cutoff = now - self._window
164
+ while self._hits and self._hits[0] <= cutoff:
165
+ self._hits.popleft()
166
+
167
+ def current_usage(self) -> int:
168
+ """Return how many acquisitions currently occupy the window.
169
+
170
+ Intended for logging and tests; the value is a snapshot and may be
171
+ stale as soon as it is returned. Must be called from inside a running
172
+ event loop, since the window is measured on that loop's clock.
173
+
174
+ Returns:
175
+ Number of hits still inside the trailing window.
176
+
177
+ Raises:
178
+ RuntimeError: If called with no running event loop.
179
+ """
180
+ self._prune(asyncio.get_running_loop().time())
181
+ return len(self._hits)
182
+
183
+ async def acquire(self, max_wait: Optional[float] = None) -> float:
184
+ """Reserve one slot, sleeping while the window is saturated.
185
+
186
+ Args:
187
+ max_wait: Maximum total seconds to wait. ``None`` waits as long as
188
+ necessary. When the projected wait exceeds this budget the
189
+ call raises instead of sleeping, and no slot is consumed.
190
+
191
+ Returns:
192
+ Total seconds spent waiting (``0.0`` when a slot was free).
193
+
194
+ Raises:
195
+ NvidiaRateLimitError: If a slot could not be reserved within
196
+ ``max_wait`` seconds.
197
+ """
198
+ waited = 0.0
199
+ while True:
200
+ async with self._lock:
201
+ now = asyncio.get_running_loop().time()
202
+ self._prune(now)
203
+ if len(self._hits) < self._limit:
204
+ self._hits.append(now)
205
+ return waited
206
+ # Window is full: the oldest hit dictates when a slot frees.
207
+ sleep_for = max(self._hits[0] + self._window - now, 0.0)
208
+ sleep_for += self._EPSILON
209
+ if max_wait is not None and waited + sleep_for > max_wait:
210
+ raise NvidiaRateLimitError(
211
+ f"Nvidia free-tier rate limit reached "
212
+ f"({self._limit} requests / {self._window:g}s); a slot would "
213
+ f"not free up within max_wait={max_wait:g}s",
214
+ retry_after=sleep_for,
215
+ )
216
+ await asyncio.sleep(sleep_for)
217
+ waited += sleep_for
218
+
219
+
220
+ class NvidiaClient(OpenAIBaseClient):
221
+ """Client for Nvidia NIM's OpenAI-compatible API gateway.
222
+
223
+ Routes all requests through ``https://integrate.api.nvidia.com/v1`` and
224
+ resolves the API key from the constructor argument or the ``NVIDIA_API_KEY``
225
+ environment variable (via ``navconfig.config``).
226
+
227
+ All inherited OpenAI machinery — ``ask``, ``ask_stream``, ``invoke``,
228
+ ``_chat_completion``, tool calling, structured output, and retry — works
229
+ without modification.
230
+
231
+ Two Nvidia-specific affordances are layered on top:
232
+
233
+ **Thinking flags.** The ``enable_thinking`` shortcut on ``ask`` /
234
+ ``ask_stream`` injects ``chat_template_kwargs`` into ``extra_body`` for
235
+ reasoning-capable models (e.g. ``nvidia/nemotron-3-nano-omni-30b-a3b-reasoning``). It is propagated to
236
+ ``_chat_completion`` via an async context variable so that no changes to
237
+ the parent's call signatures are required.
238
+
239
+ **Free-tier rate limiting.** Nvidia's free NIM endpoints cap traffic at 40
240
+ requests per minute. With ``free_tier=True`` (the default) every outbound
241
+ request first reserves a slot in a :class:`SlidingWindowRateLimiter`,
242
+ async-sleeping when the window is saturated so the quota is never
243
+ exceeded. Set ``free_tier=False`` for paid or self-hosted NIM endpoints to
244
+ remove the cap entirely. The limiter covers ``ask``, ``invoke``,
245
+ ``ask_stream``, and every other path routed through ``_chat_completion``;
246
+ each retry attempt consumes its own slot, since a retry is a real request.
247
+ Inherited OpenAI-only vision helpers (e.g. ``ask_to_image``) call the SDK
248
+ directly and are therefore not counted.
249
+
250
+ Args:
251
+ api_key: Nvidia NIM API key. Falls back to ``NVIDIA_API_KEY`` env var
252
+ (resolved via ``navconfig.config``).
253
+ free_tier: When ``True`` (default), throttle requests to
254
+ ``requests_per_minute``. When ``None``, the value is read from the
255
+ ``NVIDIA_FREE_TIER`` env var, itself defaulting to ``True``.
256
+ :data:`~parrot.clients.nvidia.models.FREE_TIER_MODELS` lists the models
257
+ NVIDIA publishes as free preview endpoints — the ones this
258
+ throttle exists for. Note that a free endpoint answers ``503
259
+ ResourceExhausted`` when it is at capacity; that is saturation,
260
+ not a bad model id, and the right response is to retry rather
261
+ than switch models.
262
+ requests_per_minute: Requests allowed per 60s window while
263
+ ``free_tier`` is active. Defaults to :data:`FREE_TIER_RPM` (40).
264
+ rate_limit_max_wait: Maximum seconds to wait for a free-tier slot.
265
+ ``None`` (default) waits as long as necessary; a number makes the
266
+ client raise :class:`NvidiaRateLimitError` rather than block past
267
+ that budget.
268
+ seed: Default sampling seed for reproducibility, overridable per call.
269
+ reasoning_budget: Default cap on the tokens a reasoning model may
270
+ spend on ``reasoning_content`` before it must answer, overridable
271
+ per call. ``None`` (default) sends nothing and lets NIM apply its
272
+ own per-model default.
273
+ **kwargs: Additional arguments passed to ``OpenAIBaseClient`` /
274
+ ``AbstractClient``. Notably ``timeout`` (defaults to
275
+ :attr:`_default_timeout`, 300s — reasoning calls are slow) and
276
+ ``max_tokens`` (defaults to :attr:`_default_max_tokens`, 65536 —
277
+ reasoning shares the answer's budget).
278
+
279
+ Example::
280
+
281
+ client = NvidiaClient(model=NvidiaModel.GLM_5_2)
282
+ response = await client.ask(
283
+ "Explain gradient descent.",
284
+ enable_thinking=True,
285
+ )
286
+
287
+ # Paid endpoint — no throttling.
288
+ paid = NvidiaClient(free_tier=False)
289
+
290
+ # Free endpoint, but fail fast instead of waiting more than 5s.
291
+ strict = NvidiaClient(rate_limit_max_wait=5.0)
292
+ """
293
+
294
+ client_type: str = "nvidia"
295
+ client_name: str = "nvidia"
296
+
297
+ # FEAT-523 folder-convention attributes (read by LLMFactory).
298
+ provider_keys: tuple[str, ...] = ("nvidia",)
299
+ models: type[Enum] = NvidiaModel
300
+ _default_model: str = NvidiaModel.MINIMAX_M3.value
301
+
302
+ # NIM's reasoning models routinely take longer than the 60s
303
+ # OpenAIBaseClient default: a measured single-turn word problem against
304
+ # ``nemotron-3-nano-omni-30b-a3b-reasoning`` took 96.8s end to end, which
305
+ # failed with APITimeoutError against a perfectly healthy endpoint.
306
+ _default_timeout: float = 300.0
307
+
308
+ # NIM accepts 65536 on its current chat models (verified against
309
+ # nemotron-3-nano-omni-30b-a3b-reasoning, minimax-m3 and gpt-oss-120b).
310
+ # ``max_tokens`` is a cap, not a reservation — nothing is billed for
311
+ # headroom that goes unused — and reasoning models need the room because
312
+ # ``reasoning_content`` is drawn from the same budget as the answer.
313
+ _default_max_tokens: int = 65536
314
+
315
+ def __init__(
316
+ self,
317
+ api_key: Optional[str] = None,
318
+ *,
319
+ free_tier: Optional[bool] = None,
320
+ requests_per_minute: Optional[int] = None,
321
+ rate_limit_max_wait: Optional[float] = None,
322
+ seed: Optional[int] = None,
323
+ reasoning_budget: Optional[int] = None,
324
+ **kwargs,
325
+ ):
326
+ resolved_key = api_key or config.get("NVIDIA_API_KEY")
327
+
328
+ # Capture an explicit top_p BEFORE super().__init__ runs: AbstractClient
329
+ # assigns self.top_p = kwargs.get('top_p', 0.2), after which an explicit
330
+ # 0.2 is indistinguishable from the default. Only a caller-supplied
331
+ # value becomes the per-instance default, so existing behaviour (no
332
+ # top_p on the wire) is preserved for everyone who never passed one.
333
+ explicit_top_p = kwargs.get("top_p")
334
+
335
+ super().__init__(
336
+ api_key=resolved_key,
337
+ base_url="https://integrate.api.nvidia.com/v1",
338
+ **kwargs,
339
+ )
340
+ # Re-set after super().__init__ because AbstractClient may overwrite
341
+ # self.api_key during its own initialisation. This mirrors the guard
342
+ # used by OpenRouterClient (openrouter.py:75).
343
+ self.api_key = resolved_key
344
+
345
+ # Rate-limit state is built after super().__init__ for the same reason:
346
+ # the parent must not be able to clobber it.
347
+ if free_tier is None:
348
+ free_tier = config.getboolean("NVIDIA_FREE_TIER", fallback=True)
349
+ self.free_tier: bool = bool(free_tier)
350
+ self.requests_per_minute: int = int(requests_per_minute if requests_per_minute is not None else FREE_TIER_RPM)
351
+ self.rate_limit_max_wait: Optional[float] = rate_limit_max_wait
352
+ self._rate_limiter: Optional[SlidingWindowRateLimiter] = (
353
+ SlidingWindowRateLimiter(self.requests_per_minute, RATE_LIMIT_WINDOW) if self.free_tier else None
354
+ )
355
+
356
+ # Per-instance sampling defaults for the parameters the parent drops.
357
+ self.seed: Optional[int] = seed
358
+ self.reasoning_budget: Optional[int] = reasoning_budget
359
+ self._default_top_p: Optional[float] = explicit_top_p
360
+
361
+ def _resolve_sampling(
362
+ self,
363
+ top_p: Optional[float],
364
+ seed: Optional[int],
365
+ reasoning_budget: Optional[int] = None,
366
+ ) -> Dict[str, Any]:
367
+ """Merge per-call sampling overrides over the per-instance defaults.
368
+
369
+ Args:
370
+ top_p: Per-call nucleus-sampling value, or ``None`` to fall back to
371
+ the value passed to the constructor.
372
+ seed: Per-call seed, or ``None`` to fall back to the constructor's.
373
+ reasoning_budget: Per-call reasoning-token cap, or ``None`` to fall
374
+ back to the constructor's.
375
+
376
+ Returns:
377
+ Only the parameters that resolved to a non-``None`` value, so
378
+ nothing extra is ever put on the wire.
379
+ """
380
+ resolved = {
381
+ "top_p": top_p if top_p is not None else self._default_top_p,
382
+ "seed": seed if seed is not None else self.seed,
383
+ "reasoning_budget": (reasoning_budget if reasoning_budget is not None else self.reasoning_budget),
384
+ }
385
+ return {key: value for key, value in resolved.items() if value is not None}
386
+
387
+ async def _acquire_rate_limit_slot(self) -> float:
388
+ """Reserve a free-tier slot before issuing a request.
389
+
390
+ A no-op returning ``0.0`` when ``free_tier`` is ``False``. Otherwise
391
+ delegates to the client's :class:`SlidingWindowRateLimiter`, sleeping
392
+ until the 40 rpm window has room.
393
+
394
+ Returns:
395
+ Seconds spent waiting for the slot.
396
+
397
+ Raises:
398
+ NvidiaRateLimitError: If ``rate_limit_max_wait`` is set and a slot
399
+ could not be reserved within that budget.
400
+ """
401
+ if self._rate_limiter is None:
402
+ return 0.0
403
+ waited = await self._rate_limiter.acquire(max_wait=self.rate_limit_max_wait)
404
+ if waited > 0:
405
+ self.logger.warning(
406
+ "Nvidia free-tier limit (%s rpm) reached; throttled for %.2fs. "
407
+ "Pass free_tier=False for paid NIM endpoints.",
408
+ self._rate_limiter.limit,
409
+ waited,
410
+ )
411
+ return waited
412
+
413
+ @staticmethod
414
+ def _merge_thinking_extra_body(
415
+ extra_body: Optional[Dict[str, Any]],
416
+ enable_thinking: bool,
417
+ clear_thinking: bool,
418
+ ) -> Optional[Dict[str, Any]]:
419
+ """Merge ``chat_template_kwargs`` reasoning flags into ``extra_body``.
420
+
421
+ When ``enable_thinking`` is ``False`` the function returns
422
+ ``extra_body`` completely unchanged (including returning ``None`` when
423
+ ``extra_body`` was ``None``).
424
+
425
+ When ``enable_thinking`` is ``True`` the function returns a new dict
426
+ that preserves every existing key in ``extra_body`` and every existing
427
+ key inside ``extra_body["chat_template_kwargs"]``, then adds
428
+ ``enable_thinking`` and ``clear_thinking`` flags.
429
+
430
+ This is an internal helper; callers should use the ``enable_thinking``
431
+ keyword on ``ask`` / ``ask_stream`` rather than calling this directly.
432
+
433
+ Args:
434
+ extra_body: Existing ``extra_body`` dict (may be ``None``).
435
+ enable_thinking: When ``True``, inject the reasoning flags.
436
+ clear_thinking: Value forwarded to ``clear_thinking`` in the
437
+ injected payload.
438
+
439
+ Returns:
440
+ Updated ``extra_body`` dict, or ``None`` when nothing was injected.
441
+ """
442
+ if not enable_thinking:
443
+ return extra_body
444
+ merged: Dict[str, Any] = dict(extra_body or {})
445
+ kwargs_block: Dict[str, Any] = dict(merged.get("chat_template_kwargs") or {})
446
+ kwargs_block["enable_thinking"] = True
447
+ kwargs_block["clear_thinking"] = clear_thinking
448
+ merged["chat_template_kwargs"] = kwargs_block
449
+ return merged
450
+
451
+ async def _chat_completion(
452
+ self,
453
+ model: str,
454
+ messages: Any,
455
+ use_tools: bool = False,
456
+ **kwargs,
457
+ ) -> Any:
458
+ """Run a chat completion against NVIDIA NIM via ``create()``.
459
+
460
+ Three NVIDIA-specific differences from ``OpenAIBaseClient._chat_completion``:
461
+
462
+ 1. Always uses ``client.chat.completions.create``. NIM rejects the
463
+ OpenAI SDK's ``parse()`` shortcut (returns 5xx / "page not found"),
464
+ so we never route through it — even when ``use_tools`` is ``False``.
465
+ 2. Reads the thinking flags from the async context variable set by
466
+ ``ask`` / ``ask_stream`` and merges them into ``extra_body`` for
467
+ reasoning-capable models (e.g. ``nvidia/nemotron-3-nano-omni-30b-a3b-reasoning``).
468
+ 3. Reserves a free-tier rate-limit slot before *each* attempt when
469
+ ``free_tier`` is active, so retries are counted against the quota
470
+ too.
471
+
472
+ Args:
473
+ model: Model identifier string.
474
+ messages: Chat messages list.
475
+ use_tools: Whether tools are enabled (kept for parity with parent).
476
+ **kwargs: Additional completion arguments forwarded to the OpenAI SDK.
477
+
478
+ Returns:
479
+ Raw OpenAI ``ChatCompletion`` response.
480
+
481
+ Raises:
482
+ NvidiaRateLimitError: If ``rate_limit_max_wait`` is set and no
483
+ free-tier slot became available within that budget.
484
+ """
485
+ from openai import (
486
+ APIConnectionError,
487
+ APIError,
488
+ APITimeoutError,
489
+ RateLimitError,
490
+ )
491
+
492
+ thinking = _thinking_ctx.get()
493
+ if thinking.get("enable_thinking"):
494
+ kwargs["extra_body"] = self._merge_thinking_extra_body(
495
+ kwargs.get("extra_body"),
496
+ True,
497
+ thinking.get("clear_thinking", False),
498
+ )
499
+
500
+ # Inject the sampling parameters the parent's fixed signature drops.
501
+ # An explicit value already in kwargs always wins, so this can never
502
+ # override something a caller managed to set by another route.
503
+ # Native OpenAI parameters go on as keyword arguments; NIM-only ones
504
+ # must be nested under extra_body or the SDK rejects the keyword.
505
+ sampling = _sampling_ctx.get()
506
+ for name, value in sampling.items():
507
+ if name in INJECTABLE_SAMPLING_PARAMS and kwargs.get(name) is None:
508
+ kwargs[name] = value
509
+
510
+ nim_params = {name: value for name, value in sampling.items() if name in INJECTABLE_EXTRA_BODY_PARAMS}
511
+ if nim_params:
512
+ extra_body = dict(kwargs.get("extra_body") or {})
513
+ # An explicit extra_body entry always wins, mirroring the rule
514
+ # applied to the keyword parameters above.
515
+ for name, value in nim_params.items():
516
+ extra_body.setdefault(name, value)
517
+ kwargs["extra_body"] = extra_body
518
+ # APITimeoutError is a *subclass* of APIConnectionError, so it used to
519
+ # be swept into the retry set. That is the worst possible thing to
520
+ # retry here: a reasoning request that outran the timeout is not a
521
+ # transient fault, and each retry pays the full timeout again *and*
522
+ # makes the endpoint regenerate the whole answer from scratch. Three
523
+ # attempts at the 300s default would block for 15 minutes before
524
+ # surfacing the same error. Fail on the first one instead, and tell the
525
+ # caller which knob actually fixes it.
526
+ retry_policy = AsyncRetrying(
527
+ retry=(
528
+ retry_if_exception_type((APIConnectionError, RateLimitError, APIError))
529
+ & retry_if_not_exception_type(APITimeoutError)
530
+ ),
531
+ wait=wait_exponential(multiplier=1, min=2, max=10),
532
+ stop=stop_after_attempt(3),
533
+ reraise=True,
534
+ )
535
+ async for attempt in retry_policy:
536
+ with attempt:
537
+ await self._acquire_rate_limit_slot()
538
+ return await self.client.chat.completions.create(
539
+ model=model,
540
+ messages=messages,
541
+ **kwargs,
542
+ )
543
+
544
+ async def ask(
545
+ self,
546
+ prompt: str,
547
+ *,
548
+ enable_thinking: bool = False,
549
+ clear_thinking: bool = False,
550
+ top_p: Optional[float] = None,
551
+ seed: Optional[int] = None,
552
+ reasoning_budget: Optional[int] = None,
553
+ **kwargs,
554
+ ) -> AIMessage:
555
+ """Submit a prompt and return the full response.
556
+
557
+ Identical to ``OpenAIBaseClient.ask`` with two additions:
558
+
559
+ 1. An ``enable_thinking`` shortcut that injects
560
+ ``chat_template_kwargs`` into ``extra_body`` for reasoning-capable
561
+ models (e.g. ``nvidia/nemotron-3-nano-omni-30b-a3b-reasoning``).
562
+ 2. ``top_p`` and ``seed`` support. ``OpenAIBaseClient.ask`` has a fixed
563
+ signature with no ``**kwargs``, so passing either to it raises
564
+ ``TypeError``; both are accepted here and merged into the request
565
+ inside ``_chat_completion``.
566
+
567
+ All of it travels via async context variables, so the parent's call
568
+ signature is preserved.
569
+
570
+ Args:
571
+ prompt: User message text.
572
+ enable_thinking: When ``True``, add
573
+ ``extra_body["chat_template_kwargs"]["enable_thinking"] = True``.
574
+ clear_thinking: Forwarded to ``clear_thinking`` in the payload
575
+ when ``enable_thinking`` is ``True``.
576
+ top_p: Nucleus-sampling value. Defaults to the ``top_p`` passed to
577
+ the constructor; when neither is given, none is sent and the
578
+ endpoint's own default applies.
579
+ seed: Sampling seed for reproducibility. Defaults to the ``seed``
580
+ passed to the constructor.
581
+ reasoning_budget: Maximum tokens a reasoning model may spend on
582
+ ``reasoning_content`` before answering. Defaults to the
583
+ constructor's value; when neither is given, NIM's own
584
+ per-model default applies.
585
+ **kwargs: All other keyword arguments delegated to
586
+ ``OpenAIBaseClient.ask`` (e.g. ``model``, ``temperature``,
587
+ ``system_prompt``, ``session_id``).
588
+
589
+ Returns:
590
+ AIMessage with the model response. For reasoning models the
591
+ thinking trace is available on ``AIMessage.reasoning``, kept
592
+ separate from the answer in ``AIMessage.output``.
593
+ """
594
+ kwargs.setdefault("model", self.model or self._default_model)
595
+ thinking_token = _thinking_ctx.set({"enable_thinking": enable_thinking, "clear_thinking": clear_thinking})
596
+ sampling_token = _sampling_ctx.set(self._resolve_sampling(top_p, seed, reasoning_budget))
597
+ try:
598
+ return await super().ask(prompt, **kwargs)
599
+ finally:
600
+ _thinking_ctx.reset(thinking_token)
601
+ _sampling_ctx.reset(sampling_token)
602
+
603
+ async def ask_stream(
604
+ self,
605
+ prompt: str,
606
+ *,
607
+ enable_thinking: bool = False,
608
+ clear_thinking: bool = False,
609
+ top_p: Optional[float] = None,
610
+ seed: Optional[int] = None,
611
+ reasoning_budget: Optional[int] = None,
612
+ **kwargs,
613
+ ) -> AsyncIterator[str]:
614
+ """Submit a prompt and stream response chunks.
615
+
616
+ Identical to ``OpenAIBaseClient.ask_stream`` with the same
617
+ ``enable_thinking`` shortcut as ``ask``. For reasoning-capable models
618
+ (e.g. ``nvidia/nemotron-3-nano-omni-30b-a3b-reasoning``) each chunk may
619
+ carry a ``delta.reasoning_content`` field in addition to
620
+ ``delta.content``.
621
+
622
+ The flags are forwarded to ``_chat_completion`` via an async context
623
+ variable, so the parent's call signature is preserved.
624
+
625
+ FEAT-438 (TASK-2298/2300): ``ask_stream`` now routes through
626
+ ``_chat_completion`` (the single completion funnel), so the
627
+ free-tier slot is reserved there — no separate reservation is made
628
+ here, which would otherwise double-count against the 40 rpm quota.
629
+
630
+ ``top_p``, ``seed`` and ``reasoning_budget`` ARE supported here. They
631
+ previously raised ``NotImplementedError`` because ``ask_stream`` did
632
+ not set the ``_sampling_ctx`` context variable that
633
+ ``_chat_completion`` reads. Once TASK-2298 routed the streaming path
634
+ through ``_chat_completion`` as well, that limitation became stale
635
+ plumbing rather than a real constraint: the injection point is shared,
636
+ so setting the same context variable here is all it took. This
637
+ matters most for reasoning models, whose natural mode is streaming.
638
+
639
+ Args:
640
+ prompt: User message text.
641
+ enable_thinking: When ``True``, inject reasoning flags into
642
+ ``extra_body``.
643
+ clear_thinking: Forwarded to ``clear_thinking`` in the payload
644
+ when ``enable_thinking`` is ``True``.
645
+ top_p: Nucleus-sampling value. Defaults to the constructor's.
646
+ seed: Sampling seed for reproducibility. Defaults to the
647
+ constructor's.
648
+ reasoning_budget: Maximum tokens spent on reasoning before
649
+ answering. Defaults to the constructor's.
650
+ **kwargs: All other keyword arguments delegated to
651
+ ``OpenAIBaseClient.ask_stream`` (e.g. ``model``, ``temperature``,
652
+ ``system_prompt``, ``session_id``).
653
+
654
+ Yields:
655
+ Response text chunks (same shape as ``OpenAIBaseClient.ask_stream``).
656
+
657
+ Raises:
658
+ NvidiaRateLimitError: If ``rate_limit_max_wait`` is set and no
659
+ free-tier slot became available within that budget.
660
+ """
661
+ kwargs.setdefault("model", self.model or self._default_model)
662
+ # FEAT-438 TASK-2300: no manual _acquire_rate_limit_slot() call here
663
+ # anymore — ask_stream() now routes through _chat_completion (the
664
+ # single completion funnel, TASK-2298), which already reserves a
665
+ # slot per attempt. A second reservation here would double-count
666
+ # every streamed call against the 40 rpm free-tier quota.
667
+ thinking_token = _thinking_ctx.set({"enable_thinking": enable_thinking, "clear_thinking": clear_thinking})
668
+ sampling_token = _sampling_ctx.set(self._resolve_sampling(top_p, seed, reasoning_budget))
669
+ try:
670
+ async for chunk in super().ask_stream(prompt, **kwargs):
671
+ yield chunk
672
+ finally:
673
+ _thinking_ctx.reset(thinking_token)
674
+ _sampling_ctx.reset(sampling_token)
@@ -0,0 +1,160 @@
1
+ """Nvidia NIM data models for AI-Parrot.
2
+
3
+ Provides model enums for Nvidia's NIM-hosted OpenAI-compatible API gateway
4
+ (https://integrate.api.nvidia.com/v1). No Pydantic wrappers are needed —
5
+ Nvidia's response shape matches the OpenAI Chat Completion shape and is
6
+ already covered by existing AIMessage / CompletionUsage models.
7
+ """
8
+
9
+ from enum import Enum
10
+
11
+
12
+ class NvidiaModel(str, Enum):
13
+ """Nvidia NIM-hosted model identifiers.
14
+
15
+ String-valued enum so members interchange with raw model strings
16
+ in OpenAI SDK calls (e.g. ``model=NvidiaModel.KIMI_K2_6.value``
17
+ or simply ``model=NvidiaModel.KIMI_K2_6`` since the class
18
+ inherits from ``str``).
19
+
20
+ Slugs use NIM's ``vendor/model`` form, where the vendor segment may itself
21
+ contain a dash (``z-ai/glm-5.2``).
22
+
23
+ .. warning::
24
+
25
+ **Several members are withdrawn upstream.** A live probe on 2026-09-02
26
+ sent a real ``POST /v1/chat/completions`` for every member then defined:
27
+
28
+ =============================== ======================================
29
+ Member Live status (2026-09-02)
30
+ =============================== ======================================
31
+ ``MINIMAX_M3`` 200 OK
32
+ ``GPT_OSS_120B`` 200 OK
33
+ ``KIMI_K2_6`` 404 — gated per account
34
+ ``DEEPSEEK_V4_PRO`` 410 Gone — EOL
35
+ ``DEEPSEEK_V4_FLASH`` 410 Gone — EOL
36
+ ``LLAMA_3_3_70B_INSTRUCT`` 410 Gone — EOL 2026-08-26
37
+ ``NEMOTRON_3_NANO_30B`` 410 Gone — EOL
38
+ ``GLM_5_2`` 410 Gone — EOL 2026-08-21
39
+ ``STEPFUN_STEP_3_7_FLASH`` 410 Gone — EOL
40
+ =============================== ======================================
41
+
42
+ The dead members are deliberately kept rather than deleted, because
43
+ removing a member breaks any caller that imports it by name. They are
44
+ documented here so nobody picks one expecting it to work, and the
45
+ replacements now live in :data:`FREE_TIER_MODELS` below.
46
+
47
+ The revision before this one claimed verification against
48
+ ``GET /v1/models`` on 2026-08-05. That is why the drift went unnoticed:
49
+ presence in the catalog listing is NOT proof a slug still serves traffic —
50
+ ``GLM_5_2`` and ``LLAMA_3_3_70B_INSTRUCT`` both vanished from the catalog
51
+ *and* began returning 410, while ``KIMI_K2_6`` is listed but returns 404
52
+ per account. Only a real completion request is evidence.
53
+
54
+ Every member of the previous revision of this enum had reached
55
+ end-of-life and been withdrawn from the catalog — requests returned
56
+ ``410 Gone`` or ``404``. The mapping applied was:
57
+
58
+ - ``moonshotai/kimi-k2-thinking`` → ``moonshotai/kimi-k2.6``
59
+ - ``moonshotai/kimi-k2-instruct-0905`` → ``moonshotai/kimi-k2.6``
60
+ - ``moonshotai/kimi-k2.5`` → ``moonshotai/kimi-k2.6``
61
+ - ``minimaxai/minimax-m2.5`` → ``minimaxai/minimax-m3``
62
+ - ``minimaxai/minimax-m2.7`` → ``minimaxai/minimax-m3``
63
+ - ``mistralai/mamba-codestral-7b-v0.1`` → ``mistralai/mistral-nemotron``
64
+ (and ``poolside/laguna-xs-2.1`` for the code-generation role)
65
+ - ``deepseek-ai/deepseek-v3.1-terminus`` → ``deepseek-ai/deepseek-v4-pro``
66
+ - ``qwen/qwen3.5-397b-a17b`` → dropped; the catalog carries no Qwen model
67
+ - ``z-ai/glm-5.1`` → ``z-ai/glm-5.2``
68
+
69
+ The three Moonshot members and the two Minimax members collapsed to a
70
+ single successor each. They are deliberately **not** kept as aliases: in a
71
+ ``str``-valued Enum, two members sharing one value silently become
72
+ aliases of the first, which would make ``NvidiaModel.KIMI_K2_THINKING``
73
+ resolve to a non-thinking model.
74
+ """
75
+
76
+ # Moonshot AI — gated per account (404 without the entitlement)
77
+ KIMI_K2_6 = "moonshotai/kimi-k2.6"
78
+
79
+ # Minimax
80
+ MINIMAX_M3 = "minimaxai/minimax-m3"
81
+
82
+ # DeepSeek
83
+ DEEPSEEK_V4_PRO = "deepseek-ai/deepseek-v4-pro"
84
+ DEEPSEEK_V4_FLASH = "deepseek-ai/deepseek-v4-flash"
85
+
86
+ # Mistral
87
+ MISTRAL_NEMOTRON = "mistralai/mistral-nemotron"
88
+
89
+ # Code generation
90
+ LAGUNA_XS_2_1 = "poolside/laguna-xs-2.1"
91
+ LLAMA_3_3_70B_INSTRUCT = "meta/llama-3.3-70b-instruct"
92
+
93
+ # OpenAI open-weights
94
+ GPT_OSS_120B = "openai/gpt-oss-120b"
95
+
96
+ # Nvidia first-party
97
+ NEMOTRON_3_NANO_30B = "nvidia/nemotron-3-nano-30b-a3b"
98
+
99
+ # -- Free preview endpoints (see FREE_TIER_MODELS below) ----------------
100
+
101
+ #: Reasoning model: emits ``reasoning_content`` beside ``content``.
102
+ #: Surfaced on ``AIMessage.reasoning``. Because the thinking is drawn from
103
+ #: the same token budget as the answer, this model needs a large
104
+ #: ``max_tokens`` and a generous timeout — both are the NvidiaClient
105
+ #: defaults (65536 / 300s). Optionally cap the thinking with
106
+ #: ``reasoning_budget``.
107
+ NEMOTRON_3_NANO_OMNI_30B_REASONING = "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning"
108
+
109
+ #: Reasoning-capable; confirmed emitting ``reasoning_content``.
110
+ NEMOTRON_3_5_LIGHTNING_30B = "nvidia/nemotron-3.5-lightning-30b-a3b"
111
+
112
+ #: Moonshot's successor to the account-gated ``kimi-k2.6``.
113
+ #: Reasoning-capable; confirmed emitting ``reasoning_content``.
114
+ KIMI_K3 = "moonshotai/kimi-k3"
115
+
116
+ #: Successor to the withdrawn ``deepseek-ai/deepseek-v4-flash``.
117
+ #: Its thinking flags use DIFFERENT ``chat_template_kwargs`` keys than the
118
+ #: ``enable_thinking``/``clear_thinking`` pair ``NvidiaClient`` injects —
119
+ #: this model documents ``{"thinking": True, "reasoning_effort": "high"}``.
120
+ #: Pass those through ``extra_body`` explicitly rather than relying on the
121
+ #: ``enable_thinking=True`` shortcut.
122
+ DEEPSEEK_V4_FLASH_0731 = "deepseek-ai/deepseek-v4-flash-0731"
123
+
124
+ GEMMA_4_31B_IT = "google/gemma-4-31b-it"
125
+
126
+ # Z-AI — WITHDRAWN upstream (410 Gone, EOL 2026-08-21). Kept only so
127
+ # existing imports keep resolving; see the warning in the class docstring.
128
+ GLM_5_2 = "z-ai/glm-5.2"
129
+
130
+ # Stepfun-ai — WITHDRAWN upstream (410 Gone).
131
+ STEPFUN_STEP_3_7_FLASH = "stepfun-ai/step-3.7-flash"
132
+
133
+
134
+ #: Models NVIDIA publishes as **free preview endpoints**
135
+ #: (https://build.nvidia.com/models?filters=nimType%3Anim_type_preview).
136
+ #:
137
+ #: These are the models the ``free_tier`` throttle in
138
+ #: :class:`~parrot.clients.nvidia.NvidiaClient` is designed for: NVIDIA caps
139
+ #: free endpoints at 40 requests per minute, which is what the client's
140
+ #: :class:`~parrot.clients.nvidia.SlidingWindowRateLimiter` enforces. A model
141
+ #: outside this set is served from a paid or otherwise-provisioned endpoint,
142
+ #: where ``free_tier=False`` removes the cap.
143
+ #:
144
+ #: Membership here is NVIDIA's published classification, not something derived
145
+ #: from the API: a free endpoint answers ``200`` when it has capacity and
146
+ #: ``503 ResourceExhausted`` ("Worker local total request limit reached") when
147
+ #: it does not — so a 503 means *busy*, never *absent*, and cannot be used to
148
+ #: infer the list. Treat saturation as the normal operating condition of a free
149
+ #: endpoint and retry rather than falling back to another model.
150
+ FREE_TIER_MODELS: frozenset[str] = frozenset(
151
+ {
152
+ NvidiaModel.KIMI_K3.value,
153
+ NvidiaModel.NEMOTRON_3_5_LIGHTNING_30B.value,
154
+ NvidiaModel.DEEPSEEK_V4_FLASH_0731.value,
155
+ NvidiaModel.LAGUNA_XS_2_1.value,
156
+ NvidiaModel.NEMOTRON_3_NANO_OMNI_30B_REASONING.value,
157
+ NvidiaModel.GEMMA_4_31B_IT.value,
158
+ NvidiaModel.MISTRAL_NEMOTRON.value,
159
+ }
160
+ )