Coverage for mcpgateway/schemas.py: 71%

457 statements  

« prev     ^ index     » next       coverage.py v7.9.1, created at 2025-06-22 12:53 +0100

1# -*- coding: utf-8 -*- 

2"""MCP Gateway Schema Definitions. 

3 

4Copyright 2025 

5SPDX-License-Identifier: Apache-2.0 

6Authors: Mihai Criveti 

7 

8This module provides Pydantic models for request/response validation in the MCP Gateway. 

9It implements schemas for: 

10- Tool registration and invocation 

11- Resource management and subscriptions 

12- Prompt templates and arguments 

13- Gateway federation 

14- RPC message formats 

15- Event messages 

16- Admin interface 

17 

18The schemas ensure proper validation according to the MCP specification while adding 

19gateway-specific extensions for federation support. 

20""" 

21 

22import base64 

23import json 

24import logging 

25from datetime import datetime 

26from typing import Any, Dict, List, Literal, Optional, Union 

27 

28from pydantic import ( 

29 AnyHttpUrl, 

30 BaseModel, 

31 Field, 

32 model_validator, 

33 root_validator, 

34 validator, 

35) 

36 

37from mcpgateway.types import ImageContent 

38from mcpgateway.types import Prompt as MCPPrompt 

39from mcpgateway.types import Resource as MCPResource 

40from mcpgateway.types import ResourceContent, TextContent 

41from mcpgateway.types import Tool as MCPTool 

42from mcpgateway.utils.services_auth import decode_auth, encode_auth 

43 

44logger = logging.getLogger(__name__) 

45 

46 

47def to_camel_case(s: str) -> str: 

48 """ 

49 Convert a string from snake_case to camelCase. 

50 

51 Args: 

52 s (str): The string to be converted, which is assumed to be in snake_case. 

53 

54 Returns: 

55 str: The string converted to camelCase. 

56 

57 Example: 

58 >>> to_camel_case("hello_world_example") 

59 'helloWorldExample' 

60 """ 

61 return "".join(word.capitalize() if i else word for i, word in enumerate(s.split("_"))) 

62 

63 

64def encode_datetime(v: datetime) -> str: 

65 """ 

66 Convert a datetime object to an ISO 8601 formatted string. 

67 

68 Args: 

69 v (datetime): The datetime object to be encoded. 

70 

71 Returns: 

72 str: The ISO 8601 formatted string representation of the datetime object. 

73 

74 Example: 

75 >>> encode_datetime(datetime(2023, 5, 22, 14, 30, 0)) 

76 '2023-05-22T14:30:00' 

77 """ 

78 return v.isoformat() 

79 

80 

81# --- Base Model --- 

82class BaseModelWithConfig(BaseModel): 

83 """Base model with common configuration. 

84 

85 Provides: 

86 - ORM mode for SQLAlchemy integration 

87 - JSON encoders for datetime handling 

88 - Automatic conversion from snake_case to camelCase for output 

89 """ 

90 

91 class Config: 

92 """ 

93 A configuration class that provides default behaviors for how to handle serialization, 

94 alias generation, enum values, and extra fields when working with models. 

95 

96 Attributes: 

97 from_attributes (bool): Flag to indicate if attributes should be taken from model fields. 

98 alias_generator (callable): Function used to generate aliases for field names (e.g., converting to camelCase). 

99 populate_by_name (bool): Flag to specify whether to populate fields by name during initialization. 

100 json_encoders (dict): Custom JSON encoders for specific types, such as datetime encoding. 

101 use_enum_values (bool): Flag to determine if enum values should be serialized or the enum type itself. 

102 extra (str): Defines behavior for extra fields in models. The "ignore" option means extra fields are ignored. 

103 json_schema_extra (dict): Additional schema information, e.g., specifying that fields can be nullable. 

104 

105 """ 

106 

107 from_attributes = True 

108 alias_generator = to_camel_case 

109 populate_by_name = True 

110 json_encoders = {datetime: encode_datetime} 

111 use_enum_values = True 

112 extra = "ignore" 

113 json_schema_extra = {"nullable": True} 

114 

115 def to_dict(self, use_alias: bool = False) -> Dict[str, Any]: 

116 """ 

117 Converts the model instance into a dictionary representation. 

118 

119 Args: 

120 use_alias (bool): Whether to use aliases for field names (default is False). If True, 

121 field names will be converted using the alias generator function. 

122 

123 Returns: 

124 Dict[str, Any]: A dictionary where keys are field names and values are corresponding field values, 

125 with any nested models recursively converted to dictionaries. 

126 """ 

127 output = {} 

128 for key, value in self.dict(by_alias=use_alias).items(): 

129 output[key] = value if not isinstance(value, BaseModel) else value.to_dict(use_alias) 

130 return output 

131 

132 

133# --- Metrics Schemas --- 

134 

135 

136class ToolMetrics(BaseModelWithConfig): 

137 """ 

138 Represents the performance and execution statistics for a tool. 

139 

140 Attributes: 

141 total_executions (int): Total number of tool invocations. 

142 successful_executions (int): Number of successful tool invocations. 

143 failed_executions (int): Number of failed tool invocations. 

144 failure_rate (float): Failure rate (failed invocations / total invocations). 

145 min_response_time (Optional[float]): Minimum response time in seconds. 

146 max_response_time (Optional[float]): Maximum response time in seconds. 

147 avg_response_time (Optional[float]): Average response time in seconds. 

148 last_execution_time (Optional[datetime]): Timestamp of the most recent invocation. 

149 """ 

150 

151 total_executions: int = Field(..., description="Total number of tool invocations") 

152 successful_executions: int = Field(..., description="Number of successful tool invocations") 

153 failed_executions: int = Field(..., description="Number of failed tool invocations") 

154 failure_rate: float = Field(..., description="Failure rate (failed invocations / total invocations)") 

155 min_response_time: Optional[float] = Field(None, description="Minimum response time in seconds") 

156 max_response_time: Optional[float] = Field(None, description="Maximum response time in seconds") 

157 avg_response_time: Optional[float] = Field(None, description="Average response time in seconds") 

158 last_execution_time: Optional[datetime] = Field(None, description="Timestamp of the most recent invocation") 

159 

160 

