From d2e23a4c161f887dc386f8f2f469dcf1eb40bc4a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BE=9D=E8=AF=BA?= Date: Wed, 10 Sep 2025 14:46:11 +0800 Subject: [PATCH] Add model information to the tool parameters --- pyproject.toml | 2 +- src/alibabacloud_dms_mcp_server/server.py | 26 ++++++++++++++--------- uv.lock | 2 +- 3 files changed, 18 insertions(+), 12 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ae0dadd..8dff19e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "alibabacloud-dms-mcp-server" -version = "0.1.13" +version = "0.1.14" description = "MCP Server for AlibabaCloud DMS" readme = "README.md" authors = [ diff --git a/src/alibabacloud_dms_mcp_server/server.py b/src/alibabacloud_dms_mcp_server/server.py index 236ef94..f26ce24 100644 --- a/src/alibabacloud_dms_mcp_server/server.py +++ b/src/alibabacloud_dms_mcp_server/server.py @@ -365,7 +365,6 @@ async def create_data_change_order( script: str = Field(description="SQL script to execute"), logic: bool = Field(default=False, description="Whether to use logical execution mode") ) -> Dict[str, Any]: - client = create_client() req = dms_enterprise_20181101_models.CreateDataCorrectOrderRequest() req.comment = "Data correct order submitted by MCP" @@ -395,7 +394,6 @@ async def create_data_change_order( async def get_order_base_info( order_id: str = Field(description="DMS order ID") ) -> Dict[str, Any]: - client = create_client() req = dms_enterprise_20181101_models.GetOrderBaseInfoRequest() req.order_id = order_id @@ -410,7 +408,6 @@ async def get_order_base_info( async def submit_order_approval( order_id: str = Field(description="DMS order ID") ) -> Dict[str, Any]: - client = create_client() req = dms_enterprise_20181101_models.SubmitOrderApprovalRequest() req.order_id = order_id @@ -421,16 +418,21 @@ async def submit_order_approval( logger.error(f"Error in submit_order_approval: {e}") raise + async def nl2sql( database_id: str = Field(description="DMS databaseId"), question: str = Field(description="Natural language question"), - knowledge: Optional[str] = Field(default=None, description="Additional context") + knowledge: Optional[str] = Field(default=None, description="Optional: additional context"), + model: Optional[str] = Field(default=None, description="Optional: if a specific model is desired, it can be specified here") ) -> SqlResult: client = create_client() req = dms_enterprise_20181101_models.GenerateSqlFromNLRequest(db_id=database_id, question=question) - if knowledge: req.knowledge = knowledge + if knowledge: + req.knowledge = knowledge if mcp.state.real_login_uid: req.real_login_user_uid = mcp.state.real_login_uid + if model: + req.model = model try: resp = client.generate_sql_from_nl(req) if not resp or not resp.body: return SqlResult(sql=None) @@ -495,13 +497,15 @@ async def execute_script_configured( async def create_data_change_order_configured( script: str = Field(description="SQL script to execute") ) -> str: - result_obj = await create_data_change_order(database_id=self.default_database_id, script=script, logic=False) + result_obj = await create_data_change_order(database_id=self.default_database_id, script=script, + logic=False) return str(result_obj) self.mcp.tool(name="getOrderInfo", description="Retrieve order information from DMS using the order ID.", annotations={"title": "获取DMS工单详情", "readOnlyHint": True})(get_order_base_info) - self.mcp.tool(name="submitOrderApproval", description="Submit the order for approval in DMS using the order ID.", + self.mcp.tool(name="submitOrderApproval", + description="Submit the order for approval in DMS using the order ID.", annotations={"title": "提交工单审批", "readOnlyHint": False})(submit_order_approval) @self.mcp.tool(name="askDatabase", @@ -511,10 +515,11 @@ async def ask_database_configured( question: str = Field( description="Your question in natural language about the pre-configured database."), knowledge: Optional[str] = Field(default=None, - description="Optional: Additional context to help formulate the SQL query.") + description="Optional: additional context to help formulate the SQL query."), + model: Optional[str] = Field(default=None, description="Optional: if a specific model is desired, it can be specified here") ) -> AskDatabaseResult: sql_result_obj = await nl2sql(database_id=self.default_database_id, question=question, - knowledge=knowledge) + knowledge=knowledge, model=model) generated_sql = "" if not sql_result_obj or not sql_result_obj.sql: logger.warning(f"Failed to generate SQL for question: {question} on preconfigured DB.") @@ -600,7 +605,8 @@ async def create_data_change_order_wrapper( self.mcp.tool(name="getOrderInfo", description="Retrieve order information from DMS using the order ID.", annotations={"title": "获取DMS工单详情", "readOnlyHint": True})(get_order_base_info) - self.mcp.tool(name="submitOrderApproval", description="Submit the order for approval in DMS using the order ID.", + self.mcp.tool(name="submitOrderApproval", + description="Submit the order for approval in DMS using the order ID.", annotations={"title": "提交工单审批", "readOnlyHint": False})(submit_order_approval) self.mcp.tool(name="generateSql", description="Generate SELECT-type SQL queries from natural language input.", diff --git a/uv.lock b/uv.lock index 3fa0a54..c96b007 100644 --- a/uv.lock +++ b/uv.lock @@ -149,7 +149,7 @@ wheels = [ [[package]] name = "alibabacloud-dms-mcp-server" -version = "0.1.12" +version = "0.1.14" source = { editable = "." } dependencies = [ { name = "alibabacloud-dms-enterprise20181101" },