Split API endpoint into endpoint and requestor

trustgraph-flow/trustgraph/gateway/agent.py

@@ -4,19 +4,18 @@ from .. schema import agent_request_queue
 from .. schema import agent_response_queue
 
 from . endpoint import MultiResponseServiceEndpoint
+from . requestor import MultiResponseServiceRequestor
 
-class AgentEndpoint(MultiResponseServiceEndpoint):
+class AgentRequestor(MultiResponseServiceRequestor):
     def __init__(self, pulsar_host, timeout, auth):
 
-        super(AgentEndpoint, self).__init__(
+        super(AgentRequestor, self).__init__(
             pulsar_host=pulsar_host,
             request_queue=agent_request_queue,
             response_queue=agent_response_queue,
             request_schema=AgentRequest,
             response_schema=AgentResponse,
-            endpoint_path="/api/v1/agent",
             timeout=timeout,
-            auth=auth,
         )
 
     def to_request(self, body):
@@ -29,3 +28,15 @@ class AgentEndpoint(MultiResponseServiceEndpoint):
             return { "answer": message.answer }, True
         else:
             return {}, False
+
+class AgentEndpoint(MultiResponseServiceEndpoint):
+    def __init__(self, pulsar_host, timeout, auth):
+
+        super(AgentEndpoint, self).__init__(
+            endpoint_path="/api/v1/text-completion",
+            auth=auth,
+            requestor = AgentRequestor(
+                pulsar_host=pulsar_host, timeout=timeout, auth=auth
+            )
+        )
+

trustgraph-flow/trustgraph/gateway/command.py

@@ -0,0 +1,58 @@
+
+import asyncio
+import queue
+from pulsar.schema import JsonSchema
+import uuid
+from aiohttp import web, WSMsgType
+
+from .. schema import GraphEmbeddings
+from .. schema import graph_embeddings_store_queue
+
+from . subscriber import Subscriber
+from . socket import SocketEndpoint
+from . serialize import serialize_graph_embeddings
+
+class CommandEndpoint(SocketEndpoint):
+
+    def __init__(
+            self, pulsar_host, auth, path="/api/v1/command"
+    ):
+
+        super(CommandEndpoint, self).__init__(
+            endpoint_path=path, auth=auth,
+        )
+
+        self.pulsar_host=pulsar_host
+
+    async def start(self):
+        pass
+
+    async def async_thread(self, ws, running):
+
+        id = str(uuid.uuid4())
+
+        while running.get():
+            await asyncio.sleep(1)
+
+        running.stop()
+
+    async def listener(self, ws, running):
+        
+        async for msg in ws:
+
+            # On error, finish
+            if msg.type == WSMsgType.ERROR:
+                break
+            else:
+
+                try:
+                    data = msg.json()
+                except Exception as e:
+                    await ws.send_json({"error": str(e)})
+                    continue
+
+                if "service" not in data:
+                    await ws.send_json({"error": "Malformed message"})
+
+        running.stop()
+

trustgraph-flow/trustgraph/gateway/dbpedia.py

@@ -4,19 +4,18 @@ from .. schema import dbpedia_lookup_request_queue
 from .. schema import dbpedia_lookup_response_queue
 
 from . endpoint import ServiceEndpoint
+from . requestor import ServiceRequestor
 
-class DbpediaEndpoint(ServiceEndpoint):
+class DbpediaRequestor(ServiceRequestor):
     def __init__(self, pulsar_host, timeout, auth):
 
-        super(DbpediaEndpoint, self).__init__(
+        super(DbpediaRequestor, self).__init__(
             pulsar_host=pulsar_host,
             request_queue=dbpedia_lookup_request_queue,
             response_queue=dbpedia_lookup_response_queue,
             request_schema=LookupRequest,
             response_schema=LookupResponse,
-            endpoint_path="/api/v1/dbpedia",
             timeout=timeout,
-            auth=auth,
         )
 
     def to_request(self, body):
@@ -28,3 +27,14 @@ class DbpediaEndpoint(ServiceEndpoint):
     def from_response(self, message):
         return { "text": message.text }
 
+class DbpediaEndpoint(ServiceEndpoint):
+    def __init__(self, pulsar_host, timeout, auth):
+
+        super(DbpediaEndpoint, self).__init__(
+            endpoint_path="/api/v1/dbpedia",
+            auth=auth,
+            requestor = DbpediaRequestor(
+                pulsar_host=pulsar_host, timeout=timeout, auth=auth
+            )
+        )
+

trustgraph-flow/trustgraph/gateway/embeddings.py

@@ -4,19 +4,18 @@ from .. schema import embeddings_request_queue
 from .. schema import embeddings_response_queue
 
 from . endpoint import ServiceEndpoint
+from . requestor import ServiceRequestor
 
-class EmbeddingsEndpoint(ServiceEndpoint):
+class EmbeddingsRequestor(ServiceRequestor):
     def __init__(self, pulsar_host, timeout, auth):
 
-        super(EmbeddingsEndpoint, self).__init__(
+        super(EmbeddingsRequestor, self).__init__(
             pulsar_host=pulsar_host,
             request_queue=embeddings_request_queue,
             response_queue=embeddings_response_queue,
             request_schema=EmbeddingsRequest,
             response_schema=EmbeddingsResponse,
-            endpoint_path="/api/v1/embeddings",
             timeout=timeout,
-            auth=auth,
         )
 
     def to_request(self, body):
@@ -26,3 +25,15 @@ class EmbeddingsEndpoint(ServiceEndpoint):
 
     def from_response(self, message):
         return { "vectors": message.vectors }
+
+class EmbeddingsEndpoint(ServiceEndpoint):
+    def __init__(self, pulsar_host, timeout, auth):
+
+        super(EmbeddingsEndpoint, self).__init__(
+            endpoint_path="/api/v1/embeddings",
+            auth=auth,
+            requestor = EmbeddingsRequestor(
+                pulsar_host=pulsar_host, timeout=timeout, auth=auth
+            )
+        )
+

trustgraph-flow/trustgraph/gateway/encyclopedia.py

@@ -4,19 +4,18 @@ from .. schema import encyclopedia_lookup_request_queue
 from .. schema import encyclopedia_lookup_response_queue
 
 from . endpoint import ServiceEndpoint
+from . requestor import ServiceRequestor
 
-class EncyclopediaEndpoint(ServiceEndpoint):
+class EncyclopediaRequestor(ServiceRequestor):
     def __init__(self, pulsar_host, timeout, auth):
 
-        super(EncyclopediaEndpoint, self).__init__(
+        super(EncyclopediaRequestor, self).__init__(
             pulsar_host=pulsar_host,
             request_queue=encyclopedia_lookup_request_queue,
             response_queue=encyclopedia_lookup_response_queue,
             request_schema=LookupRequest,
             response_schema=LookupResponse,
-            endpoint_path="/api/v1/encyclopedia",
             timeout=timeout,
-            auth=auth,
         )
 
     def to_request(self, body):
@@ -28,3 +27,15 @@ class EncyclopediaEndpoint(ServiceEndpoint):
     def from_response(self, message):
         return { "text": message.text }
 
+
+class EncyclopediaEndpoint(ServiceEndpoint):
+    def __init__(self, pulsar_host, timeout, auth):
+
+        super(EncyclopediaEndpoint, self).__init__(
+            endpoint_path="/api/v1/encyclopedia",
+            auth=auth,
+            requestor = EncyclopediaRequestor(
+                pulsar_host=pulsar_host, timeout=timeout, auth=auth
+            )
+        )
+

trustgraph-flow/trustgraph/gateway/endpoint.py

@@ -13,38 +13,17 @@ logger.setLevel(logging.INFO)
 
 class ServiceEndpoint:
 
-    def __init__(
+    def __init__(self, endpoint_path, auth, requestor):
-            self,
-            pulsar_host,
-            request_queue, request_schema,
-            response_queue, response_schema,
-            endpoint_path,
-            auth,
-            subscription="api-gateway", consumer_name="api-gateway",
-            timeout=600,
-    ):
-
-        self.pub = Publisher(
-            pulsar_host, request_queue,
-            schema=JsonSchema(request_schema)
-        )
-
-        self.sub = Subscriber(
-            pulsar_host, response_queue,
-            subscription, consumer_name,
-            JsonSchema(response_schema)
-        )
 
         self.path = endpoint_path
-        self.timeout = timeout
-        self.auth = auth
 
+        self.auth = auth
         self.operation = "service"
 
-    async def start(self):
+        self.requestor = requestor
 
-        self.pub.start()
+    async def start(self):
-        self.sub.start()
+        await self.requestor.start()
 
     def add_routes(self, app):
 
@@ -52,16 +31,8 @@ class ServiceEndpoint:
             web.post(self.path, self.handle),
         ])
 
-    def to_request(self, request):
-        raise RuntimeError("Not defined")
-
-    def from_response(self, response):
-        raise RuntimeError("Not defined")
-
     async def handle(self, request):
 
-        id = str(uuid.uuid4())
-
         print(request.path, "...")
 
         try:
@@ -82,28 +53,9 @@ class ServiceEndpoint:
 
             print(data)
 
-            q = self.sub.subscribe(id)
+            resp = await self.requestor.process(data)
 
-            await asyncio.to_thread(
+            return web.json_response(resp)
-                self.pub.send, id, self.to_request(data)
-            )
-
-            try:
-                resp = await asyncio.to_thread(q.get, timeout=self.timeout)
-            except Exception as e:
-                raise RuntimeError("Timeout")
-
-            print(resp)
-
-            if resp.error:
-                print("Error")
-                return web.json_response(
-                    { "error": resp.error.message }
-                )
-
-            return web.json_response(
-                self.from_response(resp)
-            )
 
         except Exception as e:
             logging.error(f"Exception: {e}")
@@ -112,10 +64,6 @@ class ServiceEndpoint:
                 { "error": str(e) }
             )
 
-        finally:
-            self.sub.unsubscribe(id)
-
-
 class MultiResponseServiceEndpoint(ServiceEndpoint):
 
     async def handle(self, request):
@@ -132,28 +80,12 @@ class MultiResponseServiceEndpoint(ServiceEndpoint):
                 self.pub.send, id, self.to_request(data)
             )
 