161class ResourceMetrics(BaseModelWithConfig): 

162 """ 

163 Represents the performance and execution statistics for a resource. 

164 

165 Attributes: 

166 total_executions (int): Total number of resource invocations. 

167 successful_executions (int): Number of successful resource invocations. 

168 failed_executions (int): Number of failed resource invocations. 

169 failure_rate (float): Failure rate (failed invocations / total invocations). 

170 min_response_time (Optional[float]): Minimum response time in seconds. 

171 max_response_time (Optional[float]): Maximum response time in seconds. 

172 avg_response_time (Optional[float]): Average response time in seconds. 

173 last_execution_time (Optional[datetime]): Timestamp of the most recent invocation. 

174 """ 

175 

176 total_executions: int = Field(..., description="Total number of resource invocations") 

177 successful_executions: int = Field(..., description="Number of successful resource invocations") 

178 failed_executions: int = Field(..., description="Number of failed resource invocations") 

179 failure_rate: float = Field(..., description="Failure rate (failed invocations / total invocations)") 

180 min_response_time: Optional[float] = Field(None, description="Minimum response time in seconds") 

181 max_response_time: Optional[float] = Field(None, description="Maximum response time in seconds") 

182 avg_response_time: Optional[float] = Field(None, description="Average response time in seconds") 

183 last_execution_time: Optional[datetime] = Field(None, description="Timestamp of the most recent invocation") 

184 

185 

186class ServerMetrics(BaseModelWithConfig): 

187 """ 

188 Represents the performance and execution statistics for a server. 

189 

190 Attributes: 

191 total_executions (int): Total number of server invocations. 

192 successful_executions (int): Number of successful server invocations. 

193 failed_executions (int): Number of failed server invocations. 

194 failure_rate (float): Failure rate (failed invocations / total invocations). 

195 min_response_time (Optional[float]): Minimum response time in seconds. 

196 max_response_time (Optional[float]): Maximum response time in seconds. 

197 avg_response_time (Optional[float]): Average response time in seconds. 

198 last_execution_time (Optional[datetime]): Timestamp of the most recent invocation. 

199 """ 

200 

201 total_executions: int = Field(..., description="Total number of server invocations") 

202 successful_executions: int = Field(..., description="Number of successful server invocations") 

203 failed_executions: int = Field(..., description="Number of failed server invocations") 

204 failure_rate: float = Field(..., description="Failure rate (failed invocations / total invocations)") 

205 min_response_time: Optional[float] = Field(None, description="Minimum response time in seconds") 

206 max_response_time: Optional[float] = Field(None, description="Maximum response time in seconds") 

207 avg_response_time: Optional[float] = Field(None, description="Average response time in seconds") 

208 last_execution_time: Optional[datetime] = Field(None, description="Timestamp of the most recent invocation") 

209 

210 

211class PromptMetrics(BaseModelWithConfig): 

212 """ 

213 Represents the performance and execution statistics for a prompt. 

214 

215 Attributes: 

216 total_executions (int): Total number of prompt invocations. 

217 successful_executions (int): Number of successful prompt invocations. 

218 failed_executions (int): Number of failed prompt invocations. 

219 failure_rate (float): Failure rate (failed invocations / total invocations). 

220 min_response_time (Optional[float]): Minimum response time in seconds. 

221 max_response_time (Optional[float]): Maximum response time in seconds. 

222 avg_response_time (Optional[float]): Average response time in seconds. 

223 last_execution_time (Optional[datetime]): Timestamp of the most recent invocation. 

224 """ 

225 

226 total_executions: int = Field(..., description="Total number of prompt invocations") 

227 successful_executions: int = Field(..., description="Number of successful prompt invocations") 

228 failed_executions: int = Field(..., description="Number of failed prompt invocations") 

229 failure_rate: float = Field(..., description="Failure rate (failed invocations / total invocations)") 

230 min_response_time: Optional[float] = Field(None, description="Minimum response time in seconds") 

231 max_response_time: Optional[float] = Field(None, description="Maximum response time in seconds") 

232 avg_response_time: Optional[float] = Field(None, description="Average response time in seconds") 

233 last_execution_time: Optional[datetime] = Field(None, description="Timestamp of the most recent invocation") 

234 

235 

236# --- JSON Path API modifier Schema 

237 

238 

239class JsonPathModifier(BaseModelWithConfig): 

240 """Schema for JSONPath queries. 

241 

242 Provides the structure for parsing JSONPath queries and optional mapping. 

243 """ 

244 

245 jsonpath: Optional[str] = Field(None, description="JSONPath expression for querying JSON data.") 

246 mapping: Optional[Dict[str, str]] = Field(None, description="Mapping of fields from original data to output.") 

247 

248 

249# --- Tool Schemas --- 

250# Authentication model 

251class AuthenticationValues(BaseModelWithConfig): 

252 """Schema for all Authentications. 

253 Provides the authentication values for different types of authentication. 

254 """ 

255 

256 auth_type: Optional[str] = Field(None, description="Type of authentication: basic, bearer, headers or None") 

257 auth_value: Optional[str] = Field(None, description="Encoded Authentication values") 

258 

259 # Only For tool read and view tool 

260 username: str = Field("", description="Username for basic authentication") 

261 password: str = Field("", description="Password for basic authentication") 

262 token: str = Field("", description="Bearer token for authentication") 

263 auth_header_key: str = Field("", description="Key for custom headers authentication") 

264 auth_header_value: str = Field("", description="Value for custom headers authentication") 

265 

266 

267class ToolCreate(BaseModelWithConfig): 

268 """Schema for creating a new tool. 

269 

270 Supports both MCP-compliant tools and REST integrations. Validates: 

271 - Unique tool name 

272 - Valid endpoint URL 

273 - Valid JSON Schema for input validation 

274 - Integration type: 'MCP' for MCP-compliant tools or 'REST' for REST integrations 

275 - Request type (For REST-GET,POST,PUT,DELETE and for MCP-SSE,STDIO,STREAMABLEHTTP) 

276 - Optional authentication credentials: BasicAuth or BearerTokenAuth or HeadersAuth. 

277 """ 

278 

279 name: str = Field(..., description="Unique name for the tool") 

280 url: Union[str, AnyHttpUrl] = Field(None, description="Tool endpoint URL") 

