Skip to content

Commit 127a307

Browse files
PicoNVIDIAPicoNVIDIA
committed
Add AIDP OpenAI Demo example
This example demonstrates the NVIDIA AI Data Platform (AIDP) Retrieval API following the OpenAI Vector Store Search specification, exposed via MCP. Features: - OpenAI-compatible search_vector_store tool via MCP - NVIDIA NIMs for embeddings and LLM reasoning - Milvus vector database integration - Synthetic data generation using Nemotron LLM - Full ReAct agent workflow with NeMo Agent Toolkit Signed-off-by: Patrick Moorhead <pmoorhead@nvidia.com>
1 parent 7406837 commit 127a307

6 files changed

Lines changed: 990 additions & 0 deletions

File tree

‎examples/aidp_openai_demo/README.md‎

Lines changed: 323 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,323 @@
1+
<!--
2+
SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3+
SPDX-License-Identifier: Apache-2.0
4+
-->
5+
6+
# AIDP Retrieval API Demo with NVIDIA NIMs
7+
8+
## Overview
9+
10+
### The Problem
11+
12+
Enterprise AI applications need a standardized way to access enterprise data stored in storage platforms. Without a unified interface:
13+
14+
- **Custom integrations** - Each AI agent requires custom code to access storage
15+
- **Inconsistent APIs** - Different storage vendors expose different interfaces
16+
- **Tool fragmentation** - Tools built for one platform don't work with others
17+
- **Security complexity** - Each integration needs its own authentication handling
18+
19+
### The Solution
20+
21+
This demo implements the **NVIDIA AI Data Platform (AIDP) Retrieval API** following the [OpenAI Vector Store Search specification](https://platform.openai.com/docs/api-reference/vector_stores/search). By exposing the Retrieval API via **Model Context Protocol (MCP)**, any MCP-compatible AI agent can seamlessly search enterprise data with a standardized interface.
22+
23+
### How It Works
24+
25+
The demo implements an **Agentic RAG (Retrieval-Augmented Generation)** system for searching support tickets:
26+
27+
1. **User asks a question** via the chat UI or CLI (for example, "Find GPU memory issues")
28+
2. **ReAct Agent reasons** about which tools to use
29+
3. **MCP Tool executes** - `search_vector_store` performs semantic search
30+
4. **NVIDIA NIMs process** the request using GPU-accelerated embeddings
31+
5. **Agent synthesizes** the results into a coherent response
32+
33+
### Component Selection
34+
35+
| Component | Technology | Why This Choice |
36+
|-----------|------------|-----------------|
37+
| **Protocol** | MCP (`Streamable HTTP`) | Open standard with auth support, works with any MCP client |
38+
| **Agent Framework** | NeMo Agent Toolkit | Native MCP server/client, YAML config, production-ready |
39+
| **Vector Database** | Milvus | GPU-accelerated, scales to billions of vectors |
40+
| **Embeddings** | `nvidia/nv-embedqa-e5-v5` | High-quality 1024-dim embeddings optimized for Q&A retrieval |
41+
| **LLM** | `meta/llama-3.1-70b-instruct` | Strong reasoning for agent orchestration and response generation |
42+
| **API Spec** | OpenAI Vector Store Search | Industry standard for AI platform APIs |
43+
44+
---
45+
46+
## Table of Contents
47+
48+
- [Overview](#overview)
49+
- [Key Features](#key-features)
50+
- [Architecture](#architecture)
51+
- [Prerequisites](#prerequisites)
52+
- [Installation and Setup](#installation-and-setup)
53+
- [Running the Demo](#running-the-demo)
54+
- [NVIDIA NIMs Used](#nvidia-nims-used)
55+
- [The Tool](#the-tool)
56+
- [Sample Queries](#sample-queries)
57+
- [OpenAI API Alignment](#openai-api-alignment)
58+
- [Customization Guide](#customization-guide)
59+
60+
---
61+
62+
## Key Features
63+
64+
- **OpenAI-Compatible API**: Implements the OpenAI Vector Store Search specification
65+
- **MCP Protocol**: Tools exposed via standardized Model Context Protocol for interoperability
66+
- **NVIDIA NIMs Integration**: Uses NVIDIA NIMs for embedding and LLM reasoning
67+
- **Agentic RAG**: ReAct agent orchestrating search operations with tool calling
68+
- **Vector Search**: Semantic similarity search using Milvus vector database
69+
- **YAML-based Configuration**: Fully configurable workflow through YAML files
70+
71+
---
72+
73+
## Architecture
74+
75+
This demo uses a 3-terminal architecture:
76+
77+
1. **AIDP MCP Server** (`python src/nat_aidp_openai_demo/server.py`): Exposes `search_vector_store` via MCP
78+
2. **NAT UI Server** (`nat serve`): Acts as MCP client, provides API for the UI
79+
3. **NAT UI**: Frontend that users interact with
80+
81+
```
82+
┌─────────────┐ REST ┌─────────────────┐
83+
│ NAT UI │ ◄──────────────────► │ NAT UI Server │
84+
│ (Browser) │ Port 3000 │ (MCP Client) │
85+
└─────────────┘ └────────┬────────┘
86+
│ Port 8000
87+
MCP Protocol
88+
(Streamable-HTTP)
89+
│
90+
┌────────▼────────┐
91+
│ AIDP MCP Server│
92+
│ Port 8081 │
93+
│ search_vector_ │
94+
│ store │
95+
└────────┬────────┘
96+
│
97+
┌────────────────┼────────────────┐
98+
│ │ │
99+
┌───────▼──────┐ ┌──────▼──────┐ ┌──────▼──────┐
100+
│ Embedding NIM│ │ LLM NIM │ │ Milvus │
101+
│ (API.NVIDIA)│ │ (API.NVIDIA)│ │ Port 19530 │
102+
└──────────────┘ └─────────────┘ └─────────────┘
103+
```
104+
105+
---
106+
107+
## Prerequisites
108+
109+
- Docker (for Milvus vector database)
110+
- Python 3.11+
111+
- NVIDIA API key from [build.nvidia.com](https://build.nvidia.com)
112+
- Node.js (for UI)
113+
114+
---
115+
116+
## Installation and Setup
117+
118+
### Set Up API Keys
119+
120+
```bash
121+
export NVIDIA_API_KEY=<YOUR_API_KEY>
122+
```
123+
124+
### Start Milvus Vector Database
125+
126+
```bash
127+
# Download the Milvus standalone docker-compose file
128+
curl -sfL https://github.com/milvus-io/milvus/releases/download/v2.4.0/milvus-standalone-docker-compose.yml -o docker-compose.yml
129+
130+
# Start Milvus
131+
docker compose up -d
132+
```
133+
134+
### Load Sample Data
135+
136+
```bash
137+
python scripts/load_support_tickets.py
138+
```
139+
140+
Expected output:
141+
```
142+
Creating collection: support_tickets with explicit schema
143+
Collection 'support_tickets' created successfully
144+
Inserted 10 tickets with NIM embeddings
145+
Test search for 'GPU memory' returned 3 results
146+
```
147+
148+
---
149+
150+
## Running the Demo
151+
152+
### Terminal 1: Start AIDP MCP Server
153+
154+
```bash
155+
export NVIDIA_API_KEY=<YOUR_API_KEY>
156+
python src/nat_aidp_openai_demo/server.py
157+
```
158+
159+
### Terminal 2: Start NAT UI Server
160+
161+
```bash
162+
export NVIDIA_API_KEY=<YOUR_API_KEY>
163+
nat serve --config_file src/nat_aidp_openai_demo/configs/workflow.yml --port 8000
164+
```
165+
166+
### Terminal 3: Start UI
167+
168+
```bash
169+
cd external/nat-ui
170+
npm run dev
171+
```
172+
173+
### Open Browser
174+
175+
Navigate to: http://localhost:3000
176+
177+
**Alternative: Command Line**
178+
179+
```bash
180+
nat run --config_file src/nat_aidp_openai_demo/configs/workflow.yml --input "Find GPU memory issues"
181+
```
182+
183+
---
184+
185+
## NVIDIA NIMs Used
186+
187+
| NIM | Purpose | Model |
188+
|-----|---------|-------|
189+
| **Embedding** | Generate vector embeddings for semantic search | `nvidia/nv-embedqa-e5-v5` |
190+
| **LLM** | Agent reasoning and response generation | `meta/llama-3.1-70b-instruct` |
191+
192+
---
193+
194+
## The Tool
195+
196+
### `search_vector_store`
197+
198+
Semantic search following the AIDP Retrieval API (OpenAI specification).
199+
200+
| Parameter | Type | Description | Default |
201+
|-----------|------|-------------|---------|
202+
| `query` | string | Search query (required) | - |
203+
| `vector_store_id` | string | Vector store name | `support_tickets` |
204+
| `max_num_results` | integer | Results limit (1-50) | `10` |
205+
| `filter_key` | string | Attribute to filter by | `null` |
206+
| `filter_type` | string | Filter type: `eq`, `ne`, `contains` | `null` |
207+
| `filter_value` | string | Value to match | `null` |
208+
| `score_threshold` | float | Minimum similarity score | `null` |
209+
210+
---
211+
212+
## Sample Queries
213+
214+
Try these queries in the UI:
215+
216+
- "Find GPU memory issues"
217+
- "Show me critical severity tickets"
218+
- "What CUDA errors have been reported?"
219+
- "Find driver crash issues"
220+
- "Show resolved tickets about performance"
221+
222+
---
223+
224+
## OpenAI API Alignment
225+
226+
The AIDP Retrieval API follows the [OpenAI Vector Store Search specification](https://platform.openai.com/docs/api-reference/vector_stores/search):
227+
228+
### Endpoint
229+
230+
```
231+
POST /v1/vector_stores/{vector_store_id}/search
232+
```
233+
234+
### Response Format
235+
236+
```json
237+
{
238+
"object": "vector_store.search_results.page",
239+
"search_query": "GPU memory issues",
240+
"data": [
241+
{
242+
"file_id": "d1649d77-e043-45e5-b426-9b8b7c2856f2",
243+
"filename": "CUDA out of memory error.txt",
244+
"score": 0.4976,
245+
"attributes": {
246+
"category": "Memory Problems",
247+
"severity": "high",
248+
"title": "CUDA out of memory error with large batch sizes"
249+
},
250+
"content": [
251+
{
252+
"type": "text",
253+
"text": "Training transformer model with batch size 64..."
254+
}
255+
]
256+
}
257+
],
258+
"has_more": false,
259+
"next_page": null
260+
}
261+
```
262+
263+
### Alignment Table
264+
265+
| OpenAI Spec | AIDP Implementation |
266+
|-------------|---------------------|
267+
| `query` (required) | ✅ Implemented |
268+
| `filters` (key/type/value) | ✅ Implemented |
269+
| `max_num_results` (1-50) | ✅ Implemented |
270+
| `ranking_options` | ✅ Implemented |
271+
| Response: `file_id`, `filename`, `score` | ✅ Identical |
272+
| Response: `attributes`, `content[]` | ✅ Identical |
273+
| Bearer token authentication | ✅ Implemented |
274+
275+
---
276+
277+
## Customization Guide
278+
279+
### Adding New Fields
280+
281+
1. Update the Milvus schema in `scripts/load_support_tickets.py`
282+
2. Add the field to `output_fields` in `src/nat_aidp_openai_demo/server.py`
283+
3. Include the field in the response `attributes` object
284+
285+
### Using Different Models
286+
287+
Update `src/nat_aidp_openai_demo/configs/workflow.yml`:
288+
289+
```yaml
290+
llms:
291+
nim_llm:
292+
_type: nim
293+
model_name: meta/llama-3.3-70b-instruct # Change model here
294+
temperature: 0
295+
max_tokens: 512
296+
```
297+
298+
### Connecting to Different Vector Stores
299+
300+
Set the environment variable:
301+
302+
```bash
303+
export MILVUS_URI="http://your-milvus-host:19530"
304+
```
305+
306+
---
307+
308+
## Files
309+
310+
| File | Purpose |
311+
|------|---------|
312+
| `src/nat_aidp_openai_demo/server.py` | MCP server exposing `search_vector_store` tool |
313+
| `src/nat_aidp_openai_demo/configs/workflow.yml` | NeMo Agent Toolkit workflow configuration |
314+
| `scripts/load_support_tickets.py` | Data loading script for Milvus |
315+
316+
---
317+
318+
## References
319+
320+
- [OpenAI Vector Store Search API](https://platform.openai.com/docs/api-reference/vector_stores/search)
321+
- [Model Context Protocol (MCP)](https://modelcontextprotocol.io/)
322+
- [NeMo Agent Toolkit](https://github.com/NVIDIA/NeMo-Agent-Toolkit)
323+

‎examples/aidp_openai_demo/pyproject.toml‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
[project]
5+
name = "nat_aidp_openai_demo"
6+
version = "0.1.0"
7+
description = "AIDP Retrieval API Demo with NVIDIA NIMs - OpenAI Vector Store Search specification via MCP"
8+
readme = "README.md"
9+
requires-python = ">=3.11"
10+
11+
dependencies = [
12+
"nvidia-nat[langchain,mcp]>=1.4.0a0,<1.5.0",
13+
"pymilvus~=2.6",
14+
"fastmcp",
15+
"requests",
16+
]
17+
18+
[project.optional-dependencies]
19+
dev = [
20+
"pytest",
21+
"ruff",
22+
]
23+

0 commit comments

Comments
 (0)