langchain/libs/community/tests/unit_tests/llms/test_yandex.py
Nikita Pakunov c776471ac6
community: fix AttributeError: 'YandexGPT' object has no attribute '_grpc_metadata' (#24432)
Fixes #24049

---------

Co-authored-by: Erick Friis <erick@langchain.dev>
2024-07-31 21:18:33 +00:00

39 lines
1.3 KiB
Python

import pytest
from langchain_community.llms.yandex import YandexGPT
def test_yandexgpt_initialization() -> None:
llm = YandexGPT(
iam_token="your_iam_token", # type: ignore[arg-type]
api_key="your_api_key", # type: ignore[arg-type]
folder_id="your_folder_id",
)
assert llm.model_name == "yandexgpt-lite"
assert llm.model_uri.startswith("gpt://your_folder_id/yandexgpt-lite/")
def test_yandexgpt_model_params() -> None:
llm = YandexGPT(
model_name="custom-model",
model_version="v1",
iam_token="your_iam_token", # type: ignore[arg-type]
api_key="your_api_key", # type: ignore[arg-type]
folder_id="your_folder_id",
)
assert llm.model_name == "custom-model"
assert llm.model_version == "v1"
assert llm.iam_token.get_secret_value() == "your_iam_token"
assert llm.model_uri == "gpt://your_folder_id/custom-model/v1"
def test_yandexgpt_invalid_model_params() -> None:
with pytest.raises(ValueError):
YandexGPT(model_uri="", iam_token="your_iam_token") # type: ignore[arg-type]
with pytest.raises(ValueError):
YandexGPT(
iam_token="", # type: ignore[arg-type]
api_key="your_api_key", # type: ignore[arg-type]
model_uri="",
)