119 lines
3.4 KiB
Python
119 lines
3.4 KiB
Python
"""Search Gateway event bridge.
|
|
|
|
Publishes normalized search lifecycle events into the AIOS Agent Event Bus when a
|
|
bus is available, with an in-memory fallback for tests and offline runs.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import asdict, dataclass, field
|
|
from datetime import datetime, timezone
|
|
from typing import Any, Protocol
|
|
from uuid import uuid4
|
|
|
|
|
|
def utc_now() -> str:
|
|
return datetime.now(timezone.utc).isoformat()
|
|
|
|
|
|
class SearchEventPublisher(Protocol):
|
|
def publish(
|
|
self,
|
|
kind: str,
|
|
source: str,
|
|
target: str,
|
|
summary: str,
|
|
payload: dict[str, Any],
|
|
correlation_id: str | None = None,
|
|
) -> dict[str, Any]:
|
|
...
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SearchEvent:
|
|
id: str
|
|
kind: str
|
|
source: str
|
|
target: str
|
|
summary: str
|
|
payload: dict[str, Any]
|
|
correlation_id: str | None = None
|
|
created_at: str = field(default_factory=utc_now)
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return asdict(self)
|
|
|
|
|
|
class InMemorySearchEventBus:
|
|
def __init__(self):
|
|
self.events: list[dict[str, Any]] = []
|
|
|
|
def publish(
|
|
self,
|
|
kind: str,
|
|
source: str,
|
|
target: str,
|
|
summary: str,
|
|
payload: dict[str, Any],
|
|
correlation_id: str | None = None,
|
|
) -> dict[str, Any]:
|
|
event = SearchEvent(
|
|
id=f"search_evt_{uuid4().hex[:12]}",
|
|
kind=kind,
|
|
source=source,
|
|
target=target,
|
|
summary=summary,
|
|
payload=payload,
|
|
correlation_id=correlation_id,
|
|
).to_dict()
|
|
self.events.append(event)
|
|
return event
|
|
|
|
|
|
class SearchEventBridge:
|
|
def __init__(self, bus: SearchEventPublisher | None = None):
|
|
self.bus = bus or InMemorySearchEventBus()
|
|
|
|
def search_started(self, request: dict[str, Any]) -> dict[str, Any]:
|
|
return self.bus.publish(
|
|
kind="search_started",
|
|
source=request.get("requester", "unknown"),
|
|
target="search-gateway",
|
|
summary=request.get("query", "search started"),
|
|
payload=request,
|
|
correlation_id=request.get("correlation_id"),
|
|
)
|
|
|
|
def provider_attempted(self, request_id: str, attempt: dict[str, Any]) -> dict[str, Any]:
|
|
return self.bus.publish(
|
|
kind="search_provider_attempted",
|
|
source="search-gateway",
|
|
target=attempt.get("provider", "unknown"),
|
|
summary=f"provider {attempt.get('provider', 'unknown')} {attempt.get('status', 'attempted')}",
|
|
payload=attempt,
|
|
correlation_id=request_id,
|
|
)
|
|
|
|
def fallback_triggered(self, request_id: str, reason: str) -> dict[str, Any]:
|
|
return self.bus.publish(
|
|
kind="search_fallback_triggered",
|
|
source="search-gateway",
|
|
target="search-provider-router",
|
|
summary=reason,
|
|
payload={"reason": reason},
|
|
correlation_id=request_id,
|
|
)
|
|
|
|
def search_completed(self, response: dict[str, Any]) -> dict[str, Any]:
|
|
return self.bus.publish(
|
|
kind="search_completed",
|
|
source="search-gateway",
|
|
target=response.get("requester", "requester"),
|
|
summary=f"search {response.get('status', 'completed')}",
|
|
payload=response,
|
|
correlation_id=response.get("request_id"),
|
|
)
|
|
|
|
|
|
__all__ = ["SearchEvent", "InMemorySearchEventBus", "SearchEventBridge"]
|