Socket structure working

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),

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 DocumentEmbeddings
-from .. schema import document_embeddings_store_queue
+from ... schema import document_embeddings_store_queue
-from .. base import Subscriber
+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,
         )

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),
         ])

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:

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([

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()
+