fix: lint

Author: asafgardin

libs/partners/ai21/langchain_ai21/ai21_base.py

@@ -5,12 +5,14 @@
 from langchain_core.pydantic_v1 import BaseModel, Field, SecretStr, root_validator
 from langchain_core.utils import convert_to_secret_str
 
+_DEFAULT_TIMEOUT_SEC = 300
+
 
 class AI21Base(BaseModel):
     class Config:
         arbitrary_types_allowed = True
 
-    client: Optional[AI21Client] = Field(default=None)
+    client: AI21Client = Field(default=None)
     api_key: Optional[SecretStr] = None
     api_host: Optional[str] = None
     timeout_sec: Optional[float] = None
@@ -30,14 +32,16 @@ def validate_environment(cls, values: Dict) -> Dict:
         )
         values["api_host"] = api_host
 
-        timeout_sec = values.get("timeout_sec") or os.getenv("AI21_TIMEOUT_SEC")
+        timeout_sec = values.get("timeout_sec") or float(
+            os.getenv("AI21_TIMEOUT_SEC", _DEFAULT_TIMEOUT_SEC)
+        )
         values["timeout_sec"] = timeout_sec
 
         if values.get("client") is None:
             values["client"] = AI21Client(
                 api_key=api_key.get_secret_value(),
                 api_host=api_host,
-                timeout_sec=timeout_sec,
+                timeout_sec=None if timeout_sec is None else float(timeout_sec),
                 via="langchain",
             )
 

libs/partners/ai21/langchain_ai21/chat_models.py

@@ -19,6 +19,15 @@
 from langchain_ai21.ai21_base import AI21Base
 
 
+def _get_system_message_from_message(message: BaseMessage) -> str:
+    if not isinstance(message.content, str):
+        raise ValueError(
+            f"System Message must be of type str. Got {type(message.content)}"
+        )
+
+    return message.content
+
+
 def _convert_messages_to_ai21_messages(
     messages: List[BaseMessage],
 ) -> Tuple[Optional[str], List[ChatMessage]]:
@@ -30,7 +39,7 @@ def _convert_messages_to_ai21_messages(
             if i != 0:
                 raise ValueError("System message must be at beginning of message list.")
             else:
-                system_message = message.content
+                system_message = _get_system_message_from_message(message)
         else:
             converted_message = _convert_message_to_ai21_message(message)
             converted_messages.append(converted_message)
@@ -108,6 +117,11 @@ class ChatAI21(BaseChatModel, AI21Base):
     """A penalty applied to tokens based on their frequency 
     in the generated responses."""
 
+    class Config:
+        """Configuration for this pydantic object."""
+
+        arbitrary_types_allowed = True
+
     @property
     def _llm_type(self) -> str:
         """Return type of chat model."""

libs/partners/ai21/langchain_ai21/llms.py

@@ -2,8 +2,6 @@
 from functools import partial
 from typing import (
     Any,
-    AsyncIterator,
-    Iterator,
     List,
     Optional,
 )
@@ -14,7 +12,7 @@
     CallbackManagerForLLMRun,
 )
 from langchain_core.language_models import BaseLLM
-from langchain_core.outputs import Generation, GenerationChunk, LLMResult
+from langchain_core.outputs import Generation, LLMResult
 
 from langchain_ai21.ai21_base import AI21Base
 
@@ -64,6 +62,11 @@ class AI21(BaseLLM, AI21Base):
     custom_model: Optional[str] = None
     epoch: Optional[int] = None
 
+    class Config:
+        """Configuration for this pydantic object."""
+
+        allow_population_by_field_name = True
+
     @property
     def _llm_type(self) -> str:
         """Return type of LLM."""
@@ -102,30 +105,12 @@ async def _agenerate(
             None, partial(self._generate, **kwargs), prompts, stop, run_manager
         )
 
