Spaces:
Running
Running
File size: 1,677 Bytes
6918275 | 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 | import asyncio
import os
import time
from fasta2a.task_manager import TaskManager
from fasta2a.schema import SendMessageResponse
try:
from fasta2a.schema import SendMessageResult
except ImportError:
def SendMessageResult(**kwargs):
return kwargs
class BlockingTaskManager(TaskManager):
def __init__(self, broker, storage, **kwargs):
super().__init__(broker, storage, **kwargs)
self.storage = storage
async def send_message(self, request):
response = await super().send_message(request)
config = request["params"].get("configuration") or {}
if not config.get("blocking", False):
return response
task = response["result"]
task_id = task["id"]
timeout = min(
max(float(os.getenv("A2A_BLOCKING_TIMEOUT", "60")), 1),
90,
)
deadline = time.monotonic() + timeout
terminal_states = {
"completed",
"failed",
"canceled",
"rejected",
"input-required",
"auth-required",
}
while True:
task = await self.storage.load_task(task_id)
if task and task["status"]["state"] in terminal_states:
return SendMessageResponse(
jsonrpc="2.0",
id=request["id"],
result=task,
)
if time.monotonic() >= deadline:
return SendMessageResponse(
jsonrpc="2.0",
id=request["id"],
result=task or response["result"],
)
await asyncio.sleep(0.5)
|