diff --git a/llama-index-integrations/llms/llama-index-llms-nvidia/llama_index/llms/nvidia/base.py b/llama-index-integrations/llms/llama-index-llms-nvidia/llama_index/llms/nvidia/base.py index 173e92d921d65cee663960b9c6bb5ddac4e8fbc8..61d28b47d2787c2990174727c5c7dde653921331 100644 --- a/llama-index-integrations/llms/llama-index-llms-nvidia/llama_index/llms/nvidia/base.py +++ b/llama-index-integrations/llms/llama-index-llms-nvidia/llama_index/llms/nvidia/base.py @@ -32,6 +32,7 @@ class NVIDIA(OpenAILike): model: str = DEFAULT_MODEL, nvidia_api_key: Optional[str] = None, api_key: Optional[str] = None, + max_tokens: Optional[int] = 1024, **kwargs: Any, ) -> None: api_key = get_from_param_or_env( @@ -45,6 +46,7 @@ class NVIDIA(OpenAILike): model=model, api_key=api_key, api_base=BASE_URL, + max_tokens=max_tokens, is_chat_model=True, default_headers={"User-Agent": "llama-index-llms-nvidia"}, **kwargs, 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 a90fb7fcf040ac414f02e69b895e0abc9e4eebb1..e6a493855f714712b867e4f8bcf2f5a140c15235 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.0" +version = "0.1.1" [tool.poetry.dependencies] python = ">=3.8.1,<4.0"