281 description: Optional[str] = Field(None, description="Tool description") 

282 request_type: Literal["GET", "POST", "PUT", "DELETE", "SSE", "STDIO", "STREAMABLEHTTP"] = Field("SSE", description="HTTP method to be used for invoking the tool") 

283 integration_type: Literal["MCP", "REST"] = Field("MCP", description="Tool integration type: 'MCP' for MCP-compliant tools, 'REST' for REST integrations") 

284 headers: Optional[Dict[str, str]] = Field(None, description="Additional headers to send when invoking the tool") 

285 input_schema: Optional[Dict[str, Any]] = Field( 

286 default_factory=lambda: {"type": "object", "properties": {}}, 

287 description="JSON Schema for validating tool parameters", 

288 ) 

289 jsonpath_filter: Optional[str] = Field(default="", description="JSON modification filter") 

290 auth: Optional[AuthenticationValues] = Field(None, description="Authentication credentials (Basic or Bearer Token or custom headers) if required") 

291 gateway_id: Optional[int] = Field(None, description="id of gateway for the tool") 

292 

293 @root_validator(pre=True) 

294 def assemble_auth(cls, values: Dict[str, Any]) -> Dict[str, Any]: 

295 """ 

296 Assemble authentication information from separate keys if provided. 

297 

298 Looks for keys "auth_type", "auth_username", "auth_password", "auth_token", "auth_header_key" and "auth_header_value". 

299 Constructs the "auth" field as a dictionary suitable for BasicAuth or BearerTokenAuth or HeadersAuth. 

300 

301 Args: 

302 values: Dict with authentication information 

303 

304 Returns: 

305 Dict: Reformatedd values dict 

306 """ 

307 logger.debug( 

308 "Assembling auth in ToolCreate with raw values", 

309 extra={ 

310 "auth_type": values.get("auth_type"), 

311 "auth_username": values.get("auth_username"), 

312 "auth_password": values.get("auth_password"), 

313 "auth_token": values.get("auth_token"), 

314 "auth_header_key": values.get("auth_header_key"), 

315 "auth_header_value": values.get("auth_header_value"), 

316 }, 

317 ) 

318 

319 auth_type = values.get("auth_type") 

320 if auth_type: 320 ↛ 321line 320 didn't jump to line 321 because the condition on line 320 was never true

321 if auth_type.lower() == "basic": 

322 creds = base64.b64encode(f"{values.get('auth_username', '')}:{values.get('auth_password', '')}".encode("utf-8")).decode() 

323 encoded_auth = encode_auth({"Authorization": f"Basic {creds}"}) 

324 values["auth"] = {"auth_type": "basic", "auth_value": encoded_auth} 

325 elif auth_type.lower() == "bearer": 

326 encoded_auth = encode_auth({"Authorization": f"Bearer {values.get('auth_token', '')}"}) 

327 values["auth"] = {"auth_type": "bearer", "auth_value": encoded_auth} 

328 elif auth_type.lower() == "authheaders": 

329 encoded_auth = encode_auth({values.get("auth_header_key", ""): values.get("auth_header_value", "")}) 

330 values["auth"] = {"auth_type": "authheaders", "auth_value": encoded_auth} 

331 return values 

332 

333 

334class ToolUpdate(BaseModelWithConfig): 

335 """Schema for updating an existing tool. 

336 

337 Similar to ToolCreate but all fields are optional to allow partial updates. 

338 """ 

339 

340 name: Optional[str] = Field(None, description="Unique name for the tool") 

341 url: Optional[Union[str, AnyHttpUrl]] = Field(None, description="Tool endpoint URL") 

342 description: Optional[str] = Field(None, description="Tool description") 

343 request_type: Optional[Literal["GET", "POST", "PUT", "DELETE", "SSE", "STDIO", "STREAMABLEHTTP"]] = Field(None, description="HTTP method to be used for invoking the tool") 

344 integration_type: Optional[Literal["MCP", "REST"]] = Field(None, description="Tool integration type") 

345 headers: Optional[Dict[str, str]] = Field(None, description="Additional headers to send when invoking the tool") 

346 input_schema: Optional[Dict[str, Any]] = Field(None, description="JSON Schema for validating tool parameters") 

347 jsonpath_filter: Optional[str] = Field(None, description="JSON path filter for rpc tool calls") 

348 auth: Optional[AuthenticationValues] = Field(None, description="Authentication credentials (Basic or Bearer Token or custom headers) if required") 

349 gateway_id: Optional[int] = Field(None, description="id of gateway for the tool") 

350 

351 @root_validator(pre=True) 

352 def assemble_auth(cls, values: Dict[str, Any]) -> Dict[str, Any]: 

353 """ 

354 Assemble authentication information from separate keys if provided. 

355 

356 Looks for keys "auth_type", "auth_username", "auth_password", "auth_token", "auth_header_key" and "auth_header_value". 

357 Constructs the "auth" field as a dictionary suitable for BasicAuth or BearerTokenAuth or HeadersAuth. 

358 

359 Args: 

360 values: Dict with authentication information 

361 

362 Returns: 

363 Dict: Reformatedd values dict 

364 """ 

365 

366 logger.debug( 

367 "Assembling auth in ToolCreate with raw values", 

368 extra={ 

369 "auth_type": values.get("auth_type"), 

370 "auth_username": values.get("auth_username"), 

371 "auth_password": values.get("auth_password"), 

372 "auth_token": values.get("auth_token"), 

373 "auth_header_key": values.get("auth_header_key"), 

374 "auth_header_value": values.get("auth_header_value"), 

375 }, 

376 ) 

377 

378 auth_type = values.get("auth_type") 

379 if auth_type: 379 ↛ 380line 379 didn't jump to line 380 because the condition on line 379 was never true

380 if auth_type.lower() == "basic": 

381 creds = base64.b64encode(f"{values.get('auth_username', '')}:{values.get('auth_password', '')}".encode("utf-8")).decode() 

382 encoded_auth = encode_auth({"Authorization": f"Basic {creds}"}) 

383 values["auth"] = {"auth_type": "basic", "auth_value": encoded_auth} 

384 elif auth_type.lower() == "bearer": 

