ob-metaflow-extensions 1.1.156__tar.gz → 1.1.157__tar.gz

{ob-metaflow-extensions-1.1.156 → ob-metaflow-extensions-1.1.157}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.156
+Version: 1.1.157
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob-metaflow-extensions-1.1.156 → ob-metaflow-extensions-1.1.157}/metaflow_extensions/outerbounds/plugins/nim/nim_manager.py

@@ -186,18 +186,61 @@ class NimManager(object):
         ]
         self.models = {}
 
-        for each_model in models:
-            if each_model in nvcf_models:
-                self.models[each_model] = NimChatCompletion(
-                    model=each_model,
+        # Convert models to a standard format
+        standardized_models = []
+        # If models is a single string, convert it to a list with a dict
+        if isinstance(models, str):
+            standardized_models = [{"name": models}]
+        # If models is a list, process each item
+        elif isinstance(models, list):
+            for model_item in models:
+                # If the item is a string, convert it to a dict
+                if isinstance(model_item, str):
+                    standardized_models.append({"name": model_item})
+                # If it's already a dict, use it as is
+                elif isinstance(model_item, dict):
+                    standardized_models.append(model_item)
+                else:
+                    raise ValueError(
+                        f"Model specification must be a string or dictionary, got {type(model_item)}"
+                    )
+        else:
+            raise ValueError(
+                f"Models must be a string or a list of strings/dictionaries, got {type(models)}"
+            )
+
+        # Process each standardized model
+        for each_model_dict in standardized_models:
+            model_name = each_model_dict.get("name", "")
+            nvcf_id = each_model_dict.get("nvcf_id", "")
+            nvcf_version = each_model_dict.get("nvcf_version", "")
+
+            if model_name and not (nvcf_id and nvcf_version):
+                if model_name in nvcf_models:
+                    self.models[model_name] = NimChatCompletion(
+                        model=model_name,
+                        nvcf_id=nvcf_id,
+                        nvcf_version=nvcf_version,
+                        nim_metadata=nim_metadata,
+                        monitor=monitor,
+                    )
+                else:
+                    raise ValueError(
+                        f"Model {model_name} not supported by the Outerbounds @nim offering."
+                        f"\nYou can choose from these options: {nvcf_models}\n\n"
+                        "Reach out to Outerbounds if there are other models you'd like supported."
+                    )
+            elif nvcf_id and nvcf_version:
+                self.models[model_name] = NimChatCompletion(
+                    model=model_name,
+                    nvcf_id=nvcf_id,
+                    nvcf_version=nvcf_version,
                     nim_metadata=nim_metadata,
                     monitor=monitor,
                 )
             else:
                 raise ValueError(
-                    f"Model {each_model} not supported by the Outerbounds @nim offering."
-                    f"\nYou can choose from these options: {nvcf_models}\n\n"
-                    "Reach out to Outerbounds if there are other models you'd like supported."
+                    "You must provide either a valid 'name' or a custom 'name' along with both 'nvcf_id' and 'nvcf_version'."
                 )
 
 
@@ -205,6 +248,8 @@ class NimChatCompletion(object):
     def __init__(
         self,
         model: str = "meta/llama3-8b-instruct",
+        nvcf_id: str = "",
+        nvcf_version: str = "",
         nim_metadata: NimMetadata = None,
         monitor: bool = False,
         **kwargs,
@@ -217,18 +262,34 @@ class NimChatCompletion(object):
         self.model_name = model
         self.nim_metadata = nim_metadata
         self.monitor = monitor
-
         all_nvcf_models = self.nim_metadata.get_nvcf_chat_completion_models()
-        all_nvcf_model_names = [m["name"] for m in all_nvcf_models]
 
-        if self.model_name not in all_nvcf_model_names:
-            raise ValueError(
-                f"Model {self.model_name} not found in available NVCF models"
-            )
+        if nvcf_id and nvcf_version:
+            matching_models = [
+                m
+                for m in all_nvcf_models
+                if m["function-id"] == nvcf_id and m["version-id"] == nvcf_version
+            ]
+            if matching_models:
+                self.model = matching_models[0]
+                self.function_id = self.model["function-id"]
+                self.version_id = self.model["version-id"]
+                self.model_name = self.model["name"]
+            else:
+                raise ValueError(
+                    f"Function {self.function_id} with version {self.version_id} not found on NVCF"
+                )
+        else:
+            all_nvcf_model_names = [m["name"] for m in all_nvcf_models]
+
+            if self.model_name not in all_nvcf_model_names:
+                raise ValueError(
+                    f"Model {self.model_name} not found in available NVCF models"
+                )
 
-        self.model = all_nvcf_models[all_nvcf_model_names.index(self.model_name)]
-        self.function_id = self.model["function-id"]
-        self.version_id = self.model["version-id"]
+            self.model = all_nvcf_models[all_nvcf_model_names.index(self.model_name)]
+            self.function_id = self.model["function-id"]
+            self.version_id = self.model["version-id"]
 
         self.first_request = True
 

{ob-metaflow-extensions-1.1.156 → ob-metaflow-extensions-1.1.157}/ob_metaflow_extensions.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.156
+Version: 1.1.157
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob-metaflow-extensions-1.1.156 → ob-metaflow-extensions-1.1.157}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_namespace_packages
 from pathlib import Path
 
 
-version = "1.1.156"
+version = "1.1.157"
 this_directory = Path(__file__).parent
 long_description = (this_directory / "README.md").read_text()