fix: Merge pull request #2 from AI21Labs/ai21_prepare_pr_for_merge

fix: AI21 prepare pr for merge

Author: asafgardin

libs/partners/ai21/langchain_ai21/ai21_base.py

@@ -1,17 +1,18 @@
 import os
-from typing import Optional, Dict
+from typing import Dict, Optional
 
 from ai21 import AI21Client
-
-from langchain_core.pydantic_v1 import BaseModel, Field, root_validator, SecretStr
+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
@@ -25,20 +26,22 @@ def validate_environment(cls, values: Dict) -> Dict:
         values["api_key"] = api_key
 
         api_host = (
-                values.get("api_host")
-                or os.getenv("AI21_API_URL")
-                or "https://api.ai21.com"
+            values.get("api_host")
+            or os.getenv("AI21_API_URL")
+            or "https://api.ai21.com"
         )
         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:
+        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

@@ -1,10 +1,8 @@
 import asyncio
 from functools import partial
-from typing import Any, AsyncIterator, Iterator, List, Optional, cast, Tuple
+from typing import Any, List, Optional, Tuple, cast
 
-from ai21.models import ChatMessage, RoleType, Penalty
-
-from langchain_ai21.ai21_base import AI21Base
+from ai21.models import ChatMessage, Penalty, RoleType
 from langchain_core.callbacks import (
     AsyncCallbackManagerForLLMRun,
     CallbackManagerForLLMRun,
@@ -16,7 +14,18 @@
     HumanMessage,
     SystemMessage,
 )
-from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
+from langchain_core.outputs import ChatGeneration, ChatResult
+
+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(
@@ -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)
@@ -105,7 +114,13 @@ class ChatAI21(BaseChatModel, AI21Base):
     """ A penalty applied to tokens that are already present in the prompt."""
 
     count_penalty: Optional[Penalty] = None
-    """A penalty applied to tokens based on their frequency in the generated responses."""
+    """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:

libs/partners/ai21/langchain_ai21/embeddings.py

@@ -1,9 +1,9 @@
-from typing import List, Any
+from typing import Any, List
 
 from ai21.models import EmbedType
+from langchain_core.embeddings import Embeddings
 
 from langchain_ai21.ai21_base import AI21Base
-from langchain_core.embeddings import Embeddings
 
 
 class AI21Embeddings(Embeddings, AI21Base):
@@ -19,6 +19,7 @@ class AI21Embeddings(Embeddings, AI21Base):
             embeddings = AI21Embeddings()
             query_result = embeddings.embed_query("Hello embeddings world!")
     """
+
     def embed_documents(self, texts: List[str], **kwargs: Any) -> List[List[float]]:
         """Embed search docs."""
         response = self.client.embed.create(

libs/partners/ai21/langchain_ai21/llms.py

@@ -2,21 +2,19 @@
 from functools import partial
 from typing import (
     Any,
-    AsyncIterator,
-    Iterator,
     List,
     Optional,
 )
 
 from ai21.models import CompletionsResponse, Penalty
-
-from langchain_ai21.ai21_base import AI21Base
 from langchain_core.callbacks import (
     AsyncCallbackManagerForLLMRun,
     CallbackManagerForLLMRun,
 )
 from langchain_core.language_models import BaseLLM
-from langchain_core.outputs import GenerationChunk, LLMResult, Generation, RunInfo
+from langchain_core.outputs import Generation, LLMResult
+
+from langchain_ai21.ai21_base import AI21Base
 
 
 class AI21(BaseLLM, AI21Base):
@@ -58,11 +56,17 @@ class AI21(BaseLLM, AI21Base):
     """ A penalty applied to tokens that are already present in the prompt."""
 
     count_penalty: Optional[Penalty] = None
-    """A penalty applied to tokens based on their frequency in the generated responses."""
+    """A penalty applied to tokens based on their frequency 
+    in the generated responses."""
 
     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."""
@@ -101,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_chat_models.py

@@ -1,10 +1,10 @@
 """Test ChatAI21 chat model."""
 import pytest
-
-from langchain_ai21.chat_models import ChatAI21
 from langchain_core.messages import HumanMessage
 from langchain_core.outputs import ChatGeneration
 
+from langchain_ai21.chat_models import ChatAI21
+
 
 @pytest.mark.requires("ai21")
 def test_invoke() -> None:

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

@@ -3,15 +3,15 @@
 import pytest
 from ai21 import AI21Client
 from ai21.models import (
-    CompletionsResponse,
+    ChatOutput,
+    ChatResponse,
     Completion,
     CompletionData,
     CompletionFinishReason,
-    ChatResponse,
+    CompletionsResponse,
+    FinishReason,
     Penalty,
-    ChatOutput,
     RoleType,
-    FinishReason,
 )
 from pytest_mock import MockerFixture
 
@@ -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()