fmt

Author: baskaryan

libs/aws/langchain_aws/chat_models/bedrock_converse.py

@@ -22,6 +22,7 @@
 from langchain_core._api import beta
 from langchain_core.callbacks import CallbackManagerForLLMRun
 from langchain_core.language_models import BaseChatModel, LanguageModelInput
+from langchain_core.language_models.chat_models import LangSmithParams
 from langchain_core.messages import (
     AIMessage,
     BaseMessage,
@@ -476,6 +477,23 @@ def _converse_params(
             }
         )
 
+    def _get_ls_params(
+        self, stop: Optional[List[str]] = None, **kwargs: Any
+    ) -> LangSmithParams:
+        """Get standard params for tracing."""
+        params = self._get_invocation_params(stop=stop, **kwargs)
+        ls_params = LangSmithParams(
+            ls_provider="amazon_bedrock",
+            ls_model_name=self.model_id,
+            ls_model_type="chat",
+            ls_temperature=params.get("temperature", self.temperature),
+        )
+        if ls_max_tokens := params.get("max_tokens", self.max_tokens):
+            ls_params["ls_max_tokens"] = ls_max_tokens
+        if ls_stop := stop or params.get("stop", None):
+            ls_params["ls_stop"] = ls_stop
+        return ls_params
+
     @property
     def _llm_type(self) -> str:
         """Return type of chat model."""

libs/aws/tests/integration_tests/chat_models/test_bedrock_converse_standard.py renamed to libs/aws/tests/integration_tests/chat_models/test_bedrock_converse.py

@@ -1,13 +1,40 @@
 """Test chat model integration."""
 
-from typing import cast
+from typing import Type, cast
 
+import pytest
+from langchain_core.language_models import BaseChatModel
 from langchain_core.pydantic_v1 import BaseModel, Field
 from langchain_core.runnables import RunnableBinding
+from langchain_standard_tests.unit_tests import ChatModelUnitTests
 
 from langchain_aws import ChatBedrockConverse
 
 
+class TestBedrockStandard(ChatModelUnitTests):
+    @pytest.fixture
+    def chat_model_class(self) -> Type[BaseChatModel]:
+        return ChatBedrockConverse
+
+    @pytest.fixture
+    def chat_model_params(self) -> dict:
+        return {
+            "model_id": "anthropic.claude-3-sonnet-20240229-v1:0",
+        }
+
+    @pytest.mark.xfail()
+    def test_chat_model_init_api_key(
+        self, chat_model_class: Type[BaseChatModel], chat_model_params: dict
+    ) -> None:
+        super().test_chat_model_init_api_key(chat_model_class, chat_model_params)
+
+    @pytest.mark.xfail()
+    def test_chat_model_init_streaming(
+        self, chat_model_class: Type[BaseChatModel], chat_model_params: dict
+    ) -> None:
+        super().test_chat_model_init_streaming(chat_model_class, chat_model_params)
+
+
 class GetWeather(BaseModel):
     """Get the current weather in a given location"""