fix linting errors after adding tests

libs/aws/tests/integration_tests/chat_models/test_bedrock.py

@@ -9,10 +9,10 @@
     SystemMessage,
 )
 from langchain_core.outputs import ChatGeneration, LLMResult
+from langchain_core.pydantic_v1 import BaseModel, Field
 
 from langchain_aws.chat_models.bedrock import ChatBedrock
 from tests.callbacks import FakeCallbackHandler
-from langchain_core.pydantic_v1 import BaseModel, Field
 
 
 @pytest.fixture
@@ -157,80 +157,65 @@ def test_bedrock_invoke(chat: ChatBedrock) -> None:
     """Test invoke tokens from BedrockChat."""
     result = chat.invoke("I'm Pickle Rick", config=dict(tags=["foo"]))
     assert isinstance(result.content, str)
-    assert "usage" in result.additional_kwargs 
-    assert result.additional_kwargs["usage"]["prompt_tokens"] == 13 
+    assert "usage" in result.additional_kwargs
+    assert result.additional_kwargs["usage"]["prompt_tokens"] == 13
+
 
 @pytest.mark.scheduled
-def test_bedrock_anthropic_function_call_invoke_with_system(chat: ChatBedrock) -> None:
+def test_function_call_invoke_with_system(chat: ChatBedrock) -> None:
     class GetWeather(BaseModel):
-        location: str = Field(..., description="The city and state, e.g. San Francisco, CA")
+        location: str = Field(..., description="The city and state")
 
     llm_with_tools = chat.bind_tools([GetWeather])
 
     messages = [
-        SystemMessage(
-            content="anwser only in french"
-        ),
-        HumanMessage(
-            content="what is the weather like in San Francisco"
-        )
+        SystemMessage(content="anwser only in french"),
+        HumanMessage(content="what is the weather like in San Francisco"),
     ]
 
     response = llm_with_tools.invoke(messages)
     assert isinstance(response, BaseMessage)
     assert isinstance(response.content, str)
-
+
+
 @pytest.mark.scheduled
-def test_bedrock_anthropic_function_call_invoke_without_system(chat: ChatBedrock) -> None:
+def test_function_call_invoke_without_system(chat: ChatBedrock) -> None:
     class GetWeather(BaseModel):
-        location: str = Field(..., description="The city and state, e.g. San Francisco, CA")
+        location: str = Field(..., description="The city and state")
 
     llm_with_tools = chat.bind_tools([GetWeather])
 
-    messages = [
-        HumanMessage(
-            content="what is the weather like in San Francisco"
-        )
-    ]
+    messages = [HumanMessage(content="what is the weather like in San Francisco")]
 
     response = llm_with_tools.invoke(messages)
     assert isinstance(response, BaseMessage)
     assert isinstance(response.content, str)
 
+
 @pytest.mark.scheduled
-async def test_bedrock_anthropic_function_call_invoke_with_system_astream(chat: ChatBedrock) -> None:
+async def test_function_call_invoke_with_system_astream(chat: ChatBedrock) -> None:
     class GetWeather(BaseModel):
-        location: str = Field(..., description="The city and state, e.g. San Francisco, CA")
+        location: str = Field(..., description="The city and state")
 
     llm_with_tools = chat.bind_tools([GetWeather])
 
     messages = [
-        SystemMessage(
-            content="anwser only in french"
-        ),
-        HumanMessage(
-            content="what is the weather like in San Francisco"
-        )
+        SystemMessage(content="anwser only in french"),
+        HumanMessage(content="what is the weather like in San Francisco"),
     ]
 
     for chunk in llm_with_tools.stream(messages):
         assert isinstance(chunk.content, str)
 
+
 @pytest.mark.scheduled
-async def test_bedrock_anthropic_function_call_invoke_without_system_astream(chat: ChatBedrock) -> None:
+async def test_function_call_invoke_without_system_astream(chat: ChatBedrock) -> None:
     class GetWeather(BaseModel):
-        location: str = Field(..., description="The city and state, e.g. San Francisco, CA")
+        location: str = Field(..., description="The city and state")
 
     llm_with_tools = chat.bind_tools([GetWeather])
 
-    messages = [
-        HumanMessage(
-            content="what is the weather like in San Francisco"
-        )
-    ]
+    messages = [HumanMessage(content="what is the weather like in San Francisco")]
 
-    full = None
-    for token in llm_with_tools.stream("I'm Pickle Rick"):
-        full = token if full is None else full + token 
-        assert isinstance(token.content, str)
-    assert isinstance(cast(AIMessageChunk, full).content, str)
+    for chunk in llm_with_tools.stream(messages):
+        assert isinstance(chunk.content, str)