nim

nim

NVIDIA NIM inference engine — self-hosted inference microservices.

Classes

NIMEngine

NIMEngine(host: str | None = None, *, timeout: float = 600.0)

Bases: InferenceEngine

NVIDIA NIM inference engine.

Supports local and remote NIM deployments with OpenAI-compatible API. Requires NIM_API_KEY for NVIDIA-hosted NIM endpoints, optional for self-hosted.

Source code in src/openjarvis/engine/nim.py

def __init__(self, host: str | None = None, *, timeout: float = 600.0) -> None:
    env_host = os.environ.get("NIM_HOST")
    self._host = (host or env_host or self._default_host).rstrip("/")

    api_key = os.environ.get("NIM_API_KEY", "")
    self._api_key = api_key if api_key else None

    headers = {}
    if self._api_key:
        headers["Authorization"] = f"Bearer {self._api_key}"

    self._timeout = timeout
    self._client = httpx.Client(
        base_url=self._host,
        timeout=timeout,
        headers=headers,
    )

Functions

stream_full async

stream_full(messages: Sequence[Message], *, model: str, temperature: float = 0.7, max_tokens: int = 1024, **kwargs: Any) -> AsyncIterator['StreamChunk']

Yield StreamChunks with content, tool_calls, and finish_reason.

Source code in src/openjarvis/engine/nim.py

async def stream_full(
    self,
    messages: Sequence[Message],
    *,
    model: str,
    temperature: float = 0.7,
    max_tokens: int = 1024,
    **kwargs: Any,
) -> AsyncIterator["StreamChunk"]:
    """Yield StreamChunks with content, tool_calls, and finish_reason."""
    msg_dicts = messages_to_dicts(messages)
    payload: Dict[str, Any] = {
        "model": model,
        "messages": msg_dicts,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "stream": True,
        **kwargs,
    }
    if "tools" in payload and "tool_choice" not in payload:
        payload["tool_choice"] = "auto"
    try:
        url = f"{self._api_prefix}/chat/completions"
        async with httpx.AsyncClient(
            base_url=self._host,
            timeout=self._timeout,
            headers=self._get_headers(),
        ) as client:
            async with client.stream("POST", url, json=payload) as resp:
                resp.raise_for_status()
                async for line in resp.aiter_lines():
                    if not line.startswith("data:"):
                        continue
                    data_str = line[len("data:") :].strip()
                    if data_str == "[DONE]":
                        break
                    try:
                        chunk = json.loads(data_str)
                    except json.JSONDecodeError:
                        continue
                    choices = chunk.get("choices", [])
                    if not choices:
                        continue
                    choice = choices[0]
                    delta = choice.get("delta", {})
                    finish = choice.get("finish_reason")
                    content = delta.get("content")
                    tool_calls = delta.get("tool_calls")
                    usage = chunk.get("usage")

                    if content or tool_calls or finish or usage:
                        yield StreamChunk(
                            content=content,
                            tool_calls=tool_calls,
                            finish_reason=finish,
                            usage=usage,
                        )
    except (httpx.ConnectError, httpx.TimeoutException) as exc:
        raise EngineConnectionError(
            f"NIM engine not reachable at {self._host}"
        ) from exc

Functions