From e7b12b060055a81582ef3c8d6fac181e83bf76e4 Mon Sep 17 00:00:00 2001 From: Robin Date: Tue, 15 Sep 2026 10:39:43 -0700 Subject: [PATCH] nim --- README.md | 1 + TRACING.md | 1 + agentx/integrations/__init__.py | 1 + agentx/integrations/nvidia_nim.py | 70 ++++++++++++ agentx/integrations/openai.py | 6 +- setup.py | 2 + tests/test_integrations.py | 179 ++++++++++++++++++++++++++++++ 7 files changed, 259 insertions(+), 1 deletion(-) create mode 100644 agentx/integrations/nvidia_nim.py diff --git a/README.md b/README.md index ca93a73..699e1a5 100644 --- a/README.md +++ b/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` | diff --git a/TRACING.md b/TRACING.md index 457dfb3..36f7bee 100644 --- a/TRACING.md +++ b/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` | diff --git a/agentx/integrations/__init__.py b/agentx/integrations/__init__.py index f790c6f..2092b25 100644 --- a/agentx/integrations/__init__.py +++ b/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 diff --git a/agentx/integrations/nvidia_nim.py b/agentx/integrations/nvidia_nim.py new file mode 100644 index 0000000..1ac71a3 --- /dev/null +++ b/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) diff --git a/agentx/integrations/openai.py b/agentx/integrations/openai.py index 0d43fc2..74ad839 100644 --- a/agentx/integrations/openai.py +++ b/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, diff --git a/setup.py b/setup.py index 3f1931f..bcb52ed 100644 --- a/setup.py +++ b/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"], diff --git a/tests/test_integrations.py b/tests/test_integrations.py index 8005c7f..e3b1057 100644 --- a/tests/test_integrations.py +++ b/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 # ---------------------------------------------------------------------------