File size: 1,331 Bytes
d0f3cbe | 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 | from __future__ import annotations
import logging
from typing import Any, List
from schemas.dispatch import DispatchMessage
from schemas.incident import Incident
from schemas.resource import ResourceRecommendation
from skills.generate_dispatch_message import generate_dispatch_message
logger = logging.getLogger(__name__)
class DispatchAgent:
def __init__(self, vllm_client: Any = None) -> None:
self.vllm_client = vllm_client
async def run(
self,
incidents: List[Incident],
all_resources: List[ResourceRecommendation],
) -> List[DispatchMessage]:
logger.info("DispatchAgent: generating messages for %d incidents", len(incidents))
messages: List[DispatchMessage] = []
resource_by_incident = {}
for r in all_resources:
resource_by_incident.setdefault(r.incident_id, []).append(r)
for incident in incidents:
incident_resources = resource_by_incident.get(incident.id, [])
msg = await generate_dispatch_message(
incident,
incident_resources,
channel="radio",
vllm_client=self.vllm_client,
)
messages.append(msg)
logger.info("DispatchAgent: generated %d dispatch messages", len(messages))
return messages
|