nim

README.md

@@ -185,6 +185,7 @@ extra:
 | CrewAI                | `pip install "agentx-python[crewai]"`        | `AgentXCrewObserver`      |
 | OpenAI Agents SDK     | `pip install "agentx-python[openai-agents]"` | `AgentXTracingProcessor`  |
 | OpenAI (raw client)   | `pip install "agentx-python[openai]"`        | `patch_openai_client`     |
+| NVIDIA NIM            | `pip install "agentx-python[nvidia-nim]"`    | `patch_nim_client`        |
 | Anthropic             | `pip install "agentx-python[anthropic]"`     | `patch_anthropic_client`  |
 | Google ADK            | `pip install "agentx-python[google-adk]"`    | `AgentXADKPlugin`         |
 | Google GenAI (Gemini) | `pip install "agentx-python[google-genai]"`  | `patch_genai_client`      |

TRACING.md

@@ -192,6 +192,7 @@ works. The label resolves in priority order:
    | `AgentXCrewObserver` | `crewai` |
    | `AgentXTracingProcessor` (OpenAI Agents SDK) | `openai-agents` |
    | `patch_openai_client` | `openai` |
+   | `patch_nim_client` (NVIDIA NIM) | `nvidia-nim` |
    | `patch_anthropic_client` | `anthropic` |
    | `patch_genai_client` | `google-genai` |
    | `AgentXADKPlugin` | `google-adk` |

agentx/integrations/__init__.py

@@ -7,4 +7,5 @@
 #   from agentx.integrations.anthropic import patch_anthropic_client
 #   from agentx.integrations.google_adk import AgentXADKPlugin
 #   from agentx.integrations.google_genai import patch_genai_client
+#   from agentx.integrations.nvidia_nim import patch_nim_client
 #   from agentx.integrations.moveworks import MoveworksImporter  # Data API pull sync, not in-process

agentx/integrations/nvidia_nim.py

@@ -0,0 +1,70 @@
+"""
+NVIDIA NIM integration for AgentX production tracing.
+
+NIM (NVIDIA Inference Microservices) serves models behind an OpenAI-compatible
+``/v1/chat/completions`` API, so the client you patch is the ordinary ``openai``
+Python client pointed at a NIM endpoint - a local NIM container
+(``http://localhost:8000/v1``) or NVIDIA's hosted API
+(``https://integrate.api.nvidia.com/v1``). This module reuses the OpenAI patch
+machinery verbatim and differs in exactly one way: traces are stamped
+``framework="nvidia-nim"``, so NIM traffic gets its own row in Monitor's
+Platforms chart and the framework filters instead of blending into "openai".
+
+Usage::
+
+    from agentx.integrations.nvidia_nim import patch_nim_client
+    import openai
+
+    nim = openai.OpenAI(
+        base_url="http://localhost:8000/v1",  # or https://integrate.api.nvidia.com/v1
+        api_key=os.environ.get("NVIDIA_API_KEY", "not-needed-for-local-nim"),
+    )
+    patch_nim_client(nim, agentx.tracer, name="nim-agent")
+
+    # All subsequent nim.chat.completions.create() calls are now traced.
+
+Works with both ``openai.OpenAI`` and ``openai.AsyncOpenAI`` clients. Token
+usage comes straight off the response's OpenAI-shaped ``usage`` block; NIM
+reports no prompt-cache fields, so cache token counts stay unset.
+
+Streaming calls (``stream=True``) are passed through untouched and are not
+currently traced - same posture as ``patch_openai_client``, see its docstring.
+
+Requires: ``pip install "agentx-python[nvidia-nim]"`` (installs the ``openai``
+client package; there is no separate NIM SDK dependency).
+"""
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+from agentx.tracing.tracer import Tracer
+from agentx.integrations.openai import _patch_chat_completions_create
+
+NIM_FRAMEWORK = "nvidia-nim"
+
+
+def patch_nim_client(
+    client: Any,
+    tracer: Tracer,
+    name: str = "nim-agent",
+    metadata: Optional[Dict[str, Any]] = None,
+    session_id: Optional[str] = None,
+) -> None:
+    """
+    Monkey-patch ``client.chat.completions.create`` on an OpenAI-compatible
+    client pointed at a NIM endpoint, sending a trace for every non-streaming
+    call with ``framework="nvidia-nim"``.
+
+    The original method is still called and its return value passed through
+    unchanged. Sync and async clients both work; ``stream=True`` calls pass
+    through untraced. Patching is idempotent - and because it shares the guard
+    with ``patch_openai_client``, whichever of the two patched a given client
+    first wins (patch each client with the integration that matches where its
+    ``base_url`` actually points).
+    """
+    chat = getattr(client, "chat", None)
+    completions = getattr(chat, "completions", None) if chat is not None else None
+    if completions is None:
+        raise ValueError("Provided client does not have a .chat.completions attribute")
+
+    _patch_chat_completions_create(completions, tracer, name, metadata, session_id, framework=NIM_FRAMEWORK)

