forked from themanojdesai/python-a2a
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathhttp.py
More file actions
172 lines (141 loc) · 5.33 KB
/
Copy pathhttp.py
File metadata and controls
172 lines (141 loc) · 5.33 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
"""
HTTP server implementation for the A2A protocol.
"""
import json
import traceback
from typing import Type, Optional, Dict, Any, Callable, Union
try:
from flask import Flask, request, jsonify, Response
except ImportError:
Flask = None
from ..models.message import Message, MessageRole
from ..models.conversation import Conversation
from ..models.content import TextContent, ErrorContent
from .base import BaseA2AServer
from ..exceptions import A2AImportError, A2ARequestError
class A2AServer(BaseA2AServer):
"""
A customizable A2A-compatible agent server.
This server can be configured with custom message handlers to create
specialized agents for different purposes.
"""
def __init__(self, message_handler: Optional[Callable[[Message], Message]] = None):
"""
Initialize the server with an optional message handler
Args:
message_handler: Optional function that processes messages
If provided, this function will be called to handle incoming messages
instead of the handle_message method
"""
self.message_handler = message_handler
def handle_message(self, message: Message) -> Message:
"""
Process an incoming A2A message and generate a response
If a message handler was provided in the constructor, it will be used.
Otherwise, this method should be overridden by subclasses.
Args:
message: The incoming A2A message
Returns:
The agent's response message
"""
if self.message_handler:
return self.message_handler(message)
# Default implementation for when no handler is provided
# Just echo back the input with a prefix
if message.content.type == "text":
return Message(
content=TextContent(text=f"Echo: {message.content.text}"),
role=MessageRole.AGENT,
parent_message_id=message.message_id,
conversation_id=message.conversation_id
)
else:
return Message(
content=ErrorContent(message=f"Unsupported message type: {message.content.type}"),
role=MessageRole.AGENT,
parent_message_id=message.message_id,
conversation_id=message.conversation_id
)
def get_metadata(self) -> Dict[str, Any]:
"""
Get metadata about this agent server
Returns:
A dictionary of metadata about this agent
"""
metadata = super().get_metadata()
metadata.update({
"agent_type": "A2AServer",
"capabilities": ["text"],
"has_custom_handler": self.message_handler is not None
})
return metadata
def create_flask_app(agent: BaseA2AServer) -> Flask:
"""
Create a Flask application that serves an A2A agent
Args:
agent: The A2A agent server
Returns:
A Flask application
Raises:
A2AImportError: If Flask is not installed
"""
if Flask is None:
raise A2AImportError(
"Flask is not installed. "
"Install it with 'pip install flask'"
)
app = Flask(__name__)
@app.route("/a2a", methods=["POST"])
def handle_a2a_request() -> Union[Response, tuple]:
"""Handle A2A protocol requests"""
try:
data = request.json
# Check if this is a single message or a conversation
if "messages" in data:
# This is a conversation
conversation = Conversation.from_dict(data)
response = agent.handle_conversation(conversation)
return jsonify(response.to_dict())
else:
# This is a single message
message = Message.from_dict(data)
response = agent.handle_message(message)
return jsonify(response.to_dict())
except Exception as e:
# Return an error response for any exceptions
error_dict = {
"content": {
"type": "error",
"message": f"Error processing request: {str(e)}"
},
"role": "system"
}
return jsonify(error_dict), 500
@app.route("/a2a/metadata", methods=["GET"])
def get_agent_metadata() -> Response:
"""Return metadata about the agent"""
return jsonify(agent.get_metadata())
@app.route("/a2a/health", methods=["GET"])
def health_check() -> Response:
"""Health check endpoint"""
return jsonify({"status": "ok"})
return app
def run_server(
agent: BaseA2AServer,
host: str = "0.0.0.0",
port: int = 5000,
debug: bool = False
) -> None:
"""
Run an A2A agent as a Flask server
Args:
agent: The A2A agent server
host: Host to bind to (default: "0.0.0.0")
port: Port to listen on (default: 5000)
debug: Enable debug mode (default: False)
Raises:
A2AImportError: If Flask is not installed
"""
app = create_flask_app(agent)
print(f"Starting A2A server on http://{host}:{port}/a2a")
app.run(host=host, port=port, debug=debug)