diff --git a/libs/community/langchain_community/chat_models/writer.py b/libs/community/langchain_community/chat_models/writer.py index 4d76a6a5f4dea..6538bdd1491d1 100644 --- a/libs/community/langchain_community/chat_models/writer.py +++ b/libs/community/langchain_community/chat_models/writer.py @@ -110,7 +110,8 @@ def _default_params(self) -> Dict[str, Any]: } @model_validator(mode="before") - def validate_environment(self, values: Dict) -> Any: + @classmethod + def validate_environment(cls, values: Dict) -> Any: """Validates that api key is passed and creates Writer clients.""" try: from writerai import AsyncClient, Client @@ -120,10 +121,19 @@ def validate_environment(self, values: Dict) -> Any: "Please install it with `pip install writerai`." ) from e - if not (values["client"] and values["async_client"]): + if not (values.get("client") and values.get("async_client")): api_key = get_from_dict_or_env(values, "api_key", "WRITER_API_KEY") - values["client"] = Client(api_key=api_key) - values["async_client"] = AsyncClient(api_key=api_key) + values.update({"client": Client(api_key=api_key)}) + values.update({"async_client": AsyncClient(api_key=api_key)}) + + if not ( + type(values.get("client")) is Client + and type(values.get("async_client")) is AsyncClient + ): + raise ValueError( + "'client' attribute must be with type 'Client' and " + "'async_client' must be with type 'AsyncClient' from 'writerai' package" + ) return values diff --git a/libs/community/tests/unit_tests/chat_models/test_writer.py b/libs/community/tests/unit_tests/chat_models/test_writer.py index 8b20a10ef7adf..8a0c833794ca6 100644 --- a/libs/community/tests/unit_tests/chat_models/test_writer.py +++ b/libs/community/tests/unit_tests/chat_models/test_writer.py @@ -106,25 +106,31 @@ def __init__( self.choices = choices -@pytest.mark.requires("writer-sdk") +@pytest.mark.requires("writerai") class TestChatWriterCustom: """Test case for ChatWriter""" + from writerai import AsyncClient, Client + def test_writer_model_param(self) -> None: """Test different ways to initialize the chat model.""" test_cases: List[dict] = [ { "model_name": "palmyra-x-004", + "api_key": "key", }, { "model": "palmyra-x-004", + "api_key": "key", }, { "model_name": "palmyra-x-004", + "api_key": "key", }, { "model": "palmyra-x-004", "temperature": 0.5, + "api_key": "key", }, ] @@ -280,10 +286,12 @@ def test_sync_completion( self, mock_unstreaming_completion: List[ChatCompletionChunk] ) -> None: """Test basic chat completion with mocked response.""" - mock_client = MagicMock() - mock_client.chat.chat.return_value = mock_unstreaming_completion + mock_client = self.Client(api_key="key") + mock_client.chat.chat = MagicMock(return_value=mock_unstreaming_completion) - chat = ChatWriter(client=mock_client, async_client=AsyncMock()) + async_client = self.AsyncClient(api_key="key") + + chat = ChatWriter(client=mock_client, async_client=async_client) message = HumanMessage(content="Hi there!") response = chat.invoke([message]) @@ -416,7 +424,7 @@ class GetWeather(BaseModel): assert response.tool_calls[0]["args"]["location"] == "London" -@pytest.mark.requires("writer-sdk") +@pytest.mark.requires("writerai") class TestChatWriterStandart(ChatModelUnitTests): """Test case for ChatWriter that inherits from standard LangChain tests."""