From a98aee7596864dba642722d6a89db33b06e1d3c8 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Thu, 1 May 2025 23:49:39 +0100 Subject: [PATCH] Dispatcher stuff is in place --- .../gateway/{ => dispatch}/triples_stream.py | 44 +++++----- .../trustgraph/gateway/endpoint/flows.py | 80 ++++++++++--------- .../trustgraph/gateway/endpoint/socket.py | 50 +++++++++--- 3 files changed, 99 insertions(+), 75 deletions(-) rename trustgraph-flow/trustgraph/gateway/{ => dispatch}/triples_stream.py (51%) diff --git a/trustgraph-flow/trustgraph/gateway/triples_stream.py b/trustgraph-flow/trustgraph/gateway/dispatch/triples_stream.py similarity index 51% rename from trustgraph-flow/trustgraph/gateway/triples_stream.py rename to trustgraph-flow/trustgraph/gateway/dispatch/triples_stream.py index a660591e..ec1b0347 100644 --- a/trustgraph-flow/trustgraph/gateway/triples_stream.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/triples_stream.py @@ -3,47 +3,38 @@ import asyncio import queue import uuid -from .. schema import Triples -from .. schema import triples_store_queue -from .. base import Subscriber +from ... schema import Triples +from ... base import Subscriber -from . socket import SocketEndpoint from . serialize import serialize_triples -class TriplesStreamEndpoint(SocketEndpoint): +class TriplesStream: - def __init__(self, pulsar_client, auth, path="/api/v1/stream/triples"): + def __init__(self, ws, running, pulsar_client, queue): - super(TriplesStreamEndpoint, self).__init__( - endpoint_path=path, auth=auth, - ) + self.ws = ws + self.running = runnning + self.pulsar_client = pulsar_client + self.queue = queue - self.pulsar_client=pulsar_client + async def destroy(self): + self.running.stop() + self.ws.close() + + async def receive(self, msg): + print(msg.data) + + async def run(self, ws, running): self.subscriber = Subscriber( - self.pulsar_client, triples_store_queue, + pulsar_client, queue, "api-gateway", "api-gateway", schema=Triples ) - async def listener(self, ws, running): - - worker = asyncio.create_task( - self.async_thread(ws, running) - ) - - await super(TriplesStreamEndpoint, self).listener(ws, running) - - await worker - - async def start(self): - await self.subscriber.start() - async def async_thread(self, ws, running): - id = str(uuid.uuid4()) - q = self.subscriber.subscribe_all(id) while running.get(): @@ -63,5 +54,6 @@ class TriplesStreamEndpoint(SocketEndpoint): self.subscriber.unsubscribe_all(id) + await ws.close() running.stop() diff --git a/trustgraph-flow/trustgraph/gateway/endpoint/flows.py b/trustgraph-flow/trustgraph/gateway/endpoint/flows.py index b0114219..a7b90765 100644 --- a/trustgraph-flow/trustgraph/gateway/endpoint/flows.py +++ b/trustgraph-flow/trustgraph/gateway/endpoint/flows.py @@ -1,6 +1,8 @@ import asyncio +from aiohttp import web + from . endpoint import ServiceEndpoint from . flow_endpoint import FlowEndpoint @@ -14,7 +16,22 @@ from .. dispatch.triples_query import TriplesQueryRequestor from .. dispatch.embeddings import EmbeddingsRequestor from .. dispatch.graph_embeddings_query import GraphEmbeddingsQueryRequestor from .. dispatch.prompt import PromptRequestor -from . stream import StreamEndpoint + +from .. dispatch.triples_stream import TriplesStream + +from . socket import SocketEndpoint + +class Dispatcher: + def __init__( + self, pulsar_client, timeout, queue, schema, + consumer, subscriber + ): + self.pulsar_client = pulsar_client + self.timeout = timeout + self.queue = queue + self.schema = schema + self.consumer = consumer + self.subscriber = subscriber class FlowEndpointManager: @@ -27,47 +44,36 @@ class FlowEndpointManager: self.services = { } - class DispInst: - def __init__(self, ws, running): - self.ws = ws - self.running = running - self.num = 1 - async def destroy(self): - print("Destroy..") - await self.ws.close() - self.running.stop() - async def run(self): - while self.running.get(): - await asyncio.sleep(1) - await self.ws.send_json({"number": self.num}) - self.num += 1 - print("Tick") - async def receive(self, msg): - print("Message...") - print(msg.data) - - class Dispatcher: - def __init__(self): - pass - async def create(self, ws, running, request): - print("Create") - return DispInst(ws, running) - self.endpoints = [ FlowEndpoint( endpoint_path = "/api/v1/flow/{flow}/{kind}", auth = auth, requestors = self.services, ), - StreamEndpoint( + SocketEndpoint( endpoint_path = "/api/v1/flow/{flow}/stream/{kind}", auth = auth, - dispatcher = Dispatcher(), + dispatcher = self.create_stream_dispatch ), ] self.config_receiver.add_handler(self) + async def create_stream_dispatch(self, ws, running, request): + + flow_id = request.match_info['flow'] + kind = request.match_info['kind'] + k = (flow_id, kind) + + print("Service", k) + + print(self.services) + + if k not in self.services: + raise web.HTTPBadRequest() + + raise RuntimeError("Not impl") + def add_routes(self, app): for ep in self.endpoints: ep.add_routes(app) @@ -102,6 +108,7 @@ class FlowEndpointManager: await self.services[k].stop() del self.services[k] + self.services[k] = requestor( pulsar_client=self.pulsar_client, timeout = self.timeout, request_queue = intf[api_kind]["request"], @@ -113,7 +120,7 @@ class FlowEndpointManager: kinds = { # "document-embeddings-stream": DocumentEmbeddingsStreamEndpoint, -# "triples-store" + "triples-stream": TriplesStream, # "bunch": } @@ -127,12 +134,13 @@ class FlowEndpointManager: await self.services[k].stop() del self.services[k] - self.services[k] = streamer( -# pulsar_client=self.pulsar_client, -# timeout = self.timeout, -# input_queue = intf[api_kind], -# consumer = f"api-gateway-{id}-{api_kind}-stream", -# subscriber = f"api-gateway-{id}-{api_kind}-stream", + self.services[k] = Dispatcher( + pulsar_client=self.pulsar_client, + timeout = self.timeout, + input_queue = intf[api_kind], + consumer = f"api-gateway-{id}-{api_kind}-stream", + subscriber = f"api-gateway-{id}-{api_kind}-stream", + impl=streamer, ) await self.services[k].start() diff --git a/trustgraph-flow/trustgraph/gateway/endpoint/socket.py b/trustgraph-flow/trustgraph/gateway/endpoint/socket.py index c44ecaad..543d54c2 100644 --- a/trustgraph-flow/trustgraph/gateway/endpoint/socket.py +++ b/trustgraph-flow/trustgraph/gateway/endpoint/socket.py @@ -11,32 +11,34 @@ logger.setLevel(logging.INFO) class SocketEndpoint: def __init__( - self, endpoint_path, auth, + self, endpoint_path, auth, dispatcher, ): self.path = endpoint_path self.auth = auth self.operation = "socket" - self.running = Running() + self.dispatcher = dispatcher - async def dispatcher(self, ws): - pass + async def worker(self, ws, dispatcher, running): - async def listener(self, ws): + await dispatcher.run() + + async def listener(self, ws, dispatcher, running): async for msg in ws: + # On error, finish if msg.type == WSMsgType.TEXT: - # Ignore incoming message + await dispatcher.receive(msg) continue elif msg.type == WSMsgType.BINARY: - # Ignore incoming message + await dispatcher.receive(msg) continue else: break - self.running.stop() + running.stop() await ws.close() async def handle(self, request): @@ -57,17 +59,39 @@ class SocketEndpoint: try: async with asyncio.TaskGroup() as tg: - - worker = tg.create_task( - self.dispatcher(ws) + + running = Running() + + print("Create...") + print(self.dispatcher) + dispinst = await self.dispatcher(ws, running, request) + + print("Create worker...") + worker_task = tg.create_task( + self.worker(ws, dispinst, running) ) - worker = tg.create_task( - self.listener(ws) + print("Create listener") + lsnr_task = tg.create_task( + self.listener(ws, dispinst, running) ) + print("Created taskgroup, waiting...") + # Wait for threads to complete + print("Task group closed") + + # Finally? + await dispinst.destroy() + + except ExceptionGroup as e: + + print("Exception group:", flush=True) + + for se in e.exceptions: + print(" Type:", type(se), flush=True) + print(f" Exception: {se}", flush=True) except Exception as e: print("Socket exception:", e, flush=True)