Gemini MCP Server

""" Schema builders for workflow MCP tools. This module provides workflow-specific schema generation functionality, keeping workflow concerns separated from simple tool concerns. """ from typing import Any from ..shared.base_models import WORKFLOW_FIELD_DESCRIPTIONS from ..shared.schema_builders import SchemaBuilder class WorkflowSchemaBuilder: """ Schema builder for workflow MCP tools. This class extends the base SchemaBuilder with workflow-specific fields and schema generation logic, maintaining separation of concerns. """ # Workflow-specific field schemas WORKFLOW_FIELD_SCHEMAS = { "step": { "type": "string", "description": WORKFLOW_FIELD_DESCRIPTIONS["step"], }, "step_number": { "type": "integer", "minimum": 1, "description": WORKFLOW_FIELD_DESCRIPTIONS["step_number"], }, "total_steps": { "type": "integer", "minimum": 1, "description": WORKFLOW_FIELD_DESCRIPTIONS["total_steps"], }, "next_step_required": { "type": "boolean", "description": WORKFLOW_FIELD_DESCRIPTIONS["next_step_required"], }, "findings": { "type": "string", "description": WORKFLOW_FIELD_DESCRIPTIONS["findings"], }, "files_checked": { "type": "array", "items": {"type": "string"}, "description": WORKFLOW_FIELD_DESCRIPTIONS["files_checked"], }, "relevant_files": { "type": "array", "items": {"type": "string"}, "description": WORKFLOW_FIELD_DESCRIPTIONS["relevant_files"], }, "relevant_context": { "type": "array", "items": {"type": "string"}, "description": WORKFLOW_FIELD_DESCRIPTIONS["relevant_context"], }, "issues_found": { "type": "array", "items": {"type": "object"}, "description": WORKFLOW_FIELD_DESCRIPTIONS["issues_found"], }, "confidence": { "type": "string", "enum": ["exploring", "low", "medium", "high", "very_high", "almost_certain", "certain"], "description": WORKFLOW_FIELD_DESCRIPTIONS["confidence"], }, "hypothesis": { "type": "string", "description": WORKFLOW_FIELD_DESCRIPTIONS["hypothesis"], }, "backtrack_from_step": { "type": "integer", "minimum": 1, "description": WORKFLOW_FIELD_DESCRIPTIONS["backtrack_from_step"], }, "use_assistant_model": { "type": "boolean", "default": True, "description": WORKFLOW_FIELD_DESCRIPTIONS["use_assistant_model"], }, } @staticmethod def build_schema( tool_specific_fields: dict[str, dict[str, Any]] = None, required_fields: list[str] = None, model_field_schema: dict[str, Any] = None, auto_mode: bool = False, tool_name: str = None, excluded_workflow_fields: list[str] = None, excluded_common_fields: list[str] = None, require_model: bool = False, ) -> dict[str, Any]: """ Build complete schema for workflow tools. Args: tool_specific_fields: Additional fields specific to the tool required_fields: List of required field names (beyond workflow defaults) model_field_schema: Schema for the model field auto_mode: Whether the tool is in auto mode (affects model requirement) tool_name: Name of the tool (for schema title) excluded_workflow_fields: Workflow fields to exclude from schema (e.g., for planning tools) excluded_common_fields: Common fields to exclude from schema Returns: Complete JSON schema for the workflow tool """ properties = {} # Add workflow fields first, excluding any specified fields workflow_fields = WorkflowSchemaBuilder.WORKFLOW_FIELD_SCHEMAS.copy() if excluded_workflow_fields: for field in excluded_workflow_fields: workflow_fields.pop(field, None) properties.update(workflow_fields) # Add common fields (temperature, thinking_mode, etc.) from base builder, excluding any specified fields common_fields = SchemaBuilder.COMMON_FIELD_SCHEMAS.copy() if excluded_common_fields: for field in excluded_common_fields: common_fields.pop(field, None) properties.update(common_fields) # Add model field if provided if model_field_schema: properties["model"] = model_field_schema # Add tool-specific fields if provided if tool_specific_fields: properties.update(tool_specific_fields) # Build required fields list - workflow tools have standard required fields standard_required = ["step", "step_number", "total_steps", "next_step_required", "findings"] # Filter out excluded fields from required fields if excluded_workflow_fields: standard_required = [field for field in standard_required if field not in excluded_workflow_fields] required = standard_required + (required_fields or []) if (auto_mode or require_model) and "model" not in required: required.append("model") # Build the complete schema schema = { "$schema": "http://json-schema.org/draft-07/schema#", "type": "object", "properties": properties, "required": required, "additionalProperties": False, } if tool_name: schema["title"] = f"{tool_name.capitalize()}Request" return schema @staticmethod def get_workflow_fields() -> dict[str, dict[str, Any]]: """Get the standard field schemas for workflow tools.""" combined = {} combined.update(WorkflowSchemaBuilder.WORKFLOW_FIELD_SCHEMAS) combined.update(SchemaBuilder.COMMON_FIELD_SCHEMAS) return combined @staticmethod def get_workflow_only_fields() -> dict[str, dict[str, Any]]: """Get only the workflow-specific field schemas.""" return WorkflowSchemaBuilder.WORKFLOW_FIELD_SCHEMAS.copy()