From 3196e54f03e76cd90719d916c96f9ff0763103cf Mon Sep 17 00:00:00 2001
From: raspawar <raspawar@nvidia.com>
Date: Thu, 2 Jan 2025 13:38:19 +0530
Subject: [PATCH] add local nim changes

---
 nvidia_models/intro/README.md                 |  3 ++-
 nvidia_models/intro/main.py                   |  3 ++-
 .../marketing_strategy/marketing_posts.ipynb  | 27 ++++++++++++++++---
 .../src/marketing_posts/crew.py               |  3 ++-
 4 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/nvidia_models/intro/README.md b/nvidia_models/intro/README.md
index 2cedbb56..9eb6bd40 100644
--- a/nvidia_models/intro/README.md
+++ b/nvidia_models/intro/README.md
@@ -4,10 +4,11 @@
 This is a simple example using the CrewAI framework with an NVIDIA endpoint and langchain-nvidia-ai-endpoints integration.
 
 ## Running the Script
-This example uses the Azure OpenAI API to call a model. 
+This example show cases the NVIDIA NIM endpoint integration with CrewAI.
 
 - **Configure Environment**: Set NVIDIA_API_KEY to appropriate api key.
                              Set MODEL to select appropriate model
+                             Set NVIDIA_API_URL to select the endpoint(Catalogue/local endpoint)
 - **Install Dependencies**: Run `make install`.
 - **Execute the Script**: Run `python main.py` to see a list of recommended changes to this document.
 
diff --git a/nvidia_models/intro/main.py b/nvidia_models/intro/main.py
index 676d5ebb..9814f62a 100644
--- a/nvidia_models/intro/main.py
+++ b/nvidia_models/intro/main.py
@@ -116,7 +116,8 @@ class nvllm(LLM):
 
 
 model = os.environ.get("MODEL", "meta/llama-3.1-8b-instruct")
-llm = ChatNVIDIA(model=model)
+api_base = os.environ.get("NVIDIA_API_URL", "https://integrate.api.nvidia.com/v1")
+llm = ChatNVIDIA(model=model, base_url=api_base)
 default_llm = nvllm(model_str="nvidia_nim/" + model, llm=llm)
 
 os.environ["NVIDIA_NIM_API_KEY"] = os.environ.get("NVIDIA_API_KEY")
diff --git a/nvidia_models/marketing_strategy/marketing_posts.ipynb b/nvidia_models/marketing_strategy/marketing_posts.ipynb
index 80122737..b7829646 100644
--- a/nvidia_models/marketing_strategy/marketing_posts.ipynb
+++ b/nvidia_models/marketing_strategy/marketing_posts.ipynb
@@ -51,7 +51,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -69,6 +69,18 @@
     "    os.environ[\"NVIDIA_API_KEY\"] = nvapi_key"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# set API Endoipoint\n",
+    "# to call local model set NVIDIA_API_URL to local NIM endpoint\n",
+    "os.environ[\"NVIDIA_API_URL\"] = \"http://localhost:8000/v1\" # for local NIM container\n",
+    "# os.environ[\"NVIDIA_API_URL\"] = \"https://integrate.api.nvidia.com/v1\""
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -78,12 +90,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
     "#set model\n",
-    "os.environ[\"MODEL\"] = \"meta/llama-3.1-8b-instruct\""
+    "os.environ[\"MODEL\"] = \"meta/llama-2-7b-chat\""
    ]
   },
   {
@@ -95,7 +107,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -110,6 +122,13 @@
    "source": [
     "run()"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
diff --git a/nvidia_models/marketing_strategy/src/marketing_posts/crew.py b/nvidia_models/marketing_strategy/src/marketing_posts/crew.py
index 515a15ba..44ddc855 100644
--- a/nvidia_models/marketing_strategy/src/marketing_posts/crew.py
+++ b/nvidia_models/marketing_strategy/src/marketing_posts/crew.py
@@ -16,7 +16,8 @@ from langchain_nvidia_ai_endpoints import ChatNVIDIA
 load_dotenv()
 
 model = os.getenv("MODEL", "meta/llama-3.1-8b-instruct")
-llm = ChatNVIDIA(model=model)
+api_base = os.environ.get("NVIDIA_API_URL", "https://integrate.api.nvidia.com/v1")
+llm = ChatNVIDIA(model=model, base_url=api_base)
 default_llm = nvllm(model_str="nvidia_nim/" + model, llm=llm)
 
 os.environ["NVIDIA_API_KEY"] = os.getenv("NVIDIA_API_KEY")