diff --git a/runtime/src/orion/events/base.py b/runtime/src/orion/events/base.py index 71e1397..527dda1 100644 --- a/runtime/src/orion/events/base.py +++ b/runtime/src/orion/events/base.py @@ -4,6 +4,8 @@ from pydantic import BaseModel, Field +from orion.transport.messages import MessageType + # ============================================================ # Event Status @@ -49,5 +51,8 @@ class Event(BaseModel): #: Human-readable description of the event. message: str = "" + #: Message type + type: MessageType | None = None + #: Severity of the event. status: EventStatus = EventStatus.INFO diff --git a/runtime/src/orion/events/events.py b/runtime/src/orion/events/events.py index c064147..5a48f12 100644 --- a/runtime/src/orion/events/events.py +++ b/runtime/src/orion/events/events.py @@ -1,4 +1,5 @@ from orion.events.base import Event, EventStatus +from orion.transport.messages import MessageType # ============================================================ # Pipeline Events @@ -16,6 +17,7 @@ class VoicePipelineStartEvent(PipelineStartEvent): class ChatPipelineStartEvent(PipelineStartEvent): """Published when a chat processing pipeline is started.""" + type: MessageType = MessageType.SUBMIT_PROMPT text: str @@ -44,10 +46,13 @@ class PipelineRestartEvent(Event): class VoiceRecordingStartEvent(Event): """Published when voice recording begins.""" + type: MessageType = MessageType.VOICE_START + class VoiceRecordingCompletedEvent(Event): """Published when voice recording has completed.""" + type: MessageType = MessageType.VOICE_END audio_path: str | None = None @@ -98,24 +103,28 @@ class AgentProcessingStartEvent(Event): class ResponseStartedEvent(Event): """Published when the assistant starts generating a response.""" + type: MessageType = MessageType.ASSISTANT_START + class ResponseChunkEvent(Event): """Published for each streamed response chunk.""" + type: MessageType = MessageType.ASSISTANT_CHUNK text: str class ResponseCompletedEvent(Event): """Published when the assistant has finished generating a response.""" + type: MessageType = MessageType.ASSISTANT_END status: EventStatus = EventStatus.SUCCESS - text: str class ResponseGenerationFailedEvent(Event): """Published when response generation fails.""" + type: MessageType = MessageType.ERROR status: EventStatus = EventStatus.ERROR error: str diff --git a/runtime/src/orion/orchestrator/config.py b/runtime/src/orion/orchestrator/config.py index 233945d..42125c9 100644 --- a/runtime/src/orion/orchestrator/config.py +++ b/runtime/src/orion/orchestrator/config.py @@ -5,6 +5,7 @@ from langchain_groq import ChatGroq from orion.memory.module import MemoryModule +from orion.transport.bridge import IPCBridge @dataclass(slots=True) @@ -15,3 +16,4 @@ class OrchestratorConfig: llm: ChatGroq memory: MemoryModule + bridge: IPCBridge diff --git a/runtime/src/orion/orchestrator/orchestrator.py b/runtime/src/orion/orchestrator/orchestrator.py index ac8b68d..38c49f3 100644 --- a/runtime/src/orion/orchestrator/orchestrator.py +++ b/runtime/src/orion/orchestrator/orchestrator.py @@ -1,9 +1,12 @@ -from orion.bus.event_bus import EventBus +from __future__ import annotations +from orion.bus.event_bus import EventBus from orion.orchestrator.config import OrchestratorConfig -from orion.services.logging import LoggingService -from orion.services.setup import ServiceContext, setup_services from orion.runtime.lifecycle import Lifecycle +from orion.services.base import BaseService +from orion.services.ipc_publisher import IPCPublisherService +from orion.services.logging import LoggingService +from orion.services.setup import ServiceContext, setup_runtime_services class Orchestrator(Lifecycle): @@ -11,8 +14,10 @@ class Orchestrator(Lifecycle): Coordinates the ORION runtime. Responsibilities: - - Startup / shutdown services - - Wire global observers + - Create runtime services + - Create global services + - Start and stop all services + - Register global observers """ def __init__( @@ -23,9 +28,12 @@ def __init__( self.bus = bus self.config = config - self.logger = LoggingService() + self.runtime_services: list[BaseService] = [] + self.global_services: list[BaseService] = [] + self.services: list[BaseService] = [] + + self.bridge = self.config.bridge - self.services = [] self._started = False async def startup(self) -> None: @@ -36,33 +44,49 @@ async def startup(self) -> None: if self._started: return - service_context = ServiceContext( + context = ServiceContext( llm=self.config.llm, memory=self.config.memory, ) + self.runtime_services = setup_runtime_services(context) + + self.global_services = [ + LoggingService(), + IPCPublisherService(bridge=self.bridge), + # MetricsService(...), + # TracingService(...), + ] + self.services = [ - *setup_services(service_context), - self.logger, + *self.runtime_services, + *self.global_services, ] for service in self.services: await service.startup() - self.bus.subscribe_all(self.logger.handle) + for service in self.global_services: + self.bus.subscribe_all(service.handle) self._started = True async def shutdown(self) -> None: """ - Gracefully shutdown the runtime. + Gracefully shutdown the ORION runtime. """ + if not self._started: + return + for service in reversed(self.services): try: await service.shutdown() except Exception as exc: - print(f"Failed to shutdown {service}: {exc}") + print(f"Failed to shutdown {service.__class__.__name__}: {exc}") + self.runtime_services.clear() + self.global_services.clear() self.services.clear() + self._started = False diff --git a/runtime/src/orion/runtime/run.py b/runtime/src/orion/runtime/run.py index 1b8d96f..0e8d57c 100644 --- a/runtime/src/orion/runtime/run.py +++ b/runtime/src/orion/runtime/run.py @@ -80,15 +80,17 @@ async def run() -> None: planner=RetrievalPlanner(llm=llm), ) + bridge = IPCBridge(bus) + orchestrator = Orchestrator( bus=bus, config=OrchestratorConfig( llm=llm, memory=memory, + bridge=bridge, ), ) - bridge = IPCBridge(bus) server = IPCServer( socket_path="/tmp/orion.sock", diff --git a/runtime/src/orion/services/ipc_publisher.py b/runtime/src/orion/services/ipc_publisher.py new file mode 100644 index 0000000..c60c638 --- /dev/null +++ b/runtime/src/orion/services/ipc_publisher.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from orion.events.base import Event +from orion.services.base import BaseService +from orion.transport.bridge import IPCBridge +from orion.transport.messages import Envelope + + +class IPCPublisherService(BaseService): + """ + Publishes runtime events to connected IPC clients. + + Events that do not define an IPC message type are considered + internal runtime events and are not forwarded to clients. + """ + + service_name = "ipc" + + def __init__(self, bridge: IPCBridge) -> None: + super().__init__() + self._bridge = bridge + + async def handle(self, event: Event) -> None: + """ + Publish a runtime event over IPC. + """ + + if event.type is None: + return + + envelope = Envelope( + correlation_id=event.correlation_id, + type=event.type, + payload=event.model_dump( + exclude={ + "event_id", + "correlation_id", + "session_id", + "timestamp", + "source", + "type", + }, + exclude_none=True, + ), + ) + + await self._bridge.send( + session_id=event.session_id, + envelope=envelope, + ) diff --git a/runtime/src/orion/services/setup.py b/runtime/src/orion/services/setup.py index d9a534d..de1f181 100644 --- a/runtime/src/orion/services/setup.py +++ b/runtime/src/orion/services/setup.py @@ -1,4 +1,3 @@ -from collections.abc import Sequence from dataclasses import dataclass from langchain_groq import ChatGroq @@ -18,8 +17,8 @@ class ServiceContext: memory: MemoryModule -def setup_services(ctx: ServiceContext) -> Sequence[BaseService]: - services = [ +def setup_runtime_services(ctx: ServiceContext) -> list[BaseService]: + services: list[BaseService] = [ VoiceRecordingService(), TranscriptGenerationService(), AgentService( diff --git a/runtime/src/orion/transport/bridge.py b/runtime/src/orion/transport/bridge.py index a285cb4..aa5e36d 100644 --- a/runtime/src/orion/transport/bridge.py +++ b/runtime/src/orion/transport/bridge.py @@ -1,27 +1,28 @@ """ IPC bridge. -The bridge connects the transport layer to Orion's internal event system. +The bridge translates between the IPC transport protocol and Orion's +internal event system. -Incoming protocol messages are translated into domain events and -published to the EventBus. +Incoming protocol messages are converted into domain events and published +to the EventBus. -Outgoing domain events are translated into protocol messages and sent -back to the appropriate client session. +Outgoing protocol messages are routed to the appropriate connected client. """ from __future__ import annotations + from uuid import UUID from orion.bus.event_bus import EventBus from orion.events.events import ChatPipelineStartEvent - -from .messages import ( +from orion.transport.messages import ( Envelope, MessageType, + PongPayload, SubmitPromptPayload, ) -from .session import ClientSession +from orion.transport.session import ClientSession class IPCBridge: @@ -29,61 +30,66 @@ class IPCBridge: Bridges IPC protocol messages and Orion domain events. """ - def __init__( + def __init__(self, event_bus: EventBus) -> None: + self._event_bus = event_bus + self._sessions: dict[UUID, ClientSession] = {} + + # ------------------------------------------------------------------ + # Session Management + # ------------------------------------------------------------------ + + def register_session(self, session: ClientSession) -> None: + """Register a connected client.""" + self._sessions[session.id] = session + + def unregister_session(self, session: ClientSession) -> None: + """Remove a disconnected client.""" + self._sessions.pop(session.id, None) + + async def send( self, - event_bus: EventBus, + session_id: UUID, + envelope: Envelope, ) -> None: """ - Initialise the IPC bridge. - - Args: - event_bus: - Orion's event bus. + Send an IPC message to a connected client. """ - self._event_bus = event_bus - self._session: dict[UUID, ClientSession] = {} + session = self._sessions.get(session_id) - def register_session(self, session: ClientSession) -> None: - """ - Register a new client session. + if session is None: + return - Args: - session: - Client session to register. - """ - self._session[session.id] = session + await session.send(envelope) - def unregister_session(self, session: ClientSession) -> None: + async def broadcast( + self, + envelope: Envelope, + ) -> None: """ - Deregister a client session. - - Args: - session: - Client session to deregister. + Broadcast an IPC message to every connected client. """ - self._session.pop(session.id, None) + for session in self._sessions.values(): + await session.send(envelope) + + # ------------------------------------------------------------------ + # Incoming Messages + # ------------------------------------------------------------------ async def handle( self, session: ClientSession, - message: Envelope, + envelope: Envelope, ) -> None: """ - Handle one incoming protocol message. - - Args: - session: - Client that sent the message. - - message: - Incoming protocol message. + Handle one incoming IPC message. """ - match message.type: + + match envelope.type: case MessageType.PING: - await self._handle_ping(session) + await self._handle_ping(session, envelope) case MessageType.SUBMIT_PROMPT: - await self._handle_prompt(session, message) + await self._handle_submit_prompt(session, envelope) case MessageType.VOICE_START: ... @@ -95,63 +101,61 @@ async def handle( ... case _: - raise ValueError(f"Unsupported message: {message.type}") + raise ValueError(f"Unsupported IPC message: {envelope.type}") async def _handle_ping( self, session: ClientSession, + envelope: Envelope, ) -> None: """ Respond to a ping request. """ - ... - async def _handle_prompt( + await session.send( + Envelope( + correlation_id=envelope.correlation_id, + type=MessageType.PONG, + payload=PongPayload().model_dump(), + ) + ) + + async def _handle_submit_prompt( self, session: ClientSession, - message: Envelope, + envelope: Envelope, ) -> None: """ - Handle a prompt submission from a client. - - The incoming protocol message is translated into a - ChatPipelineStartEvent and published to the EventBus. + Translate a prompt submission into a domain event. """ - payload = SubmitPromptPayload.model_validate(message.payload) - - event = ChatPipelineStartEvent( - correlation_id=message.correlation_id, - session_id=session.id, - source="ipc", - message="Prompt submitted via IPC.", - text=payload.text, - ) - await self._event_bus.publish(event) + payload = SubmitPromptPayload.model_validate(envelope.payload) - await session.send( - Envelope( - correlation_id=message.correlation_id, - type=MessageType.STATUS, - payload={ - "message": "Prompt accepted.", - }, + await self._event_bus.publish( + ChatPipelineStartEvent( + correlation_id=envelope.correlation_id, + session_id=session.id, + source="ipc", + message="Prompt submitted via IPC.", + text=payload.text, ) ) + # ------------------------------------------------------------------ + # Session Lifecycle + # ------------------------------------------------------------------ + async def serve(self, session: ClientSession) -> None: - print(f"[IPC] Client connected: {session.id}") + """ + Process messages from a connected client until it disconnects. + """ self.register_session(session) try: while True: - print("[IPC] Waiting for message...") - message = await session.receive() - print(f"[IPC] Received: {message.type}") - - await self.handle(session, message) + envelope = await session.receive() + await self.handle(session, envelope) finally: - print(f"[IPC] Client disconnected: {session.id}") self.unregister_session(session)