Stream kinda working

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
 
-

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

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

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

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(