From db832a8b5fa322da44eaf612d56e2137ce992865 Mon Sep 17 00:00:00 2001
From: Bagatur <22008038+baskaryan@users.noreply.github.com>
Date: Fri, 2 Aug 2024 20:59:39 -0700
Subject: [PATCH] aws[patch]: add with_structured_output pydantic v2 support
 bedrock (#138)

---
 libs/aws/langchain_aws/chat_models/bedrock.py          | 8 ++++----
 libs/aws/langchain_aws/chat_models/bedrock_converse.py | 9 ++++-----
 2 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/libs/aws/langchain_aws/chat_models/bedrock.py b/libs/aws/langchain_aws/chat_models/bedrock.py
index e5ef55ae..ba5dbe36 100644
--- a/libs/aws/langchain_aws/chat_models/bedrock.py
+++ b/libs/aws/langchain_aws/chat_models/bedrock.py
@@ -11,7 +11,6 @@
     Optional,
     Sequence,
     Tuple,
-    Type,
     Union,
     cast,
 )
@@ -34,6 +33,7 @@
 from langchain_core.pydantic_v1 import BaseModel, Extra
 from langchain_core.runnables import Runnable, RunnableMap, RunnablePassthrough
 from langchain_core.tools import BaseTool
+from langchain_core.utils.pydantic import TypeBaseModel, is_basemodel_subclass
 
 from langchain_aws.chat_models.bedrock_converse import ChatBedrockConverse
 from langchain_aws.function_calling import (
@@ -595,7 +595,7 @@ def set_system_prompt_with_tools(self, xml_tools_system_prompt: str) -> None:
 
     def bind_tools(
         self,
-        tools: Sequence[Union[Dict[str, Any], Type[BaseModel], Callable, BaseTool]],
+        tools: Sequence[Union[Dict[str, Any], TypeBaseModel, Callable, BaseTool]],
         *,
         tool_choice: Optional[Union[dict, str, Literal["auto", "none"], bool]] = None,
         **kwargs: Any,
@@ -650,7 +650,7 @@ def bind_tools(
 
     def with_structured_output(
         self,
-        schema: Union[Dict, Type[BaseModel]],
+        schema: Union[Dict, TypeBaseModel],
         *,
         include_raw: bool = False,
         **kwargs: Any,
@@ -771,7 +771,7 @@ class AnswerWithJustification(BaseModel):
 
         tool_name = convert_to_anthropic_tool(schema)["name"]
         llm = self.bind_tools([schema], tool_choice=tool_name)
-        if isinstance(schema, type) and issubclass(schema, BaseModel):
+        if isinstance(schema, type) and is_basemodel_subclass(schema):
             output_parser = ToolsOutputParser(
                 first_tool_only=True, pydantic_schemas=[schema]
             )
diff --git a/libs/aws/langchain_aws/chat_models/bedrock_converse.py b/libs/aws/langchain_aws/chat_models/bedrock_converse.py
index 9bc8e59f..e133ed23 100644
--- a/libs/aws/langchain_aws/chat_models/bedrock_converse.py
+++ b/libs/aws/langchain_aws/chat_models/bedrock_converse.py
@@ -12,7 +12,6 @@
     Optional,
     Sequence,
     Tuple,
-    Type,
     TypeVar,
     Union,
     cast,
@@ -43,7 +42,7 @@
 from langchain_core.tools import BaseTool
 from langchain_core.utils import get_from_dict_or_env
 from langchain_core.utils.function_calling import convert_to_openai_function
-from langchain_core.utils.pydantic import is_basemodel_subclass
+from langchain_core.utils.pydantic import TypeBaseModel, is_basemodel_subclass
 
 from langchain_aws.function_calling import ToolsOutputParser
 
@@ -436,7 +435,7 @@ def _stream(
 
     def bind_tools(
         self,
-        tools: Sequence[Union[Dict[str, Any], Type[BaseModel], Callable, BaseTool]],
+        tools: Sequence[Union[Dict[str, Any], TypeBaseModel, Callable, BaseTool]],
         *,
         tool_choice: Optional[Union[dict, str, Literal["auto", "any"]]] = None,
         **kwargs: Any,
@@ -447,7 +446,7 @@ def bind_tools(
 
     def with_structured_output(
         self,
-        schema: Union[Dict, Type[BaseModel]],
+        schema: Union[Dict, TypeBaseModel],
         *,
         include_raw: bool = False,
         **kwargs: Any,
@@ -770,7 +769,7 @@ def _bedrock_to_anthropic(content: List[Dict[str, Any]]) -> List[Dict[str, Any]]
 
 
 def _format_tools(
-    tools: Sequence[Union[Dict[str, Any], Type[BaseModel], Callable, BaseTool],],
+    tools: Sequence[Union[Dict[str, Any], TypeBaseModel, Callable, BaseTool],],
 ) -> List[Dict[Literal["toolSpec"], Dict[str, Union[Dict[str, Any], str]]]]:
     formatted_tools: List = []
     for tool in tools: