jaf-py 2.5.9__py3-none-any.whl → 2.5.11__py3-none-any.whl

jaf/a2a/protocol.py

@@ -12,29 +12,23 @@ from .types import A2AAgent, A2AErrorCodes
 def validate_jsonrpc_request(request: Dict[str, Any]) -> bool:
     """Pure function to validate JSON-RPC request"""
     return (
-        isinstance(request, dict) and
-        request.get("jsonrpc") == "2.0" and
-        "id" in request and
-        isinstance(request.get("method"), str)
+        isinstance(request, dict)
+        and request.get("jsonrpc") == "2.0"
+        and "id" in request
+        and isinstance(request.get("method"), str)
     )
 
 
 def create_jsonrpc_success_response_dict(id: Union[str, int, None], result: Any) -> Dict[str, Any]:
     """Pure function to create JSON-RPC success response as dict"""
-    return {
-        "jsonrpc": "2.0",
-        "id": id,
-        "result": result
-    }
+    return {"jsonrpc": "2.0", "id": id, "result": result}
 
 
-def create_jsonrpc_error_response_dict(id: Union[str, int, None], error: Dict[str, Any]) -> Dict[str, Any]:
+def create_jsonrpc_error_response_dict(
+    id: Union[str, int, None], error: Dict[str, Any]
+) -> Dict[str, Any]:
     """Pure function to create JSON-RPC error response as dict"""
-    return {
-        "jsonrpc": "2.0",
-        "id": id,
-        "error": error
-    }
+    return {"jsonrpc": "2.0", "id": id, "error": error}
 
 
 def map_error_to_a2a_error(error: Exception) -> Dict[str, Any]:
@@ -43,13 +37,10 @@ def map_error_to_a2a_error(error: Exception) -> Dict[str, Any]:
         return {
             "code": A2AErrorCodes.INTERNAL_ERROR.value,
             "message": str(error),
-            "data": {"type": type(error).__name__}
+            "data": {"type": type(error).__name__},
         }
 
-    return {
-        "code": A2AErrorCodes.INTERNAL_ERROR.value,
-        "message": "Unknown error occurred"
-    }
+    return {"code": A2AErrorCodes.INTERNAL_ERROR.value, "message": "Unknown error occurred"}
 
 
 def validate_send_message_request(request: Dict[str, Any]) -> Dict[str, Any]:
@@ -61,8 +52,8 @@ def validate_send_message_request(request: Dict[str, Any]) -> Dict[str, Any]:
                 "is_valid": False,
                 "error": {
                     "code": A2AErrorCodes.INVALID_REQUEST.value,
-                    "message": "Invalid JSON-RPC request"
-                }
+                    "message": "Invalid JSON-RPC request",
+                },
             }
 
         if request.get("method") not in ["message/send", "message/stream"]:
@@ -70,8 +61,8 @@ def validate_send_message_request(request: Dict[str, Any]) -> Dict[str, Any]:
                 "is_valid": False,
                 "error": {
                     "code": A2AErrorCodes.METHOD_NOT_FOUND.value,
-                    "message": f"Method {request.get('method')} not supported"
-                }
+                    "message": f"Method {request.get('method')} not supported",
+                },
             }
 
         params = request.get("params", {})
@@ -81,8 +72,8 @@ def validate_send_message_request(request: Dict[str, Any]) -> Dict[str, Any]:
                 "error": {
                     "code": A2AErrorCodes.INVALID_PARAMS.value,
                     "message": "Invalid params format - must be an object",
-                    "data": {"expected": "object", "received": type(params).__name__}
-                }
+                    "data": {"expected": "object", "received": type(params).__name__},
+                },
             }
 
         message = params.get("message")
@@ -92,8 +83,11 @@ def validate_send_message_request(request: Dict[str, Any]) -> Dict[str, Any]:
                 "error": {
                     "code": A2AErrorCodes.INVALID_PARAMS.value,
                     "message": "Invalid message format - message must be an object",
-                    "data": {"expected": "object", "received": type(message).__name__ if message is not None else "null"}
-                }
+                    "data": {
+                        "expected": "object",
+                        "received": type(message).__name__ if message is not None else "null",
+                    },
+                },
             }
 
         # Validate required message fields
@@ -105,8 +99,8 @@ def validate_send_message_request(request: Dict[str, Any]) -> Dict[str, Any]:
                 "error": {
                     "code": A2AErrorCodes.INVALID_PARAMS.value,
                     "message": f"Missing required message fields: {', '.join(missing_fields)}",
-                    "data": {"missing_fields": missing_fields}
-                }
+                    "data": {"missing_fields": missing_fields},
+                },
             }
 
         # Validate message structure
@@ -116,8 +110,8 @@ def validate_send_message_request(request: Dict[str, Any]) -> Dict[str, Any]:
                 "error": {
                     "code": A2AErrorCodes.INVALID_PARAMS.value,
                     "message": "Message kind must be 'message'",
-                    "data": {"expected": "message", "received": message.get("kind")}
-                }
+                    "data": {"expected": "message", "received": message.get("kind")},
+                },
             }
 
         if message.get("role") not in ["user", "agent"]:
@@ -126,8 +120,8 @@ def validate_send_message_request(request: Dict[str, Any]) -> Dict[str, Any]:
                 "error": {
                     "code": A2AErrorCodes.INVALID_PARAMS.value,
                     "message": "Message role must be 'user' or 'agent'",
-                    "data": {"expected": ["user", "agent"], "received": message.get("role")}
-                }
+                    "data": {"expected": ["user", "agent"], "received": message.get("role")},
+                },
             }
 
         if not isinstance(message.get("parts"), list):
@@ -136,8 +130,8 @@ def validate_send_message_request(request: Dict[str, Any]) -> Dict[str, Any]:
                 "error": {
                     "code": A2AErrorCodes.INVALID_PARAMS.value,
                     "message": "Message parts must be a list",
-                    "data": {"expected": "array", "received": type(message.get("parts")).__name__}
-                }
+                    "data": {"expected": "array", "received": type(message.get("parts")).__name__},
+                },
             }
 
         if len(message.get("parts", [])) == 0:
@@ -146,8 +140,8 @@ def validate_send_message_request(request: Dict[str, Any]) -> Dict[str, Any]:
                 "error": {
                     "code": A2AErrorCodes.INVALID_PARAMS.value,
                     "message": "Message parts cannot be empty",
-                    "data": {"minimum_parts": 1}
-                }
+                    "data": {"minimum_parts": 1},
+                },
             }
 
         # Validate parts structure
@@ -158,8 +152,12 @@ def validate_send_message_request(request: Dict[str, Any]) -> Dict[str, Any]:
                     "error": {
                         "code": A2AErrorCodes.INVALID_PARAMS.value,
                         "message": f"Message part {i} must be an object",
-                        "data": {"part_index": i, "expected": "object", "received": type(part).__name__}
-                    }
+                        "data": {
+                            "part_index": i,
+                            "expected": "object",
+                            "received": type(part).__name__,
+                        },
+                    },
                 }
 
             if "kind" not in part:
@@ -168,8 +166,8 @@ def validate_send_message_request(request: Dict[str, Any]) -> Dict[str, Any]:
                     "error": {
                         "code": A2AErrorCodes.INVALID_PARAMS.value,
                         "message": f"Message part {i} missing 'kind' field",
-                        "data": {"part_index": i, "missing_field": "kind"}
-                    }
+                        "data": {"part_index": i, "missing_field": "kind"},
+                    },
                 }
 
             kind = part.get("kind")
@@ -179,8 +177,8 @@ def validate_send_message_request(request: Dict[str, Any]) -> Dict[str, Any]:
                     "error": {
                         "code": A2AErrorCodes.INVALID_PARAMS.value,
                         "message": f"Text part {i} missing 'text' field",
-                        "data": {"part_index": i, "part_kind": "text", "missing_field": "text"}
-                    }
+                        "data": {"part_index": i, "part_kind": "text", "missing_field": "text"},
+                    },
                 }
             elif kind == "data" and "data" not in part:
                 return {
@@ -188,8 +186,8 @@ def validate_send_message_request(request: Dict[str, Any]) -> Dict[str, Any]:
                     "error": {
                         "code": A2AErrorCodes.INVALID_PARAMS.value,
                         "message": f"Data part {i} missing 'data' field",
-                        "data": {"part_index": i, "part_kind": "data", "missing_field": "data"}
-                    }
+                        "data": {"part_index": i, "part_kind": "data", "missing_field": "data"},
+                    },
                 }
             elif kind == "file" and "file" not in part:
                 return {
@@ -197,8 +195,8 @@ def validate_send_message_request(request: Dict[str, Any]) -> Dict[str, Any]:
                     "error": {
                         "code": A2AErrorCodes.INVALID_PARAMS.value,
                         "message": f"File part {i} missing 'file' field",
-                        "data": {"part_index": i, "part_kind": "file", "missing_field": "file"}
-                    }
+                        "data": {"part_index": i, "part_kind": "file", "missing_field": "file"},
+                    },
                 }
             elif kind not in ["text", "data", "file"]:
                 return {
@@ -206,14 +204,15 @@ def validate_send_message_request(request: Dict[str, Any]) -> Dict[str, Any]:
                     "error": {
                         "code": A2AErrorCodes.INVALID_PARAMS.value,
                         "message": f"Unknown part kind '{kind}' in part {i}",
-                        "data": {"part_index": i, "supported_kinds": ["text", "data", "file"], "received": kind}
-                    }
+                        "data": {
+                            "part_index": i,
+                            "supported_kinds": ["text", "data", "file"],
+                            "received": kind,
+                        },
+                    },
                 }
 
-        return {
-            "is_valid": True,
-            "data": request
-        }
+        return {"is_valid": True, "data": request}
 
     except Exception as e:
         return {
@@ -221,16 +220,16 @@ def validate_send_message_request(request: Dict[str, Any]) -> Dict[str, Any]:
             "error": {
                 "code": A2AErrorCodes.INTERNAL_ERROR.value,
                 "message": f"Request validation failed: {e!s}",
-                "data": {"error_type": type(e).__name__}
-            }
+                "data": {"error_type": type(e).__name__},
+            },
         }
 
 
 async def handle_message_send(
     request: Dict[str, Any],
-    agent: 'A2AAgent',  # Forward reference to avoid circular imports
+    agent: "A2AAgent",  # Forward reference to avoid circular imports
     model_provider: Any,
-    executor_func: Callable
+    executor_func: Callable,
 ) -> Dict[str, Any]:
     """Pure function to handle message/send method"""
     try:
