Added google tool calling with gemini 2.0 Flash

src/airflow_wingman/tools/conversion.py

@@ -66,7 +66,7 @@ def convert_to_openai_tools(airflow_tools: list) -> list:
                     # Add default value if available
                     if "default" in param_info and param_info["default"] is not None:
                         param_def["default"] = param_info["default"]
-                        
+
                     # Add items property for array types
                     if param_def.get("type") == "array" and "items" not in param_def:
                         # If items is defined in the original schema, use it
@@ -116,6 +116,84 @@ def convert_to_anthropic_tools(airflow_tools: list) -> list:
     return anthropic_tools
 
 
+def convert_to_google_tools(airflow_tools: list) -> list:
+    """
+    Convert Airflow tools to Google Gemini format.
+
+    Args:
+        airflow_tools: List of Airflow tools from MCP server
+
+    Returns:
+        List of Google Gemini tool definitions wrapped in correct SDK structure
+    """
+    logger = logging.getLogger("airflow.plugins.wingman")
+    logger.info(f"Converting {len(airflow_tools)} Airflow tools to Google Gemini format")
+
+    # This will hold our list of function declarations
+    function_declarations = []
+
+    for tool in airflow_tools:
+        # Create function declaration for Google's format
+        function_declaration = {
+            "name": tool.name if hasattr(tool, "name") else str(tool),
+            "description": tool.description if hasattr(tool, "description") else "",
+            "parameters": {"type": "object", "properties": {}, "required": []},
+        }
+
+        # Extract parameters from inputSchema if available
+        if hasattr(tool, "inputSchema") and tool.inputSchema:
+            # Add required parameters if specified
+            if "required" in tool.inputSchema:
+                function_declaration["parameters"]["required"] = tool.inputSchema["required"]
+
+            # Add properties from the input schema
+            if "properties" in tool.inputSchema:
+                for param_name, param_info in tool.inputSchema["properties"].items():
+                    # Create parameter definition
+                    param_def = {}
+
+                    # Handle different schema constructs
+                    if "anyOf" in param_info:
+                        _handle_schema_construct(param_def, param_info, "anyOf")
+                    elif "oneOf" in param_info:
+                        _handle_schema_construct(param_def, param_info, "oneOf")
+                    elif "allOf" in param_info:
+                        _handle_schema_construct(param_def, param_info, "allOf")
+                    elif "type" in param_info:
+                        param_def["type"] = param_info["type"]
+                        # Add format if available
+                        if "format" in param_info:
+                            param_def["format"] = param_info["format"]
+                    else:
+                        param_def["type"] = "string"  # Default type
+
+                    # Add description
+                    param_def["description"] = param_info.get("description", param_info.get("title", param_name))
+
+                    # Add enum values if available
+                    if "enum" in param_info:
+                        param_def["enum"] = param_info["enum"]
+
+                    # Add items property for array types
+                    if param_def.get("type") == "array" and "items" not in param_def:
+                        if "items" in param_info:
+                            param_def["items"] = param_info["items"]
+                        else:
+                            param_def["items"] = {"type": "string"}
+
+                    # Add to properties
+                    function_declaration["parameters"]["properties"][param_name] = param_def
+
+        function_declarations.append(function_declaration)
+
+    # For Google API, we need to wrap the function declarations in a specific structure
+    # The correct structure is [{'function_declarations': [func1, func2, ...]}]
+    google_tools = [{"function_declarations": function_declarations}]
+
+    logger.info(f"Converted {len(function_declarations)} tools to Google Gemini format with correct SDK structure")
+    return google_tools
+
+
 def _handle_schema_construct(param_def: dict[str, Any], param_info: dict[str, Any], construct_type: str) -> None:
     """
     Helper function to handle JSON Schema constructs like anyOf, oneOf, allOf.
@@ -150,7 +228,7 @@ def _handle_schema_construct(param_def: dict[str, Any], param_info: dict[str, An
     # If no type was found, default to string
     if "type" not in param_def:
         param_def["type"] = "string"
-        
+
     # Add items property for array types
     if param_def.get("type") == "array" and "items" not in param_def:
         param_def["items"] = {"type": "string"}