-    def _stream(
-        self,
-        prompt: str,
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> Iterator[GenerationChunk]:
-        raise NotImplementedError
-
-    async def _astream(
-        self,
-        prompt: str,
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> AsyncIterator[GenerationChunk]:
-        raise NotImplementedError
-
     def _invoke_completion(
         self,
         prompt: str,
         model: str,
         stop_sequences: Optional[List[str]] = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> CompletionsResponse:
         return self.client.completion.create(
             prompt=prompt,

libs/partners/ai21/tests/integration_tests/test_llms.py

@@ -1,4 +1,5 @@
 """Test AI21LLM llm."""
+
 import pytest
 from ai21.models import Penalty
 
@@ -45,9 +46,8 @@ def test_stream() -> None:
     """Test streaming tokens from AI21."""
     llm = AI21()
 
-    with pytest.raises(NotImplementedError):
-        for token in llm.stream("I'm Pickle Rick"):
-            assert isinstance(token, str)
+    for token in llm.stream("I'm Pickle Rick"):
+        assert isinstance(token, str)
 
 
 @pytest.mark.requires("ai21")
@@ -109,7 +109,7 @@ def test__generate() -> None:
     )
 
     assert len(llm_result.generations) > 0
-    assert llm_result.llm_output["token_count"] != 0
+    assert llm_result.llm_output["token_count"] != 0  # type: ignore
 
 
 @pytest.mark.requires("ai21")
@@ -121,4 +121,4 @@ async def test__agenerate() -> None:
     )
 
     assert len(llm_result.generations) > 0
-    assert llm_result.llm_output["token_count"] != 0
+    assert llm_result.llm_output["token_count"] != 0  # type: ignore

libs/partners/ai21/tests/unit_tests/conftest.py

@@ -34,7 +34,7 @@
 
 
 @pytest.fixture
-def mocked_completion_response(mocker: MockerFixture):
+def mocked_completion_response(mocker: MockerFixture) -> Mock:
     mocked_response = mocker.MagicMock(spec=CompletionsResponse)
     mocked_response.prompt = "this is a test prompt"
     mocked_response.completions = [
@@ -48,7 +48,7 @@ def mocked_completion_response(mocker: MockerFixture):
 
 @pytest.fixture
 def mock_client_with_completion(
-    mocker: MockerFixture, mocked_completion_response
+    mocker: MockerFixture, mocked_completion_response: Mock
 ) -> Mock:
     mock_client = mocker.MagicMock(spec=AI21Client)
     mock_client.completion = mocker.MagicMock()

libs/partners/ai21/tests/unit_tests/test_chat_models.py

@@ -1,6 +1,6 @@
 """Test chat model integration."""
 from typing import List, Optional
-from unittest.mock import call
+from unittest.mock import Mock, call
 
 import pytest
 from ai21 import MissingApiKeyError
@@ -40,7 +40,7 @@ def test_initialization__when_default_parameters_in_init() -> None:
 
 
 @pytest.mark.requires("ai21")
-def test_initialization__when_custom_parameters_in_init():
+def test_initialization__when_custom_parameters_in_init() -> None:
     model = "j2-mid"
     num_results = 1
     max_tokens = 10
@@ -97,7 +97,7 @@ def test_initialization__when_custom_parameters_in_init():
 )
 def test_convert_message_to_ai21_message(
     message: BaseMessage, expected_ai21_message: ChatMessage
-):
+) -> None:
     ai21_message = _convert_message_to_ai21_message(message)
     assert ai21_message == expected_ai21_message
 
@@ -115,8 +115,8 @@ def test_convert_message_to_ai21_message(
     ],
 )
 def test_convert_message_to_ai21_message__when_invalid_role__should_raise_exception(
-    message,
-):
+    message: BaseMessage,
+) -> None:
     with pytest.raises(ValueError) as e:
         _convert_message_to_ai21_message(message)
     assert e.value.args[0] == (
@@ -157,15 +157,17 @@ def test_convert_message_to_ai21_message__when_invalid_role__should_raise_except
     ],
 )
 def test_convert_messages(
-    messages, expected_system: Optional[str], expected_messages: List[ChatMessage]
-):
+    messages: List[BaseMessage],
+    expected_system: Optional[str],
+    expected_messages: List[ChatMessage],
+) -> None:
     system, ai21_messages = _convert_messages_to_ai21_messages(messages)
     assert ai21_messages == expected_messages
     assert system == expected_system
 
 
 @pytest.mark.requires("ai21")
-def test_convert_messages_when_system_is_not_first__should_raise_value_error():
+def test_convert_messages_when_system_is_not_first__should_raise_value_error() -> None:
     messages = [
         HumanMessage(content="Human Message Content 1"),
         SystemMessage(content="System Message Content 1"),
@@ -175,7 +177,7 @@ def test_convert_messages_when_system_is_not_first__should_raise_value_error():
 
 
 @pytest.mark.requires("ai21")
-def test_invoke(mock_client_with_chat):
+def test_invoke(mock_client_with_chat: Mock) -> None:
     chat_input = "I'm Pickle Rick"
 
     llm = ChatAI21(
@@ -195,7 +197,7 @@ def test_invoke(mock_client_with_chat):
 
 
 @pytest.mark.requires("ai21")
-def test_generate(mock_client_with_chat):
+def test_generate(mock_client_with_chat: Mock) -> None:
     messages0 = [
         HumanMessage(content="I'm Pickle Rick"),
         AIMessage(content="Hello Pickle Rick! I am your AI Assistant"),
@@ -216,9 +218,14 @@ def test_generate(mock_client_with_chat):
             call(
                 model="j2-ultra",
                 messages=[
-                    ChatMessage(role=RoleType.USER, text=messages0[0].content),
-                    ChatMessage(role=RoleType.ASSISTANT, text=messages0[1].content),
-                    ChatMessage(role=RoleType.USER, text=messages0[2].content),
+                    ChatMessage(
+                        role=RoleType.USER,
+                        text=str(messages0[0].content),
+                    ),
+                    ChatMessage(
+                        role=RoleType.ASSISTANT, text=str(messages0[1].content)
+                    ),
+                    ChatMessage(role=RoleType.USER, text=str(messages0[2].content)),
                 ],
                 system="",
                 stop_sequences=None,
@@ -227,7 +234,7 @@ def test_generate(mock_client_with_chat):
             call(
                 model="j2-ultra",
                 messages=[
-                    ChatMessage(role=RoleType.USER, text=messages1[1].content),
+                    ChatMessage(role=RoleType.USER, text=str(messages1[1].content)),
                 ],
                 system="system message",
                 stop_sequences=None,

libs/partners/ai21/tests/unit_tests/test_llms.py

@@ -1,5 +1,5 @@
 """Test AI21 Chat API wrapper."""
-from unittest.mock import call
+from unittest.mock import Mock, call
 
 import pytest
 from ai21 import MissingApiKeyError
@@ -51,7 +51,7 @@ def test_initialization__when_custom_parameters_to_init() -> None:
 
 
 @pytest.mark.requires("ai21")
-def test_generate(mock_client_with_completion):
+def test_generate(mock_client_with_completion: Mock) -> None:
     # Setup test
     prompt0 = "Hi, my name is what?"
     prompt1 = "My name is who?"