385 encoded_auth = encode_auth({"Authorization": f"Bearer {values.get('auth_token', '')}"}) 

386 values["auth"] = {"auth_type": "bearer", "auth_value": encoded_auth} 

387 elif auth_type.lower() == "authheaders": 

388 encoded_auth = encode_auth({values.get("auth_header_key", ""): values.get("auth_header_value", "")}) 

389 values["auth"] = {"auth_type": "authheaders", "auth_value": encoded_auth} 

390 return values 

391 

392 

393class ToolRead(BaseModelWithConfig): 

394 """Schema for reading tool information. 

395 

396 Includes all tool fields plus: 

397 - Database ID 

398 - Creation/update timestamps 

399 - Active status 

400 - Gateway ID for federation 

401 - Execution count indicating the number of times the tool has been executed. 

402 - Metrics: Aggregated metrics for the tool invocations. 

403 - Request type and authentication settings. 

404 """ 

405 

406 id: int 

407 name: str 

408 url: Optional[str] 

409 description: Optional[str] 

410 request_type: str 

411 integration_type: str 

412 headers: Optional[Dict[str, str]] 

413 input_schema: Dict[str, Any] 

414 jsonpath_filter: Optional[str] 

415 auth: Optional[AuthenticationValues] 

416 created_at: datetime 

417 updated_at: datetime 

418 is_active: bool 

419 gateway_id: Optional[int] 

420 execution_count: int 

421 metrics: ToolMetrics 

422 

423 class Config(BaseModelWithConfig.Config): 

424 """ 

425 A configuration class that inherits from BaseModelWithConfig.Config. 

426 This class may be used to define specific configurations, extending 

427 the base functionality of BaseModelWithConfig. 

428 """ 

429 

430 

431class ToolInvocation(BaseModelWithConfig): 

432 """Schema for tool invocation requests. 

433 

434 Captures: 

435 - Tool name to invoke 

436 - Arguments matching tool's input schema 

437 """ 

438 

439 name: str = Field(..., description="Name of tool to invoke") 

440 arguments: Dict[str, Any] = Field(default_factory=dict, description="Arguments matching tool's input schema") 

441 

442 

443class ToolResult(BaseModelWithConfig): 

444 """Schema for tool invocation results. 

445 

446 Supports: 

447 - Multiple content types (text/image) 

448 - Error reporting 

449 - Optional error messages 

450 """ 

451 

452 content: List[Union[TextContent, ImageContent]] 

453 is_error: bool = False 

454 error_message: Optional[str] = None 

455 

456 

457class ResourceCreate(BaseModelWithConfig): 

458 """Schema for creating a new resource. 

459 

460 Supports: 

461 - Static resources with URI 

462 - Template resources with parameters 

463 - Both text and binary content 

464 """ 

465 

466 uri: str = Field(..., description="Unique URI for the resource") 

467 name: str = Field(..., description="Human-readable resource name") 

468 description: Optional[str] = Field(None, description="Resource description") 

469 mime_type: Optional[str] = Field(None, description="Resource MIME type") 

470 template: Optional[str] = Field(None, description="URI template for parameterized resources") 

471 content: Union[str, bytes] = Field(..., description="Resource content (text or binary)") 

472 

473 

474class ResourceUpdate(BaseModelWithConfig): 

475 """Schema for updating an existing resource. 

476 

477 Similar to ResourceCreate but URI is not required and all fields are optional. 

478 """ 

479 

480 name: Optional[str] = Field(None, description="Human-readable resource name") 

481 description: Optional[str] = Field(None, description="Resource description") 

482 mime_type: Optional[str] = Field(None, description="Resource MIME type") 

483 template: Optional[str] = Field(None, description="URI template for parameterized resources") 

484 content: Optional[Union[str, bytes]] = Field(None, description="Resource content (text or binary)") 

485 

486 

487class ResourceRead(BaseModelWithConfig): 

488 """Schema for reading resource information. 

489 

490 Includes all resource fields plus: 

491 - Database ID 

492 - Content size 

493 - Creation/update timestamps 

494 - Active status 

495 - Metrics: Aggregated metrics for the resource invocations. 

496 """ 

497 

498 id: int 

499 uri: str 

500 name: str 

501 description: Optional[str] 

502 mime_type: Optional[str] 

503 size: Optional[int] 

504 created_at: datetime 

505 updated_at: datetime 

506 is_active: bool 

507 metrics: ResourceMetrics 

508 

509 

510class ResourceSubscription(BaseModelWithConfig): 

511 """Schema for resource subscriptions. 

512 

513 Tracks: 

514 - Resource URI being subscribed to 

515 - Unique subscriber identifier 

516 """ 

517 

518 uri: str = Field(..., description="URI of resource to subscribe to") 

519 subscriber_id: str = Field(..., description="Unique subscriber identifier") 

520 

521 

522class ResourceNotification(BaseModelWithConfig): 

523 """Schema for resource update notifications. 

524 

525 Contains: 

526 - Resource URI 

527 - Updated content 

528 - Update timestamp 

529 """ 

530 

531 uri: str 

532 content: ResourceContent 

533 timestamp: datetime = Field(default_factory=datetime.utcnow) 

534 

535 

536# --- Prompt Schemas --- 

537 

538 

539class PromptArgument(BaseModelWithConfig): 

540 """Schema for prompt template arguments. 

541 

542 Defines: 

543 - Argument name 

544 - Optional description 

545 - Required flag 

546 """ 

547 

548 name: str = Field(..., description="Argument name") 

549 description: Optional[str] = Field(None, description="Argument description") 

550 required: bool = Field(default=False, description="Whether argument is required") 

551 

552 class Config(BaseModelWithConfig.Config): 

553 """ 

554 A configuration class that inherits from BaseModelWithConfig.Config. 

555 

556 This class defines an example schema for configuration, which includes: 

557 - 'name': A string representing the name of the configuration (e.g., "language"). 

558 - 'description': A brief description of the configuration (e.g., "Programming language"). 

559 - 'required': A boolean indicating if the configuration is mandatory (e.g., True). 

560 

561 The `schema_extra` attribute provides an example of how the configuration should be structured. 

562 """ 

563 

564 schema_extra = {"example": {"name": "language", "description": "Programming language", "required": True}} 

565 

566 

