phowell/pocketflow
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

finish a2a

Browse Source
This commit is contained in:
zachary62
2025-04-12 17:42:08 -04:00
parent 91dc7bbc99
commit 5610fbe9d2
23 changed files with 602 additions and 276 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cookbook/pocketflow-a2a/README.md
+228
View File
@@ -0,0 +1,228 @@
# PocketFlow Agent with A2A Protocol
This project demonstrates how to take an existing agent built with the PocketFlow library and make it accessible to other agents using the **Agent-to-Agent (A2A) communication protocol**.
## How it Works: A2A Integration
This project combines two main parts:
1. **PocketFlow Agent Logic:** The original agent code ([`nodes.py`](nodes.py), [`utils.py`](utils.py), [`flow.py`](flow.py)) defines the internal workflow (Decide -> Search -> Answer). This code is taken directly from the [PocketFlow Agent Tutorial](https://github.com/The-Pocket/PocketFlow/tree/main/cookbook/pocketflow-agent).
2. **A2A Server Wrapper:** Code from the [google/A2A samples repository](https://github.com/google/A2A/tree/main/samples/python) (`common/` directory) provides the necessary infrastructure to host the agent as an A2A-compliant server. *Note: Minor modifications were made to the common server/client code to add detailed logging for educational purposes.*
3. **The Bridge ([`task_manager.py`](task_manager.py)):** A custom `PocketFlowTaskManager` class acts as the bridge. It receives A2A requests (like `tasks/send`), extracts the user query, runs the PocketFlow `agent_flow`, takes the final result from the flow's shared state, and packages it back into an A2A `Task` object with the answer as an `Artifact`.
This demonstrates how a non-A2A agent framework can be exposed over the A2A protocol by implementing a specific `TaskManager`.
## Simplified Interaction Sequence
```mermaid
sequenceDiagram
participant Client as "Client ([minimal_a2a_client.py](a2a_client.py))"
participant Server as "Server (localhost:10003)"
Note over Client: User enters question
Client->>+Server: POST / (JSON-RPC Request: tasks/send)
Note over Server: Processes request internally (runs PocketFlow)
Server-->>-Client: HTTP 200 OK (JSON-RPC Response: result=Task)
Note over Client: Displays final answer
```
## Getting Started
### Prerequisites
* Python 3.10+ (due to type hinting used in the A2A `common` code)
* An OpenAI API Key
### Installation
1. Install dependencies:
```bash
pip install -r requirements.txt
```
2. Set your OpenAI API key as an environment variable:
```bash
export OPENAI_API_KEY="your-api-key-here"
```
Let's do a quick check to make sure your API key is working properly:
```bash
python utils.py
```
3. Run the server from this directory:
```bash
python a2a_server.py --port 10003
```
You should see logs indicating the server has started on `http://localhost:10003`.
4. Run the Client in a *separate terminal*
```bash
python a2a_client.py --agent-url http://localhost:10003
```
5. Follow the instructions in the client terminal to ask questions. Type `:q` or `quit` to exit the client.
## Example Interaction Logs
**(Server Log - showing internal PocketFlow steps)**
```
2025-04-12 17:20:40,893 - __main__ - INFO - Starting PocketFlow A2A server on http://localhost:10003
INFO: Started server process [677223]
INFO: Waiting for application startup.
INFO: Application startup complete.
INFO: Uvicorn running on http://localhost:10003 (Press CTRL+C to quit)
2025-04-12 17:20:57,647 - A2AServer - INFO - <- Received Request (ID: d3f3fb93350d47d9a94ca12bb62b656b):
{
"jsonrpc": "2.0",
"id": "d3f3fb93350d47d9a94ca12bb62b656b",
"method": "tasks/send",
"params": {
"id": "46c3ce7b941a4fff9b8e3b644d6db5f4",
"sessionId": "f3e12b8424c44241be881cd4bb8a269f",
"message": {
"role": "user",
"parts": [
{
"type": "text",
"text": "Who won the Nobel Prize in Physics 2024?"
}
]
},
"acceptedOutputModes": [
"text",
"text/plain"
]
}
}
2025-04-12 17:20:57,647 - task_manager - INFO - Received task send request: 46c3ce7b941a4fff9b8e3b644d6db5f4
2025-04-12 17:20:57,647 - common.server.task_manager - INFO - Upserting task 46c3ce7b941a4fff9b8e3b644d6db5f4
2025-04-12 17:20:57,647 - task_manager - INFO - Running PocketFlow for task 46c3ce7b941a4fff9b8e3b644d6db5f4...
🤔 Agent deciding what to do next...
2025-04-12 17:20:59,213 - httpx - INFO - HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 200 OK"
🔍 Agent decided to search for: 2024 Nobel Prize in Physics winner
🌐 Searching the web for: 2024 Nobel Prize in Physics winner
2025-04-12 17:20:59,974 - primp - INFO - response: https://lite.duckduckgo.com/lite/ 200
📚 Found information, analyzing results...
🤔 Agent deciding what to do next...
2025-04-12 17:21:01,619 - httpx - INFO - HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 200 OK"
💡 Agent decided to answer the question
✍️ Crafting final answer...
2025-04-12 17:21:03,833 - httpx - INFO - HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 200 OK"
✅ Answer generated successfully
2025-04-12 17:21:03,834 - task_manager - INFO - PocketFlow completed for task 46c3ce7b941a4fff9b8e3b644d6db5f4
2025-04-12 17:21:03,834 - A2AServer - INFO - -> Response (ID: d3f3fb93350d47d9a94ca12bb62b656b):
{
"jsonrpc": "2.0",
"id": "d3f3fb93350d47d9a94ca12bb62b656b",
"result": {
"id": "46c3ce7b941a4fff9b8e3b644d6db5f4",
"sessionId": "f3e12b8424c44241be881cd4bb8a269f",
"status": {
"state": "completed",
"timestamp": "2025-04-12T17:21:03.834542"
},
"artifacts": [
{
"parts": [
{
"type": "text",
"text": "The 2024 Nobel Prize in Physics was awarded to John J. Hopfield and Geoffrey Hinton for their foundational discoveries and inventions that have significantly advanced the field of machine learning through the use of artificial neural networks. Their pioneering work has been crucial in the development and implementation of algorithms that enable machines to learn and process information in a manner that mimics human cognitive functions. This advancement in artificial intelligence technology has had a profound impact on numerous industries, facilitating innovations across various applications, from image and speech recognition to self-driving cars."
}
],
"index": 0
}
],
"history": []
}
}
```
**(Client Log - showing request/response)**
```
Connecting to agent at: http://localhost:10003
Using Session ID: f3e12b8424c44241be881cd4bb8a269f
Enter your question (:q or quit to exit) > Who won the Nobel Prize in Physics 2024?
Sending task 46c3ce7b941a4fff9b8e3b644d6db5f4...
2025-04-12 17:20:57,643 - A2AClient - INFO - -> Sending Request (ID: d3f3fb93350d47d9a94ca12bb62b656b, Method: tasks/send):
{
"jsonrpc": "2.0",
"id": "d3f3fb93350d47d9a94ca12bb62b656b",
"method": "tasks/send",
"params": {
"id": "46c3ce7b941a4fff9b8e3b644d6db5f4",
"sessionId": "f3e12b8424c44241be881cd4bb8a269f",
"message": {
"role": "user",
"parts": [
{
"type": "text",
"text": "Who won the Nobel Prize in Physics 2024?"
}
]
},
"acceptedOutputModes": [
"text",
"text/plain"
]
}
}
2025-04-12 17:21:03,835 - httpx - INFO - HTTP Request: POST http://localhost:10003 "HTTP/1.1 200 OK"
2025-04-12 17:21:03,836 - A2AClient - INFO - <- Received HTTP Status 200 for Request (ID: d3f3fb93350d47d9a94ca12bb62b656b)
2025-04-12 17:21:03,836 - A2AClient - INFO - <- Received Success Response (ID: d3f3fb93350d47d9a94ca12bb62b656b):
{
"jsonrpc": "2.0",
"id": "d3f3fb93350d47d9a94ca12bb62b656b",
"result": {
"id": "46c3ce7b941a4fff9b8e3b644d6db5f4",
"sessionId": "f3e12b8424c44241be881cd4bb8a269f",
"status": {
"state": "completed",
"timestamp": "2025-04-12T17:21:03.834542"
},
"artifacts": [
{
"parts": [
{
"type": "text",
"text": "The 2024 Nobel Prize in Physics was awarded to John J. Hopfield and Geoffrey Hinton for their foundational discoveries and inventions that have significantly advanced the field of machine learning through the use of artificial neural networks. Their pioneering work has been crucial in the development and implementation of algorithms that enable machines to learn and process information in a manner that mimics human cognitive functions. This advancement in artificial intelligence technology has had a profound impact on numerous industries, facilitating innovations across various applications, from image and speech recognition to self-driving cars."
}
],
"index": 0
}
],
"history": []
}
}
Task 46c3ce7b941a4fff9b8e3b644d6db5f4 finished with state: completed
Agent Response:
The 2024 Nobel Prize in Physics was awarded to John J. Hopfield and Geoffrey Hinton for their foundational discoveries and inventions that have significantly advanced the field of machine learning through the use of artificial neural networks. Their pioneering work has been crucial in the development and implementation of algorithms that enable machines to learn and process information in a manner that mimics human cognitive functions. This advancement in artificial intelligence technology has had a profound impact on numerous industries, facilitating innovations across various applications, from image and speech recognition to self-driving cars.
```
## Key A2A Integration Points
To make the PocketFlow agent A2A-compatible, the following were essential:
1. **A2A Server ([`common/server/server.py`](common/server/server.py)):** An ASGI application (using Starlette/Uvicorn) that listens for HTTP POST requests, parses JSON-RPC, and routes requests based on the `method` field.
2. **A2A Data Types ([`common/types.py`](common/types.py)):** Pydantic models defining the structure of A2A messages, tasks, artifacts, errors, and the agent card, ensuring compliance with the `a2a.json` specification.
3. **Task Manager ([`task_manager.py`](task_manager.py)):** A custom class (`PocketFlowTaskManager`) inheriting from the common `InMemoryTaskManager`. Its primary role is implementing the `on_send_task` method (and potentially others like `on_send_task_subscribe` if streaming were supported). This method:
* Receives the validated A2A `SendTaskRequest`.
* Extracts the user's query (`TextPart`) from the request's `message`.
* Initializes the PocketFlow `shared_data` dictionary.
* Creates and runs the PocketFlow `agent_flow`.
* Retrieves the final answer from the `shared_data` dictionary *after* the flow completes.
* Updates the task's state (e.g., to `COMPLETED` or `FAILED`) in the `InMemoryTaskManager`'s store.
* Packages the final answer into an A2A `Artifact` containing a `TextPart`.
* Constructs the final A2A `Task` object for the response.
4. **Agent Card ([`a2a_server.py`](a2a_server.py)):** A Pydantic model (`AgentCard`) defining the agent's metadata (name, description, URL, capabilities, skills) served at `/.well-known/agent.json`.
5. **Server Entry Point ([`a2a_server.py`](a2a_server.py)):** A script that initializes the `AgentCard`, the `PocketFlowTaskManager`, and the `A2AServer`, then starts the Uvicorn server process.
cookbook/pocketflow-a2a/a2a_client.py
+159
View File
@@ -0,0 +1,159 @@
import asyncio
import asyncclick as click # Using asyncclick for async main
from uuid import uuid4
import json # For potentially inspecting raw errors
import anyio
import functools
import logging
# Import from the common directory placed alongside this script
from common.client import A2AClient
from common.types import (
TaskState,
A2AClientError,
TextPart, # Used to construct the message
JSONRPCResponse # Potentially useful for error checking
)
# --- Configure logging ---
# Set level to INFO to see client requests and responses
# Set level to DEBUG to see raw response bodies and SSE data lines
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
# Optionally silence overly verbose libraries
# logging.getLogger("httpx").setLevel(logging.WARNING)
# logging.getLogger("httpcore").setLevel(logging.WARNING)
# --- ANSI Colors (Optional but helpful) ---
C_RED = "\x1b[31m"
C_GREEN = "\x1b[32m"
C_YELLOW = "\x1b[33m"
C_BLUE = "\x1b[34m"
C_MAGENTA = "\x1b[35m"
C_CYAN = "\x1b[36m"
C_WHITE = "\x1b[37m"
C_GRAY = "\x1b[90m"
C_BRIGHT_MAGENTA = "\x1b[95m"
C_RESET = "\x1b[0m"
C_BOLD = "\x1b[1m"
def colorize(color, text):
return f"{color}{text}{C_RESET}"
@click.command()
@click.option(
"--agent-url",
default="http://localhost:10003", # Default to the port used in server __main__
help="URL of the PocketFlow A2A agent server.",
)
async def cli(agent_url: str):
"""Minimal CLI client to interact with an A2A agent."""
print(colorize(C_BRIGHT_MAGENTA, f"Connecting to agent at: {agent_url}"))
# Instantiate the client - only URL is needed if not fetching card first
# Note: The PocketFlow wrapper doesn't expose much via the AgentCard,
# so we skip fetching it for this minimal client.
client = A2AClient(url=agent_url)
sessionId = uuid4().hex # Generate a new session ID for this run
print(colorize(C_GRAY, f"Using Session ID: {sessionId}"))
while True:
taskId = uuid4().hex # Generate a new task ID for each interaction
try:
# Use functools.partial to prepare the prompt function call
prompt_func = functools.partial(
click.prompt,
colorize(C_CYAN, "\nEnter your question (:q or quit to exit)"),
prompt_suffix=" > ",
type=str
)
# Run the synchronous prompt function in a worker thread
prompt = await anyio.to_thread.run_sync(prompt_func)
except (EOFError, RuntimeError, KeyboardInterrupt):
# Catch potential errors during input or if stdin closes
print(colorize(C_RED, "\nInput closed or interrupted. Exiting."))
break
if prompt.lower() in [":q", "quit"]:
print(colorize(C_YELLOW, "Exiting client."))
break
# --- Construct A2A Request Payload ---
payload = {
"id": taskId,
"sessionId": sessionId,
"message": {
"role": "user",
"parts": [
{
"type": "text", # Explicitly match TextPart structure
"text": prompt,
}
],
},
"acceptedOutputModes": ["text", "text/plain"], # What the client wants back
# historyLength could be added if needed
}
print(colorize(C_GRAY, f"Sending task {taskId}..."))
try:
# --- Send Task (Non-Streaming) ---
response = await client.send_task(payload)
# --- Process Response ---
if response.error:
print(colorize(C_RED, f"Error from agent (Code: {response.error.code}): {response.error.message}"))
if response.error.data:
print(colorize(C_GRAY, f"Error Data: {response.error.data}"))
elif response.result:
task_result = response.result
print(colorize(C_GREEN, f"Task {task_result.id} finished with state: {task_result.status.state}"))
final_answer = "Agent did not provide a final artifact."
# Extract answer from artifacts (as implemented in PocketFlowTaskManager)
if task_result.artifacts:
try:
# Find the first text part in the first artifact
first_artifact = task_result.artifacts[0]
first_text_part = next(
(p for p in first_artifact.parts if isinstance(p, TextPart)),
None
)
if first_text_part:
final_answer = first_text_part.text
else:
final_answer = f"(Non-text artifact received: {first_artifact.parts})"
except (IndexError, AttributeError, TypeError) as e:
final_answer = f"(Error parsing artifact: {e})"
elif task_result.status.message and task_result.status.message.parts:
# Fallback to status message if no artifact
try:
first_text_part = next(
(p for p in task_result.status.message.parts if isinstance(p, TextPart)),
None
)
if first_text_part:
final_answer = f"(Final status message: {first_text_part.text})"
except (AttributeError, TypeError) as e:
final_answer = f"(Error parsing status message: {e})"
print(colorize(C_BOLD + C_WHITE, f"\nAgent Response:\n{final_answer}"))
else:
# Should not happen if error is None
print(colorize(C_YELLOW, "Received response with no result and no error."))
except A2AClientError as e:
print(colorize(C_RED, f"\nClient Error: {e}"))
except Exception as e:
print(colorize(C_RED, f"\nAn unexpected error occurred: {e}"))
if __name__ == "__main__":
asyncio.run(cli())
cookbook/pocketflow-a2a/a2a_server.py
+91
View File
@@ -0,0 +1,91 @@
import click
import logging
import os
# Import from the common code you copied
from common.server import A2AServer
from common.types import AgentCard, AgentCapabilities, AgentSkill, MissingAPIKeyError
# Import your custom TaskManager (which now imports from your original files)
from task_manager import PocketFlowTaskManager
# --- Configure logging ---
# Set level to INFO to see server start, requests, responses
# Set level to DEBUG to see raw response bodies from client
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
# Optionally silence overly verbose libraries
# logging.getLogger("httpx").setLevel(logging.WARNING)
# logging.getLogger("httpcore").setLevel(logging.WARNING)
# logging.getLogger("uvicorn.access").setLevel(logging.WARNING)
logger = logging.getLogger(__name__)
@click.command()
@click.option("--host", "host", default="localhost")
@click.option("--port", "port", default=10003) # Use a different port from other agents
def main(host, port):
"""Starts the PocketFlow A2A Agent server."""
try:
# Check for necessary API keys (add others if needed)
if not os.getenv("OPENAI_API_KEY"):
raise MissingAPIKeyError("OPENAI_API_KEY environment variable not set.")
# --- Define the Agent Card ---
capabilities = AgentCapabilities(
streaming=False, # This simple implementation is synchronous
pushNotifications=False,
stateTransitionHistory=False # PocketFlow state isn't exposed via A2A history
)
skill = AgentSkill(
id="web_research_qa",
name="Web Research and Answering",
description="Answers questions using web search results when necessary.",
tags=["research", "qa", "web search"],
examples=[
"Who won the Nobel Prize in Physics 2024?",
"What is quantum computing?",
"Summarize the latest news about AI.",
],
# Input/Output modes defined in the TaskManager
inputModes=PocketFlowTaskManager.SUPPORTED_CONTENT_TYPES,
outputModes=PocketFlowTaskManager.SUPPORTED_CONTENT_TYPES,
)
agent_card = AgentCard(
name="PocketFlow Research Agent (A2A Wrapped)",
description="A simple research agent based on PocketFlow, made accessible via A2A.",
url=f"http://{host}:{port}/", # The endpoint A2A clients will use
version="0.1.0-a2a",
capabilities=capabilities,
skills=[skill],
# Assuming no specific provider or auth for this example
provider=None,
authentication=None,
defaultInputModes=PocketFlowTaskManager.SUPPORTED_CONTENT_TYPES,
defaultOutputModes=PocketFlowTaskManager.SUPPORTED_CONTENT_TYPES,
)
# --- Initialize and Start Server ---
task_manager = PocketFlowTaskManager() # Instantiate your custom manager
server = A2AServer(
agent_card=agent_card,
task_manager=task_manager,
host=host,
port=port,
)
logger.info(f"Starting PocketFlow A2A server on http://{host}:{port}")
server.start()
except MissingAPIKeyError as e:
logger.error(f"Configuration Error: {e}")
exit(1)
except Exception as e:
logger.error(f"An error occurred during server startup: {e}", exc_info=True)
exit(1)
if __name__ == "__main__":
main()
cookbook/pocketflow-a2a/common/__init__.py
View File
cookbook/pocketflow-a2a/common/client/__init__.py
+4
View File
@@ -0,0 +1,4 @@
from .client import A2AClient
from .card_resolver import A2ACardResolver
__all__ = ["A2AClient", "A2ACardResolver"]
cookbook/pocketflow-a2a/common/client/card_resolver.py
+21
View File
@@ -0,0 +1,21 @@
import httpx
from common.types import (
AgentCard,
A2AClientJSONError,
)
import json
class A2ACardResolver:
def __init__(self, base_url, agent_card_path="/.well-known/agent.json"):
self.base_url = base_url.rstrip("/")
self.agent_card_path = agent_card_path.lstrip("/")
def get_agent_card(self) -> AgentCard:
with httpx.Client() as client:
response = client.get(self.base_url + "/" + self.agent_card_path)
response.raise_for_status()
try:
return AgentCard(**response.json())
except json.JSONDecodeError as e:
raise A2AClientJSONError(str(e)) from e
cookbook/pocketflow-a2a/common/client/client.py
+182
View File
@@ -0,0 +1,182 @@
import httpx
from httpx_sse import connect_sse
from typing import Any, AsyncIterable
from common.types import (
AgentCard,
GetTaskRequest,
SendTaskRequest,
SendTaskResponse,
JSONRPCRequest,
JSONRPCResponse,
JSONRPCError,
GetTaskResponse,
CancelTaskResponse,
CancelTaskRequest,
SetTaskPushNotificationRequest,
SetTaskPushNotificationResponse,
GetTaskPushNotificationRequest,
GetTaskPushNotificationResponse,
A2AClientHTTPError,
A2AClientJSONError,
SendTaskStreamingRequest,
SendTaskStreamingResponse,
Task,
TaskPushNotificationConfig,
TaskStatusUpdateEvent,
TaskArtifactUpdateEvent,
)
import json
import logging
# Configure a logger specific to the client
logger = logging.getLogger("A2AClient")
class A2AClientError(Exception):
"""Base class for A2A client errors"""
def __init__(self, message):
super().__init__(message)
class RpcError(Exception):
code: int
data: Any = None
def __init__(self, code: int, message: str, data: Any = None):
super().__init__(message)
self.name = "RpcError"
self.code = code
self.data = data
class A2AClient:
def __init__(self, agent_card: AgentCard = None, url: str = None):
if agent_card:
self.url = agent_card.url.rstrip("/")
elif url:
self.url = url.rstrip("/")
else:
raise ValueError("Must provide either agent_card or url")
self.fetchImpl = httpx.AsyncClient(timeout=None)
def _generateRequestId(self):
import time
return int(time.time() * 1000)
async def _send_request(self, request: JSONRPCRequest) -> dict[str, Any]:
req_id = request.id
req_method = request.method
req_dump = request.model_dump(exclude_none=True)
logger.info(f"-> Sending Request (ID: {req_id}, Method: {req_method}):\n{json.dumps(req_dump, indent=2)}")
try:
response = await self.fetchImpl.post(
self.url, json=req_dump, timeout=60.0
)
logger.info(f"<- Received HTTP Status {response.status_code} for Request (ID: {req_id})")
response_text = await response.aread()
logger.debug(f"Raw Response Body (ID: {req_id}):\n{response_text.decode('utf-8', errors='replace')}")
response.raise_for_status()
try:
json_response = json.loads(response_text)
except json.JSONDecodeError as e:
logger.error(f"Failed to decode JSON response (ID: {req_id}): {e}")
raise A2AClientJSONError(f"Failed to decode JSON: {e}") from e
if "error" in json_response and json_response["error"] is not None:
rpc_error = json_response["error"]
logger.warning(f"<- Received JSON-RPC Error (ID: {req_id}): Code={rpc_error.get('code')}, Msg='{rpc_error.get('message')}'")
raise RpcError(rpc_error.get("code", -32000), rpc_error.get("message", "Unknown RPC Error"), rpc_error.get("data"))
logger.info(f"<- Received Success Response (ID: {req_id}):\n{json.dumps(json_response, indent=2)}")
return json_response
except httpx.HTTPStatusError as e:
logger.error(f"HTTP Error for Request (ID: {req_id}): {e.response.status_code} - {e.request.url}")
error_body = await e.response.aread()
raise A2AClientHTTPError(e.response.status_code, f"{e}. Body: {error_body.decode('utf-8', errors='replace')}") from e
except httpx.RequestError as e:
logger.error(f"Request Error for (ID: {req_id}): {e}")
raise A2AClientError(f"Network or request error: {e}") from e
except RpcError:
raise
except Exception as e:
logger.error(f"Unexpected error during request (ID: {req_id}): {e}", exc_info=True)
raise A2AClientError(f"Unexpected error: {e}") from e
async def send_task(self, payload: dict[str, Any]) -> SendTaskResponse:
request = SendTaskRequest(params=payload)
response_dict = await self._send_request(request)
return SendTaskResponse(**response_dict)
async def send_task_streaming(
self, payload: dict[str, Any]
) -> AsyncIterable[SendTaskStreamingResponse]:
request = SendTaskStreamingRequest(params=payload)
req_id = request.id
req_dump = request.model_dump(exclude_none=True)
logger.info(f"-> Sending Streaming Request (ID: {req_id}, Method: {request.method}):\n{json.dumps(req_dump, indent=2)}")
try:
async with self.fetchImpl.stream("POST", self.url, json=req_dump, timeout=None) as response:
logger.info(f"<- Received HTTP Status {response.status_code} for Streaming Request (ID: {req_id})")
response.raise_for_status()
buffer = ""
async for line in response.aiter_lines():
if not line:
if buffer.startswith("data:"):
data_str = buffer[len("data:"):].strip()
logger.debug(f"Received SSE Data Line (ID: {req_id}): {data_str}")
try:
sse_data_dict = json.loads(data_str)
yield SendTaskStreamingResponse(**sse_data_dict)
except json.JSONDecodeError as e:
logger.error(f"Failed to decode SSE JSON (ID: {req_id}): {e}. Data: '{data_str}'")
except Exception as e:
logger.error(f"Error processing SSE data (ID: {req_id}): {e}. Data: '{data_str}'", exc_info=True)
elif buffer:
logger.debug(f"Received non-data SSE line (ID: {req_id}): {buffer}")
buffer = ""
else:
buffer += line + "\n"
if buffer:
logger.warning(f"SSE stream ended with partial data in buffer (ID: {req_id}): {buffer}")
logger.info(f"SSE Stream ended for request ID: {req_id}")
except httpx.HTTPStatusError as e:
logger.error(f"HTTP Error during streaming connection (ID: {req_id}): {e.response.status_code} - {e.request.url}")
error_body = await e.response.aread()
raise A2AClientHTTPError(e.response.status_code, f"{e}. Body: {error_body.decode('utf-8', errors='replace')}") from e
except httpx.RequestError as e:
logger.error(f"Request Error during streaming (ID: {req_id}): {e}")
raise A2AClientError(f"Network or request error during streaming: {e}") from e
except Exception as e:
logger.error(f"Unexpected error during streaming (ID: {req_id}): {e}", exc_info=True)
raise A2AClientError(f"Unexpected streaming error: {e}") from e
async def get_task(self, payload: dict[str, Any]) -> GetTaskResponse:
request = GetTaskRequest(params=payload)
response_dict = await self._send_request(request)
return GetTaskResponse(**response_dict)
async def cancel_task(self, payload: dict[str, Any]) -> CancelTaskResponse:
request = CancelTaskRequest(params=payload)
response_dict = await self._send_request(request)
return CancelTaskResponse(**response_dict)
async def set_task_callback(
self, payload: dict[str, Any]
) -> SetTaskPushNotificationResponse:
request = SetTaskPushNotificationRequest(params=payload)
response_dict = await self._send_request(request)
return SetTaskPushNotificationResponse(**response_dict)
async def get_task_callback(
self, payload: dict[str, Any]
) -> GetTaskPushNotificationResponse:
request = GetTaskPushNotificationRequest(params=payload)
response_dict = await self._send_request(request)
return GetTaskPushNotificationResponse(**response_dict)
cookbook/pocketflow-a2a/common/server/__init__.py
+4
View File
@@ -0,0 +1,4 @@
from .server import A2AServer
from .task_manager import TaskManager, InMemoryTaskManager
__all__ = ["A2AServer", "TaskManager", "InMemoryTaskManager"]
cookbook/pocketflow-a2a/common/server/server.py
+168
View File
@@ -0,0 +1,168 @@
from starlette.applications import Starlette
from starlette.responses import JSONResponse
from sse_starlette.sse import EventSourceResponse
from starlette.requests import Request
from common.types import (
A2ARequest,
JSONRPCResponse,
InvalidRequestError,
JSONParseError,
GetTaskRequest,
CancelTaskRequest,
SendTaskRequest,
SetTaskPushNotificationRequest,
GetTaskPushNotificationRequest,
InternalError,
AgentCard,
TaskResubscriptionRequest,
SendTaskStreamingRequest,
Message,
)
from pydantic import ValidationError
import json
from typing import AsyncIterable, Any
from common.server.task_manager import TaskManager
import logging
# Configure a logger specific to the server
logger = logging.getLogger("A2AServer")
class A2AServer:
def __init__(
self,
host="0.0.0.0",
port=5000,
endpoint="/",
agent_card: AgentCard = None,
task_manager: TaskManager = None,
):
self.host = host
self.port = port
self.endpoint = endpoint
self.task_manager = task_manager
self.agent_card = agent_card
self.app = Starlette()
self.app.add_route(self.endpoint, self._process_request, methods=["POST"])
self.app.add_route(
"/.well-known/agent.json", self._get_agent_card, methods=["GET"]
)
def start(self):
if self.agent_card is None:
raise ValueError("agent_card is not defined")
if self.task_manager is None:
raise ValueError("request_handler is not defined")
import uvicorn
# Basic logging config moved to __main__.py for application-level control
uvicorn.run(self.app, host=self.host, port=self.port)
def _get_agent_card(self, request: Request) -> JSONResponse:
logger.info("Serving Agent Card request")
return JSONResponse(self.agent_card.model_dump(exclude_none=True))
async def _process_request(self, request: Request):
request_id_for_log = "N/A" # Default if parsing fails early
raw_body = b""
try:
# Log raw body first
raw_body = await request.body()
body = json.loads(raw_body) # Attempt parsing
request_id_for_log = body.get("id", "N/A") # Get ID if possible
logger.info(f"<- Received Request (ID: {request_id_for_log}):\n{json.dumps(body, indent=2)}")
json_rpc_request = A2ARequest.validate_python(body)
# Route based on method (same as before)
if isinstance(json_rpc_request, GetTaskRequest):
result = await self.task_manager.on_get_task(json_rpc_request)
elif isinstance(json_rpc_request, SendTaskRequest):
result = await self.task_manager.on_send_task(json_rpc_request)
elif isinstance(json_rpc_request, SendTaskStreamingRequest):
result = await self.task_manager.on_send_task_subscribe(
json_rpc_request
)
elif isinstance(json_rpc_request, CancelTaskRequest):
result = await self.task_manager.on_cancel_task(json_rpc_request)
elif isinstance(json_rpc_request, SetTaskPushNotificationRequest):
result = await self.task_manager.on_set_task_push_notification(json_rpc_request)
elif isinstance(json_rpc_request, GetTaskPushNotificationRequest):
result = await self.task_manager.on_get_task_push_notification(json_rpc_request)
elif isinstance(json_rpc_request, TaskResubscriptionRequest):
result = await self.task_manager.on_resubscribe_to_task(
json_rpc_request
)
else:
logger.warning(f"Unexpected request type: {type(json_rpc_request)}")
raise ValueError(f"Unexpected request type: {type(request)}")
return self._create_response(result) # Pass result to response creation
except json.decoder.JSONDecodeError as e:
logger.error(f"JSON Parse Error for Request body: <<<{raw_body.decode('utf-8', errors='replace')}>>>\nError: {e}")
return self._handle_exception(e, request_id_for_log) # Pass ID if known
except ValidationError as e:
logger.error(f"Request Validation Error (ID: {request_id_for_log}): {e.json()}")
return self._handle_exception(e, request_id_for_log)
except Exception as e:
logger.error(f"Unhandled Exception processing request (ID: {request_id_for_log}): {e}", exc_info=True)
return self._handle_exception(e, request_id_for_log) # Pass ID if known
def _handle_exception(self, e: Exception, req_id=None) -> JSONResponse: # Accept req_id
if isinstance(e, json.decoder.JSONDecodeError):
json_rpc_error = JSONParseError()
elif isinstance(e, ValidationError):
json_rpc_error = InvalidRequestError(data=json.loads(e.json()))
else:
# Log the full exception details
logger.error(f"Internal Server Error (ReqID: {req_id}): {e}", exc_info=True)
json_rpc_error = InternalError(message=f"Internal Server Error: {type(e).__name__}")
response = JSONRPCResponse(id=req_id, error=json_rpc_error)
response_dump = response.model_dump(exclude_none=True)
logger.info(f"-> Sending Error Response (ReqID: {req_id}):\n{json.dumps(response_dump, indent=2)}")
# A2A errors are still sent with HTTP 200
return JSONResponse(response_dump, status_code=200)
def _create_response(self, result: Any) -> JSONResponse | EventSourceResponse:
if isinstance(result, AsyncIterable):
# Streaming response
async def event_generator(result_stream) -> AsyncIterable[dict[str, str]]:
stream_request_id = None # Capture ID from the first event if possible
try:
async for item in result_stream:
# Log each streamed item
response_json = item.model_dump_json(exclude_none=True)
stream_request_id = item.id # Update ID
logger.info(f"-> Sending SSE Event (ID: {stream_request_id}):\n{json.dumps(json.loads(response_json), indent=2)}")
yield {"data": response_json}
logger.info(f"SSE Stream ended for request ID: {stream_request_id}")
except Exception as e:
logger.error(f"Error during SSE generation (ReqID: {stream_request_id}): {e}", exc_info=True)
# Optionally yield an error event if the protocol allows/requires it
# error_payload = JSONRPCResponse(id=stream_request_id, error=InternalError(message=f"SSE Error: {e}"))
# yield {"data": error_payload.model_dump_json(exclude_none=True)}
logger.info("Starting SSE stream...") # Log stream start
return EventSourceResponse(event_generator(result))
elif isinstance(result, JSONRPCResponse):
# Standard JSON response
response_dump = result.model_dump(exclude_none=True)
log_id = result.id if result.id is not None else "N/A (Notification?)"
log_prefix = "->"
log_type = "Response"
if result.error:
log_prefix = "-> Sending Error"
log_type = "Error Response"
logger.info(f"{log_prefix} {log_type} (ID: {log_id}):\n{json.dumps(response_dump, indent=2)}")
return JSONResponse(response_dump)
else:
# This should ideally not happen if task manager returns correctly
logger.error(f"Task manager returned unexpected type: {type(result)}")
err_resp = JSONRPCResponse(id=None, error=InternalError(message="Invalid internal response type"))
return JSONResponse(err_resp.model_dump(exclude_none=True), status_code=500)
cookbook/pocketflow-a2a/common/server/task_manager.py
+277
View File
@@ -0,0 +1,277 @@
from abc import ABC, abstractmethod
from typing import Union, AsyncIterable, List
from common.types import Task
from common.types import (
JSONRPCResponse,
TaskIdParams,
TaskQueryParams,
GetTaskRequest,
TaskNotFoundError,
SendTaskRequest,
CancelTaskRequest,
TaskNotCancelableError,
SetTaskPushNotificationRequest,
GetTaskPushNotificationRequest,
GetTaskResponse,
CancelTaskResponse,
SendTaskResponse,
SetTaskPushNotificationResponse,
GetTaskPushNotificationResponse,
PushNotificationNotSupportedError,
TaskSendParams,
TaskStatus,
TaskState,
TaskResubscriptionRequest,
SendTaskStreamingRequest,
SendTaskStreamingResponse,
Artifact,
PushNotificationConfig,
TaskStatusUpdateEvent,
JSONRPCError,
TaskPushNotificationConfig,
InternalError,
)
from common.server.utils import new_not_implemented_error
import asyncio
import logging
logger = logging.getLogger(__name__)
class TaskManager(ABC):
@abstractmethod
async def on_get_task(self, request: GetTaskRequest) -> GetTaskResponse:
pass
@abstractmethod
async def on_cancel_task(self, request: CancelTaskRequest) -> CancelTaskResponse:
pass
@abstractmethod
async def on_send_task(self, request: SendTaskRequest) -> SendTaskResponse:
pass
@abstractmethod
async def on_send_task_subscribe(
self, request: SendTaskStreamingRequest
) -> Union[AsyncIterable[SendTaskStreamingResponse], JSONRPCResponse]:
pass
@abstractmethod
async def on_set_task_push_notification(
self, request: SetTaskPushNotificationRequest
) -> SetTaskPushNotificationResponse:
pass
@abstractmethod
async def on_get_task_push_notification(
self, request: GetTaskPushNotificationRequest
) -> GetTaskPushNotificationResponse:
pass
@abstractmethod
async def on_resubscribe_to_task(
self, request: TaskResubscriptionRequest
) -> Union[AsyncIterable[SendTaskResponse], JSONRPCResponse]:
pass
class InMemoryTaskManager(TaskManager):
def __init__(self):
self.tasks: dict[str, Task] = {}
self.push_notification_infos: dict[str, PushNotificationConfig] = {}
self.lock = asyncio.Lock()
self.task_sse_subscribers: dict[str, List[asyncio.Queue]] = {}
self.subscriber_lock = asyncio.Lock()
async def on_get_task(self, request: GetTaskRequest) -> GetTaskResponse:
logger.info(f"Getting task {request.params.id}")
task_query_params: TaskQueryParams = request.params
async with self.lock:
task = self.tasks.get(task_query_params.id)
if task is None:
return GetTaskResponse(id=request.id, error=TaskNotFoundError())
task_result = self.append_task_history(
task, task_query_params.historyLength
)
return GetTaskResponse(id=request.id, result=task_result)
async def on_cancel_task(self, request: CancelTaskRequest) -> CancelTaskResponse:
logger.info(f"Cancelling task {request.params.id}")
task_id_params: TaskIdParams = request.params
async with self.lock:
task = self.tasks.get(task_id_params.id)
if task is None:
return CancelTaskResponse(id=request.id, error=TaskNotFoundError())
return CancelTaskResponse(id=request.id, error=TaskNotCancelableError())
@abstractmethod
async def on_send_task(self, request: SendTaskRequest) -> SendTaskResponse:
pass
@abstractmethod
async def on_send_task_subscribe(
self, request: SendTaskStreamingRequest
) -> Union[AsyncIterable[SendTaskStreamingResponse], JSONRPCResponse]:
pass
async def set_push_notification_info(self, task_id: str, notification_config: PushNotificationConfig):
async with self.lock:
task = self.tasks.get(task_id)
if task is None:
raise ValueError(f"Task not found for {task_id}")
self.push_notification_infos[task_id] = notification_config
return
async def get_push_notification_info(self, task_id: str) -> PushNotificationConfig:
async with self.lock:
task = self.tasks.get(task_id)
if task is None:
raise ValueError(f"Task not found for {task_id}")
return self.push_notification_infos[task_id]
return
async def has_push_notification_info(self, task_id: str) -> bool:
async with self.lock:
return task_id in self.push_notification_infos
async def on_set_task_push_notification(
self, request: SetTaskPushNotificationRequest
) -> SetTaskPushNotificationResponse:
logger.info(f"Setting task push notification {request.params.id}")
task_notification_params: TaskPushNotificationConfig = request.params
try:
await self.set_push_notification_info(task_notification_params.id, task_notification_params.pushNotificationConfig)
except Exception as e:
logger.error(f"Error while setting push notification info: {e}")
return JSONRPCResponse(
id=request.id,
error=InternalError(
message="An error occurred while setting push notification info"
),
)
return SetTaskPushNotificationResponse(id=request.id, result=task_notification_params)
async def on_get_task_push_notification(
self, request: GetTaskPushNotificationRequest
) -> GetTaskPushNotificationResponse:
logger.info(f"Getting task push notification {request.params.id}")
task_params: TaskIdParams = request.params
try:
notification_info = await self.get_push_notification_info(task_params.id)
except Exception as e:
logger.error(f"Error while getting push notification info: {e}")
return GetTaskPushNotificationResponse(
id=request.id,
error=InternalError(
message="An error occurred while getting push notification info"
),
)
return GetTaskPushNotificationResponse(id=request.id, result=TaskPushNotificationConfig(id=task_params.id, pushNotificationConfig=notification_info))
async def upsert_task(self, task_send_params: TaskSendParams) -> Task:
logger.info(f"Upserting task {task_send_params.id}")
async with self.lock:
task = self.tasks.get(task_send_params.id)
if task is None:
task = Task(
id=task_send_params.id,
sessionId = task_send_params.sessionId,
messages=[task_send_params.message],
status=TaskStatus(state=TaskState.SUBMITTED),
history=[task_send_params.message],
)
self.tasks[task_send_params.id] = task
else:
task.history.append(task_send_params.message)
return task
async def on_resubscribe_to_task(
self, request: TaskResubscriptionRequest
) -> Union[AsyncIterable[SendTaskStreamingResponse], JSONRPCResponse]:
return new_not_implemented_error(request.id)
async def update_store(
self, task_id: str, status: TaskStatus, artifacts: list[Artifact]
) -> Task:
async with self.lock:
try:
task = self.tasks[task_id]
except KeyError:
logger.error(f"Task {task_id} not found for updating the task")
raise ValueError(f"Task {task_id} not found")
task.status = status
if status.message is not None:
task.history.append(status.message)
if artifacts is not None:
if task.artifacts is None:
task.artifacts = []
task.artifacts.extend(artifacts)
return task
def append_task_history(self, task: Task, historyLength: int | None):
new_task = task.model_copy()
if historyLength is not None and historyLength > 0:
new_task.history = new_task.history[-historyLength:]
else:
new_task.history = []
return new_task
async def setup_sse_consumer(self, task_id: str, is_resubscribe: bool = False):
async with self.subscriber_lock:
if task_id not in self.task_sse_subscribers:
if is_resubscribe:
raise ValueError("Task not found for resubscription")
else:
self.task_sse_subscribers[task_id] = []
sse_event_queue = asyncio.Queue(maxsize=0) # <=0 is unlimited
self.task_sse_subscribers[task_id].append(sse_event_queue)
return sse_event_queue
async def enqueue_events_for_sse(self, task_id, task_update_event):
async with self.subscriber_lock:
if task_id not in self.task_sse_subscribers:
return
current_subscribers = self.task_sse_subscribers[task_id]
for subscriber in current_subscribers:
await subscriber.put(task_update_event)
async def dequeue_events_for_sse(
self, request_id, task_id, sse_event_queue: asyncio.Queue
) -> AsyncIterable[SendTaskStreamingResponse] | JSONRPCResponse:
try:
while True:
event = await sse_event_queue.get()
if isinstance(event, JSONRPCError):
yield SendTaskStreamingResponse(id=request_id, error=event)
break
yield SendTaskStreamingResponse(id=request_id, result=event)
if isinstance(event, TaskStatusUpdateEvent) and event.final:
break
finally:
async with self.subscriber_lock:
if task_id in self.task_sse_subscribers:
self.task_sse_subscribers[task_id].remove(sse_event_queue)
cookbook/pocketflow-a2a/common/server/utils.py
+28
View File
@@ -0,0 +1,28 @@
from common.types import (
JSONRPCResponse,
ContentTypeNotSupportedError,
UnsupportedOperationError,
)
from typing import List
def are_modalities_compatible(
server_output_modes: List[str], client_output_modes: List[str]
):
"""Modalities are compatible if they are both non-empty
and there is at least one common element."""
if client_output_modes is None or len(client_output_modes) == 0:
return True
if server_output_modes is None or len(server_output_modes) == 0:
return True
return any(x in server_output_modes for x in client_output_modes)
def new_incompatible_types_error(request_id):
return JSONRPCResponse(id=request_id, error=ContentTypeNotSupportedError())
def new_not_implemented_error(request_id):
return JSONRPCResponse(id=request_id, error=UnsupportedOperationError())
cookbook/pocketflow-a2a/common/types.py
+365
View File
@@ -0,0 +1,365 @@
from typing import Union, Any
from pydantic import BaseModel, Field, TypeAdapter
from typing import Literal, List, Annotated, Optional
from datetime import datetime
from pydantic import model_validator, ConfigDict, field_serializer
from uuid import uuid4
from enum import Enum
from typing_extensions import Self
class TaskState(str, Enum):
SUBMITTED = "submitted"
WORKING = "working"
INPUT_REQUIRED = "input-required"
COMPLETED = "completed"
CANCELED = "canceled"
FAILED = "failed"
UNKNOWN = "unknown"
class TextPart(BaseModel):
type: Literal["text"] = "text"
text: str
metadata: dict[str, Any] | None = None
class FileContent(BaseModel):
name: str | None = None
mimeType: str | None = None
bytes: str | None = None
uri: str | None = None
@model_validator(mode="after")
def check_content(self) -> Self:
if not (self.bytes or self.uri):
raise ValueError("Either 'bytes' or 'uri' must be present in the file data")
if self.bytes and self.uri:
raise ValueError(
"Only one of 'bytes' or 'uri' can be present in the file data"
)
return self
class FilePart(BaseModel):
type: Literal["file"] = "file"
file: FileContent
metadata: dict[str, Any] | None = None
class DataPart(BaseModel):
type: Literal["data"] = "data"
data: dict[str, Any]
metadata: dict[str, Any] | None = None
Part = Annotated[Union[TextPart, FilePart, DataPart], Field(discriminator="type")]
class Message(BaseModel):
role: Literal["user", "agent"]
parts: List[Part]
metadata: dict[str, Any] | None = None
class TaskStatus(BaseModel):
state: TaskState
message: Message | None = None
timestamp: datetime = Field(default_factory=datetime.now)
@field_serializer("timestamp")
def serialize_dt(self, dt: datetime, _info):
return dt.isoformat()
class Artifact(BaseModel):
name: str | None = None
description: str | None = None
parts: List[Part]
metadata: dict[str, Any] | None = None
index: int = 0
append: bool | None = None
lastChunk: bool | None = None
class Task(BaseModel):
id: str
sessionId: str | None = None
status: TaskStatus
artifacts: List[Artifact] | None = None
history: List[Message] | None = None
metadata: dict[str, Any] | None = None
class TaskStatusUpdateEvent(BaseModel):
id: str
status: TaskStatus
final: bool = False
metadata: dict[str, Any] | None = None
class TaskArtifactUpdateEvent(BaseModel):
id: str
artifact: Artifact
metadata: dict[str, Any] | None = None
class AuthenticationInfo(BaseModel):
model_config = ConfigDict(extra="allow")
schemes: List[str]
credentials: str | None = None
class PushNotificationConfig(BaseModel):
url: str
token: str | None = None
authentication: AuthenticationInfo | None = None
class TaskIdParams(BaseModel):
id: str
metadata: dict[str, Any] | None = None
class TaskQueryParams(TaskIdParams):
historyLength: int | None = None
class TaskSendParams(BaseModel):
id: str
sessionId: str = Field(default_factory=lambda: uuid4().hex)
message: Message
acceptedOutputModes: Optional[List[str]] = None
pushNotification: PushNotificationConfig | None = None
historyLength: int | None = None
metadata: dict[str, Any] | None = None
class TaskPushNotificationConfig(BaseModel):
id: str
pushNotificationConfig: PushNotificationConfig
## RPC Messages
class JSONRPCMessage(BaseModel):
jsonrpc: Literal["2.0"] = "2.0"
id: int | str | None = Field(default_factory=lambda: uuid4().hex)
class JSONRPCRequest(JSONRPCMessage):
method: str
params: dict[str, Any] | None = None
class JSONRPCError(BaseModel):
code: int
message: str
data: Any | None = None
class JSONRPCResponse(JSONRPCMessage):
result: Any | None = None
error: JSONRPCError | None = None
class SendTaskRequest(JSONRPCRequest):
method: Literal["tasks/send"] = "tasks/send"
params: TaskSendParams
class SendTaskResponse(JSONRPCResponse):
result: Task | None = None
class SendTaskStreamingRequest(JSONRPCRequest):
method: Literal["tasks/sendSubscribe"] = "tasks/sendSubscribe"
params: TaskSendParams
class SendTaskStreamingResponse(JSONRPCResponse):
result: TaskStatusUpdateEvent | TaskArtifactUpdateEvent | None = None
class GetTaskRequest(JSONRPCRequest):
method: Literal["tasks/get"] = "tasks/get"
params: TaskQueryParams
class GetTaskResponse(JSONRPCResponse):
result: Task | None = None
class CancelTaskRequest(JSONRPCRequest):
method: Literal["tasks/cancel",] = "tasks/cancel"
params: TaskIdParams
class CancelTaskResponse(JSONRPCResponse):
result: Task | None = None
class SetTaskPushNotificationRequest(JSONRPCRequest):
method: Literal["tasks/pushNotification/set",] = "tasks/pushNotification/set"
params: TaskPushNotificationConfig
class SetTaskPushNotificationResponse(JSONRPCResponse):
result: TaskPushNotificationConfig | None = None
class GetTaskPushNotificationRequest(JSONRPCRequest):
method: Literal["tasks/pushNotification/get",] = "tasks/pushNotification/get"
params: TaskIdParams
class GetTaskPushNotificationResponse(JSONRPCResponse):
result: TaskPushNotificationConfig | None = None
class TaskResubscriptionRequest(JSONRPCRequest):
method: Literal["tasks/resubscribe",] = "tasks/resubscribe"
params: TaskIdParams
A2ARequest = TypeAdapter(
Annotated[
Union[
SendTaskRequest,
GetTaskRequest,
CancelTaskRequest,
SetTaskPushNotificationRequest,
GetTaskPushNotificationRequest,
TaskResubscriptionRequest,
SendTaskStreamingRequest,
],
Field(discriminator="method"),
]
)
## Error types
class JSONParseError(JSONRPCError):
code: int = -32700
message: str = "Invalid JSON payload"
data: Any | None = None
class InvalidRequestError(JSONRPCError):
code: int = -32600
message: str = "Request payload validation error"
data: Any | None = None
class MethodNotFoundError(JSONRPCError):
code: int = -32601
message: str = "Method not found"
data: None = None
class InvalidParamsError(JSONRPCError):
code: int = -32602
message: str = "Invalid parameters"
data: Any | None = None
class InternalError(JSONRPCError):
code: int = -32603
message: str = "Internal error"
data: Any | None = None
class TaskNotFoundError(JSONRPCError):
code: int = -32001
message: str = "Task not found"
data: None = None
class TaskNotCancelableError(JSONRPCError):
code: int = -32002
message: str = "Task cannot be canceled"
data: None = None
class PushNotificationNotSupportedError(JSONRPCError):
code: int = -32003
message: str = "Push Notification is not supported"
data: None = None
class UnsupportedOperationError(JSONRPCError):
code: int = -32004
message: str = "This operation is not supported"
data: None = None
class ContentTypeNotSupportedError(JSONRPCError):
code: int = -32005
message: str = "Incompatible content types"
data: None = None
class AgentProvider(BaseModel):
organization: str
url: str | None = None
class AgentCapabilities(BaseModel):
streaming: bool = False
pushNotifications: bool = False
stateTransitionHistory: bool = False
class AgentAuthentication(BaseModel):
schemes: List[str]
credentials: str | None = None
class AgentSkill(BaseModel):
id: str
name: str
description: str | None = None
tags: List[str] | None = None
examples: List[str] | None = None
inputModes: List[str] | None = None
outputModes: List[str] | None = None
class AgentCard(BaseModel):
name: str
description: str | None = None
url: str
provider: AgentProvider | None = None
version: str
documentationUrl: str | None = None
capabilities: AgentCapabilities
authentication: AgentAuthentication | None = None
defaultInputModes: List[str] = ["text"]
defaultOutputModes: List[str] = ["text"]
skills: List[AgentSkill]
class A2AClientError(Exception):
pass
class A2AClientHTTPError(A2AClientError):
def __init__(self, status_code: int, message: str):
self.status_code = status_code
self.message = message
super().__init__(f"HTTP Error {status_code}: {message}")
class A2AClientJSONError(A2AClientError):
def __init__(self, message: str):
self.message = message
super().__init__(f"JSON Error: {message}")
class MissingAPIKeyError(Exception):
"""Exception for missing API key."""
pass
cookbook/pocketflow-a2a/common/utils/in_memory_cache.py
+109
View File
@@ -0,0 +1,109 @@
"""In Memory Cache utility."""
import threading
import time
from typing import Any, Dict, Optional
class InMemoryCache:
"""A thread-safe Singleton class to manage cache data.
Ensures only one instance of the cache exists across the application.
"""
_instance: Optional["InMemoryCache"] = None
_lock: threading.Lock = threading.Lock()
_initialized: bool = False
def __new__(cls):
"""Override __new__ to control instance creation (Singleton pattern).
Uses a lock to ensure thread safety during the first instantiation.
Returns:
The singleton instance of InMemoryCache.
"""
if cls._instance is None:
with cls._lock:
if cls._instance is None:
cls._instance = super().__new__(cls)
return cls._instance
def __init__(self):
"""Initialize the cache storage.
Uses a flag (_initialized) to ensure this logic runs only on the very first
creation of the singleton instance.
"""
if not self._initialized:
with self._lock:
if not self._initialized:
# print("Initializing SessionCache storage")
self._cache_data: Dict[str, Dict[str, Any]] = {}
self._ttl: Dict[str, float] = {}
self._data_lock: threading.Lock = threading.Lock()
self._initialized = True
def set(self, key: str, value: Any, ttl: Optional[int] = None) -> None:
"""Set a key-value pair.
Args:
key: The key for the data.
value: The data to store.
ttl: Time to live in seconds. If None, data will not expire.
"""
with self._data_lock:
self._cache_data[key] = value
if ttl is not None:
self._ttl[key] = time.time() + ttl
else:
if key in self._ttl:
del self._ttl[key]
def get(self, key: str, default: Any = None) -> Any:
"""Get the value associated with a key.
Args:
key: The key for the data within the session.
default: The value to return if the session or key is not found.
Returns:
The cached value, or the default value if not found.
"""
with self._data_lock:
if key in self._ttl and time.time() > self._ttl[key]:
del self._cache_data[key]
del self._ttl[key]
return default
return self._cache_data.get(key, default)
def delete(self, key: str) -> None:
"""Delete a specific key-value pair from a cache.
Args:
key: The key to delete.
Returns:
True if the key was found and deleted, False otherwise.
"""
with self._data_lock:
if key in self._cache_data:
del self._cache_data[key]
if key in self._ttl:
del self._ttl[key]
return True
return False
def clear(self) -> bool:
"""Remove all data.
Returns:
True if the data was cleared, False otherwise.
"""
with self._data_lock:
self._cache_data.clear()
self._ttl.clear()
return True
return False
cookbook/pocketflow-a2a/common/utils/push_notification_auth.py
+135
View File
@@ -0,0 +1,135 @@
from jwcrypto import jwk
import uuid
from starlette.responses import JSONResponse
from starlette.requests import Request
from typing import Any
import jwt
import time
import json
import hashlib
import httpx
import logging
from jwt import PyJWK, PyJWKClient
logger = logging.getLogger(__name__)
AUTH_HEADER_PREFIX = 'Bearer '
class PushNotificationAuth:
def _calculate_request_body_sha256(self, data: dict[str, Any]):
"""Calculates the SHA256 hash of a request body.
This logic needs to be same for both the agent who signs the payload and the client verifier.
"""
body_str = json.dumps(
data,
ensure_ascii=False,
allow_nan=False,
indent=None,
separators=(",", ":"),
)
return hashlib.sha256(body_str.encode()).hexdigest()
class PushNotificationSenderAuth(PushNotificationAuth):
def __init__(self):
self.public_keys = []
self.private_key_jwk: PyJWK = None
@staticmethod
async def verify_push_notification_url(url: str) -> bool:
async with httpx.AsyncClient(timeout=10) as client:
try:
validation_token = str(uuid.uuid4())
response = await client.get(
url,
params={"validationToken": validation_token}
)
response.raise_for_status()
is_verified = response.text == validation_token
logger.info(f"Verified push-notification URL: {url} => {is_verified}")
return is_verified
except Exception as e:
logger.warning(f"Error during sending push-notification for URL {url}: {e}")
return False
def generate_jwk(self):
key = jwk.JWK.generate(kty='RSA', size=2048, kid=str(uuid.uuid4()), use="sig")
self.public_keys.append(key.export_public(as_dict=True))
self.private_key_jwk = PyJWK.from_json(key.export_private())
def handle_jwks_endpoint(self, _request: Request):
"""Allow clients to fetch public keys.
"""
return JSONResponse({
"keys": self.public_keys
})
def _generate_jwt(self, data: dict[str, Any]):
"""JWT is generated by signing both the request payload SHA digest and time of token generation.
Payload is signed with private key and it ensures the integrity of payload for client.
Including iat prevents from replay attack.
"""
iat = int(time.time())
return jwt.encode(
{"iat": iat, "request_body_sha256": self._calculate_request_body_sha256(data)},
key=self.private_key_jwk,
headers={"kid": self.private_key_jwk.key_id},
algorithm="RS256"
)
async def send_push_notification(self, url: str, data: dict[str, Any]):
jwt_token = self._generate_jwt(data)
headers = {'Authorization': f"Bearer {jwt_token}"}
async with httpx.AsyncClient(timeout=10) as client:
try:
response = await client.post(
url,
json=data,
headers=headers
)
response.raise_for_status()
logger.info(f"Push-notification sent for URL: {url}")
except Exception as e:
logger.warning(f"Error during sending push-notification for URL {url}: {e}")
class PushNotificationReceiverAuth(PushNotificationAuth):
def __init__(self):
self.public_keys_jwks = []
self.jwks_client = None
async def load_jwks(self, jwks_url: str):
self.jwks_client = PyJWKClient(jwks_url)
async def verify_push_notification(self, request: Request) -> bool:
auth_header = request.headers.get("Authorization")
if not auth_header or not auth_header.startswith(AUTH_HEADER_PREFIX):
print("Invalid authorization header")
return False
token = auth_header[len(AUTH_HEADER_PREFIX):]
signing_key = self.jwks_client.get_signing_key_from_jwt(token)
decode_token = jwt.decode(
token,
signing_key,
options={"require": ["iat", "request_body_sha256"]},
algorithms=["RS256"],
)
actual_body_sha256 = self._calculate_request_body_sha256(await request.json())
if actual_body_sha256 != decode_token["request_body_sha256"]:
# Payload signature does not match the digest in signed token.
raise ValueError("Invalid request body")
if time.time() - decode_token["iat"] > 60 * 5:
# Do not allow push-notifications older than 5 minutes.
# This is to prevent replay attack.
raise ValueError("Token is expired")
return True
cookbook/pocketflow-a2a/flow.py
+33
View File
@@ -0,0 +1,33 @@
from pocketflow import Flow
from nodes import DecideAction, SearchWeb, AnswerQuestion
def create_agent_flow():
"""
Create and connect the nodes to form a complete agent flow.
The flow works like this:
1. DecideAction node decides whether to search or answer
2. If search, go to SearchWeb node
3. If answer, go to AnswerQuestion node
4. After SearchWeb completes, go back to DecideAction
Returns:
Flow: A complete research agent flow
"""
# Create instances of each node
decide = DecideAction()
search = SearchWeb()
answer = AnswerQuestion()
# Connect the nodes
# If DecideAction returns "search", go to SearchWeb
decide - "search" >> search
# If DecideAction returns "answer", go to AnswerQuestion
decide - "answer" >> answer
# After SearchWeb completes and returns "decide", go back to DecideAction
search - "decide" >> decide
# Create and return the flow, starting with the DecideAction node
return Flow(start=decide)
cookbook/pocketflow-a2a/main.py
+27
View File
@@ -0,0 +1,27 @@
import sys
from flow import create_agent_flow
def main():
"""Simple function to process a question."""
# Default question
default_question = "Who won the Nobel Prize in Physics 2024?"
# Get question from command line if provided with --
question = default_question
for arg in sys.argv[1:]:
if arg.startswith("--"):
question = arg[2:]
break
# Create the agent flow
agent_flow = create_agent_flow()
# Process the question
shared = {"question": question}
print(f"🤔 Processing question: {question}")
agent_flow.run(shared)
print("\n🎯 Final Answer:")
print(shared.get("answer", "No answer found"))
if __name__ == "__main__":
main()
cookbook/pocketflow-a2a/nodes.py
+135
View File
@@ -0,0 +1,135 @@
from pocketflow import Node
from utils import call_llm, search_web
import yaml
class DecideAction(Node):
def prep(self, shared):
"""Prepare the context and question for the decision-making process."""
# Get the current context (default to "No previous search" if none exists)
context = shared.get("context", "No previous search")
# Get the question from the shared store
question = shared["question"]
# Return both for the exec step
return question, context
def exec(self, inputs):
"""Call the LLM to decide whether to search or answer."""
question, context = inputs
print(f"🤔 Agent deciding what to do next...")
# Create a prompt to help the LLM decide what to do next with proper yaml formatting
prompt = f"""
### CONTEXT
You are a research assistant that can search the web.
Question: {question}
Previous Research: {context}
### ACTION SPACE
[1] search
Description: Look up more information on the web
Parameters:
- query (str): What to search for
[2] answer
Description: Answer the question with current knowledge
Parameters:
- answer (str): Final answer to the question
## NEXT ACTION
Decide the next action based on the context and available actions.
Return your response in this format:
```yaml
thinking: |
<your step-by-step reasoning process>
action: search OR answer
reason: <why you chose this action>
answer: <if action is answer>
search_query: <specific search query if action is search>
```
IMPORTANT: Make sure to:
1. Use proper indentation (4 spaces) for all multi-line fields
2. Use the | character for multi-line text fields
3. Keep single-line fields without the | character
"""
# Call the LLM to make a decision
response = call_llm(prompt)
# Parse the response to get the decision
yaml_str = response.split("```yaml")[1].split("```")[0].strip()
decision = yaml.safe_load(yaml_str)
return decision
def post(self, shared, prep_res, exec_res):
"""Save the decision and determine the next step in the flow."""
# If LLM decided to search, save the search query
if exec_res["action"] == "search":
shared["search_query"] = exec_res["search_query"]
print(f"🔍 Agent decided to search for: {exec_res['search_query']}")
else:
shared["context"] = exec_res["answer"] #save the context if LLM gives the answer without searching.
print(f"💡 Agent decided to answer the question")
# Return the action to determine the next node in the flow
return exec_res["action"]
class SearchWeb(Node):
def prep(self, shared):
"""Get the search query from the shared store."""
return shared["search_query"]
def exec(self, search_query):
"""Search the web for the given query."""
# Call the search utility function
print(f"🌐 Searching the web for: {search_query}")
results = search_web(search_query)
return results
def post(self, shared, prep_res, exec_res):
"""Save the search results and go back to the decision node."""
# Add the search results to the context in the shared store
previous = shared.get("context", "")
shared["context"] = previous + "\n\nSEARCH: " + shared["search_query"] + "\nRESULTS: " + exec_res
print(f"📚 Found information, analyzing results...")
# Always go back to the decision node after searching
return "decide"
class AnswerQuestion(Node):
def prep(self, shared):
"""Get the question and context for answering."""
return shared["question"], shared.get("context", "")
def exec(self, inputs):
"""Call the LLM to generate a final answer."""
question, context = inputs
print(f"✍️ Crafting final answer...")
# Create a prompt for the LLM to answer the question
prompt = f"""
### CONTEXT
Based on the following information, answer the question.
Question: {question}
Research: {context}
## YOUR ANSWER:
Provide a comprehensive answer using the research results.
"""
# Call the LLM to generate an answer
answer = call_llm(prompt)
return answer
def post(self, shared, prep_res, exec_res):
"""Save the final answer and complete the flow."""
# Save the answer in the shared store
shared["answer"] = exec_res
print(f"✅ Answer generated successfully")
# We're done - no need to continue the flow
return "done"
cookbook/pocketflow-a2a/requirements.txt
+22
View File
@@ -0,0 +1,22 @@
# For PocketFlow Agent Logic
pocketflow>=0.0.1
openai>=1.0.0
duckduckgo-search>=7.5.2
pyyaml>=5.1
# For A2A Server Infrastructure (from common)
starlette>=0.37.2,<0.38.0
uvicorn[standard]>=0.29.0,<0.30.0
sse-starlette>=1.8.2,<2.0.0
pydantic>=2.0.0,<3.0.0
httpx>=0.27.0,<0.28.0
anyio>=3.0.0,<5.0.0 # Dependency of starlette/httpx
# For running __main__.py
click>=8.0.0,<9.0.0
# For A2A Client
httpx>=0.27.0,<0.28.0
httpx-sse>=0.4.0
asyncclick>=8.1.8 # Or just 'click' if you prefer asyncio.run
pydantic>=2.0.0,<3.0.0 # For common.types
cookbook/pocketflow-a2a/task_manager.py
+113
View File
@@ -0,0 +1,113 @@
# FILE: pocketflow_a2a_agent/task_manager.py
import logging
from typing import AsyncIterable, Union
import asyncio
# Import from the common code you copied
from common.server.task_manager import InMemoryTaskManager
from common.types import (
JSONRPCResponse, SendTaskRequest, SendTaskResponse,
SendTaskStreamingRequest, SendTaskStreamingResponse, Task, TaskSendParams,
TaskState, TaskStatus, TextPart, Artifact, UnsupportedOperationError,
InternalError, InvalidParamsError,
Message
)
import common.server.utils as server_utils
# Import directly from your original PocketFlow files
from flow import create_agent_flow
logger = logging.getLogger(__name__)
class PocketFlowTaskManager(InMemoryTaskManager):
""" TaskManager implementation that runs the PocketFlow agent. """
SUPPORTED_CONTENT_TYPES = ["text", "text/plain"] # Define what the agent accepts/outputs
async def on_send_task(self, request: SendTaskRequest) -> SendTaskResponse:
"""Handles non-streaming task requests."""
logger.info(f"Received task send request: {request.params.id}")
# Validate output modes
if not server_utils.are_modalities_compatible(
request.params.acceptedOutputModes, self.SUPPORTED_CONTENT_TYPES
):
logger.warning(
"Unsupported output mode. Received %s, Support %s",
request.params.acceptedOutputModes, self.SUPPORTED_CONTENT_TYPES
)
return SendTaskResponse(id=request.id, error=server_utils.new_incompatible_types_error(request.id).error)
# Upsert the task in the store (initial state: submitted)
# We create the task first so its state can be tracked, even if the sync execution fails
await self.upsert_task(request.params)
# Update state to working before running
await self.update_store(request.params.id, TaskStatus(state=TaskState.WORKING), [])
# --- Run the PocketFlow logic ---
task_params: TaskSendParams = request.params
query = self._get_user_query(task_params)
if query is None:
fail_status = TaskStatus(state=TaskState.FAILED, message=Message(role="agent", parts=[TextPart(text="No text query found")]))
await self.update_store(task_params.id, fail_status, [])
return SendTaskResponse(id=request.id, error=InvalidParamsError(message="No text query found in message parts"))
shared_data = {"question": query}
agent_flow = create_agent_flow() # Create the flow instance
try:
# Run the synchronous PocketFlow
# In a real async server, you might run this in a separate thread/process
# executor to avoid blocking the event loop. For simplicity here, we run it directly.
# Consider adding a timeout if flows can hang.
logger.info(f"Running PocketFlow for task {task_params.id}...")
agent_flow.run(shared_data) # Run the flow, modifying shared_data in place
logger.info(f"PocketFlow completed for task {task_params.id}")
# Access the original shared_data dictionary, which was modified by the flow
answer_text = shared_data.get("answer", "Agent did not produce a final answer text.")
# --- Package result into A2A Task ---
final_task_status = TaskStatus(state=TaskState.COMPLETED)
# Package the answer as an artifact
final_artifact = Artifact(parts=[TextPart(text=answer_text)])
# Update the task in the store with final status and artifact
final_task = await self.update_store(
task_params.id, final_task_status, [final_artifact]
)
# Prepare and return the A2A response
task_result = self.append_task_history(final_task, task_params.historyLength)
return SendTaskResponse(id=request.id, result=task_result)
except Exception as e:
logger.error(f"Error executing PocketFlow for task {task_params.id}: {e}", exc_info=True)
# Update task state to FAILED
fail_status = TaskStatus(
state=TaskState.FAILED,
message=Message(role="agent", parts=[TextPart(text=f"Agent execution failed: {e}")])
)
await self.update_store(task_params.id, fail_status, [])
return SendTaskResponse(id=request.id, error=InternalError(message=f"Agent error: {e}"))
async def on_send_task_subscribe(
self, request: SendTaskStreamingRequest
) -> Union[AsyncIterable[SendTaskStreamingResponse], JSONRPCResponse]:
"""Handles streaming requests - Not implemented for this synchronous agent."""
logger.warning(f"Streaming requested for task {request.params.id}, but not supported by this PocketFlow agent implementation.")
# Return an error indicating streaming is not supported
return JSONRPCResponse(id=request.id, error=UnsupportedOperationError(message="Streaming not supported by this agent"))
def _get_user_query(self, task_send_params: TaskSendParams) -> str | None:
"""Extracts the first text part from the user message."""
if not task_send_params.message or not task_send_params.message.parts:
logger.warning(f"No message parts found for task {task_send_params.id}")
return None
for part in task_send_params.message.parts:
# Ensure part is treated as a dictionary if it came from JSON
part_dict = part if isinstance(part, dict) else part.model_dump()
if part_dict.get("type") == "text" and "text" in part_dict:
return part_dict["text"]
logger.warning(f"No text part found in message for task {task_send_params.id}")
return None # No text part found
cookbook/pocketflow-a2a/utils.py
+30
View File
@@ -0,0 +1,30 @@
from openai import OpenAI
import os
from duckduckgo_search import DDGS
def call_llm(prompt):
client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY", "your-api-key"))
r = client.chat.completions.create(
model="gpt-4o",
messages=[{"role": "user", "content": prompt}]
)
return r.choices[0].message.content
def search_web(query):
results = DDGS().text(query, max_results=5)
# Convert results to a string
results_str = "\n\n".join([f"Title: {r['title']}\nURL: {r['href']}\nSnippet: {r['body']}" for r in results])
return results_str
if __name__ == "__main__":
print("## Testing call_llm")
prompt = "In a few words, what is the meaning of life?"
print(f"## Prompt: {prompt}")
response = call_llm(prompt)
print(f"## Response: {response}")
print("## Testing search_web")
query = "Who won the Nobel Prize in Physics 2024?"
print(f"## Query: {query}")
results = search_web(query)
print(f"## Results: {results}")