agentx/integrations/openai.py

@@ -113,7 +113,11 @@ def _patch_chat_completions_create(
     name: str,
     metadata: Optional[Dict[str, Any]],
     session_id: Optional[str],
+    framework: str = "openai",
 ) -> None:
+    # `framework` exists for OpenAI-compatible endpoints served by other vendors
+    # (agentx.integrations.nvidia_nim stamps "nvidia-nim" through here) - the request/response
+    # shapes are identical, so they share this machinery instead of duplicating it.
     original = completions_resource.create
     if getattr(original, "_agentx_patched", False):
         return  # already patched
@@ -148,7 +152,7 @@ def on_finish(response: Optional[Any], error: Optional[str]) -> None:
             finish_llm_call(
                 tracer,
                 name=name,
-                framework="openai",
+                framework=framework,
                 metadata=metadata,
                 session_id=session_id,
                 start_t=start_t,

setup.py

@@ -51,6 +51,8 @@ def get_long_description():
         "crewai": ["crewai>=0.80.0"],
         "openai-agents": ["openai-agents>=0.0.3"],
         "openai": ["openai>=1.0.0"],
+        # NIM endpoints speak the OpenAI-compatible API; the client package IS openai.
+        "nvidia-nim": ["openai>=1.0.0"],
         "anthropic": ["anthropic>=0.25.0"],
         "google-adk": ["google-adk>=1.0.0"],
         "google-genai": ["google-genai>=1.0.0"],

tests/test_integrations.py

@@ -513,6 +513,185 @@ def create(self, **kwargs):
     assert kwargs["error"] == "rate limited"
 
 
+def test_openai_patch_stamps_openai_framework():
+    # Regression guard for the shared-machinery refactor: _patch_chat_completions_create
+    # grew a framework parameter for nvidia_nim.py, and the OpenAI default must stay "openai".
+    from agentx.integrations.openai import patch_openai_client
+
+    class FakeCompletions:
+        def create(self, **kwargs):
+            return _FakeOpenAIChatCompletion("hello")
+
+    client = _fake_openai_client(FakeCompletions())
+    tracer = make_tracer()
+    patch_openai_client(client, tracer, name="gpt-agent")
+    client.chat.completions.create(model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}])
+
+    _, kwargs = tracer._send.call_args
+    assert kwargs["framework"] == "openai"
+
+
+# ---------------------------------------------------------------------------
+# 6b. nvidia_nim.py — the OpenAI-compatible patch with the NIM framework label
+# ---------------------------------------------------------------------------
+
+def test_nim_sync_client_traces_call_with_nim_framework():
+    from agentx.integrations.nvidia_nim import patch_nim_client
+
+    response = _FakeOpenAIChatCompletion("hello from nim")
+
+    class FakeCompletions:
+        def create(self, **kwargs):
+            return response
+
+    client = _fake_openai_client(FakeCompletions())
+    tracer = make_tracer()
+    patch_nim_client(client, tracer, name="nim-agent")
+
+    result = client.chat.completions.create(
+        model="meta/llama-3.1-8b-instruct", messages=[{"role": "user", "content": "hi"}]
+    )
+
+    assert result is response
+    tracer._send.assert_called_once()
+    _, kwargs = tracer._send.call_args
+    assert kwargs["framework"] == "nvidia-nim"
+    assert kwargs["name"] == "nim-agent"
+    assert kwargs["output"] == "hello from nim"
+    assert kwargs["model"] == "meta/llama-3.1-8b-instruct"
+    assert kwargs["input_tokens"] == 12
+    assert kwargs["output_tokens"] == 6
+    # NIM reports no prompt-cache fields; the counts must stay unset, not become 0.
+    assert not kwargs.get("cache_read_tokens")
+
+
+def test_nim_async_client_traces_the_real_response():
+    from agentx.integrations.nvidia_nim import patch_nim_client
+
+    response = _FakeOpenAIChatCompletion("hello from async nim")
+
+    class FakeAsyncCompletions:
+        async def create(self, **kwargs):
+            await asyncio.sleep(0.01)
+            return response
+
+    client = _fake_openai_client(FakeAsyncCompletions())
+    tracer = make_tracer()
+    patch_nim_client(client, tracer, name="nim-agent")
+
+    result = asyncio.run(
+        client.chat.completions.create(model="meta/llama-3.1-8b-instruct", messages=[{"role": "user", "content": "hi"}])
+    )
+
+    assert result is response
+    tracer._send.assert_called_once()
+    _, kwargs = tracer._send.call_args
+    assert kwargs["framework"] == "nvidia-nim"
+    assert kwargs["output"] == "hello from async nim"
+    assert kwargs["input_tokens"] == 12
+
+
+def test_nim_streaming_calls_are_passed_through_untraced():
+    from agentx.integrations.nvidia_nim import patch_nim_client
+
+    sentinel_stream = object()
+
+    class FakeCompletions:
+        def create(self, **kwargs):
+            assert kwargs.get("stream") is True
+            return sentinel_stream
+
+    client = _fake_openai_client(FakeCompletions())
+    tracer = make_tracer()
+    patch_nim_client(client, tracer, name="nim-agent")
+
+    result = client.chat.completions.create(
+        model="meta/llama-3.1-8b-instruct", messages=[{"role": "user", "content": "hi"}], stream=True
+    )
+
+    assert result is sentinel_stream
+    tracer._send.assert_not_called()
+
+
+def test_nim_sync_client_records_errors():
+    from agentx.integrations.nvidia_nim import patch_nim_client
+
+    class FakeCompletions:
+        def create(self, **kwargs):
+            raise ValueError("nim endpoint unavailable")
+
+    client = _fake_openai_client(FakeCompletions())
+    tracer = make_tracer()
+    patch_nim_client(client, tracer, name="nim-agent")
+
+    with pytest.raises(ValueError, match="nim endpoint unavailable"):
+        client.chat.completions.create(model="meta/llama-3.1-8b-instruct", messages=[{"role": "user", "content": "hi"}])
+
+    tracer._send.assert_called_once()
+    _, kwargs = tracer._send.call_args
+    assert kwargs["error"] == "nim endpoint unavailable"
+    assert kwargs["framework"] == "nvidia-nim"
+
+
+def test_nim_patch_is_idempotent_and_first_patch_wins():
+    from agentx.integrations.nvidia_nim import patch_nim_client
+    from agentx.integrations.openai import patch_openai_client
+
+    class FakeCompletions:
+        def create(self, **kwargs):
+            return _FakeOpenAIChatCompletion("once")
+
+    client = _fake_openai_client(FakeCompletions())
+    tracer = make_tracer()
+    patch_nim_client(client, tracer, name="nim-agent")
+    # Double NIM patch and a later OpenAI patch are both no-ops (shared _agentx_patched guard):
+    # exactly one trace per call, and the first patch's framework label stays.
+    patch_nim_client(client, tracer, name="nim-agent")
+    patch_openai_client(client, tracer, name="gpt-agent")
+
+    client.chat.completions.create(model="meta/llama-3.1-8b-instruct", messages=[{"role": "user", "content": "hi"}])
+
+    tracer._send.assert_called_once()
+    _, kwargs = tracer._send.call_args
+    assert kwargs["framework"] == "nvidia-nim"
+
+
+def test_nim_rejects_client_without_chat_completions():
+    from agentx.integrations.nvidia_nim import patch_nim_client
+
+    with pytest.raises(ValueError, match="chat.completions"):
+        patch_nim_client(object(), make_tracer())
+
+
+def test_nim_request_tools_land_in_trace_metadata():
+    # The docs promise the request's tools=[...] definitions feed the unregistered-tool
+    # listing via metadata.tools - pin the shared capture path (openai.py machinery) here.
+    from agentx.integrations.nvidia_nim import patch_nim_client
+
+    class FakeCompletions:
+        def create(self, **kwargs):
+            return _FakeOpenAIChatCompletion("used a tool")
+
+    tool_def = {
+        "type": "function",
+        "function": {"name": "lookup_order", "parameters": {"type": "object", "properties": {}}},
+    }
+    client = _fake_openai_client(FakeCompletions())
+    tracer = make_tracer()
+    patch_nim_client(client, tracer, name="nim-agent", metadata={"env": "test"})
+
+    client.chat.completions.create(
+        model="meta/llama-3.1-8b-instruct",
+        messages=[{"role": "user", "content": "hi"}],
+        tools=[tool_def],
+    )
+
+    _, kwargs = tracer._send.call_args
+    assert kwargs["metadata"]["tools"] == [tool_def]
+    # The caller's own static metadata must survive the tools merge.
+    assert kwargs["metadata"]["env"] == "test"
+
+
 # ---------------------------------------------------------------------------
 # 7. langchain.py — nested-run state cleanup + TTL safety net
 # ---------------------------------------------------------------------------