567class PromptCreate(BaseModelWithConfig): 

568 """Schema for creating a new prompt template. 

569 

570 Includes: 

571 - Template name and description 

572 - Template text 

573 - Expected arguments 

574 """ 

575 

576 name: str = Field(..., description="Unique name for the prompt") 

577 description: Optional[str] = Field(None, description="Prompt description") 

578 template: str = Field(..., description="Prompt template text") 

579 arguments: List[PromptArgument] = Field(default_factory=list, description="List of arguments for the template") 

580 

581 

582class PromptUpdate(BaseModelWithConfig): 

583 """Schema for updating an existing prompt. 

584 

585 Similar to PromptCreate but all fields are optional to allow partial updates. 

586 """ 

587 

588 name: Optional[str] = Field(None, description="Unique name for the prompt") 

589 description: Optional[str] = Field(None, description="Prompt description") 

590 template: Optional[str] = Field(None, description="Prompt template text") 

591 arguments: Optional[List[PromptArgument]] = Field(None, description="List of arguments for the template") 

592 

593 

594class PromptRead(BaseModelWithConfig): 

595 """Schema for reading prompt information. 

596 

597 Includes all prompt fields plus: 

598 - Database ID 

599 - Creation/update timestamps 

600 - Active status 

601 - Metrics: Aggregated metrics for the prompt invocations. 

602 """ 

603 

604 id: int 

605 name: str 

606 description: Optional[str] 

607 template: str 

608 arguments: List[PromptArgument] 

609 created_at: datetime 

610 updated_at: datetime 

611 is_active: bool 

612 metrics: PromptMetrics 

613 

614 

615class PromptInvocation(BaseModelWithConfig): 

616 """Schema for prompt invocation requests. 

617 

618 Contains: 

619 - Prompt name to use 

620 - Arguments for template rendering 

621 """ 

622 

623 name: str = Field(..., description="Name of prompt to use") 

624 arguments: Dict[str, str] = Field(default_factory=dict, description="Arguments for template rendering") 

625 

626 

627# --- Gateway Schemas --- 

628 

629 

630class GatewayCreate(BaseModelWithConfig): 

631 """Schema for registering a new federation gateway. 

632 

633 Captures: 

634 - Gateway name 

635 - Endpoint URL 

636 - Optional description 

637 - Authentication type: basic, bearer, headers 

638 - Authentication credentials: username/password or token or custom headers 

639 """ 

640 

641 name: str = Field(..., description="Unique name for the gateway") 

642 url: Union[str, AnyHttpUrl] = Field(..., description="Gateway endpoint URL") 

643 description: Optional[str] = Field(None, description="Gateway description") 

644 transport: str = Field(default="SSE", description="Transport used by MCP server: SSE or STREAMABLEHTTP") 

645 

646 # Authorizations 

647 auth_type: Optional[str] = Field(None, description="Type of authentication: basic, bearer, headers, or none") 

648 # Fields for various types of authentication 

649 auth_username: Optional[str] = Field(None, description="Username for basic authentication") 

650 auth_password: Optional[str] = Field(None, description="Password for basic authentication") 

651 auth_token: Optional[str] = Field(None, description="Token for bearer authentication") 

652 auth_header_key: Optional[str] = Field(None, description="Key for custom headers authentication") 

653 auth_header_value: Optional[str] = Field(None, description="Value for custom headers authentication") 

654 

655 # Adding `auth_value` as an alias for better access post-validation 

656 auth_value: Optional[str] = None 

657 

658 @validator("url", pre=True) 

659 def ensure_url_scheme(cls, v: str) -> str: 

660 """ 

661 Ensure URL has an http/https scheme. 

662 

663 Args: 

664 v: Input url 

665 

666 Returns: 

667 str: URL with correct schema 

668 

669 """ 

670 if isinstance(v, str) and not (v.startswith("http://") or v.startswith("https://")): 670 ↛ 671line 670 didn't jump to line 671 because the condition on line 670 was never true

671 return f"http://{v}" 

672 return v 

673 

674 @validator("auth_value", pre=True, always=True) 

675 def create_auth_value(cls, v, values): 

676 """ 

677 This validator will run before the model is fully instantiated (pre=True) 

678 It will process the auth fields based on auth_type and generate auth_value. 

679 

680 Args: 

681 v: Input url 

682 values: Dict containing auth_type 

683 

684 Returns: 

685 str: Auth value 

686 """ 

687 auth_type = values.get("auth_type") 

688 

689 if (auth_type is None) or (auth_type == ""): 689 ↛ 693line 689 didn't jump to line 693 because the condition on line 689 was always true

690 return v # If no auth_type is provided, no need to create auth_value 

691 

692 # Process the auth fields and generate auth_value based on auth_type 

693 auth_value = cls._process_auth_fields(values) 

694 

695 return auth_value 

696 

697 @staticmethod 

698 def _process_auth_fields(values: Dict[str, Any]) -> Optional[Dict[str, Any]]: 

699 """ 

700 Processes the input authentication fields and returns the correct auth_value. 

701 This method is called based on the selected auth_type. 

702 

703 Args: 

704 values: Dict containing auth fields 

705 

706 Returns: 

707 Dict with encoded auth 

708 

709 Raises: 

710 ValueError: If auth_type is invalid 

711 """ 

712 auth_type = values.get("auth_type") 

713 

714 if auth_type == "basic": 

715 # For basic authentication, both username and password must be present 

716 username = values.get("auth_username") 

717 password = values.get("auth_password") 

718 

719 if not username or not password: 

720 raise ValueError("For 'basic' auth, both 'auth_username' and 'auth_password' must be provided.") 

721 

722 creds = base64.b64encode(f"{username}:{password}".encode("utf-8")).decode() 

723 return encode_auth({"Authorization": f"Basic {creds}"}) 

724 

725 if auth_type == "bearer": 

726 # For bearer authentication, only token is required 

727 token = values.get("auth_token") 

728 

729 if not token: 

730 raise ValueError("For 'bearer' auth, 'auth_token' must be provided.") 

731 

732 return encode_auth({"Authorization": f"Bearer {token}"}) 

733 

734 if auth_type == "authheaders": 

735 # For headers authentication, both key and value must be present 

736 header_key = values.get("auth_header_key") 

