From f5c398cd41f7c291ce47a2a985bf7cb276b80ff0 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Thu, 1 May 2025 16:58:38 +0100 Subject: [PATCH] Socket structure working --- .../trustgraph/gateway/dispatch/librarian.py | 4 -- .../document_embeddings_stream.py | 11 ++--- .../gateway/endpoint/flow_endpoint.py | 1 - .../trustgraph/gateway/endpoint/flows.py | 34 +++++++++++++- .../gateway/{ => endpoint}/socket.py | 33 ++++++++++---- .../trustgraph/gateway/endpoint/stream.py | 45 +++++++++++++++++++ 6 files changed, 108 insertions(+), 20 deletions(-) rename trustgraph-flow/trustgraph/gateway/{ => endpoint}/document_embeddings_stream.py (85%) rename trustgraph-flow/trustgraph/gateway/{ => endpoint}/socket.py (70%) create mode 100644 trustgraph-flow/trustgraph/gateway/endpoint/stream.py diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/librarian.py b/trustgraph-flow/trustgraph/gateway/dispatch/librarian.py index 81b809f3..81de0427 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/librarian.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/librarian.py @@ -21,20 +21,16 @@ class LibrarianRequestor(ServiceRequestor): def to_request(self, body): - print("TRR") if "document" in body: dp = to_document_package(body["document"]) else: dp = None - print("GOT") if "criteria" in body: criteria = to_criteria(body["criteria"]) else: criteria = None - print("ASLDKJ") - return LibrarianRequest( operation = body.get("operation", None), id = body.get("id", None), diff --git a/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py b/trustgraph-flow/trustgraph/gateway/endpoint/document_embeddings_stream.py similarity index 85% rename from trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py rename to trustgraph-flow/trustgraph/gateway/endpoint/document_embeddings_stream.py index e59a0370..7613ed95 100644 --- a/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py +++ b/trustgraph-flow/trustgraph/gateway/endpoint/document_embeddings_stream.py @@ -3,9 +3,9 @@ import asyncio import queue import uuid -from .. schema import DocumentEmbeddings -from .. schema import document_embeddings_store_queue -from .. base import Subscriber +from ... schema import DocumentEmbeddings +from ... schema import document_embeddings_store_queue +from ... base import Subscriber from . socket import SocketEndpoint from . serialize import serialize_document_embeddings @@ -14,7 +14,8 @@ class DocumentEmbeddingsStreamEndpoint(SocketEndpoint): def __init__( self, pulsar_client, auth, - path="/api/v1/stream/document-embeddings" + queue, + path="/api/v1/stream/document-embeddings", ): super(DocumentEmbeddingsStreamEndpoint, self).__init__( @@ -24,7 +25,7 @@ class DocumentEmbeddingsStreamEndpoint(SocketEndpoint): self.pulsar_client=pulsar_client self.subscriber = Subscriber( - self.pulsar_client, document_embeddings_store_queue, + self.pulsar_client, queue, "api-gateway", "api-gateway", schema=DocumentEmbeddings, ) diff --git a/trustgraph-flow/trustgraph/gateway/endpoint/flow_endpoint.py b/trustgraph-flow/trustgraph/gateway/endpoint/flow_endpoint.py index 8c69af76..e061ab2a 100644 --- a/trustgraph-flow/trustgraph/gateway/endpoint/flow_endpoint.py +++ b/trustgraph-flow/trustgraph/gateway/endpoint/flow_endpoint.py @@ -23,7 +23,6 @@ class FlowEndpoint: def add_routes(self, app): - pass app.add_routes([ web.post(self.path, self.handle), ]) diff --git a/trustgraph-flow/trustgraph/gateway/endpoint/flows.py b/trustgraph-flow/trustgraph/gateway/endpoint/flows.py index 9fed0336..89ef76ed 100644 --- a/trustgraph-flow/trustgraph/gateway/endpoint/flows.py +++ b/trustgraph-flow/trustgraph/gateway/endpoint/flows.py @@ -12,6 +12,7 @@ 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 class FlowEndpointManager: @@ -30,6 +31,11 @@ class FlowEndpointManager: auth = auth, requestors = self.services, ), + StreamEndpoint( + endpoint_path = "/api/v1/flow/{flow}/stream/{kind}", + auth = auth, + dispatchers = self.services, + ), ] self.config_receiver.add_handler(self) @@ -62,6 +68,7 @@ class FlowEndpointManager: for api_kind, requestor in kinds.items(): if api_kind in intf: + k = (id, api_kind) if k in self.services: await self.services[k].stop() @@ -76,12 +83,35 @@ class FlowEndpointManager: ) await self.services[k].start() + kinds = { +# "document-embeddings-stream": DocumentEmbeddingsStreamEndpoint, +# "triples-store" +# "bunch": + } + + for api_kind, streamer in kinds.items(): + +# if api_kind in intf: + if True: + + k = (id, api_kind) + if k in self.services: + await self.services[k].stop() + del self.services[k] + + self.services[k] = steamer( +# 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", + ) + await self.services[k].start() + async def stop_flow(self, id, flow): print("STOP FLOW", id) - intf = flow["interfaces"] - svc_list = list(self.services.keys()) for k in svc_list: diff --git a/trustgraph-flow/trustgraph/gateway/socket.py b/trustgraph-flow/trustgraph/gateway/endpoint/socket.py similarity index 70% rename from trustgraph-flow/trustgraph/gateway/socket.py rename to trustgraph-flow/trustgraph/gateway/endpoint/socket.py index c32a28af..c44ecaad 100644 --- a/trustgraph-flow/trustgraph/gateway/socket.py +++ b/trustgraph-flow/trustgraph/gateway/endpoint/socket.py @@ -3,7 +3,7 @@ import asyncio from aiohttp import web, WSMsgType import logging -from . running import Running +from .. running import Running logger = logging.getLogger("socket") logger.setLevel(logging.INFO) @@ -18,7 +18,12 @@ class SocketEndpoint: self.auth = auth self.operation = "socket" - async def listener(self, ws, running): + self.running = Running() + + async def dispatcher(self, ws): + pass + + async def listener(self, ws): async for msg in ws: # On error, finish @@ -31,7 +36,8 @@ class SocketEndpoint: else: break - running.stop() + self.running.stop() + await ws.close() async def handle(self, request): @@ -43,20 +49,28 @@ class SocketEndpoint: if not self.auth.permitted(token, self.operation): return web.HTTPUnauthorized() - running = Running() - # 50MB max message size ws = web.WebSocketResponse(max_msg_size=52428800) await ws.prepare(request) try: - await self.listener(ws, running) + + async with asyncio.TaskGroup() as tg: + + worker = tg.create_task( + self.dispatcher(ws) + ) + + worker = tg.create_task( + self.listener(ws) + ) + + # Wait for threads to complete + except Exception as e: print("Socket exception:", e, flush=True) - running.stop() - await ws.close() return ws @@ -64,6 +78,9 @@ class SocketEndpoint: async def start(self): pass + async def stop(self): + self.running.stop() + def add_routes(self, app): app.add_routes([ diff --git a/trustgraph-flow/trustgraph/gateway/endpoint/stream.py b/trustgraph-flow/trustgraph/gateway/endpoint/stream.py new file mode 100644 index 00000000..9cdd56fb --- /dev/null +++ b/trustgraph-flow/trustgraph/gateway/endpoint/stream.py @@ -0,0 +1,45 @@ + +import asyncio +from aiohttp import web, WSMsgType +import logging + +from .. running import Running +from . socket import SocketEndpoint + +logger = logging.getLogger("socket") +logger.setLevel(logging.INFO) + +class StreamEndpoint(SocketEndpoint): + + def __init__( + self, endpoint_path, auth, dispatchers, + ): + + super(StreamEndpoint, self).__init__(endpoint_path, auth) + + async def listener(self, ws): + + async for msg in ws: + # On error, finish + if msg.type == WSMsgType.TEXT: + print("Received:", msg) + continue + elif msg.type == WSMsgType.BINARY: + # Ignore incoming message + continue + else: + break + + self.running.stop() + + async def dispatcher(self, ws): + + for i in range(0, 10): + + await asyncio.sleep(1) + + await ws.send_json({"number": i}) + + await ws.close() + self.running.stop() +