update

Author: h3xxit

plugins/communication_protocols/gql/old_tests/test_graphql_transport.py

@@ -1,9 +1,9 @@
 from utcp.plugins.discovery import register_communication_protocol, register_call_template
 
 from .gql_communication_protocol import GraphQLCommunicationProtocol
-from .gql_call_template import GraphQLProvider, GraphQLProviderSerializer
+from .gql_call_template import GraphQLCallTemplate, GraphQLCallTemplateSerializer
 
 
 def register():
     register_communication_protocol("graphql", GraphQLCommunicationProtocol())
-    register_call_template("graphql", GraphQLProviderSerializer())
+    register_call_template("graphql", GraphQLCallTemplateSerializer())

plugins/communication_protocols/gql/src/utcp_gql/__init__.py

@@ -6,13 +6,17 @@
 from typing import Dict, List, Optional, Literal
 from pydantic import Field, field_serializer, field_validator
 
-class GraphQLProvider(CallTemplate):
+class GraphQLCallTemplate(CallTemplate):
     """Provider configuration for GraphQL-based tools.
 
     Enables communication with GraphQL endpoints supporting queries, mutations,
     and subscriptions. Provides flexible query execution with custom headers
     and authentication.
 
+    For maximum flexibility, use the `query` field to provide a complete GraphQL
+    query string with proper selection sets and variable types. This allows agents
+    to call any existing GraphQL endpoint without limitations.
+
     Attributes:
         call_template_type: Always "graphql" for GraphQL providers.
         url: The GraphQL endpoint URL.
@@ -21,6 +25,23 @@ class GraphQLProvider(CallTemplate):
         auth: Optional authentication configuration.
         headers: Optional static headers to include in requests.
         header_fields: List of tool argument names to map to HTTP request headers.
+        query: Custom GraphQL query string with full control over selection sets
+            and variable types. Example: 'query GetUser($id: ID!) { user(id: $id) { id name } }'
+        variable_types: Map of variable names to GraphQL types for auto-generated queries.
+            Example: {'id': 'ID!', 'limit': 'Int'}. Defaults to 'String' if not specified.
+
+    Example:
+        # Full flexibility with custom query
+        template = GraphQLCallTemplate(
+            url="https://api.example.com/graphql",
+            query="query GetUser($id: ID!) { user(id: $id) { id name email } }",
+        )
+
+        # Auto-generation with proper types
+        template = GraphQLCallTemplate(
+            url="https://api.example.com/graphql",
+            variable_types={"limit": "Int", "active": "Boolean"},
+        )
     """
 
     call_template_type: Literal["graphql"] = "graphql"
@@ -30,6 +51,18 @@ class GraphQLProvider(CallTemplate):
     auth: Optional[Auth] = None
     headers: Optional[Dict[str, str]] = None
     header_fields: Optional[List[str]] = Field(default=None, description="List of input fields to be sent as request headers for the initial connection.")
+    query: Optional[str] = Field(
+        default=None,
+        description="Custom GraphQL query/mutation string. Use $varName syntax for variables. "
+                    "If provided, this takes precedence over auto-generation. "
+                    "Example: 'query GetUser($id: ID!) { user(id: $id) { id name email } }'"
+    )
+    variable_types: Optional[Dict[str, str]] = Field(
+        default=None,
+        description="Map of variable names to GraphQL types for auto-generated queries. "
+                    "Example: {'id': 'ID!', 'limit': 'Int', 'active': 'Boolean'}. "
+                    "Defaults to 'String' if not specified."
+    )
 
     @field_serializer("auth")
     def serialize_auth(self, auth: Optional[Auth]):
@@ -47,14 +80,14 @@ def validate_auth(cls, v: Optional[Auth | dict]):
         return AuthSerializer().validate_dict(v)
 
 
-class GraphQLProviderSerializer(Serializer[GraphQLProvider]):
-    def to_dict(self, obj: GraphQLProvider) -> dict:
+class GraphQLCallTemplateSerializer(Serializer[GraphQLCallTemplate]):
+    def to_dict(self, obj: GraphQLCallTemplate) -> dict:
         return obj.model_dump()
 
-    def validate_dict(self, data: dict) -> GraphQLProvider:
+    def validate_dict(self, data: dict) -> GraphQLCallTemplate:
         try:
-            return GraphQLProvider.model_validate(data)
+            return GraphQLCallTemplate.model_validate(data)
         except Exception as e:
             raise UtcpSerializerValidationError(
-                f"Invalid GraphQLProvider: {e}\n{traceback.format_exc()}"
-            )
+                f"Invalid GraphQLCallTemplate: {e}\n{traceback.format_exc()}"
+            )

plugins/communication_protocols/gql/src/utcp_gql/gql_call_template.py

@@ -14,7 +14,7 @@
 from utcp.data.auth_implementations.basic_auth import BasicAuth
 from utcp.data.auth_implementations.oauth2_auth import OAuth2Auth
 
-from utcp_gql.gql_call_template import GraphQLProvider
+from utcp_gql.gql_call_template import GraphQLCallTemplate
 
 if TYPE_CHECKING:
     from utcp.utcp_client import UtcpClient
@@ -70,25 +70,25 @@ async def _handle_oauth2(self, auth: OAuth2Auth) -> str:
                 return token_response["access_token"]
 
     async def _prepare_headers(
-        self, provider: GraphQLProvider, tool_args: Optional[Dict[str, Any]] = None
+        self, call_template: GraphQLCallTemplate, tool_args: Optional[Dict[str, Any]] = None
     ) -> Dict[str, str]:
-        headers: Dict[str, str] = provider.headers.copy() if provider.headers else {}
-        if provider.auth:
-            if isinstance(provider.auth, ApiKeyAuth):
-                if provider.auth.api_key and provider.auth.location == "header":
-                    headers[provider.auth.var_name] = provider.auth.api_key
-            elif isinstance(provider.auth, BasicAuth):
+        headers: Dict[str, str] = call_template.headers.copy() if call_template.headers else {}
+        if call_template.auth:
+            if isinstance(call_template.auth, ApiKeyAuth):
+                if call_template.auth.api_key and call_template.auth.location == "header":
+                    headers[call_template.auth.var_name] = call_template.auth.api_key
+            elif isinstance(call_template.auth, BasicAuth):
                 import base64
 
-                userpass = f"{provider.auth.username}:{provider.auth.password}"
+                userpass = f"{call_template.auth.username}:{call_template.auth.password}"
                 headers["Authorization"] = "Basic " + base64.b64encode(userpass.encode()).decode()
-            elif isinstance(provider.auth, OAuth2Auth):
-                token = await self._handle_oauth2(provider.auth)
+            elif isinstance(call_template.auth, OAuth2Auth):
+                token = await self._handle_oauth2(call_template.auth)
                 headers["Authorization"] = f"Bearer {token}"
 
         # Map selected tool_args into headers if requested
-        if tool_args and provider.header_fields:
-            for field in provider.header_fields:
+        if tool_args and call_template.header_fields:
+            for field in call_template.header_fields:
                 if field in tool_args and isinstance(tool_args[field], str):
                     headers[field] = tool_args[field]
 
@@ -97,8 +97,8 @@ async def _prepare_headers(
     async def register_manual(
         self, caller: "UtcpClient", manual_call_template: CallTemplate
     ) -> RegisterManualResult:
-        if not isinstance(manual_call_template, GraphQLProvider):
-            raise ValueError("GraphQLCommunicationProtocol requires a GraphQLProvider call template")
+        if not isinstance(manual_call_template, GraphQLCallTemplate):
+            raise ValueError("GraphQLCommunicationProtocol requires a GraphQLCallTemplate call template")
         self._enforce_https_or_localhost(manual_call_template.url)
 
         try:
@@ -176,26 +176,40 @@ async def call_tool(
         tool_args: Dict[str, Any],
         tool_call_template: CallTemplate,
     ) -> Any:
-        if not isinstance(tool_call_template, GraphQLProvider):
-            raise ValueError("GraphQLCommunicationProtocol requires a GraphQLProvider call template")
+        if not isinstance(tool_call_template, GraphQLCallTemplate):
+            raise ValueError("GraphQLCommunicationProtocol requires a GraphQLCallTemplate call template")
         self._enforce_https_or_localhost(tool_call_template.url)
 
         headers = await self._prepare_headers(tool_call_template, tool_args)
         transport = AIOHTTPTransport(url=tool_call_template.url, headers=headers)
         async with GqlClient(transport=transport, fetch_schema_from_transport=True) as session:
-            op_type = getattr(tool_call_template, "operation_type", "query")
-            # Strip manual prefix if present (client prefixes at save time)
-            base_tool_name = tool_name.split(".", 1)[-1] if "." in tool_name else tool_name
             # Filter out header fields from GraphQL variables; these are sent via HTTP headers
             header_fields = tool_call_template.header_fields or []
             filtered_args = {k: v for k, v in tool_args.items() if k not in header_fields}
 
-            arg_str = ", ".join(f"${k}: String" for k in filtered_args.keys())
-            var_defs = f"({arg_str})" if arg_str else ""
-            arg_pass = ", ".join(f"{k}: ${k}" for k in filtered_args.keys())
-            arg_pass = f"({arg_pass})" if arg_pass else ""
+            # Use custom query if provided (highest flexibility for agents)
+            if tool_call_template.query:
+                gql_str = tool_call_template.query
+            else:
+                # Auto-generate query - use variable_types for proper typing
+                op_type = getattr(tool_call_template, "operation_type", "query")
+                base_tool_name = tool_name.split(".", 1)[-1] if "." in tool_name else tool_name
+                variable_types = tool_call_template.variable_types or {}
+
+                # Build variable definitions with proper types (default to String)
+                arg_str = ", ".join(
+                    f"${k}: {variable_types.get(k, 'String')}"
+                    for k in filtered_args.keys()
+                )
+                var_defs = f"({arg_str})" if arg_str else ""
+                arg_pass = ", ".join(f"{k}: ${k}" for k in filtered_args.keys())
+                arg_pass = f"({arg_pass})" if arg_pass else ""
+
+                # Note: Auto-generated queries for object-returning fields will still fail
+                # without a selection set. Use the `query` field for full control.
+                gql_str = f"{op_type} {var_defs} {{ {base_tool_name}{arg_pass} }}"
+                logger.debug(f"Auto-generated GraphQL: {gql_str}")
 
-            gql_str = f"{op_type} {var_defs} {{ {base_tool_name}{arg_pass} }}"
             document = gql_query(gql_str)
             result = await session.execute(document, variable_values=filtered_args)
             return result