737 header_value = values.get("auth_header_value") 

738 

739 if not header_key or not header_value: 

740 raise ValueError("For 'headers' auth, both 'auth_header_key' and 'auth_header_value' must be provided.") 

741 

742 return encode_auth({header_key: header_value}) 

743 

744 raise ValueError("Invalid 'auth_type'. Must be one of: basic, bearer, or headers.") 

745 

746 

747class GatewayUpdate(BaseModelWithConfig): 

748 """Schema for updating an existing federation gateway. 

749 

750 Similar to GatewayCreate but all fields are optional to allow partial updates. 

751 """ 

752 

753 name: Optional[str] = Field(None, description="Unique name for the gateway") 

754 url: Optional[Union[str, AnyHttpUrl]] = Field(None, description="Gateway endpoint URL") 

755 description: Optional[str] = Field(None, description="Gateway description") 

756 transport: str = Field(default="SSE", description="Transport used by MCP server: SSE or STREAMABLEHTTP") 

757 

758 name: Optional[str] = Field(None, description="Unique name for the prompt") 

759 # Authorizations 

760 auth_type: Optional[str] = Field(None, description="auth_type: basic, bearer, headers or None") 

761 auth_username: Optional[str] = Field(None, description="username for basic authentication") 

762 auth_password: Optional[str] = Field(None, description="password for basic authentication") 

763 auth_token: Optional[str] = Field(None, description="token for bearer authentication") 

764 auth_header_key: Optional[str] = Field(None, description="key for custom headers authentication") 

765 auth_header_value: Optional[str] = Field(None, description="vallue for custom headers authentication") 

766 

767 # Adding `auth_value` as an alias for better access post-validation 

768 auth_value: Optional[str] = None 

769 

770 @validator("url", pre=True) 

771 def ensure_url_scheme(cls, v: Optional[str]) -> Optional[str]: 

772 """ 

773 Ensure URL has an http/https scheme. 

774 

775 Args: 

776 v: Input URL 

777 

778 Returns: 

779 str: Validated URL 

780 """ 

781 if isinstance(v, str) and not (v.startswith("http://") or v.startswith("https://")): 781 ↛ 782line 781 didn't jump to line 782 because the condition on line 781 was never true

782 return f"http://{v}" 

783 return v 

784 

785 @validator("auth_value", pre=True, always=True) 

786 def create_auth_value(cls, v, values): 

787 """ 

788 This validator will run before the model is fully instantiated (pre=True) 

789 It will process the auth fields based on auth_type and generate auth_value. 

790 

791 Args: 

792 v: Input URL 

793 values: Dict containing auth_type 

794 

795 Returns: 

796 str: Auth value or URL 

797 """ 

798 auth_type = values.get("auth_type") 

799 

800 if (auth_type is None) or (auth_type == ""): 800 ↛ 804line 800 didn't jump to line 804 because the condition on line 800 was always true

801 return v # If no auth_type is provided, no need to create auth_value 

802 

803 # Process the auth fields and generate auth_value based on auth_type 

804 auth_value = cls._process_auth_fields(values) 

805 

806 return auth_value 

807 

808 @staticmethod 

809 def _process_auth_fields(values: Dict[str, Any]) -> Optional[Dict[str, Any]]: 

810 """ 

811 Processes the input authentication fields and returns the correct auth_value. 

812 This method is called based on the selected auth_type. 

813 

814 Args: 

815 values: Dict container auth information auth_type, auth_username, auth_password, auth_token, auth_header_key and auth_header_value 

816 

817 Returns: 

818 dict: Encoded auth information 

819 

820 Raises: 

821 ValueError: If auth type is invalid 

822 """ 

823 auth_type = values.get("auth_type") 

824 

825 if auth_type == "basic": 

826 # For basic authentication, both username and password must be present 

827 username = values.get("auth_username") 

828 password = values.get("auth_password") 

829 

830 if not username or not password: 

831 raise ValueError("For 'basic' auth, both 'auth_username' and 'auth_password' must be provided.") 

832 

833 creds = base64.b64encode(f"{username}:{password}".encode("utf-8")).decode() 

834 return encode_auth({"Authorization": f"Basic {creds}"}) 

835 

836 if auth_type == "bearer": 

837 # For bearer authentication, only token is required 

838 token = values.get("auth_token") 

839 

840 if not token: 

841 raise ValueError("For 'bearer' auth, 'auth_token' must be provided.") 

842 

843 return encode_auth({"Authorization": f"Bearer {token}"}) 

844 

845 if auth_type == "authheaders": 

846 # For headers authentication, both key and value must be present 

847 header_key = values.get("auth_header_key") 

848 header_value = values.get("auth_header_value") 

849 

850 if not header_key or not header_value: 

851 raise ValueError("For 'headers' auth, both 'auth_header_key' and 'auth_header_value' must be provided.") 

852 

853 return encode_auth({header_key: header_value}) 

854 

855 raise ValueError("Invalid 'auth_type'. Must be one of: basic, bearer, or headers.") 

856 

857 

858class GatewayRead(BaseModelWithConfig): 

859 """Schema for reading gateway information. 

860 

861 Includes all gateway fields plus: 

862 - Database ID 

863 - Capabilities dictionary 

864 - Creation/update timestamps 

865 - Active status 

866 - Last seen timestamp 

867 - Authentication type: basic, bearer, headers 

868 - Authentication value: username/password or token or custom headers 

869 

870 Auto Populated fields: 

871 - Authentication username: for basic auth 

872 - Authentication password: for basic auth 

873 - Authentication token: for bearer auth 

874 - Authentication header key: for headers auth 

875 - Authentication header value: for headers auth 

876 """ 

877 

878 id: int = Field(None, description="Unique ID of the gateway") 

879 name: str = Field(..., description="Unique name for the gateway") 

880 url: str = Field(..., description="Gateway endpoint URL") 

881 description: Optional[str] = Field(None, description="Gateway description") 

882 transport: str = Field(default="SSE", description="Transport used by MCP server: SSE or STREAMABLEHTTP") 

883 capabilities: Dict[str, Any] = Field(default_factory=dict, description="Gateway capabilities") 

884 created_at: datetime = Field(default_factory=datetime.utcnow, description="Creation timestamp") 