@@ -240,7 +239,7 @@ async def handle_message_send(
         context = {
             "message": message,
             "session_id": message.get("contextId", f"session_{id(request)}"),
-            "metadata": params.get("metadata")
+            "metadata": params.get("metadata"),
         }
 
         result = await executor_func(context, agent, model_provider)
@@ -248,29 +247,22 @@ async def handle_message_send(
         if result.get("error"):
             return create_jsonrpc_error_response_dict(
                 request.get("id"),
-                {
-                    "code": A2AErrorCodes.INTERNAL_ERROR.value,
-                    "message": result["error"]
-                }
+                {"code": A2AErrorCodes.INTERNAL_ERROR.value, "message": result["error"]},
             )
 
         return create_jsonrpc_success_response_dict(
-            request.get("id"),
-            result.get("final_task", {"message": "No result available"})
+            request.get("id"), result.get("final_task", {"message": "No result available"})
         )
 
     except Exception as error:
-        return create_jsonrpc_error_response_dict(
-            request.get("id"),
-            map_error_to_a2a_error(error)
-        )
+        return create_jsonrpc_error_response_dict(request.get("id"), map_error_to_a2a_error(error))
 
 
 async def handle_message_stream(
     request: Dict[str, Any],
-    agent: 'A2AAgent',  # Forward reference
+    agent: "A2AAgent",  # Forward reference
     model_provider: Any,
-    executor_func: Callable
+    executor_func: Callable,
 ) -> AsyncGenerator[Dict[str, Any], None]:
     """Pure function to handle message/stream method"""
     try:
@@ -280,22 +272,18 @@ async def handle_message_stream(
         context = {
             "message": message,
             "session_id": message.get("contextId", f"session_{id(request)}"),
-            "metadata": params.get("metadata")
+            "metadata": params.get("metadata"),
         }
 
         async for event in executor_func(context, agent, model_provider):
             yield create_jsonrpc_success_response_dict(request.get("id"), event)
 
     except Exception as error:
-        yield create_jsonrpc_error_response_dict(
-            request.get("id"),
-            map_error_to_a2a_error(error)
-        )
+        yield create_jsonrpc_error_response_dict(request.get("id"), map_error_to_a2a_error(error))
 
 
 async def handle_tasks_get(
-    request: Dict[str, Any],
-    task_storage: Dict[str, Dict[str, Any]]
+    request: Dict[str, Any], task_storage: Dict[str, Dict[str, Any]]
 ) -> Dict[str, Any]:
     """Pure function to handle tasks/get method"""
     try:
@@ -305,10 +293,7 @@ async def handle_tasks_get(
         if not task_id:
             return create_jsonrpc_error_response_dict(
                 request.get("id"),
-                {
-                    "code": A2AErrorCodes.INVALID_PARAMS.value,
-                    "message": "Task ID is required"
-                }
+                {"code": A2AErrorCodes.INVALID_PARAMS.value, "message": "Task ID is required"},
             )
 
         task = task_storage.get(task_id)
@@ -318,8 +303,8 @@ async def handle_tasks_get(
                 request.get("id"),
                 {
                     "code": A2AErrorCodes.TASK_NOT_FOUND.value,
-                    "message": f"Task with id {task_id} not found"
-                }
+                    "message": f"Task with id {task_id} not found",
+                },
             )
 
         # Apply history length limit if specified
@@ -331,15 +316,11 @@ async def handle_tasks_get(
         return create_jsonrpc_success_response_dict(request.get("id"), result_task)
 
     except Exception as error:
-        return create_jsonrpc_error_response_dict(
-            request.get("id"),
-            map_error_to_a2a_error(error)
-        )
+        return create_jsonrpc_error_response_dict(request.get("id"), map_error_to_a2a_error(error))
 
 
 async def handle_tasks_cancel(
-    request: Dict[str, Any],
-    task_storage: Dict[str, Dict[str, Any]]
+    request: Dict[str, Any], task_storage: Dict[str, Dict[str, Any]]
 ) -> Dict[str, Any]:
     """Pure function to handle tasks/cancel method"""
     try:
@@ -349,10 +330,7 @@ async def handle_tasks_cancel(
         if not task_id:
             return create_jsonrpc_error_response_dict(
                 request.get("id"),
-                {
-                    "code": A2AErrorCodes.INVALID_PARAMS.value,
-                    "message": "Task ID is required"
-                }
+                {"code": A2AErrorCodes.INVALID_PARAMS.value, "message": "Task ID is required"},
             )
 
         task = task_storage.get(task_id)
@@ -362,8 +340,8 @@ async def handle_tasks_cancel(
                 request.get("id"),
                 {
                     "code": A2AErrorCodes.TASK_NOT_FOUND.value,
-                    "message": f"Task with id {task_id} not found"
-                }
+                    "message": f"Task with id {task_id} not found",
+                },
             )
 
         # Check if task can be canceled
@@ -373,29 +351,25 @@ async def handle_tasks_cancel(
                 request.get("id"),
                 {
                     "code": A2AErrorCodes.TASK_NOT_CANCELABLE.value,
-                    "message": f"Task {task_id} cannot be canceled in state {current_state}"
-                }
+                    "message": f"Task {task_id} cannot be canceled in state {current_state}",
+                },
             )
 
         # Create canceled task
         canceled_task = task.copy()
         canceled_task["status"] = {
             "state": "canceled",
-            "timestamp": None  # Would be set by the system
+            "timestamp": None,  # Would be set by the system
         }
 
         return create_jsonrpc_success_response_dict(request.get("id"), canceled_task)
 
     except Exception as error:
-        return create_jsonrpc_error_response_dict(
-            request.get("id"),
-            map_error_to_a2a_error(error)
-        )
+        return create_jsonrpc_error_response_dict(request.get("id"), map_error_to_a2a_error(error))
 
 
 async def handle_get_authenticated_extended_card(
-    request: Dict[str, Any],
-    agent_card: Dict[str, Any]
+    request: Dict[str, Any], agent_card: Dict[str, Any]
 ) -> Dict[str, Any]:
     """Pure function to handle agent/getAuthenticatedExtendedCard method"""
     try:
@@ -404,20 +378,17 @@ async def handle_get_authenticated_extended_card(
         return create_jsonrpc_success_response_dict(request.get("id"), agent_card)
 
     except Exception as error:
-        return create_jsonrpc_error_response_dict(
-            request.get("id"),
-            map_error_to_a2a_error(error)
-        )
+        return create_jsonrpc_error_response_dict(request.get("id"), map_error_to_a2a_error(error))
 
 
 def route_a2a_request(
     request: Dict[str, Any],
-    agent: 'A2AAgent',  # Forward reference
+    agent: "A2AAgent",  # Forward reference
     model_provider: Any,
     task_storage: Dict[str, Dict[str, Any]],
     agent_card: Dict[str, Any],
     executor_func: Callable,
-    streaming_executor_func: Callable
+    streaming_executor_func: Callable,
 ) -> Union[Dict[str, Any], AsyncGenerator[Dict[str, Any], None]]:
     """Pure function to route A2A requests"""
 
@@ -427,8 +398,8 @@ def route_a2a_request(
                 request.get("id"),
                 {
                     "code": A2AErrorCodes.INVALID_REQUEST.value,
-                    "message": "Invalid JSON-RPC request"
-                }
+                    "message": "Invalid JSON-RPC request",
+                },
             )
 
         method = request.get("method")
@@ -436,17 +407,16 @@ def route_a2a_request(
         if method == "message/send":
             validation = validate_send_message_request(request)
             if not validation["is_valid"]:
-                return create_jsonrpc_error_response_dict(
-                    request.get("id"),
-                    validation["error"]
-                )
+                return create_jsonrpc_error_response_dict(request.get("id"), validation["error"])
             return await handle_message_send(request, agent, model_provider, executor_func)
 
         elif method == "message/stream":
             validation = validate_send_message_request(request)
             if not validation["is_valid"]:
+
                 async def error_generator():
                     yield create_jsonrpc_error_response_dict(request.get("id"), validation["error"])
+
                 return error_generator()
             # Return the async generator directly
             return handle_message_stream(request, agent, model_provider, streaming_executor_func)
@@ -465,8 +435,8 @@ def route_a2a_request(
                 request.get("id"),
                 {
                     "code": A2AErrorCodes.METHOD_NOT_FOUND.value,
-                    "message": f"Method {method} not found"
-                }
+                    "message": f"Method {method} not found",
+                },
             )
 
     # Handle streaming vs non-streaming
@@ -479,11 +449,11 @@ def route_a2a_request(
 
 
 def create_protocol_handler_config(
-    agents: Dict[str, 'A2AAgent'],  # Forward reference
+    agents: Dict[str, "A2AAgent"],  # Forward reference
     model_provider: Any,
     agent_card: Dict[str, Any],
     executor_func: Callable,
-    streaming_executor_func: Callable
+    streaming_executor_func: Callable,
 ) -> Dict[str, Any]:
     """Pure function to create protocol handler configuration"""
 
@@ -502,8 +472,8 @@ def create_protocol_handler_config(
                 request.get("id"),
                 {
                     "code": A2AErrorCodes.INVALID_PARAMS.value,
-                    "message": f"Agent {agent_name or 'default'} not found"
-                }
+                    "message": f"Agent {agent_name or 'default'} not found",
+                },
             )
 
         return route_a2a_request(
@@ -513,7 +483,7 @@ def create_protocol_handler_config(
             task_storage,
             agent_card,
             executor_func,
-            streaming_executor_func
+            streaming_executor_func,
         )
 
     return {
@@ -521,5 +491,5 @@ def create_protocol_handler_config(
         "model_provider": model_provider,
         "agent_card": agent_card,
         "task_storage": task_storage,
-        "handle_request": handle_request
+        "handle_request": handle_request,
     }