diff --git a/trustgraph-flow/trustgraph/rev_gateway/__init__.py b/trustgraph-flow/trustgraph/rev_gateway/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/trustgraph-flow/trustgraph/rev_gateway/__main__.py b/trustgraph-flow/trustgraph/rev_gateway/__main__.py new file mode 100644 index 00000000..69fa0895 --- /dev/null +++ b/trustgraph-flow/trustgraph/rev_gateway/__main__.py @@ -0,0 +1,23 @@ +import asyncio +import logging +import sys +from .service import ReverseGateway + +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' +) + +async def main(): + gateway = ReverseGateway() + + try: + await gateway.run() + except KeyboardInterrupt: + print("\nShutdown requested by user") + except Exception as e: + print(f"Fatal error: {e}") + sys.exit(1) + +if __name__ == "__main__": + asyncio.run(main()) \ No newline at end of file diff --git a/trustgraph-flow/trustgraph/rev_gateway/dispatcher.py b/trustgraph-flow/trustgraph/rev_gateway/dispatcher.py new file mode 100644 index 00000000..266d5b80 --- /dev/null +++ b/trustgraph-flow/trustgraph/rev_gateway/dispatcher.py @@ -0,0 +1,47 @@ +import asyncio +import logging +import uuid +from typing import Dict, Any, Optional + +logger = logging.getLogger("dispatcher") +logger.setLevel(logging.INFO) + +class MessageDispatcher: + + def __init__(self, max_workers: int = 10): + self.max_workers = max_workers + self.semaphore = asyncio.Semaphore(max_workers) + self.active_tasks = set() + + async def handle_message(self, message: Dict[Any, Any]) -> Optional[Dict[Any, Any]]: + async with self.semaphore: + task = asyncio.create_task(self._process_message(message)) + self.active_tasks.add(task) + + try: + result = await task + return result + finally: + self.active_tasks.discard(task) + + async def _process_message(self, message: Dict[Any, Any]) -> Dict[Any, Any]: + request_id = message.get('id', str(uuid.uuid4())) + + logger.info(f"Processing message {request_id}") + + await asyncio.sleep(2.0) + + response = { + 'id': request_id, + 'response': 'hello world', + 'status': 'success' + } + + logger.info(f"Completed processing message {request_id}") + return response + + async def shutdown(self): + if self.active_tasks: + logger.info(f"Waiting for {len(self.active_tasks)} active tasks to complete") + await asyncio.gather(*self.active_tasks, return_exceptions=True) + logger.info("Dispatcher shutdown complete") \ No newline at end of file diff --git a/trustgraph-flow/trustgraph/rev_gateway/service.py b/trustgraph-flow/trustgraph/rev_gateway/service.py new file mode 100644 index 00000000..8ed4ecca --- /dev/null +++ b/trustgraph-flow/trustgraph/rev_gateway/service.py @@ -0,0 +1,117 @@ +import asyncio +import logging +import json +import sys +from aiohttp import ClientSession, WSMsgType, ClientWebSocketResponse +from typing import Optional + +from .dispatcher import MessageDispatcher + +logger = logging.getLogger("rev_gateway") +logger.setLevel(logging.INFO) + +class ReverseGateway: + + def __init__(self, host: str = "api.trustgraph.ai", max_workers: int = 10): + self.host = host + self.url = f"wss://{host}/ws" + self.max_workers = max_workers + self.ws: Optional[ClientWebSocketResponse] = None + self.session: Optional[ClientSession] = None + self.dispatcher = MessageDispatcher(max_workers) + self.running = False + self.reconnect_delay = 3.0 + + async def connect(self) -> bool: + try: + if self.session is None: + self.session = ClientSession() + + logger.info(f"Connecting to {self.url}") + self.ws = await self.session.ws_connect(self.url) + logger.info("WebSocket connection established") + return True + + except Exception as e: + logger.error(f"Failed to connect: {e}") + return False + + async def disconnect(self): + if self.ws and not self.ws.closed: + await self.ws.close() + if self.session and not self.session.closed: + await self.session.close() + self.ws = None + self.session = None + + async def send_message(self, message: dict): + if self.ws and not self.ws.closed: + try: + await self.ws.send_str(json.dumps(message)) + except Exception as e: + logger.error(f"Failed to send message: {e}") + + async def handle_message(self, message: str): + try: + print(f"Received: {message}", flush=True) + + msg_data = json.loads(message) + response = await self.dispatcher.handle_message(msg_data) + + if response: + await self.send_message(response) + + except Exception as e: + logger.error(f"Error handling message: {e}") + + async def listen(self): + while self.running and self.ws and not self.ws.closed: + try: + msg = await self.ws.receive() + + if msg.type == WSMsgType.TEXT: + await self.handle_message(msg.data) + elif msg.type == WSMsgType.BINARY: + await self.handle_message(msg.data.decode('utf-8')) + elif msg.type in (WSMsgType.CLOSE, WSMsgType.ERROR): + logger.warning("WebSocket closed or error occurred") + break + + except Exception as e: + logger.error(f"Error in listen loop: {e}") + break + + async def run(self): + self.running = True + logger.info("Starting reverse gateway") + + while self.running: + try: + if await self.connect(): + await self.listen() + else: + logger.warning(f"Connection failed, retrying in {self.reconnect_delay} seconds") + + await self.disconnect() + + if self.running: + await asyncio.sleep(self.reconnect_delay) + + except KeyboardInterrupt: + logger.info("Shutdown requested") + break + except Exception as e: + logger.error(f"Unexpected error: {e}") + if self.running: + await asyncio.sleep(self.reconnect_delay) + + await self.shutdown() + + async def shutdown(self): + logger.info("Shutting down reverse gateway") + self.running = False + await self.dispatcher.shutdown() + await self.disconnect() + + def stop(self): + self.running = False \ No newline at end of file