From ac71ff603c56fa36148c49e297e55882b6631641 Mon Sep 17 00:00:00 2001 From: Soumili Nandi Date: Mon, 22 Sep 2025 15:53:54 -0700 Subject: [PATCH 1/3] feat: Add NVIDIA NIM model integration example via LiteLLM --- .../samples/hello_world_nvidia/README.md | 77 ++++++++++++++++ .../samples/hello_world_nvidia/__init__.py | 16 ++++ .../samples/hello_world_nvidia/agent.py | 90 +++++++++++++++++++ .../samples/hello_world_nvidia/main.py | 86 ++++++++++++++++++ 4 files changed, 269 insertions(+) create mode 100644 contributing/samples/hello_world_nvidia/README.md create mode 100644 contributing/samples/hello_world_nvidia/__init__.py create mode 100644 contributing/samples/hello_world_nvidia/agent.py create mode 100644 contributing/samples/hello_world_nvidia/main.py diff --git a/contributing/samples/hello_world_nvidia/README.md b/contributing/samples/hello_world_nvidia/README.md new file mode 100644 index 0000000000..e892000197 --- /dev/null +++ b/contributing/samples/hello_world_nvidia/README.md @@ -0,0 +1,77 @@ +# Using NVIDIA Models with ADK via LiteLLM + +This example demonstrates how to use NVIDIA NIM (NVIDIA Inference Microservices) models with ADK through LiteLLM integration. + +For comprehensive information about using NVIDIA models with LiteLLM, refer to the [official LiteLLM NVIDIA NIM documentation](https://github.com/BerriAI/litellm/blob/main/docs/my-website/docs/providers/nvidia_nim.md?plain=1). + +## Setup + +### 1. Get NVIDIA NIM API Key + +1. Sign up at [NVIDIA AI Foundation](https://www.nvidia.com/en-us/ai-data-science/cloud/) +2. Navigate to API section and generate an API key +3. Copy your API key + +### 2. Ensure LiteLLM is Installed + +Install LiteLLM if you haven't already: + +```bash +pip install litellm +``` + + +## Using NVIDIA Models in ADK + +### Environment Variables + +Set the required environment variables: + +```bash +export NVIDIA_NIM_API_KEY="your-nvidia-api-key" +export NVIDIA_NIM_API_BASE="https://integrate.api.nvidia.com/v1/" # Optional +``` + +### Code Examples + +#### Basic Agent Creation + +```python +from google.adk import Agent +from google.adk.models.lite_llm import LiteLlm + +# Create agent with NVIDIA NIM model +agent = Agent( + model=LiteLlm(model="nvidia_nim/meta/llama3-8b-instruct"), + name="nvidia_agent", + instruction="You are a helpful assistant.", + description="Agent using NVIDIA model", +) +``` + +#### Available Models + +All NVIDIA NIM models are supported. Use the `nvidia_nim/` prefix: + +```python +# Examples of available models +models = [ + "nvidia_nim/meta/llama3-8b-instruct", + "nvidia_nim/meta/llama3-70b-instruct", + "nvidia_nim/microsoft/phi-3-medium-4k-instruct", + "nvidia_nim/mistralai/mistral-7b-instruct", + "nvidia_nim/google/gemma-7b", + "nvidia_nim/nvidia/nemotron-4-340b-instruct", + # ... and many more +] +``` + +## Integration Details + +ADK uses LiteLLM as a wrapper to access NVIDIA NIM models. The integration: + +1. **Model Format**: Uses `nvidia_nim//` format +2. **Authentication**: Requires `NVIDIA_NIM_API_KEY` environment variable +3. **Base URL**: Optional `NVIDIA_NIM_API_BASE` for custom endpoints +4. **Compatibility**: Works with all ADK features (tools, sessions, etc.) +5. **Supported Endpoints**: `/chat/completions`, `/completions`, `/embeddings` \ No newline at end of file diff --git a/contributing/samples/hello_world_nvidia/__init__.py b/contributing/samples/hello_world_nvidia/__init__.py new file mode 100644 index 0000000000..b19b9fbce2 --- /dev/null +++ b/contributing/samples/hello_world_nvidia/__init__.py @@ -0,0 +1,16 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from . import agent \ No newline at end of file diff --git a/contributing/samples/hello_world_nvidia/agent.py b/contributing/samples/hello_world_nvidia/agent.py new file mode 100644 index 0000000000..b56b11d1d2 --- /dev/null +++ b/contributing/samples/hello_world_nvidia/agent.py @@ -0,0 +1,90 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import random + +from google.adk import Agent +from google.adk.models.lite_llm import LiteLlm + + +def roll_die(sides: int) -> int: + """Roll a die and return the rolled result. + + Args: + sides: The integer number of sides the die has. + + Returns: + An integer of the result of rolling the die. + """ + return random.randint(1, sides) + + +def check_prime(nums: list[int]) -> str: + """Check if a given list of numbers are prime. + + Args: + nums: The list of numbers to check. + + Returns: + A str indicating which number is prime. + """ + primes = set() + for number in nums: + if number <= 1: + continue + is_prime = True + for i in range(2, int(number**0.5) + 1): + if number % i == 0: + is_prime = False + break + if is_prime: + primes.add(number) + return ( + "No prime numbers found." + if not primes + else f"{', '.join(str(num) for num in primes)} are prime numbers." + ) + + +root_agent = Agent( + model=LiteLlm(model="nvidia_nim/meta/llama3-8b-instruct"), + name="hello_world_nvidia_nim_agent", + description=( + "hello world agent powered by NVIDIA NIM that can roll a dice of 8 sides and check prime" + " numbers." + ), + instruction=""" + You are a helpful assistant powered by NVIDIA AI models through NVIDIA NIM. + You roll dice and answer questions about the outcome of the dice rolls. + You can roll dice of different sizes. + You can use multiple tools in parallel by calling functions in parallel(in one request and in one round). + It is ok to discuss previous dice rolls, and comment on the dice rolls. + When you are asked to roll a die, you must call the roll_die tool with the number of sides. Be sure to pass in an integer. Do not pass in a string. + You should never roll a die on your own. + When checking prime numbers, call the check_prime tool with a list of integers. Be sure to pass in a list of integers. You should never pass in a string. + You should not check prime numbers before calling the tool. + When you are asked to roll a die and check prime numbers, you should always make the following two function calls: + 1. You should first call the roll_die tool to get a roll. Wait for the function response before calling the check_prime tool. + 2. After you get the function response from roll_die tool, you should call the check_prime tool with the roll_die result. + 2.1 If user asks you to check primes based on previous rolls, make sure you include the previous rolls in the list. + 3. When you respond, you must include the roll_die result from step 1. + You should always perform the previous 3 steps when asking for a roll and checking prime numbers. + You should not rely on the previous history on prime results. + """, + tools=[ + roll_die, + check_prime, + ], +) \ No newline at end of file diff --git a/contributing/samples/hello_world_nvidia/main.py b/contributing/samples/hello_world_nvidia/main.py new file mode 100644 index 0000000000..63a905ab54 --- /dev/null +++ b/contributing/samples/hello_world_nvidia/main.py @@ -0,0 +1,86 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import asyncio +import os +import time + +import agent +from dotenv import load_dotenv +from google.adk import Runner +from google.adk.artifacts import InMemoryArtifactService +from google.adk.cli.utils import logs +from google.adk.sessions import InMemorySessionService +from google.adk.sessions import Session +from google.genai import types + +load_dotenv(override=True) +logs.log_to_tmp_folder() + + +async def main(): + # Check for NVIDIA NIM API key + if not os.environ.get("NVIDIA_NIM_API_KEY"): + print(" NVIDIA_NIM_API_KEY not set. Please set it before running:") + print(" export NVIDIA_NIM_API_KEY='your-nvidia-api-key'") + return + + app_name = 'nvidia_nim_app' + user_id_1 = 'user1' + session_service = InMemorySessionService() + artifact_service = InMemoryArtifactService() + runner = Runner( + app_name=app_name, + agent=agent.root_agent, + artifact_service=artifact_service, + session_service=session_service, + ) + session_11 = await session_service.create_session( + app_name=app_name, user_id=user_id_1 + ) + + async def run_prompt(session: Session, new_message: str): + content = types.Content( + role='user', parts=[types.Part.from_text(text=new_message)] + ) + print('** User says:', content.model_dump(exclude_none=True)) + async for event in runner.run_async( + user_id=user_id_1, + session_id=session.id, + new_message=content, + ): + if event.content.parts and event.content.parts[0].text: + print(f'** {event.author}: {event.content.parts[0].text}') + + start_time = time.time() + print('Start time:', start_time) + print('------------------------------------') + print('Testing NVIDIA NIM integration with ADK') + print('Model: nvidia_nim/meta/llama3-8b-instruct') + print('------------------------------------') + await run_prompt(session_11, 'Hi, introduce yourself.') + await run_prompt( + session_11, + 'Run the following request 5 times: roll a die with 20 sides and check' + ' if it is prime', + ) + end_time = time.time() + print('------------------------------------') + print('End time:', end_time) + print('Total time:', end_time - start_time) + + +if __name__ == '__main__': + asyncio.run(main()) \ No newline at end of file From b94a69c7cb929e86f67b461d415af288fce36e5b Mon Sep 17 00:00:00 2001 From: Soumili Nandi Date: Wed, 29 Oct 2025 12:26:14 -0700 Subject: [PATCH 2/3] docs: Improve README and main.py documentation --- contributing/samples/hello_world_nvidia/README.md | 13 ++++++++----- contributing/samples/hello_world_nvidia/main.py | 4 ++-- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/contributing/samples/hello_world_nvidia/README.md b/contributing/samples/hello_world_nvidia/README.md index e892000197..a8c282ab1f 100644 --- a/contributing/samples/hello_world_nvidia/README.md +++ b/contributing/samples/hello_world_nvidia/README.md @@ -1,20 +1,23 @@ -# Using NVIDIA Models with ADK via LiteLLM +# Using NVIDIA Models with ADK and LiteLLM This example demonstrates how to use NVIDIA NIM (NVIDIA Inference Microservices) models with ADK through LiteLLM integration. -For comprehensive information about using NVIDIA models with LiteLLM, refer to the [official LiteLLM NVIDIA NIM documentation](https://github.com/BerriAI/litellm/blob/main/docs/my-website/docs/providers/nvidia_nim.md?plain=1). +For comprehensive information about using NVIDIA models with LiteLLM, refer to the [official LiteLLM documentation](https://github.com/BerriAI/litellm/blob/main/docs/my-website/docs/providers/nvidia_nim.md). ## Setup -### 1. Get NVIDIA NIM API Key +Use the following procedure to get an NVIDIA NIM API key. + 1. Sign up at [NVIDIA AI Foundation](https://www.nvidia.com/en-us/ai-data-science/cloud/) 2. Navigate to API section and generate an API key 3. Copy your API key -### 2. Ensure LiteLLM is Installed +### 2. Install LiteLLM + +Install LiteLLM by running the following code. + -Install LiteLLM if you haven't already: ```bash pip install litellm diff --git a/contributing/samples/hello_world_nvidia/main.py b/contributing/samples/hello_world_nvidia/main.py index 63a905ab54..6a35e45990 100644 --- a/contributing/samples/hello_world_nvidia/main.py +++ b/contributing/samples/hello_world_nvidia/main.py @@ -33,7 +33,7 @@ async def main(): # Check for NVIDIA NIM API key if not os.environ.get("NVIDIA_NIM_API_KEY"): - print(" NVIDIA_NIM_API_KEY not set. Please set it before running:") + print(" Your NVIDIA_NIM_API_KEY is not set. Set it before you continue by running the following code:") print(" export NVIDIA_NIM_API_KEY='your-nvidia-api-key'") return @@ -74,7 +74,7 @@ async def run_prompt(session: Session, new_message: str): await run_prompt( session_11, 'Run the following request 5 times: roll a die with 20 sides and check' - ' if it is prime', + 'if the result is prime', ) end_time = time.time() print('------------------------------------') From af3ca77e0116cf90b5673a4e6555a9771570bc23 Mon Sep 17 00:00:00 2001 From: Soumili Nandi Date: Wed, 29 Oct 2025 12:28:15 -0700 Subject: [PATCH 3/3] minor fix in readme --- contributing/samples/hello_world_nvidia/README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/contributing/samples/hello_world_nvidia/README.md b/contributing/samples/hello_world_nvidia/README.md index a8c282ab1f..9e875da200 100644 --- a/contributing/samples/hello_world_nvidia/README.md +++ b/contributing/samples/hello_world_nvidia/README.md @@ -6,6 +6,7 @@ For comprehensive information about using NVIDIA models with LiteLLM, refer to t ## Setup +### 1. Get an NVIDIA NIM API Key Use the following procedure to get an NVIDIA NIM API key.