885 updated_at: datetime = Field(default_factory=datetime.utcnow, description="Last update timestamp") 

886 is_active: bool = Field(default=True, description="Is the gateway active?") 

887 last_seen: Optional[datetime] = Field(default_factory=datetime.utcnow, description="Last seen timestamp") 

888 

889 # Authorizations 

890 auth_type: Optional[str] = Field(None, description="auth_type: basic, bearer, headers or None") 

891 auth_value: Optional[str] = Field(None, description="auth value: username/password or token or custom headers") 

892 

893 # auth_value will populate the following fields 

894 auth_username: Optional[str] = Field(None, description="username for basic authentication") 

895 auth_password: Optional[str] = Field(None, description="password for basic authentication") 

896 auth_token: Optional[str] = Field(None, description="token for bearer authentication") 

897 auth_header_key: Optional[str] = Field(None, description="key for custom headers authentication") 

898 auth_header_value: Optional[str] = Field(None, description="vallue for custom headers authentication") 

899 

900 # This will be the main method to automatically populate fields 

901 @model_validator(mode="after") 

902 def _populate_auth(cls, values: Dict[str, Any]) -> Dict[str, Any]: 

903 auth_type = values.auth_type 

904 auth_value_encoded = values.auth_value 

905 auth_value = decode_auth(auth_value_encoded) 

906 if auth_type == "basic": 

907 u = auth_value.get("username") 

908 p = auth_value.get("password") 

909 if not u or not p: 

910 raise ValueError("basic auth requires both username and password") 

911 values.auth_username, values.auth_password = u, p 

912 

913 elif auth_type == "bearer": 

914 auth = auth_value.get("Authorization") 

915 if not (isinstance(auth, str) and auth.startswith("Bearer ")): 

916 raise ValueError("bearer auth requires an Authorization header of the form 'Bearer <token>'") 

917 values.auth_token = auth.removeprefix("Bearer ") 

918 

919 elif auth_type == "authheaders": 

920 # must be exactly one header 

921 if len(auth_value) != 1: 

922 raise ValueError("authheaders requires exactly one key/value pair") 

923 k, v = next(iter(auth_value.items())) 

924 values.auth_header_key, values.auth_header_value = k, v 

925 

926 return values 

927 

928 

929class FederatedTool(BaseModelWithConfig): 

930 """Schema for tools provided by federated gateways. 

931 

932 Contains: 

933 - Tool definition 

934 - Source gateway information 

935 """ 

936 

937 tool: MCPTool 

938 gateway_id: str 

939 gateway_name: str 

940 gateway_url: str 

941 

942 

943class FederatedResource(BaseModelWithConfig): 

944 """Schema for resources from federated gateways. 

945 

946 Contains: 

947 - Resource definition 

948 - Source gateway information 

949 """ 

950 

951 resource: MCPResource 

952 gateway_id: str 

953 gateway_name: str 

954 gateway_url: str 

955 

956 

957class FederatedPrompt(BaseModelWithConfig): 

958 """Schema for prompts from federated gateways. 

959 

960 Contains: 

961 - Prompt definition 

962 - Source gateway information 

963 """ 

964 

965 prompt: MCPPrompt 

966 gateway_id: str 

967 gateway_name: str 

968 gateway_url: str 

969 

970 

971# --- RPC Schemas --- 

972 

973 

974class RPCRequest(BaseModelWithConfig): 

975 """Schema for JSON-RPC 2.0 requests. 

976 

977 Validates: 

978 - Protocol version 

979 - Method name 

980 - Optional parameters 

981 - Optional request ID 

982 """ 

983 

984 jsonrpc: Literal["2.0"] 

985 method: str 

986 params: Optional[Dict[str, Any]] = None 

987 id: Optional[Union[int, str]] = None 

988 

989 

990class RPCResponse(BaseModelWithConfig): 

991 """Schema for JSON-RPC 2.0 responses. 

992 

993 Contains: 

994 - Protocol version 

995 - Result or error 

996 - Request ID 

997 """ 

998 

999 jsonrpc: Literal["2.0"] 

1000 result: Optional[Any] = None 

1001 error: Optional[Dict[str, Any]] = None 

1002 id: Optional[Union[int, str]] = None 

1003 

1004 

1005# --- Event and Admin Schemas --- 

1006 

1007 

1008class EventMessage(BaseModelWithConfig): 

1009 """Schema for SSE event messages. 

1010 

1011 Includes: 

1012 - Event type 

1013 - Event data payload 

1014 - Event timestamp 

1015 """ 

1016 

1017 type: str = Field(..., description="Event type (tool_added, resource_updated, etc)") 

1018 data: Dict[str, Any] = Field(..., description="Event payload") 

1019 timestamp: datetime = Field(default_factory=datetime.utcnow) 

1020 

1021 

1022class AdminToolCreate(BaseModelWithConfig): 

1023 """Schema for creating tools via admin UI. 

1024 

1025 Handles: 

1026 - Basic tool information 

1027 - JSON string inputs for headers/schema 

1028 """ 

1029 

1030 name: str 

1031 url: str 

1032 description: Optional[str] = None 

1033 integration_type: str = "MCP" 

1034 headers: Optional[str] = None # JSON string 

1035 input_schema: Optional[str] = None # JSON string 

1036 

1037 @validator("headers", "input_schema") 

1038 def validate_json(cls, v: Optional[str]) -> Optional[Dict[str, Any]]: 

1039 """ 

1040 Validate and parse JSON string inputs. 

1041 

1042 Args: 

1043 v: Input string 

1044 

1045 Returns: 

1046 dict: Output JSON version of v 

1047 

1048 Raises: 

1049 ValueError: When unable to convert to JSON 

1050 """ 

1051 if not v: 1051 ↛ 1052line 1051 didn't jump to line 1052 because the condition on line 1051 was never true

1052 return None 

1053 try: 

1054 return json.loads(v) 

1055 except json.JSONDecodeError: 

1056 raise ValueError("Invalid JSON") 

1057 

1058 

1059class AdminGatewayCreate(BaseModelWithConfig): 

1060 """Schema for creating gateways via admin UI. 

1061 

1062 Captures: 

1063 - Gateway name 

1064 - Endpoint URL 

1065 - Optional description 

1066 """ 

