GradLLM / rabbit_repo.py
johnbridges's picture
.
2c8368f
raw
history blame
1.95 kB
import uuid
from typing import Any, Optional
import aio_pika
from models import CloudEvent
from utils import to_json, json_compress_str
class RabbitRepo:
def __init__(self, external_source: str):
super().__init__(exchange_type_resolver=self._resolve_type)
self._external_source = external_source # like SystemUrl.ExternalUrl
def _resolve_type(self, exch: str) -> str:
# longest prefix wins (like your .NET mapping)
matches = [k for k in settings.EXCHANGE_TYPES.keys() if exch.lower().startswith(k.lower())]
if matches:
return settings.EXCHANGE_TYPES[max(matches, key=len)]
return settings.RABBIT_EXCHANGE_TYPE
async def publish(self, exchange: str, obj: Any, routing_key: str = ""):
ex = await self._base.ensure_exchange(exchange)
evt = CloudEvent.wrap(
event_id=str(uuid.uuid4()),
event_type=(obj.__class__.__name__ if obj is not None else "NullOrEmpty"),
source=self._source,
data=obj,
)
body = evt.model_dump_json(by_alias=False, exclude_none=True).encode("utf-8")
await ex.publish(aio_pika.Message(body=body), routing_key=routing_key)
async def publish_jsonz(self, exchange: str, obj: Any, routing_key: str = "", with_id: Optional[str] = None) -> str:
ex = await self._base.ensure_exchange(exchange)
datajson = to_json(obj)
datajsonZ = json_compress_str(datajson)
payload: Any = (datajsonZ, with_id) if with_id else datajsonZ
evt = CloudEvent.wrap(
event_id=str(uuid.uuid4()),
event_type=(obj.__class__.__name__ if obj is not None else "NullOrEmpty"),
source=self._source,
data=payload,
)
body = evt.model_dump_json(by_alias=False, exclude_none=True).encode("utf-8")
await ex.publish(aio_pika.Message(body=body), routing_key=routing_key)
return datajsonZ