lint / format

Author: streamnsight

langchain_mcp_adapters/tools.py

@@ -4,7 +4,7 @@
 tools, handle tool execution, and manage tool conversion between the two formats.
 """
 
-from typing import Any, cast, get_args, Annotated, List, Union, Literal
+from typing import Annotated, Any, Literal, Union, cast, get_args
 
 from langchain_core.tools import (
     BaseTool,
@@ -20,7 +20,6 @@
 from mcp.types import Tool as MCPTool
 from pydantic import BaseModel, create_model
 
-
 from langchain_mcp_adapters.sessions import Connection, create_session
 
 NonTextContent = ImageContent | EmbeddedResource
@@ -136,66 +135,66 @@ async def call_tool(
         else:
             call_tool_result = await session.call_tool(tool.name, arguments)
         return _convert_call_tool_result(call_tool_result)
-    
+
     # base types being mapped from JSON
     type_map = {
-        'null': None,
-        'integer': int,
-        'float': float,
-        'string': str,
-        'bool': bool,
-        'object': dict,
-        'bytes': bytes,
+        "null": None,
+        "integer": int,
+        "float": float,
+        "string": str,
+        "bool": bool,
+        "object": dict,
+        "bytes": bytes,
     }
 
-    def _parse_model_fields(args, injected_state):
-        """Parse a JSON field into a Pydantic Field, taking into account injected state
+    def _parse_model_fields(args: dict, injected_state: str | None = None) -> dict:
+        """Parse a JSON field into a Pydantic Field, taking into account injected state.
 
         :param args: the function parameter schema
         :type args: dict
         :param injected_state: the name of the key used for the InjectedState
         :type injected_state: str
-        :return: returns a dict of fields with their pydantic type and default value if any
+        :return: returns a dict of fields with their pydantic type
+        and default value if any
         :rtype: dict
         """
         model_fields = {}
 
-        def _parse_field(props):
-            if 'anyOf' in props.keys():
-                types = tuple(_parse_field(p) for p in props['anyOf'])
-                return Union[types]
-            if 'enum' in props.keys():
-                return Literal[tuple(props['enum'])]
-            if props['type'] == 'array':
-                items = props['items']
-                return List[_parse_field(items)]
-            else:
-                return type_map.get(props['type'], dict)
-
-        for field, props in args['properties'].items():
+        def _parse_field(props: dict) -> tuple[type, Any]:
+            if "anyOf" in props:
+                return Union[tuple(_parse_field(p) for p in props["anyOf"])]  # noqa: UP007
+            if "enum" in props:
+                return Literal[tuple(props["enum"])]
+            if props["type"] == "array":
+                items = props["items"]
+                return list[_parse_field(items)]
+            return type_map.get(props["type"], dict)
+
+        for field, props in args["properties"].items():
             if field == injected_state:
                 field_type = Annotated[dict, InjectedState]
             else:
                 field_type = _parse_field(props)
-            if 'default' in props.keys():
-                default = props['default']
+            if "default" in props:
+                default = props["default"]
                 model_fields[field] = (field_type, default)
             else:
                 model_fields[field] = (field_type, ...)
         return model_fields
 
-    args = tool.inputSchema   
-    # check for the `injected_state`` annotation on the MCP tool. 
-    # The injected_state value is the name of the function parameter used as the injected state
-    injected_state = tool.annotations.model_extra.get('injected_state')
+    args = tool.inputSchema
+    # check for the `injected_state`` annotation on the MCP tool.
+    # The injected_state value is the name of the function parameter used
+    # as the injected state
+    injected_state = tool.annotations.model_extra.get("injected_state")
     if injected_state:
         # import langgraph InjectedState only if we need it
         from langgraph.prebuilt import InjectedState
     model_fields = _parse_model_fields(args, injected_state)
 
-    # recreate a dynamic model based on the parsed JSON schema, which will be properly parsed
-    # with annotations for the InjectedState
-    args_schema = create_model(tool.name, **{k: v for k, v in model_fields.items()})
+    # recreate a dynamic model based on the parsed JSON schema,
+    # which will be properly parsed with annotations for the InjectedState
+    args_schema = create_model(tool.name, **model_fields)
 
     return StructuredTool(
         name=tool.name,