feat: Updated Graphql implementation to be compatible with UTCP 1.0v

Author: Salman Mohammed

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

@@ -0,0 +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
+
+
+def register():
+    register_communication_protocol("graphql", GraphQLCommunicationProtocol())
+    register_call_template("graphql", GraphQLProviderSerializer())

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

@@ -1,7 +1,10 @@
-from utcp.data.call_template import CallTemplate
-from utcp.data.auth import Auth
+from utcp.data.call_template import CallTemplate, CallTemplateSerializer
+from utcp.data.auth import Auth, AuthSerializer
+from utcp.interfaces.serializer import Serializer
+from utcp.exceptions import UtcpSerializerValidationError
+import traceback
 from typing import Dict, List, Optional, Literal
-from pydantic import Field
+from pydantic import Field, field_serializer, field_validator
 
 class GraphQLProvider(CallTemplate):
     """Provider configuration for GraphQL-based tools.
@@ -27,3 +30,31 @@ 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.")
+
+    @field_serializer("auth")
+    def serialize_auth(self, auth: Optional[Auth]):
+        if auth is None:
+            return None
+        return AuthSerializer().to_dict(auth)
+
+    @field_validator("auth", mode="before")
+    @classmethod
+    def validate_auth(cls, v: Optional[Auth | dict]):
+        if v is None:
+            return None
+        if isinstance(v, Auth):
+            return v
+        return AuthSerializer().validate_dict(v)
+
+
+class GraphQLProviderSerializer(Serializer[GraphQLProvider]):
+    def to_dict(self, obj: GraphQLProvider) -> dict:
+        return obj.model_dump()
+
+    def validate_dict(self, data: dict) -> GraphQLProvider:
+        try:
+            return GraphQLProvider.model_validate(data)
+        except Exception as e:
+            raise UtcpSerializerValidationError(
+                f"Invalid GraphQLProvider: {e}\n{traceback.format_exc()}"
+            )

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

@@ -1,36 +1,54 @@
-import sys
-from typing import Dict, Any, List, Optional, Callable
+import logging
+from typing import Dict, Any, List, Optional, AsyncGenerator, TYPE_CHECKING
+
 import aiohttp
-import asyncio
-import ssl
 from gql import Client as GqlClient, gql as gql_query
 from gql.transport.aiohttp import AIOHTTPTransport
-from utcp.client.client_transport_interface import ClientTransportInterface
-from utcp.shared.provider import Provider, GraphQLProvider
-from utcp.shared.tool import Tool, ToolInputOutputSchema
-from utcp.shared.auth import ApiKeyAuth, BasicAuth, OAuth2Auth
-import logging
+
+from utcp.interfaces.communication_protocol import CommunicationProtocol
+from utcp.data.call_template import CallTemplate
+from utcp.data.tool import Tool, JsonSchema
+from utcp.data.utcp_manual import UtcpManual
+from utcp.data.register_manual_response import RegisterManualResult
+from utcp.data.auth_implementations.api_key_auth import ApiKeyAuth
+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
+
+if TYPE_CHECKING:
+    from utcp.utcp_client import UtcpClient
+
 
 logging.basicConfig(
     level=logging.INFO,
-    format="%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d - %(message)s"
+    format="%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d - %(message)s",
 )
 
 logger = logging.getLogger(__name__)
 
-class GraphQLClientTransport(ClientTransportInterface):
-    """
-    Simple, robust, production-ready GraphQL transport using gql.
-    Stateless, per-operation. Supports all GraphQL features.
+
+class GraphQLCommunicationProtocol(CommunicationProtocol):
+    """GraphQL protocol implementation for UTCP 1.0.
+
+    - Discovers tools via GraphQL schema introspection.
+    - Executes per-call sessions using `gql` over HTTP(S).
+    - Supports `ApiKeyAuth`, `BasicAuth`, and `OAuth2Auth`.
+    - Enforces HTTPS or localhost for security.
     """
-    def __init__(self):
+
+    def __init__(self) -> None:
         self._oauth_tokens: Dict[str, Dict[str, Any]] = {}
 
