Skip to content

Commit 8fb26be

Browse files
PicoNVIDIAPicoNVIDIA
authored
Add MCP MilvusDB with NVIDIA NIMs (#7)
- Implements MCP protocol with NAT as client/server - Integrates NVIDIA NIMs as tool calls: Uses: embedding (nv-embedqa-e5-v5), LLM (llama-3.1-70b-instruct), and reranking (llama-3.2-nv-rerankqa-1b-v2) (within tool calls) - Demonstrates semantic search, filtering, and reranking - Includes 4 custom tools registered via @register_function Please add labels: category:example MCP Tools: <img width="1070" height="863" alt="image" src="https://github.com/user-attachments/assets/a5f4371b-165f-4df0-b67d-279bd785a505" /> Sample Run: <img width="1316" height="798" alt="image" src="https://github.com/user-attachments/assets/559ae4c7-fab5-48fe-838a-e64b33d658d8" /> ## Summary by CodeRabbit * **New Features** * Added a new Support Tickets MCP Demo showcasing AI-powered support ticket management with semantic search, category and priority filtering, and intelligent reranking. * **Documentation** * Added comprehensive README with prerequisites, setup instructions, and step-by-step workflow to run the demo locally. Authors: - https://github.com/PicoNVIDIA - Anuradha Karuppiah (https://github.com/AnuradhaKaruppiah) Approvers: - Anuradha Karuppiah (https://github.com/AnuradhaKaruppiah) URL: #7
1 parent 92955cc commit 8fb26be

9 files changed

Lines changed: 984 additions & 0 deletions

File tree

‎ci/vale/styles/config/vocabularies/nemo-agent-toolkit-examples/accept.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,9 @@ PDF(s?)
103103
Pydantic
104104
PyPI
105105
pytest
106+
[Rr]erank
107+
[Rr]eranker(s?)
108+
[Rr]eranking
106109
[Rr]edis
107110
[Rr]einstall(s?)
108111
[Rr]eplatform(ing)?

‎examples/mcp_rag_demo/README.md‎

Lines changed: 233 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,233 @@
1+
<!--
2+
SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3+
SPDX-License-Identifier: Apache-2.0
4+
5+
Licensed under the Apache License, Version 2.0 (the "License");
6+
you may not use this file except in compliance with the License.
7+
You may obtain a copy of the License at
8+
9+
http://www.apache.org/licenses/LICENSE-2.0
10+
11+
Unless required by applicable law or agreed to in writing, software
12+
distributed under the License is distributed on an "AS IS" BASIS,
13+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
See the License for the specific language governing permissions and
15+
limitations under the License.
16+
-->
17+
18+
# MCP RAG Demo with NVIDIA NIMs
19+
20+
This example demonstrates how to expose custom tools via the Model Context Protocol (MCP) using NVIDIA NeMo Agent toolkit with NVIDIA NIM integration. It showcases semantic search, filtering, and reranking of support tickets using NVIDIA NIMs for embedding, LLM reasoning, and reranking.
21+
22+
## Table of Contents
23+
24+
- [MCP RAG Demo with NVIDIA NIMs](#mcp-rag-demo-with-nvidia-nims)
25+
- [Table of Contents](#table-of-contents)
26+
- [Key Features](#key-features)
27+
- [Architecture](#architecture)
28+
- [Prerequisites](#prerequisites)
29+
- [Installation and Setup](#installation-and-setup)
30+
- [Install this Workflow](#install-this-workflow)
31+
- [Set Up API Keys](#set-up-api-keys)
32+
- [Start Milvus Vector Database](#start-milvus-vector-database)
33+
- [Load Sample Data](#load-sample-data)
34+
- [Running the Demo](#running-the-demo)
35+
- [Terminal 1: Start NAT MCP Server](#terminal-1-start-nat-mcp-server)
36+
- [Terminal 2: Start NAT UI Server](#terminal-2-start-nat-ui-server)
37+
- [Terminal 3: Start UI](#terminal-3-start-ui)
38+
- [Open Browser](#open-browser)
39+
- [NVIDIA NIMs Used](#nvidia-nims-used)
40+
- [The 4 Tools](#the-4-tools)
41+
- [Sample Queries](#sample-queries)
42+
- [Customization Guide](#customization-guide)
43+
- [Adding New Functions](#adding-new-functions)
44+
- [Modifying the Agent](#modifying-the-agent)
45+
- [Using Different Models](#using-different-models)
46+
47+
---
48+
49+
## Key Features
50+
51+
- **MCP Protocol**: Tools exposed via standardized Model Context Protocol for interoperability
52+
- **NVIDIA NIMs Integration**: Uses NVIDIA NIMs for embedding, LLM reasoning, and reranking
53+
- **Custom Functions**: 4 tools registered via `@register_function` decorator
54+
- **Agentic RAG**: ReAct agent orchestrating search operations with tool calling
55+
- **Vector Search**: Semantic similarity search using Milvus vector database
56+
- **YAML-based Configuration**: Fully configurable workflow through YAML files
57+
58+
---
59+
60+
## Architecture
61+
62+
This demo uses a 3-terminal architecture:
63+
64+
1. **NAT MCP Server** (`nat mcp serve`): Exposes the 4 custom tools via MCP protocol
65+
2. **NAT UI Server** (`nat serve`): Acts as MCP client, provides API for the UI
66+
3. **NAT UI**: Frontend that users interact with
67+
68+
```
69+
┌─────────────┐ REST ┌─────────────────┐
70+
│ NAT UI │ ◄──────────────────► │ NAT UI Server │
71+
│ (Browser) │ │ (MCP Client) │
72+
└─────────────┘ └────────┬────────┘
73+
│
74+
MCP Protocol
75+
(Streamable-HTTP)
76+
│
77+
┌────────▼────────┐
78+
│ NAT MCP Server │
79+
│ (Tool Server) │
80+
└────────┬────────┘
81+
│
82+
┌────────────────┼────────────────┐
83+
│ │ │
84+
┌───────▼──────┐ ┌──────▼──────┐ ┌──────▼──────┐
85+
│ Embedding NIM│ │ LLM NIM │ │ Rerank NIM │
86+
└──────────────┘ └─────────────┘ └─────────────┘
87+
```
88+
89+
---
90+
91+
## Prerequisites
92+
93+
- Docker (for Milvus vector database)
94+
- Python 3.11+
95+
- NVIDIA API key
96+
- Node.js (for UI)
97+
98+
---
99+
100+
## Installation and Setup
101+
102+
### Install this Workflow
103+
104+
From the root directory of the NeMo Agent Toolkit Examples repository, run:
105+
106+
```bash
107+
uv pip install --prerelease=allow -e examples/mcp_rag_demo
108+
```
109+
110+
### Set Up API Keys
111+
112+
```bash
113+
export NVIDIA_API_KEY=<YOUR_API_KEY>
114+
```
115+
116+
### Start Milvus Vector Database
117+
118+
```bash
119+
# If you don't have Milvus set up, create a docker-compose.yml:
120+
mkdir -p ~/milvus_setup && cd ~/milvus_setup
121+
122+
# Download the Milvus standalone docker-compose file
123+
wget https://github.com/milvus-io/milvus/releases/download/v2.4.0/milvus-standalone-docker-compose.yml -O docker-compose.yml
124+
125+
# Start Milvus
126+
docker-compose up -d
127+
```
128+
129+
### Load Sample Data
130+
131+
```bash
132+
python examples/mcp_rag_demo/scripts/load_support_tickets.py
133+
```
134+
135+
Expected output:
136+
```
137+
✓ Connected to Milvus
138+
✓ Created support_tickets collection
139+
✓ Prepared 15 support tickets
140+
✓ Calling NVIDIA NIM API to generate embeddings...
141+
✓ Generated 15 embeddings using NVIDIA NIM (`nvidia/nv-embedqa-e5-v5`)
142+
✓ Inserted 15 tickets into Milvus
143+
✓ Created vector index
144+
✓ Loaded collection into memory
145+
============================================================
146+
Successfully loaded 15 support tickets into Milvus
147+
Collection: support_tickets
148+
============================================================
149+
```
150+
151+
---
152+
153+
## Running the Demo
154+
155+
### Terminal 1: Start NAT MCP Server
156+
157+
```bash
158+
export NVIDIA_API_KEY=<YOUR_API_KEY>
159+
nat mcp serve --config_file examples/mcp_rag_demo/configs/support-ui.yml --port 9904
160+
```
161+
162+
### Terminal 2: Start NAT UI Server
163+
164+
```bash
165+
export NVIDIA_API_KEY=<YOUR_API_KEY>
166+
nat serve --config_file examples/mcp_rag_demo/configs/mcp-client-for-ui.yml --port 8000
167+
```
168+
169+
### Terminal 3: Start UI
170+
171+
```bash
172+
cd external/nat-ui
173+
export NAT_BACKEND_URL="http://localhost:8000"
174+
npm run dev
175+
```
176+
177+
### Open Browser
178+
179+
Navigate to: http://localhost:3000
180+
181+
---
182+
183+
## NVIDIA NIMs Used
184+
185+
| NIM | Purpose | Model |
186+
|-----|---------|-------|
187+
| **Embedding** | Generate vector embeddings for semantic search | `nvidia/nv-embedqa-e5-v5` |
188+
| **LLM** | Agent reasoning and response generation | `meta/llama-3.1-70b-instruct` |
189+
| **Reranking** | Improve search relevance | `nvidia/llama-3.2-nv-rerankqa-1b-v2` |
190+
191+
---
192+
193+
## The 4 Tools
194+
195+
1. `search_support_tickets`: Semantic search using NIM embeddings
196+
2. `query_by_category`: Filter tickets by category e.g. `bug_report`, `feature_request`, `question`, `incident`
197+
3. `query_by_priority`: Filter tickets by priority (critical, high, medium, low)
198+
4. `rerank_support_tickets`: Rerank results using NIM reranker for improved relevance
199+
200+
---
201+
202+
## Sample Queries
203+
204+
Try these queries in the UI:
205+
206+
- "Find tickets about GPU driver crashes"
207+
- "Show me critical incidents"
208+
- "What bugs are related to CUDA?"
209+
- "Find feature requests for the API"
210+
- "Show me resolved Milvus performance issues"
211+
212+
---
213+
214+
## Customization Guide
215+
216+
### Adding New Functions
217+
218+
1. Define a new config class in `register.py` that inherits from `FunctionBaseConfig`
219+
2. Create a function decorated with `@register_function`
220+
3. Add the function to your `support-ui.yml` config file
221+
4. Update the workflow's `tool_names` list
222+
223+
### Modifying the Agent
224+
225+
- Change the `_type` in the workflow section to use different agent types
226+
- Adjust `max_iterations`, `temperature`, or other parameters as needed
227+
- Add custom system prompts or modify the agent behavior
228+
229+
### Using Different Models
230+
231+
- Update the `model_name` in the LLM configuration
232+
- Adjust parameters such as `temperature` and `max_tokens`
233+
- Switch between different LLM providers (OpenAI, NIM, and so on)

‎examples/mcp_rag_demo/configs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
src/nat_mcp_rag_demo/configs

‎examples/mcp_rag_demo/pyproject.toml‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
[build-system]
17+
build-backend = "setuptools.build_meta"
18+
requires = ["setuptools >= 64", "setuptools-scm>=8"]
19+
20+
[tool.setuptools_scm]
21+
git_describe_command = "git describe --long --first-parent"
22+
root = "../.."
23+
24+
[project]
25+
name = "nat_mcp_rag_demo"
26+
dynamic = ["version"]
27+
dependencies = [
28+
"nvidia-nat[langchain,mcp]>=1.4.0a0,<1.5.0",
29+
"pymilvus~=2.6",
30+
"langchain-nvidia-ai-endpoints~=0.3",
31+
]
32+
requires-python = ">=3.11,<3.14"
33+
description = "MCP RAG demo with NVIDIA NIMs for support ticket search"
34+
keywords = ["ai", "rag", "mcp", "nvidia", "nim"]
35+
classifiers = ["Programming Language :: Python"]
36+
37+
[project.entry-points.'nat.components']
38+
nat_mcp_rag_demo = "nat_mcp_rag_demo.register"
39+
40+
[tool.setuptools.packages.find]
41+
where = ["src"]

0 commit comments

Comments
 (0)