diff --git a/llama-index-integrations/multi_modal_llms/llama-index-multi-modal-llms-nvidia/llama_index/multi_modal_llms/nvidia/base.py b/llama-index-integrations/multi_modal_llms/llama-index-multi-modal-llms-nvidia/llama_index/multi_modal_llms/nvidia/base.py index 02cea8ed17ed7c6378d1d04dc08e53e04932f23a..e14dfa616f510769c85a3335954641896057452a 100644 --- a/llama-index-integrations/multi_modal_llms/llama-index-multi-modal-llms-nvidia/llama_index/multi_modal_llms/nvidia/base.py +++ b/llama-index-integrations/multi_modal_llms/llama-index-multi-modal-llms-nvidia/llama_index/multi_modal_llms/nvidia/base.py @@ -39,6 +39,12 @@ from llama_index.multi_modal_llms.nvidia.utils import ( import aiohttp import json +from llama_index.core.bridge.pydantic import BaseModel + + +class Model(BaseModel): + id: str + class NVIDIAClient: def __init__( @@ -58,14 +64,14 @@ class NVIDIAClient: headers["accept"] = "text/event-stream" if stream else "application/json" return headers - def get_model_details(self) -> List[str]: + def get_model_details(self) -> List[Model]: """ Get model details. Returns: List of models """ - return list(NVIDIA_MULTI_MODAL_MODELS.keys()) + return [Model(id=model) for model in NVIDIA_MULTI_MODAL_MODELS] def request( self, @@ -198,7 +204,7 @@ class NVIDIAMultiModal(MultiModalLLM): ) @property - def available_models(self): + def available_models(self) -> List[Model]: return self._client.get_model_details() def _get_credential_kwargs(self) -> Dict[str, Any]: diff --git a/llama-index-integrations/multi_modal_llms/llama-index-multi-modal-llms-nvidia/pyproject.toml b/llama-index-integrations/multi_modal_llms/llama-index-multi-modal-llms-nvidia/pyproject.toml index 9555020026cb10a9348f5fe873ce6676eaeed519..32d8121d7eaa67d7a6368a74b9de49b72c056353 100644 --- a/llama-index-integrations/multi_modal_llms/llama-index-multi-modal-llms-nvidia/pyproject.toml +++ b/llama-index-integrations/multi_modal_llms/llama-index-multi-modal-llms-nvidia/pyproject.toml @@ -27,7 +27,7 @@ license = "MIT" name = "llama-index-multi-modal-llms-nvidia" packages = [{include = "llama_index/"}] readme = "README.md" -version = "0.1.0" +version = "0.2.0" [tool.poetry.dependencies] python = ">=3.8.1,<4.0" diff --git a/llama-index-integrations/multi_modal_llms/llama-index-multi-modal-llms-nvidia/tests/test_available_models.py b/llama-index-integrations/multi_modal_llms/llama-index-multi-modal-llms-nvidia/tests/test_available_models.py index 829622e8724401a47135a0d12c3edb51e00c94d3..26d14ad038abfb18b44f3c480bb552a2c0683be8 100644 --- a/llama-index-integrations/multi_modal_llms/llama-index-multi-modal-llms-nvidia/tests/test_available_models.py +++ b/llama-index-integrations/multi_modal_llms/llama-index-multi-modal-llms-nvidia/tests/test_available_models.py @@ -8,4 +8,4 @@ def test_available_models() -> None: models = NVIDIAMultiModal().available_models assert models assert isinstance(models, list) - assert all(isinstance(model, str) for model in models) + assert all(isinstance(model.id, str) for model in models)