-    def _enforce_https_or_localhost(self, url: str):
-        if not (url.startswith("https://") or url.startswith("http://localhost") or url.startswith("http://127.0.0.1")):
+    def _enforce_https_or_localhost(self, url: str) -> None:
+        if not (
+            url.startswith("https://")
+            or url.startswith("http://localhost")
+            or url.startswith("http://127.0.0.1")
+        ):
             raise ValueError(
-                f"Security error: URL must use HTTPS or start with 'http://localhost' or 'http://127.0.0.1'. Got: {url}. "
-                "Non-secure URLs are vulnerable to man-in-the-middle attacks."
+                "Security error: URL must use HTTPS or start with 'http://localhost' or 'http://127.0.0.1'. "
+                f"Got: {url}."
             )
 
     async def _handle_oauth2(self, auth: OAuth2Auth) -> str:
@@ -39,98 +57,155 @@ async def _handle_oauth2(self, auth: OAuth2Auth) -> str:
             return self._oauth_tokens[client_id]["access_token"]
         async with aiohttp.ClientSession() as session:
             data = {
-                'grant_type': 'client_credentials',
-                'client_id': client_id,
-                'client_secret': auth.client_secret,
-                'scope': auth.scope
+                "grant_type": "client_credentials",
+                "client_id": client_id,
+                "client_secret": auth.client_secret,
+                "scope": auth.scope,
             }
             async with session.post(auth.token_url, data=data) as resp:
                 resp.raise_for_status()
                 token_response = await resp.json()
                 self._oauth_tokens[client_id] = token_response
                 return token_response["access_token"]
 