1067 

1068 name: str 

1069 url: str 

1070 description: Optional[str] = None 

1071 

1072 

1073# --- New Schemas for Status Toggle Operations --- 

1074 

1075 

1076class StatusToggleRequest(BaseModelWithConfig): 

1077 """Request schema for toggling active status.""" 

1078 

1079 activate: bool = Field(..., description="Whether to activate (true) or deactivate (false) the item") 

1080 

1081 

1082class StatusToggleResponse(BaseModelWithConfig): 

1083 """Response schema for status toggle operations.""" 

1084 

1085 id: int 

1086 name: str 

1087 is_active: bool 

1088 message: str = Field(..., description="Success message") 

1089 

1090 

1091# --- Optional Filter Parameters for Listing Operations --- 

1092 

1093 

1094class ListFilters(BaseModelWithConfig): 

1095 """Filtering options for list operations.""" 

1096 

1097 include_inactive: bool = Field(False, description="Whether to include inactive items in the results") 

1098 

1099 

1100# --- Server Schemas --- 

1101 

1102 

1103class ServerCreate(BaseModelWithConfig): 

1104 """Schema for creating a new server. 

1105 

1106 Attributes: 

1107 name: The server's name (required). 

1108 description: Optional text description. 

1109 icon: Optional URL for the server's icon. 

1110 associated_tools: Optional list of tool IDs (as strings). 

1111 associated_resources: Optional list of resource IDs (as strings). 

1112 associated_prompts: Optional list of prompt IDs (as strings). 

1113 """ 

1114 

1115 name: str = Field(..., description="The server's name") 

1116 description: Optional[str] = Field(None, description="Server description") 

1117 icon: Optional[str] = Field(None, description="URL for the server's icon") 

1118 associated_tools: Optional[List[str]] = Field(None, description="Comma-separated tool IDs") 

1119 associated_resources: Optional[List[str]] = Field(None, description="Comma-separated resource IDs") 

1120 associated_prompts: Optional[List[str]] = Field(None, description="Comma-separated prompt IDs") 

1121 

1122 @validator("associated_tools", "associated_resources", "associated_prompts", pre=True) 

1123 def split_comma_separated(cls, v): 

1124 """ 

1125 Splits a comma-separated string into a list of strings if needed. 

1126 

1127 Args: 

1128 v: Input string 

1129 

1130 Returns: 

1131 list: Comma separated array of input string 

1132 """ 

1133 if isinstance(v, str): 

1134 return [item.strip() for item in v.split(",") if item.strip()] 

1135 return v 

1136 

1137 

1138class ServerUpdate(BaseModelWithConfig): 

1139 """Schema for updating an existing server. 

1140 

1141 All fields are optional to allow partial updates. 

1142 """ 

1143 

1144 name: Optional[str] = Field(None, description="The server's name") 

1145 description: Optional[str] = Field(None, description="Server description") 

1146 icon: Optional[str] = Field(None, description="URL for the server's icon") 

1147 associated_tools: Optional[List[str]] = Field(None, description="Comma-separated tool IDs") 

1148 associated_resources: Optional[List[str]] = Field(None, description="Comma-separated resource IDs") 

1149 associated_prompts: Optional[List[str]] = Field(None, description="Comma-separated prompt IDs") 

1150 

1151 @validator("associated_tools", "associated_resources", "associated_prompts", pre=True) 

1152 def split_comma_separated(cls, v): 

1153 """ 

1154 Splits a comma-separated string into a list of strings if needed. 

1155 

1156 Args: 

1157 v: Input string 

1158 

1159 Returns: 

1160 list: Comma separated array of input string 

1161 """ 

1162 if isinstance(v, str): 

1163 return [item.strip() for item in v.split(",") if item.strip()] 

1164 return v 

1165 

1166 

1167class ServerRead(BaseModelWithConfig): 

1168 """Schema for reading server information. 

1169 

1170 Includes all server fields plus: 

1171 - Database ID 

1172 - Associated tool, resource, and prompt IDs 

1173 - Creation/update timestamps 

1174 - Active status 

1175 - Metrics: Aggregated metrics for the server invocations. 

1176 """ 

1177 

1178 id: int 

1179 name: str 

1180 description: Optional[str] 

1181 icon: Optional[str] 

1182 created_at: datetime 

1183 updated_at: datetime 

1184 is_active: bool 

1185 associated_tools: List[int] = [] 

1186 associated_resources: List[int] = [] 

1187 associated_prompts: List[int] = [] 

1188 metrics: ServerMetrics 

1189 

1190 @root_validator(pre=True) 

1191 def populate_associated_ids(cls, values): 

1192 """ 

1193 Pre-validation method that converts associated objects to their 'id'. 

1194 

1195 This method checks 'associated_tools', 'associated_resources', and 

1196 'associated_prompts' in the input and replaces each object with its `id` 

1197 if present. 

1198 

1199 Args: 

1200 values (dict): The input values. 

1201 

1202 Returns: 

1203 dict: Updated values with object ids, or the original values if no 

1204 changes are made. 

1205 """ 

1206 # If values is not a dict (e.g. it's a Server instance), convert it 

1207 if not isinstance(values, dict): 1207 ↛ 1208line 1207 didn't jump to line 1208 because the condition on line 1207 was never true

1208 try: 

1209 values = vars(values) 

1210 except Exception: 

1211 return values 

1212 if "associated_tools" in values and values["associated_tools"]: 1212 ↛ 1214line 1212 didn't jump to line 1214 because the condition on line 1212 was always true

1213 values["associated_tools"] = [tool.id if hasattr(tool, "id") else tool for tool in values["associated_tools"]] 

1214 if "associated_resources" in values and values["associated_resources"]: 1214 ↛ 1216line 1214 didn't jump to line 1216 because the condition on line 1214 was always true

1215 values["associated_resources"] = [res.id if hasattr(res, "id") else res for res in values["associated_resources"]] 

1216 if "associated_prompts" in values and values["associated_prompts"]: 1216 ↛ 1218line 1216 didn't jump to line 1218 because the condition on line 1216 was always true

1217 values["associated_prompts"] = [prompt.id if hasattr(prompt, "id") else prompt for prompt in values["associated_prompts"]] 

1218 return values