-            # Keeps looking at responses...
+            def responder(x):
+                print("Resp:", x)
 
-            while True:
+            resp = await self.requestor.process(data, responder)
 
-                try:
+            return web.json_response(resp)
-                    resp = await asyncio.to_thread(q.get, timeout=self.timeout)
-                except Exception as e:
-                    raise RuntimeError("Timeout waiting for response")
-
-                if resp.error:
-                    return web.json_response(
-                        { "error": resp.error.message }
-                    )
-
-                # Until from_response says we have a finished answer
-                resp, fin = self.from_response(resp)
-
-
-                if fin:
-                    return web.json_response(resp)
-
-                # Not finished, so loop round and continue
 
         except Exception as e:
             logging.error(f"Exception: {e}")

trustgraph-flow/trustgraph/gateway/graph_rag.py

@@ -4,19 +4,18 @@ from .. schema import graph_rag_request_queue
 from .. schema import graph_rag_response_queue
 
 from . endpoint import ServiceEndpoint
+from . requestor import ServiceRequestor
 
-class GraphRagEndpoint(ServiceEndpoint):
+class GraphRagRequestor(ServiceRequestor):
     def __init__(self, pulsar_host, timeout, auth):
 
-        super(GraphRagEndpoint, self).__init__(
+        super(GraphRagRequestor, self).__init__(
             pulsar_host=pulsar_host,
             request_queue=graph_rag_request_queue,
             response_queue=graph_rag_response_queue,
             request_schema=GraphRagQuery,
             response_schema=GraphRagResponse,
-            endpoint_path="/api/v1/graph-rag",
             timeout=timeout,
-            auth=auth,
         )
 
     def to_request(self, body):
@@ -29,3 +28,14 @@ class GraphRagEndpoint(ServiceEndpoint):
     def from_response(self, message):
         return { "response": message.response }
 
+class GraphRagEndpoint(ServiceEndpoint):
+    def __init__(self, pulsar_host, timeout, auth):
+
+        super(GraphRagEndpoint, self).__init__(
+            endpoint_path="/api/v1/graph-rag",
+            auth=auth,
+            requestor = GraphRagRequestor(
+                pulsar_host=pulsar_host, timeout=timeout, auth=auth
+            )
+        )
+

trustgraph-flow/trustgraph/gateway/internet_search.py

@@ -4,19 +4,18 @@ from .. schema import internet_search_request_queue
 from .. schema import internet_search_response_queue
 
 from . endpoint import ServiceEndpoint
+from . requestor import ServiceRequestor
 
-class InternetSearchEndpoint(ServiceEndpoint):
+class InternetSearchRequestor(ServiceRequestor):
     def __init__(self, pulsar_host, timeout, auth):
 
-        super(InternetSearchEndpoint, self).__init__(
+        super(InternetSearchRequestor, self).__init__(
             pulsar_host=pulsar_host,
             request_queue=internet_search_request_queue,
             response_queue=internet_search_response_queue,
             request_schema=LookupRequest,
             response_schema=LookupResponse,
-            endpoint_path="/api/v1/internet-search",
             timeout=timeout,
-            auth=auth,
         )
 
     def to_request(self, body):
@@ -28,3 +27,14 @@ class InternetSearchEndpoint(ServiceEndpoint):
     def from_response(self, message):
         return { "text": message.text }
 
+class InternetSearchEndpoint(ServiceEndpoint):
+    def __init__(self, pulsar_host, timeout, auth):
+
+        super(InternetSearchEndpoint, self).__init__(
+            endpoint_path="/api/v1/internet-search",
+            auth=auth,
+            requestor = InternetSearchRequestor(
+                pulsar_host=pulsar_host, timeout=timeout, auth=auth
+            )
+        )
+

trustgraph-flow/trustgraph/gateway/prompt.py

@@ -6,19 +6,18 @@ from .. schema import prompt_request_queue
 from .. schema import prompt_response_queue
 
 from . endpoint import ServiceEndpoint
+from . requestor import ServiceRequestor
 
-class PromptEndpoint(ServiceEndpoint):
+class PromptRequestor(ServiceRequestor):
     def __init__(self, pulsar_host, timeout, auth):
 
-        super(PromptEndpoint, self).__init__(
+        super(PromptRequestor, self).__init__(
             pulsar_host=pulsar_host,
             request_queue=prompt_request_queue,
             response_queue=prompt_response_queue,
             request_schema=PromptRequest,
             response_schema=PromptResponse,
-            endpoint_path="/api/v1/prompt",
             timeout=timeout,
-            auth=auth,
         )
 
     def to_request(self, body):
@@ -40,3 +39,14 @@ class PromptEndpoint(ServiceEndpoint):
                 "text": message.text
             }
 
+class PromptEndpoint(ServiceEndpoint):
+    def __init__(self, pulsar_host, timeout, auth):
+
+        super(PromptEndpoint, self).__init__(
+            endpoint_path="/api/v1/prompt",
+            auth=auth,
+            requestor = PromptRequestor(
+                pulsar_host=pulsar_host, timeout=timeout, auth=auth
+            )
+        )
+

trustgraph-flow/trustgraph/gateway/requestor.py

@@ -0,0 +1,128 @@
+
+import asyncio
+from pulsar.schema import JsonSchema
+import uuid
+import logging
+
+from . publisher import Publisher
+from . subscriber import Subscriber
+
+logger = logging.getLogger("requestor")
+logger.setLevel(logging.INFO)
+
+class ServiceRequestor:
+
+    def __init__(
+            self,
+            pulsar_host,
+            request_queue, request_schema,
+            response_queue, response_schema,
+            subscription="api-gateway", consumer_name="api-gateway",
+            timeout=600,
+    ):
+
+        self.pub = Publisher(
+            pulsar_host, request_queue,
+            schema=JsonSchema(request_schema)
+        )
+
+        self.sub = Subscriber(
+            pulsar_host, response_queue,
+            subscription, consumer_name,
+            JsonSchema(response_schema)
+        )
+
+        self.timeout = timeout
+
+    async def start(self):
+
+        self.pub.start()
+        self.sub.start()
+
+    def to_request(self, request):
+        raise RuntimeError("Not defined")
+
+    def from_response(self, response):
+        raise RuntimeError("Not defined")
+
+    async def process(self, request):
+
+        id = str(uuid.uuid4())
+        print(id)
+
+        try:
+
+            q = self.sub.subscribe(id)
+
+            print("Pub...")
+            await asyncio.to_thread(
+                self.pub.send, id, self.to_request(request)
+            )
+
+            print("Pubd")
+            try:
+                resp = await asyncio.to_thread(q.get, timeout=self.timeout)
+            except Exception as e:
+                raise RuntimeError("Timeout")
+
+            print("Recvd...")
+
+            if resp.error:
+                return { "error": resp.error.message }
+
+            return self.from_response(resp)
+
+        except Exception as e:
+
+            logging.error(f"Exception: {e}")
+
+            return { "error": str(e) }
+
+        finally:
+            self.sub.unsubscribe(id)
+
+class MultiResponseServiceRequestor(ServiceRequestor):
+
+    async def process(self, request, responder):
+
+        id = str(uuid.uuid4())
+
+        try:
+
+            q = self.sub.subscribe(id)
+
+            await asyncio.to_thread(
+                self.pub.send, id, self.to_request(request)
+            )
+
+            # Keeps looking at responses...
+
+            while True:
+
+                try:
+                    resp = await asyncio.to_thread(q.get, timeout=self.timeout)
+                except Exception as e:
+                    raise RuntimeError("Timeout")
+
+                if resp.error:
+                    return { "error": resp.error.message }
+
+                # Until from_response says we have a finished answer
+                resp, fin = self.from_response(resp)
+
+                responder(resp)
+
+                if fin:
+                    return resp
+
+                # Not finished, so loop round and continue
+
+        except Exception as e:
+
+            logging.error(f"Exception: {e}")
+
+            return { "error": str(e) }
+
+        finally:
+            self.sub.unsubscribe(id)
+

trustgraph-flow/trustgraph/gateway/service.py

@@ -45,6 +45,8 @@ from . triples_stream import TriplesStreamEndpoint
 from . graph_embeddings_stream import GraphEmbeddingsStreamEndpoint
 from . triples_load import TriplesLoadEndpoint
 from . graph_embeddings_load import GraphEmbeddingsLoadEndpoint
+from . command import CommandEndpoint
+
 from . auth import Authenticator
 
 logger = logging.getLogger("api")
@@ -129,6 +131,10 @@ class Api:
                 pulsar_host=self.pulsar_host,
                 auth = self.auth,
             ),
+            CommandEndpoint(
+                pulsar_host=self.pulsar_host,
+                auth = self.auth,
+            ),
         ]
 
         self.document_out = Publisher(

trustgraph-flow/trustgraph/gateway/text_completion.py

@@ -4,19 +4,18 @@ from .. schema import text_completion_request_queue
 from .. schema import text_completion_response_queue
 
 from . endpoint import ServiceEndpoint
+from . requestor import ServiceRequestor
 
-class TextCompletionEndpoint(ServiceEndpoint):
+class TextCompletionRequestor(ServiceRequestor):
     def __init__(self, pulsar_host, timeout, auth):
 
-        super(TextCompletionEndpoint, self).__init__(
+        super(TextCompletionRequestor, self).__init__(
             pulsar_host=pulsar_host,
             request_queue=text_completion_request_queue,
             response_queue=text_completion_response_queue,
             request_schema=TextCompletionRequest,
             response_schema=TextCompletionResponse,
-            endpoint_path="/api/v1/text-completion",
             timeout=timeout,
-            auth=auth,
         )
 
     def to_request(self, body):
@@ -27,3 +26,15 @@ class TextCompletionEndpoint(ServiceEndpoint):
 
     def from_response(self, message):
         return { "response": message.response }
+
+class TextCompletionEndpoint(ServiceEndpoint):
+    def __init__(self, pulsar_host, timeout, auth):
+
+        super(TextCompletionEndpoint, self).__init__(
+            endpoint_path="/api/v1/text-completion",
+            auth=auth,
+            requestor = TextCompletionRequestor(
+                pulsar_host=pulsar_host, timeout=timeout, auth=auth
+            )
+        )
+

trustgraph-flow/trustgraph/gateway/triples_query.py

@@ -4,20 +4,19 @@ from .. schema import triples_request_queue
 from .. schema import triples_response_queue
 
 from . endpoint import ServiceEndpoint
+from . requestor import ServiceRequestor
 from . serialize import to_value, serialize_subgraph
 
-class TriplesQueryEndpoint(ServiceEndpoint):
+class TriplesQueryRequestor(ServiceRequestor):
     def __init__(self, pulsar_host, timeout, auth):
 
-        super(TriplesQueryEndpoint, self).__init__(
+        super(TriplesQueryRequestor, self).__init__(
             pulsar_host=pulsar_host,
             request_queue=triples_request_queue,
             response_queue=triples_response_queue,
             request_schema=TriplesQueryRequest,
             response_schema=TriplesQueryResponse,
-            endpoint_path="/api/v1/triples-query",
             timeout=timeout,
-            auth=auth,
         )
 
     def to_request(self, body):
@@ -52,3 +51,14 @@ class TriplesQueryEndpoint(ServiceEndpoint):
             "response": serialize_subgraph(message.triples)
         }
 
+class TriplesQueryEndpoint(ServiceEndpoint):
+    def __init__(self, pulsar_host, timeout, auth):
+
+        super(TriplesQueryEndpoint, self).__init__(
+            endpoint_path="/api/v1/triples-query",
+            auth=auth,
+            requestor = TriplesQueryRequestor(
+                pulsar_host=pulsar_host, timeout=timeout, auth=auth
+            )
+        )
+