From 3b187fc7f8de34cb0d9cf16b3203653980df3bae Mon Sep 17 00:00:00 2001
From: Matthew Farrellee <matt@cs.wisc.edu>
Date: Tue, 21 May 2024 13:28:49 -0400
Subject: [PATCH] add microsoft/phi-3-small-8k-instruct,
 microsoft/phi-3-small-128k-instruct, microsoft/phi-3-medium-4k-instruct and
 microsoft/phi-3-vision-128k-instruct to set of ChatNVIDIA models

---
 .../langchain_nvidia_ai_endpoints/_statics.py    | 16 ++++++++++++++++
 libs/ai-endpoints/pyproject.toml                 |  2 +-
 .../tests/integration_tests/conftest.py          |  7 +++++++
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_statics.py b/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_statics.py
index 68012232..963ea49e 100644
--- a/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_statics.py
+++ b/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_statics.py
@@ -179,6 +179,22 @@ class Model(BaseModel):
             "model_type": "embedding",
             "model_name": "snowflake/arctic-embed-l",
         },
+        "ai-phi-3-small-8k-instruct": {
+            "model_type": "chat",
+            "model_name": "microsoft/phi-3-small-8k-instruct",
+        },
+        "ai-phi-3-small-128k-instruct": {
+            "model_type": "chat",
+            "model_name": "microsoft/phi-3-small-128k-instruct",
+        },
+        "ai-phi-3-medium-4k-instruct": {
+            "model_type": "chat",
+            "model_name": "microsoft/phi-3-medium-4k-instruct",
+        },
+        "ai-phi-3-vision-128k-instruct": {
+            "model_type": "image_in",
+            "model_name": "microsoft/phi-3-vision-128k-instruct",
+        },
     }
 )
 
diff --git a/libs/ai-endpoints/pyproject.toml b/libs/ai-endpoints/pyproject.toml
index d9a0549a..e1cad2b6 100644
--- a/libs/ai-endpoints/pyproject.toml
+++ b/libs/ai-endpoints/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langchain-nvidia-ai-endpoints"
-version = "0.0.16"
+version = "0.0.17"
 description = "An integration package connecting NVIDIA AI Endpoints and LangChain"
 authors = []
 readme = "README.md"
diff --git a/libs/ai-endpoints/tests/integration_tests/conftest.py b/libs/ai-endpoints/tests/integration_tests/conftest.py
index 7a27d97c..5ae928ad 100644
--- a/libs/ai-endpoints/tests/integration_tests/conftest.py
+++ b/libs/ai-endpoints/tests/integration_tests/conftest.py
@@ -29,6 +29,11 @@ def pytest_addoption(parser: pytest.Parser) -> None:
         action="store",
         help="Run tests for a specific rerank model",
     )
+    parser.addoption(
+        "--vlm-model-id",
+        action="store",
+        help="Run tests for a specific vlm model",
+    )
     parser.addoption(
         "--all-models",
         action="store_true",
@@ -75,6 +80,8 @@ def get_all_models() -> List[Model]:
 
     if "image_in_model" in metafunc.fixturenames:
         models = ["adept/fuyu-8b"]
+        if model := metafunc.config.getoption("vlm_model_id"):
+            models = [model]
         if metafunc.config.getoption("all_models"):
             models = [
                 model.id for model in get_all_models() if model.model_type == "image_in"