From 54bd509d862e558b85a7da96dbae28747628a643 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Fri, 2 May 2025 15:17:43 +0100 Subject: [PATCH] Embeddings round --- .../trustgraph/gateway/dispatch/manager.py | 58 ++++++++++++++++++- .../trustgraph/gateway/endpoint/flows.py | 17 ++---- trustgraph-flow/trustgraph/gateway/service.py | 9 ++- 3 files changed, 69 insertions(+), 15 deletions(-) diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py index 01352524..53ed82dd 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py @@ -1,6 +1,8 @@ import asyncio +from . embeddings import EmbeddingsRequestor + class TestDispatcher: def __init__(self, pulsar_client, timeout=120): self.pulsar_client = pulsar_client @@ -66,11 +68,63 @@ class TestDispatcher3: return Runner(ws, running) class DispatcherManager: - def __init__(self, pulsar_client): + + def __init__(self, pulsar_client, config_receiver): self.pulsar_client = pulsar_client + self.config_receiver = config_receiver + self.config_receiver.add_handler(self) + + self.flows = {} + self.dispatchers = {} + + async def start_flow(self, id, flow): + print("Start flow", id) + self.flows[id] = flow + return + + async def stop_flow(self, id, flow): + print("Stop flow", id) + del self.flows[id] + return + def dispatch_test_service(self): return TestDispatcher(pulsar_client = self.pulsar_client) + def dispatch_flow_service(self): - return TestDispatcher2(pulsar_client = self.pulsar_client) +# return TestDispatcher2(pulsar_client = self.pulsar_client) + return self + def dispatch_socket_service(self): return TestDispatcher3(pulsar_client = self.pulsar_client).dispatch + + async def process(self, data, responder, params): + + flow = params.get("flow") + kind = params.get("kind") + + key = (flow, kind) + + if flow not in self.flows: + raise RuntimeError("Invalid flow") + + if key in self.dispatchers: + return await self.dispatchers[key].process(data, responder) + + qconfig = self.flows[flow]["interfaces"]["embeddings"] + + dispatcher = EmbeddingsRequestor( + pulsar_client = self.pulsar_client, + request_queue = qconfig["request"], + response_queue = qconfig["response"], + timeout = 120, + consumer = f"api-gateway-{flow}-{kind}-request", + subscriber = f"api-gateway-{flow}-{kind}-request", + ) + + await dispatcher.start() + + self.dispatchers[key] = dispatcher + + print("CREATE") + return await dispatcher.process(data, responder) + diff --git a/trustgraph-flow/trustgraph/gateway/endpoint/flows.py b/trustgraph-flow/trustgraph/gateway/endpoint/flows.py index a7a06000..09ad423a 100644 --- a/trustgraph-flow/trustgraph/gateway/endpoint/flows.py +++ b/trustgraph-flow/trustgraph/gateway/endpoint/flows.py @@ -13,37 +13,32 @@ from .. dispatch.manager import DispatcherManager class FlowEndpointManager: - def __init__(self, config_receiver, pulsar_client, auth, timeout=600): + def __init__(self, dispatcher_manager, auth, timeout=600): - self.config_receiver = config_receiver - self.pulsar_client = pulsar_client + self.dispatcher_manager = dispatcher_manager self.timeout = timeout self.services = { } - dm = DispatcherManager(pulsar_client) - self.endpoints = [ ConstantEndpoint( endpoint_path = "/api/v1/test", auth = auth, - dispatcher = dm.dispatch_test_service(), + dispatcher = dispatcher_manager.dispatch_test_service(), ), VariableEndpoint( - endpoint_path = "/api/v1/test/{thing}", + endpoint_path = "/api/v1/flow/{flow}/{kind}", auth = auth, - dispatcher = dm.dispatch_flow_service(), + dispatcher = dispatcher_manager.dispatch_flow_service(), ), SocketEndpoint( endpoint_path = "/api/v1/test2", auth = auth, - dispatcher = dm.dispatch_socket_service() + dispatcher = dispatcher_manager.dispatch_socket_service() ), ] - self.config_receiver.add_handler(self) - def add_routes(self, app): for ep in self.endpoints: ep.add_routes(app) diff --git a/trustgraph-flow/trustgraph/gateway/service.py b/trustgraph-flow/trustgraph/gateway/service.py index 16e6019d..3a44b428 100755 --- a/trustgraph-flow/trustgraph/gateway/service.py +++ b/trustgraph-flow/trustgraph/gateway/service.py @@ -13,6 +13,7 @@ from .. log_level import LogLevel from . auth import Authenticator from . config.receiver import ConfigReceiver +from . dispatch.manager import DispatcherManager from . endpoint.globals import GlobalEndpointManager from . endpoint.flows import FlowEndpointManager @@ -69,6 +70,11 @@ class Api: self.config_receiver = ConfigReceiver(self.pulsar_client) + self.dispatcher_manager = DispatcherManager( + pulsar_client = self.pulsar_client, + config_receiver = self.config_receiver, + ) + self.global_manager = GlobalEndpointManager( pulsar_client = self.pulsar_client, auth = self.auth, @@ -77,8 +83,7 @@ class Api: ) self.flow_manager = FlowEndpointManager( - config_receiver = self.config_receiver, - pulsar_client = self.pulsar_client, + dispatcher_manager = self.dispatcher_manager, auth = self.auth, timeout = self.timeout,