-    async def _prepare_headers(self, provider: GraphQLProvider) -> Dict[str, str]:
-        headers = provider.headers.copy() if provider.headers else {}
+    async def _prepare_headers(
+        self, provider: GraphQLProvider, 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:
-                    if provider.auth.location == "header":
-                        headers[provider.auth.var_name] = provider.auth.api_key
-                # (query/cookie not supported for GraphQL by default)
+                if provider.auth.api_key and provider.auth.location == "header":
+                    headers[provider.auth.var_name] = provider.auth.api_key
             elif isinstance(provider.auth, BasicAuth):
                 import base64
+
                 userpass = f"{provider.auth.username}:{provider.auth.password}"
                 headers["Authorization"] = "Basic " + base64.b64encode(userpass.encode()).decode()
             elif isinstance(provider.auth, OAuth2Auth):
                 token = await self._handle_oauth2(provider.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 field in tool_args and isinstance(tool_args[field], str):
+                    headers[field] = tool_args[field]
+
         return headers
 
-    async def register_tool_provider(self, manual_provider: Provider) -> List[Tool]:
-        if not isinstance(manual_provider, GraphQLProvider):
-            raise ValueError("GraphQLClientTransport can only be used with GraphQLProvider")
-        self._enforce_https_or_localhost(manual_provider.url)
-        headers = await self._prepare_headers(manual_provider)
-        transport = AIOHTTPTransport(url=manual_provider.url, headers=headers)
-        async with GqlClient(transport=transport, fetch_schema_from_transport=True) as session:
-            schema = session.client.schema
-            tools = []
-            # Queries
-            if hasattr(schema, 'query_type') and schema.query_type:
-                for name, field in schema.query_type.fields.items():
-                    tools.append(Tool(
-                        name=name,
-                        description=getattr(field, 'description', '') or '',
-                        inputs=ToolInputOutputSchema(required=None),
-                        tool_provider=manual_provider
-                    ))
-            # Mutations
-            if hasattr(schema, 'mutation_type') and schema.mutation_type:
-                for name, field in schema.mutation_type.fields.items():
-                    tools.append(Tool(
-                        name=name,
-                        description=getattr(field, 'description', '') or '',
-                        inputs=ToolInputOutputSchema(required=None),
-                        tool_provider=manual_provider
-                    ))
-            # Subscriptions (listed, but not called here)
-            if hasattr(schema, 'subscription_type') and schema.subscription_type:
-                for name, field in schema.subscription_type.fields.items():
-                    tools.append(Tool(
-                        name=name,
-                        description=getattr(field, 'description', '') or '',
-                        inputs=ToolInputOutputSchema(required=None),
-                        tool_provider=manual_provider
-                    ))
-            return tools
-
-    async def deregister_tool_provider(self, manual_provider: Provider) -> None:
-        # Stateless: nothing to do
-        pass
-
-    async def call_tool(self, tool_name: str, tool_args: Dict[str, Any], tool_provider: Provider, query: Optional[str] = None) -> Any:
-        if not isinstance(tool_provider, GraphQLProvider):
-            raise ValueError("GraphQLClientTransport can only be used with GraphQLProvider")
-        self._enforce_https_or_localhost(tool_provider.url)
-        headers = await self._prepare_headers(tool_provider)
-        transport = AIOHTTPTransport(url=tool_provider.url, headers=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")
+        self._enforce_https_or_localhost(manual_call_template.url)
+
+        try:
+            headers = await self._prepare_headers(manual_call_template)
+            transport = AIOHTTPTransport(url=manual_call_template.url, headers=headers)
+            async with GqlClient(transport=transport, fetch_schema_from_transport=True) as session:
+                schema = session.client.schema
+                tools: List[Tool] = []
+
+                # Queries
+                if hasattr(schema, "query_type") and schema.query_type:
+                    for name, field in schema.query_type.fields.items():
+                        tools.append(
+                            Tool(
+                                name=name,
+                                description=getattr(field, "description", "") or "",
+                                inputs=JsonSchema(type="object"),
+                                outputs=JsonSchema(type="object"),
+                                tool_call_template=manual_call_template,
+                            )
+                        )
+
+                # Mutations
+                if hasattr(schema, "mutation_type") and schema.mutation_type:
+                    for name, field in schema.mutation_type.fields.items():
+                        tools.append(
+                            Tool(
+                                name=name,
+                                description=getattr(field, "description", "") or "",
+                                inputs=JsonSchema(type="object"),
+                                outputs=JsonSchema(type="object"),
+                                tool_call_template=manual_call_template,
+                            )
+                        )
+
+                # Subscriptions (listed for completeness)
+                if hasattr(schema, "subscription_type") and schema.subscription_type:
+                    for name, field in schema.subscription_type.fields.items():
+                        tools.append(
+                            Tool(
+                                name=name,
+                                description=getattr(field, "description", "") or "",
+                                inputs=JsonSchema(type="object"),
+                                outputs=JsonSchema(type="object"),
+                                tool_call_template=manual_call_template,
+                            )
+                        )
+
+                manual = UtcpManual(tools=tools)
+                return RegisterManualResult(
+                    manual_call_template=manual_call_template,
+                    manual=manual,
+                    success=True,
+                    errors=[],
+                )
+        except Exception as e:
+            logger.error(f"GraphQL manual registration failed for '{manual_call_template.name}': {e}")
+            return RegisterManualResult(
+                manual_call_template=manual_call_template,
+                manual=UtcpManual(manual_version="0.0.0", tools=[]),
+                success=False,
+                errors=[str(e)],
+            )
+
+    async def deregister_manual(
+        self, caller: "UtcpClient", manual_call_template: CallTemplate
+    ) -> None:
+        # Stateless: nothing to clean up
+        return None
+
+    async def call_tool(
+        self,
+        caller: "UtcpClient",
+        tool_name: str,
+        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")
+        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:
-            if query is not None:
-                document = gql_query(query)
-                result = await session.execute(document, variable_values=tool_args)
-                return result
-            # If no query provided, build a simple query
-            # Default to query operation
-            op_type = getattr(tool_provider, 'operation_type', 'query')
-            arg_str = ', '.join(f"${k}: String" for k in tool_args.keys())
+            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
+
+            arg_str = ", ".join(f"${k}: String" for k in tool_args.keys())
             var_defs = f"({arg_str})" if arg_str else ""
-            arg_pass = ', '.join(f"{k}: ${k}" for k in tool_args.keys())
+            arg_pass = ", ".join(f"{k}: ${k}" for k in tool_args.keys())
             arg_pass = f"({arg_pass})" if arg_pass else ""
-            gql_str = f"{op_type} {var_defs} {{ {tool_name}{arg_pass} }}"
+
+            gql_str = f"{op_type} {var_defs} {{ {base_tool_name}{arg_pass} }}"
             document = gql_query(gql_str)
             result = await session.execute(document, variable_values=tool_args)
             return result
 
+    async def call_tool_streaming(
+        self,
+        caller: "UtcpClient",
+        tool_name: str,
+        tool_args: Dict[str, Any],
+        tool_call_template: CallTemplate,
+    ) -> AsyncGenerator[Any, None]:
+        # Basic implementation: execute non-streaming and yield once
+        result = await self.call_tool(caller, tool_name, tool_args, tool_call_template)
+        yield result
+
     async def close(self) -> None:
-        self._oauth_tokens.clear()
+        self._oauth_tokens.clear()