From 8d5ea16f68bca3b8c046f587b5729ae11065e9e1 Mon Sep 17 00:00:00 2001
From: Matthew Farrellee <matt@cs.wisc.edu>
Date: Thu, 27 Jun 2024 16:59:15 -0400
Subject: [PATCH] add known context lengths for hosted models (#14436)

---
 .../llama_index/llms/nvidia/utils.py          | 36 ++++++++++++++-----
 .../llama-index-llms-nvidia/pyproject.toml    |  2 +-
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/llama-index-integrations/llms/llama-index-llms-nvidia/llama_index/llms/nvidia/utils.py b/llama-index-integrations/llms/llama-index-llms-nvidia/llama_index/llms/nvidia/utils.py
index 3897347eb71e4..9c6af24afc047 100644
--- a/llama-index-integrations/llms/llama-index-llms-nvidia/llama_index/llms/nvidia/utils.py
+++ b/llama-index-integrations/llms/llama-index-llms-nvidia/llama_index/llms/nvidia/utils.py
@@ -1,17 +1,35 @@
 from typing import Dict, Optional
 
 API_CATALOG_MODELS: Dict[str, int] = {
-    "mistralai/mistral-7b-instruct-v0.2": 16384,
-    "mistralai/mixtral-8x7b-instruct-v0.1": 16384,
-    "mistralai/mixtral-8x22b-instruct-v0.1": 32768,
-    "mistralai/mistral-large": 16384,
-    "google/gemma-7b": 4096,
-    "google/gemma-2b": 4096,
+    "01-ai/yi-large": 16384,
+    "aisingapore/sea-lion-7b-instruct": 1024,
+    "databricks/dbrx-instruct": 2048,
+    "google/codegemma-1.1-7b": 1024,
     "google/codegemma-7b": 4096,
-    "meta/llama2-70b": 1024,
+    "google/gemma-2b": 4096,
+    "google/recurrentgemma-2b": 1024,
+    "ibm/granite-34b-code-instruct": 2048,
+    "ibm/granite-8b-code-instruct": 4096,
+    "mediatek/breeze-7b-instruct": 1024,
     "meta/codellama-70b": 1024,
-    "meta/llama3-8b-instruct": 6000,
-    "meta/llama3-70b-instruct": 6000,
+    "meta/llama3-70b-instruct": 8192,
+    "meta/llama3-8b-instruct": 8192,
+    "microsoft/phi-3-medium-4k-instruct": 1024,
+    "microsoft/phi-3-mini-128k-instruct": 2048,
+    "microsoft/phi-3-mini-4k-instruct": 2048,
+    "microsoft/phi-3-small-128k-instruct": 2048,
+    "microsoft/phi-3-small-8k-instruct": 2048,
+    "mistralai/codestral-22b-instruct-v0.1": 32768,
+    "mistralai/mistral-7b-instruct-v0.2": 16384,
+    "mistralai/mistral-7b-instruct-v0.3": 32768,
+    "mistralai/mistral-large": 16384,
+    "mistralai/mixtral-8x22b-instruct-v0.1": 65536,
+    "mistralai/mixtral-8x7b-instruct-v0.1": 16384,
+    "nvidia/llama3-chatqa-1.5-70b": 1024,
+    "nvidia/llama3-chatqa-1.5-8b": 1024,
+    "nvidia/nemotron-4-340b-instruct": 2048,
+    "seallms/seallm-7b-v2.5": 4096,
+    "upstage/solar-10.7b-instruct": 4096,
 }
 
 
diff --git a/llama-index-integrations/llms/llama-index-llms-nvidia/pyproject.toml b/llama-index-integrations/llms/llama-index-llms-nvidia/pyproject.toml
index edee1b459c7d7..8a3ae27e6a671 100644
--- a/llama-index-integrations/llms/llama-index-llms-nvidia/pyproject.toml
+++ b/llama-index-integrations/llms/llama-index-llms-nvidia/pyproject.toml
@@ -30,7 +30,7 @@ license = "MIT"
 name = "llama-index-llms-nvidia"
 packages = [{include = "llama_index/"}]
 readme = "README.md"
-version = "0.1.3"
+version = "0.1.4"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"