feat: Add handler-base library for NATS AI/ML services
- Handler base class with graceful shutdown and signal handling - NATSClient with JetStream and msgpack serialization - Pydantic Settings for environment configuration - HealthServer for Kubernetes probes - OpenTelemetry telemetry setup - Service clients: STT, TTS, LLM, Embeddings, Reranker, Milvus
This commit is contained in:
221
handler_base/handler.py
Normal file
221
handler_base/handler.py
Normal file
@@ -0,0 +1,221 @@
|
||||
"""
|
||||
Base handler class for building NATS-based services.
|
||||
"""
|
||||
import asyncio
|
||||
import logging
|
||||
import signal
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Optional
|
||||
|
||||
from nats.aio.msg import Msg
|
||||
|
||||
from handler_base.config import Settings
|
||||
from handler_base.health import HealthServer
|
||||
from handler_base.nats_client import NATSClient
|
||||
from handler_base.telemetry import setup_telemetry, create_span
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Handler(ABC):
|
||||
"""
|
||||
Base class for NATS message handlers.
|
||||
|
||||
Subclass and implement:
|
||||
- setup(): Initialize your service clients
|
||||
- handle_message(): Process incoming messages
|
||||
- teardown(): Clean up resources (optional)
|
||||
|
||||
Example:
|
||||
class MyHandler(Handler):
|
||||
async def setup(self):
|
||||
self.embeddings = EmbeddingsClient()
|
||||
|
||||
async def handle_message(self, msg: Msg, data: dict) -> Optional[dict]:
|
||||
result = await self.embeddings.embed(data["text"])
|
||||
return {"embedding": result}
|
||||
|
||||
if __name__ == "__main__":
|
||||
MyHandler(subject="my.subject").run()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
subject: str,
|
||||
settings: Optional[Settings] = None,
|
||||
queue_group: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Initialize the handler.
|
||||
|
||||
Args:
|
||||
subject: NATS subject to subscribe to
|
||||
settings: Configuration settings
|
||||
queue_group: Optional queue group for load balancing
|
||||
"""
|
||||
self.subject = subject
|
||||
self.settings = settings or Settings()
|
||||
self.queue_group = queue_group or self.settings.nats_queue_group
|
||||
|
||||
self.nats = NATSClient(self.settings)
|
||||
self.health_server = HealthServer(self.settings, self._check_ready)
|
||||
|
||||
self._running = False
|
||||
self._shutdown_event = asyncio.Event()
|
||||
|
||||
@abstractmethod
|
||||
async def setup(self) -> None:
|
||||
"""
|
||||
Initialize service clients and resources.
|
||||
|
||||
Called once before starting to handle messages.
|
||||
Override this to set up your service-specific clients.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def handle_message(self, msg: Msg, data: Any) -> Optional[Any]:
|
||||
"""
|
||||
Handle an incoming message.
|
||||
|
||||
Args:
|
||||
msg: Raw NATS message
|
||||
data: Decoded message data (msgpack unpacked)
|
||||
|
||||
Returns:
|
||||
Optional response data. If returned and msg has a reply subject,
|
||||
the response will be sent automatically.
|
||||
"""
|
||||
pass
|
||||
|
||||
async def teardown(self) -> None:
|
||||
"""
|
||||
Clean up resources.
|
||||
|
||||
Called during graceful shutdown.
|
||||
Override to add custom cleanup logic.
|
||||
"""
|
||||
pass
|
||||
|
||||
async def _check_ready(self) -> bool:
|
||||
"""Check if the service is ready to handle requests."""
|
||||
return self._running and self.nats._nc is not None
|
||||
|
||||
async def _message_handler(self, msg: Msg) -> None:
|
||||
"""Internal message handler with tracing and error handling."""
|
||||
with create_span(f"handle.{self.subject}") as span:
|
||||
try:
|
||||
# Decode message
|
||||
data = NATSClient.decode_msgpack(msg)
|
||||
|
||||
if span:
|
||||
span.set_attribute("messaging.destination", msg.subject)
|
||||
if isinstance(data, dict):
|
||||
request_id = data.get("request_id", data.get("id"))
|
||||
if request_id:
|
||||
span.set_attribute("request.id", str(request_id))
|
||||
|
||||
# Handle message
|
||||
response = await self.handle_message(msg, data)
|
||||
|
||||
# Send response if applicable
|
||||
if response is not None and msg.reply:
|
||||
await self.nats.publish(msg.reply, response)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"Error handling message on {msg.subject}")
|
||||
if span:
|
||||
span.set_attribute("error", True)
|
||||
span.set_attribute("error.message", str(e))
|
||||
|
||||
# Send error response if reply expected
|
||||
if msg.reply:
|
||||
error_response = {
|
||||
"error": True,
|
||||
"message": str(e),
|
||||
"type": type(e).__name__,
|
||||
}
|
||||
await self.nats.publish(msg.reply, error_response)
|
||||
|
||||
def _setup_signals(self) -> None:
|
||||
"""Set up signal handlers for graceful shutdown."""
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
for sig in (signal.SIGTERM, signal.SIGINT):
|
||||
loop.add_signal_handler(sig, self._handle_signal, sig)
|
||||
|
||||
def _handle_signal(self, sig: signal.Signals) -> None:
|
||||
"""Handle shutdown signal."""
|
||||
logger.info(f"Received {sig.name}, initiating graceful shutdown...")
|
||||
self._shutdown_event.set()
|
||||
|
||||
async def _run(self) -> None:
|
||||
"""Main async run loop."""
|
||||
# Setup telemetry
|
||||
setup_telemetry(self.settings)
|
||||
|
||||
# Start health server
|
||||
self.health_server.start()
|
||||
|
||||
try:
|
||||
# Connect to NATS
|
||||
await self.nats.connect()
|
||||
|
||||
# Run user setup
|
||||
logger.info("Running service setup...")
|
||||
await self.setup()
|
||||
|
||||
# Subscribe to subject
|
||||
await self.nats.subscribe(
|
||||
self.subject,
|
||||
self._message_handler,
|
||||
queue=self.queue_group,
|
||||
)
|
||||
|
||||
self._running = True
|
||||
logger.info(f"Handler ready, listening on {self.subject}")
|
||||
|
||||
# Wait for shutdown signal
|
||||
await self._shutdown_event.wait()
|
||||
|
||||
except Exception as e:
|
||||
logger.exception("Fatal error in handler")
|
||||
raise
|
||||
finally:
|
||||
self._running = False
|
||||
|
||||
# Graceful shutdown
|
||||
logger.info("Shutting down...")
|
||||
|
||||
try:
|
||||
await self.teardown()
|
||||
except Exception as e:
|
||||
logger.warning(f"Error during teardown: {e}")
|
||||
|
||||
await self.nats.close()
|
||||
self.health_server.stop()
|
||||
|
||||
logger.info("Shutdown complete")
|
||||
|
||||
def run(self) -> None:
|
||||
"""
|
||||
Run the handler.
|
||||
|
||||
This is the main entry point. It sets up signal handlers
|
||||
and runs the async event loop.
|
||||
"""
|
||||
# Configure logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
|
||||
logger.info(f"Starting {self.settings.service_name} v{self.settings.service_version}")
|
||||
|
||||
# Run the async loop
|
||||
asyncio.run(self._run_with_signals())
|
||||
|
||||
async def _run_with_signals(self) -> None:
|
||||
"""Run with signal handling."""
|
||||
self._setup_signals()
|
||||
await self._run()
|
||||
Reference in New Issue
Block a user