From 3d307cfb89c995aa4c572b6c552aebbdc55726d6 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Fri, 2 May 2025 16:19:43 +0100 Subject: [PATCH] Stream kinda working --- .../trustgraph/clients/embeddings_client.py | 8 --- .../trustgraph/gateway/dispatch/manager.py | 47 ++++++++++++++++- .../gateway/dispatch/triples_stream.py | 51 ++++++++++++------- .../trustgraph/gateway/endpoint/flows.py | 4 +- .../trustgraph/gateway/endpoint/socket.py | 7 ++- 5 files changed, 87 insertions(+), 30 deletions(-) diff --git a/trustgraph-base/trustgraph/clients/embeddings_client.py b/trustgraph-base/trustgraph/clients/embeddings_client.py index 811f6ed2..1b1c0dc8 100644 --- a/trustgraph-base/trustgraph/clients/embeddings_client.py +++ b/trustgraph-base/trustgraph/clients/embeddings_client.py @@ -1,7 +1,6 @@ from pulsar.schema import JsonSchema from .. schema import EmbeddingsRequest, EmbeddingsResponse -from .. schema import embeddings_request_queue, embeddings_response_queue from . base import BaseClient import _pulsar @@ -23,12 +22,6 @@ class EmbeddingsClient(BaseClient): pulsar_api_key=None, ): - if input_queue == None: - input_queue=embeddings_request_queue - - if output_queue == None: - output_queue=embeddings_response_queue - super(EmbeddingsClient, self).__init__( log_level=log_level, subscriber=subscriber, @@ -43,4 +36,3 @@ class EmbeddingsClient(BaseClient): def request(self, text, timeout=300): return self.call(text=text, timeout=timeout).vectors - diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py index 8ba8c658..9f83b3ca 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py @@ -11,6 +11,7 @@ from . triples_query import TriplesQueryRequestor from . embeddings import EmbeddingsRequestor from . graph_embeddings_query import GraphEmbeddingsQueryRequestor from . prompt import PromptRequestor +from . triples_stream import TriplesStream request_response_dispatchers = { "agent": AgentRequestor, @@ -23,6 +24,10 @@ request_response_dispatchers = { "triples-query": TriplesQueryRequestor, } +receive_dispatchers = { + "embeddings": TriplesStream, +} + class TestDispatcher: def __init__(self, pulsar_client, timeout=120): self.pulsar_client = pulsar_client @@ -113,8 +118,46 @@ class DispatcherManager: def dispatch_flow_service(self): return self - def dispatch_socket_service(self): - return TestDispatcher3(pulsar_client = self.pulsar_client).dispatch + def dispatch_flow_receive(self): + return self.dispatch_receive + + async def dispatch_receive(self, ws, running, params): + + print("HERE") + flow = params.get("flow") + kind = params.get("kind") + + if flow not in self.flows: + raise RuntimeError("Invalid flow") + + if kind not in receive_dispatchers: + raise RuntimeError("Invalid kind") + + key = (flow, kind) + + if key in self.dispatchers: + return self.dispatchers[key] + + intf_defs = self.flows[flow]["interfaces"] + + if kind not in intf_defs: + raise RuntimeError("This kind not supported by flow") + + qconfig = intf_defs[kind] + + dispatcher = receive_dispatchers[kind]( + pulsar_client = self.pulsar_client, + ws = ws, + running = running, + # FIXME! + queue = qconfig["response"], + consumer = f"api-gateway-{flow}-{kind}-request", + subscriber = f"api-gateway-{flow}-{kind}-request", + ) + + self.dispatchers[key] = dispatcher + + return dispatcher async def process(self, data, responder, params): diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/triples_stream.py b/trustgraph-flow/trustgraph/gateway/dispatch/triples_stream.py index ec1b0347..66d61930 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/triples_stream.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/triples_stream.py @@ -3,44 +3,59 @@ import asyncio import queue import uuid -from ... schema import Triples +from ... schema import Triples, EmbeddingsResponse from ... base import Subscriber from . serialize import serialize_triples class TriplesStream: - def __init__(self, ws, running, pulsar_client, queue): + def __init__( + self, ws, running, pulsar_client, queue, consumer, subscriber + ): self.ws = ws - self.running = runnning + self.running = running self.pulsar_client = pulsar_client self.queue = queue + self.consumer = consumer + self.subscriber = subscriber async def destroy(self): self.running.stop() - self.ws.close() + await self.ws.close() async def receive(self, msg): print(msg.data) - async def run(self, ws, running): + async def run(self): - self.subscriber = Subscriber( - pulsar_client, queue, - "api-gateway", "api-gateway", - schema=Triples + print("c", self.consumer) + print("s", self.subscriber) + print("q", self.queue) + + subs = Subscriber( + client = self.pulsar_client, topic = self.queue, + consumer_name = self.consumer, subscription = self.subscriber, +# schema = Triples + schema = EmbeddingsResponse ) - await self.subscriber.start() + await subs.start() id = str(uuid.uuid4()) - q = self.subscriber.subscribe_all(id) + q = await subs.subscribe_all(id) - while running.get(): + while self.running.get(): try: - resp = await asyncio.to_thread(q.get, timeout=0.5) - await ws.send_json(serialize_triples(resp)) + + resp = await asyncio.wait_for(q.get(), timeout=0.5) +# await self.ws.send_json(serialize_triples(resp)) + + print("GOT MESSAGE!!!") + + + await self.ws.send_json(str(resp)) except TimeoutError: continue @@ -52,8 +67,10 @@ class TriplesStream: print(f"Exception: {str(e)}", flush=True) break - self.subscriber.unsubscribe_all(id) + await subs.unsubscribe_all(id) - await ws.close() - running.stop() + await subs.stop() + + await self.ws.close() + self.running.stop() diff --git a/trustgraph-flow/trustgraph/gateway/endpoint/flows.py b/trustgraph-flow/trustgraph/gateway/endpoint/flows.py index 09ad423a..b297dece 100644 --- a/trustgraph-flow/trustgraph/gateway/endpoint/flows.py +++ b/trustgraph-flow/trustgraph/gateway/endpoint/flows.py @@ -33,9 +33,9 @@ class FlowEndpointManager: dispatcher = dispatcher_manager.dispatch_flow_service(), ), SocketEndpoint( - endpoint_path = "/api/v1/test2", + endpoint_path = "/api/v1/flow/{flow}/receive/{kind}", auth = auth, - dispatcher = dispatcher_manager.dispatch_socket_service() + dispatcher = dispatcher_manager.dispatch_flow_receive() ), ] diff --git a/trustgraph-flow/trustgraph/gateway/endpoint/socket.py b/trustgraph-flow/trustgraph/gateway/endpoint/socket.py index efe55caa..c7b194be 100644 --- a/trustgraph-flow/trustgraph/gateway/endpoint/socket.py +++ b/trustgraph-flow/trustgraph/gateway/endpoint/socket.py @@ -26,10 +26,13 @@ class SocketEndpoint: async def listener(self, ws, dispatcher, running): + print("In listener...!!!!!!!!!!!!!!!!!!!!") async for msg in ws: # On error, finish + print ("mt:", msg.type) if msg.type == WSMsgType.TEXT: + print("YER!!!!!") await dispatcher.receive(msg) continue elif msg.type == WSMsgType.BINARY: @@ -63,7 +66,9 @@ class SocketEndpoint: running = Running() print("Create...") - dispatcher = await self.dispatcher(ws, running, request) + dispatcher = await self.dispatcher( + ws, running, request.match_info + ) print("Create worker...") worker_task = tg.create_task(