diff --git a/api/core/mcp/server/streamable_http.py b/api/core/mcp/server/streamable_http.py index 278add8cc9..8de002ae55 100644 --- a/api/core/mcp/server/streamable_http.py +++ b/api/core/mcp/server/streamable_http.py @@ -1,7 +1,7 @@ import json import logging from collections.abc import Mapping -from typing import Any, cast +from typing import Any, NotRequired, TypedDict, cast from graphon.variables.input_entities import VariableEntity, VariableEntityType @@ -15,6 +15,17 @@ from services.app_generate_service import AppGenerateService logger = logging.getLogger(__name__) +class ToolParameterSchemaDict(TypedDict): + type: str + properties: dict[str, Any] + required: list[str] + + +class ToolArgumentsDict(TypedDict): + query: NotRequired[str] + inputs: dict[str, Any] + + def handle_mcp_request( app: App, request: mcp_types.ClientRequest, @@ -119,7 +130,7 @@ def handle_list_tools( mcp_types.Tool( name=app_name, description=description, - inputSchema=parameter_schema, + inputSchema=cast(dict[str, Any], parameter_schema), ) ], ) @@ -154,7 +165,7 @@ def build_parameter_schema( app_mode: str, user_input_form: list[VariableEntity], parameters_dict: dict[str, str], -) -> dict[str, Any]: +) -> ToolParameterSchemaDict: """Build parameter schema for the tool""" parameters, required = convert_input_form_to_parameters(user_input_form, parameters_dict) @@ -174,7 +185,7 @@ def build_parameter_schema( } -def prepare_tool_arguments(app: App, arguments: dict[str, Any]) -> dict[str, Any]: +def prepare_tool_arguments(app: App, arguments: dict[str, Any]) -> ToolArgumentsDict: """Prepare arguments based on app mode""" if app.mode == AppMode.WORKFLOW: return {"inputs": arguments}