From 86c26fd64c69f481312d84796ba9fa3a7370c6ab Mon Sep 17 00:00:00 2001 From: Yousif Astarabadi <6870090+yousifa@users.noreply.github.com> Date: Thu, 26 Jun 2025 19:57:58 -0700 Subject: [PATCH] moved needs_mcp_clean_schema to LLMService --- .../services/gemini_multimodal_live/gemini.py | 11 ++++ src/pipecat/services/google/llm.py | 11 ++++ src/pipecat/services/llm_service.py | 11 ++++ src/pipecat/services/mcp_service.py | 54 ++++++++++++++++--- 4 files changed, 80 insertions(+), 7 deletions(-) diff --git a/src/pipecat/services/gemini_multimodal_live/gemini.py b/src/pipecat/services/gemini_multimodal_live/gemini.py index afaf966dc..f9584df9d 100644 --- a/src/pipecat/services/gemini_multimodal_live/gemini.py +++ b/src/pipecat/services/gemini_multimodal_live/gemini.py @@ -524,6 +524,17 @@ class GeminiMultimodalLiveLLMService(LLMService): """ return True + def needs_mcp_clean_schema(self) -> bool: + """Check if this LLM service requires MCP schema cleaning. + + Google/Gemini has stricter JSON schema validation and requires + certain properties to be removed or modified for compatibility. + + Returns: + True for Google/Gemini services. + """ + return True + def set_audio_input_paused(self, paused: bool): """Set the audio input pause state. diff --git a/src/pipecat/services/google/llm.py b/src/pipecat/services/google/llm.py index 6b8f51f33..e36c6591e 100644 --- a/src/pipecat/services/google/llm.py +++ b/src/pipecat/services/google/llm.py @@ -631,6 +631,17 @@ class GoogleLLMService(LLMService): """ return True + def needs_mcp_clean_schema(self) -> bool: + """Check if this LLM service requires MCP schema cleaning. + + Google/Gemini has stricter JSON schema validation and requires + certain properties to be removed or modified for compatibility. + + Returns: + True for Google/Gemini services. + """ + return True + def _create_client(self, api_key: str): self._client = genai.Client(api_key=api_key) diff --git a/src/pipecat/services/llm_service.py b/src/pipecat/services/llm_service.py index f7779df98..24ed932d8 100644 --- a/src/pipecat/services/llm_service.py +++ b/src/pipecat/services/llm_service.py @@ -307,6 +307,17 @@ class LLMService(AIService): return True return function_name in self._functions.keys() + def needs_mcp_clean_schema(self) -> bool: + """Check if this LLM service requires MCP schema cleaning. + + Some LLM services have stricter JSON schema validation and require + certain properties to be removed or modified for compatibility. + + Returns: + True if MCP schemas should be cleaned for this service, False otherwise. + """ + return False + async def run_function_calls(self, function_calls: Sequence[FunctionCallFromLLM]): """Execute a sequence of function calls from the LLM. diff --git a/src/pipecat/services/mcp_service.py b/src/pipecat/services/mcp_service.py index c51a9bd9c..3d757e32e 100644 --- a/src/pipecat/services/mcp_service.py +++ b/src/pipecat/services/mcp_service.py @@ -51,6 +51,7 @@ class MCPClient(BaseObject): super().__init__(**kwargs) self._server_params = server_params self._session = ClientSession + self._needs_schema_cleaning = False if isinstance(server_params, StdioServerParameters): self._client = stdio_client @@ -78,18 +79,56 @@ class MCPClient(BaseObject): Returns: A ToolsSchema containing all successfully registered tools. """ + # Check once if the LLM needs schema cleaning + self._needs_schema_cleaning = llm and llm.needs_mcp_clean_schema() tools_schema = await self._register_tools(llm) return tools_schema + def _clean_schema_for_strict_validation(self, schema: Dict[str, Any]) -> Dict[str, Any]: + """Clean a JSON schema to be compatible with LLMs that have strict validation. + + Some LLMs have stricter validation and don't allow certain schema properties + that are valid in standard JSON Schema. + + Args: + schema: The JSON schema to clean + + Returns: + A cleaned schema compatible with strict validation + """ + if not isinstance(schema, dict): + return schema + + cleaned = {} + + for key, value in schema.items(): + # Skip additionalProperties as some LLMs don't like additionalProperties: false + if key == "additionalProperties": + continue + + # Recursively clean nested objects + if isinstance(value, dict): + cleaned[key] = self._clean_schema_for_strict_validation(value) + elif isinstance(value, list): + cleaned[key] = [ + self._clean_schema_for_strict_validation(item) + if isinstance(item, dict) + else item + for item in value + ] + else: + cleaned[key] = value + + return cleaned + def _convert_mcp_schema_to_pipecat( - self, tool_name: str, tool_schema: Dict[str, Any], llm=None + self, tool_name: str, tool_schema: Dict[str, Any] ) -> FunctionSchema: """Convert an mcp tool schema to Pipecat's FunctionSchema format. Args: tool_name: The name of the tool tool_schema: The mcp tool schema - llm: The LLM service instance (used to determine if we need Gemini compatibility) Returns: A FunctionSchema instance """ @@ -99,10 +138,10 @@ class MCPClient(BaseObject): properties = tool_schema["input_schema"].get("properties", {}) required = tool_schema["input_schema"].get("required", []) - # Only clean properties for Google/Gemini LLM services - if llm and self._is_google_llm(llm): - logger.debug(f"Detected Google LLM service, cleaning schema for Gemini compatibility") - properties = self._clean_schema_for_gemini(properties) + # Only clean properties for LLMs that need strict schema validation + if self._needs_schema_cleaning: + logger.debug("Cleaning schema for strict validation") + properties = self._clean_schema_for_strict_validation(properties) schema = FunctionSchema( name=tool_name, @@ -283,7 +322,8 @@ class MCPClient(BaseObject): try: # Convert the schema function_schema = self._convert_mcp_schema_to_pipecat( - tool_name, {"description": tool.description, "input_schema": tool.inputSchema} + tool_name, + {"description": tool.description, "input_schema": tool.inputSchema}, ) # Register the wrapped function