diff --git a/trustgraph-base/trustgraph/base/subscriber.py b/trustgraph-base/trustgraph/base/subscriber.py index 4f5d1455..66c1a4fe 100644 --- a/trustgraph-base/trustgraph/base/subscriber.py +++ b/trustgraph-base/trustgraph/base/subscriber.py @@ -45,6 +45,8 @@ class Subscriber: try: + # FIXME: Create consumer in start method so we know + # it is definitely running when start completes consumer = self.client.subscribe( topic = self.topic, subscription_name = self.subscription, diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/config.py b/trustgraph-flow/trustgraph/gateway/dispatch/config.py index 001a051b..4b6a0439 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/config.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/config.py @@ -6,7 +6,7 @@ from ... schema import config_response_queue from . requestor import ServiceRequestor class ConfigRequestor(ServiceRequestor): - def __init__(self, pulsar_client, timeout): + def __init__(self, pulsar_client, timeout=120): super(ConfigRequestor, self).__init__( pulsar_client=pulsar_client, diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/flow.py b/trustgraph-flow/trustgraph/gateway/dispatch/flow.py index 7b360601..7a2f8a39 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/flow.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/flow.py @@ -6,7 +6,7 @@ from ... schema import flow_response_queue from . requestor import ServiceRequestor class FlowRequestor(ServiceRequestor): - def __init__(self, pulsar_client, timeout): + def __init__(self, pulsar_client, timeout=120): super(FlowRequestor, self).__init__( pulsar_client=pulsar_client, diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/librarian.py b/trustgraph-flow/trustgraph/gateway/dispatch/librarian.py index 81de0427..8705614a 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/librarian.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/librarian.py @@ -8,7 +8,7 @@ from . serialize import serialize_document_package, serialize_document_info from . serialize import to_document_package, to_document_info, to_criteria class LibrarianRequestor(ServiceRequestor): - def __init__(self, pulsar_client, timeout): + def __init__(self, pulsar_client, timeout=120): super(LibrarianRequestor, self).__init__( pulsar_client=pulsar_client, diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py index 44305d2f..4b896718 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py @@ -2,6 +2,10 @@ import asyncio import uuid +from . config import ConfigRequestor +from . flow import FlowRequestor +from . librarian import LibrarianRequestor + from . embeddings import EmbeddingsRequestor from . agent import AgentRequestor from . text_completion import TextCompletionRequestor @@ -12,6 +16,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 from . graph_embeddings_stream import GraphEmbeddingsStream from . document_embeddings_stream import DocumentEmbeddingsStream @@ -33,69 +38,15 @@ receive_dispatchers = { "document-embeddings": DocumentEmbeddingsStream, } -class TestDispatcher: - def __init__(self, pulsar_client, timeout=120): - self.pulsar_client = pulsar_client - timeout = timeout +class DispatcherWrapper: + def __init__(self, mgr, name, impl): + self.mgr = mgr + self.name = name + self.impl = impl async def process(self, data, responder): - result = { "result": "Hello world!" } - - if responder: - await responder(result, True) - - return result - -class TestDispatcher2: - def __init__(self, pulsar_client, timeout=120): - self.pulsar_client = pulsar_client - timeout = timeout - async def process(self, data, responder, params): - - thing = params['thing'] - - result = { "result": "Hello world!!", "thing": thing } - - if responder: - await responder(result, True) - - return result - -class TestDispatcher3: - def __init__(self, pulsar_client, timeout=120): - self.pulsar_client = pulsar_client - self.timeout = timeout - - async def dispatch(self, ws, running, request): - - class Runner: - def __init__(self, ws, running): - self.ws = ws - self.running = running - - async def destroy(self): - - if self.ws: - await self.ws.close() - self.ws = None - - self.running.stop() - - async def run(self): - - i = 0 - - while self.running.get(): - await self.ws.send_json({"i": i}) - i += 1 - await asyncio.sleep(1) - - await self.ws.close() - self.ws = None - - async def receive(self, msg): - print("Receive:", msg.data) - - return Runner(ws, running) + return await self.mgr.process_impl( + data, responder, self.name, self.impl + ) class DispatcherManager: @@ -117,8 +68,31 @@ class DispatcherManager: del self.flows[id] return - def dispatch_test_service(self): - return TestDispatcher(pulsar_client = self.pulsar_client) + def dispatch_config(self): + return DispatcherWrapper(self, "config", ConfigRequestor) + + def dispatch_flow(self): + return DispatcherWrapper(self, "flow", FlowRequestor) + + def dispatch_librarian(self): + return DispatcherWrapper(self, "librarian", LibrarianRequestor) + + async def process_impl(self, data, responder, name, impl): + + key = (None, name) + + if key in self.dispatchers: + return await self.dispatchers[key].process(data, responder) + + dispatcher = impl( + pulsar_client = self.pulsar_client + ) + + await dispatcher.start() + + self.dispatchers[key] = dispatcher + + return await dispatcher.process(data, responder) def dispatch_flow_service(self): return self @@ -152,7 +126,6 @@ class DispatcherManager: pulsar_client = self.pulsar_client, ws = ws, running = running, - # FIXME! queue = qconfig, consumer = f"api-gateway-{id}", subscriber = f"api-gateway-{id}", diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/requestor.py b/trustgraph-flow/trustgraph/gateway/dispatch/requestor.py index 226fc9a9..1ce5ac68 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/requestor.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/requestor.py @@ -36,9 +36,9 @@ class ServiceRequestor: self.running = True async def start(self): - await self.pub.start() - await self.sub.start() self.running = True + await self.sub.start() + await self.pub.start() async def stop(self): await self.pub.stop() @@ -68,6 +68,7 @@ class ServiceRequestor: q.get(), timeout=self.timeout ) except Exception as e: + print("Exception", e) raise RuntimeError("Timeout") if resp.error: diff --git a/trustgraph-flow/trustgraph/gateway/endpoint/flows.py b/trustgraph-flow/trustgraph/gateway/endpoint/flows.py deleted file mode 100644 index b297dece..00000000 --- a/trustgraph-flow/trustgraph/gateway/endpoint/flows.py +++ /dev/null @@ -1,128 +0,0 @@ - -import asyncio - -from aiohttp import web - -#from . endpoint import ServiceEndpoint - -from . constant_endpoint import ConstantEndpoint -from . variable_endpoint import VariableEndpoint -from . socket import SocketEndpoint - -from .. dispatch.manager import DispatcherManager - -class FlowEndpointManager: - - def __init__(self, dispatcher_manager, auth, timeout=600): - - self.dispatcher_manager = dispatcher_manager - self.timeout = timeout - - self.services = { - } - - self.endpoints = [ - ConstantEndpoint( - endpoint_path = "/api/v1/test", - auth = auth, - dispatcher = dispatcher_manager.dispatch_test_service(), - ), - VariableEndpoint( - endpoint_path = "/api/v1/flow/{flow}/{kind}", - auth = auth, - dispatcher = dispatcher_manager.dispatch_flow_service(), - ), - SocketEndpoint( - endpoint_path = "/api/v1/flow/{flow}/receive/{kind}", - auth = auth, - dispatcher = dispatcher_manager.dispatch_flow_receive() - ), - ] - - def add_routes(self, app): - for ep in self.endpoints: - ep.add_routes(app) - - async def start(self): - for ep in self.endpoints: - await ep.start() - - async def start_flow(self, id, flow): - - print("START FLOW", id) - - return - - intf = flow["interfaces"] - - kinds = { - "agent": AgentRequestor, - "text-completion": TextCompletionRequestor, - "prompt": PromptRequestor, - "graph-rag": GraphRagRequestor, - "document-rag": DocumentRagRequestor, - "embeddings": EmbeddingsRequestor, - "graph-embeddings": GraphEmbeddingsQueryRequestor, - "triples-query": TriplesQueryRequestor, - } - - 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() - del self.services[k] - - - self.services[k] = requestor( - pulsar_client=self.pulsar_client, timeout = self.timeout, - request_queue = intf[api_kind]["request"], - response_queue = intf[api_kind]["response"], - consumer = f"api-gateway-{id}-{api_kind}-request", - subscriber = f"api-gateway-{id}-{api_kind}-request", - ) - await self.services[k].start() - - kinds = { -# "document-embeddings-stream": DocumentEmbeddingsStreamEndpoint, - "triples-stream": TriplesStream, -# "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] = 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() - - async def stop_flow(self, id, flow): - - print("STOP FLOW", id) - - return - - svc_list = list(self.services.keys()) - - for k in svc_list: - - kid, kkind = k - - if id == kid: - await self.services[k].stop() - del self.services[k] diff --git a/trustgraph-flow/trustgraph/gateway/service.py b/trustgraph-flow/trustgraph/gateway/service.py index 3a44b428..e83b79d1 100755 --- a/trustgraph-flow/trustgraph/gateway/service.py +++ b/trustgraph-flow/trustgraph/gateway/service.py @@ -14,8 +14,8 @@ 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 + +from . endpoint.manager import EndpointManager import pulsar from prometheus_client import start_http_server @@ -75,16 +75,10 @@ class Api: config_receiver = self.config_receiver, ) - self.global_manager = GlobalEndpointManager( - pulsar_client = self.pulsar_client, - auth = self.auth, - prometheus_url = self.prometheus_url, - timeout = self.timeout, - ) - - self.flow_manager = FlowEndpointManager( + self.endpoint_manager = EndpointManager( dispatcher_manager = self.dispatcher_manager, auth = self.auth, + prometheus_url = self.prometheus_url, timeout = self.timeout, ) @@ -108,11 +102,8 @@ class Api: for ep in self.endpoints: await ep.start() - self.global_manager.add_routes(self.app) - await self.global_manager.start() - - self.flow_manager.add_routes(self.app) - await self.flow_manager.start() + self.endpoint_manager.add_routes(self.app) + await self.endpoint_manager.start() return self.app