Dispatcher stuff is in place

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
-from .. schema import triples_store_queue
+from ... base import Subscriber
-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__(
+        self.ws = ws
-            endpoint_path=path, auth=auth,
+        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()
 

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(
+                self.services[k] = Dispatcher(
-#                    pulsar_client=self.pulsar_client,
+                    pulsar_client=self.pulsar_client,
-#                    timeout = self.timeout,
+                    timeout = self.timeout,
-#                    input_queue = intf[api_kind],
+                    input_queue = intf[api_kind],
-#                    consumer = f"api-gateway-{id}-{api_kind}-stream",
+                    consumer = f"api-gateway-{id}-{api_kind}-stream",
-#                    subscriber = f"api-gateway-{id}-{api_kind}-stream",
+                    subscriber = f"api-gateway-{id}-{api_kind}-stream",
+                    impl=streamer,
                 )
                 await self.services[k].start()
 

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):
+    async def worker(self, ws, dispatcher, running):
-        pass
 
-    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):
@@ -58,16 +60,38 @@ class SocketEndpoint:
 
             async with asyncio.TaskGroup() as tg:
 
-                worker = tg.create_task(
+                running = Running()
-                    self.dispatcher(ws)
+
+                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(
+                print("Create listener")
-                    self.listener(ws)
+                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)