More API stuff

trustgraph-flow/trustgraph/gateway/agent.py

@@ -1,20 +1,23 @@
 
 from .. schema import AgentRequest, AgentResponse
-from .. schema import agent_request_queue
-from .. schema import agent_response_queue
 
 from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 
 class AgentRequestor(ServiceRequestor):
-    def __init__(self, pulsar_client, timeout, auth):
+    def __init__(
+            self, pulsar_client, request_queue, response_queue, timeout, auth,
+            consumer, subscriber,
+    ):
 
         super(AgentRequestor, self).__init__(
             pulsar_client=pulsar_client,
-            request_queue=agent_request_queue,
+            request_queue=request_queue,
-            response_queue=agent_response_queue,
+            response_queue=response_queue,
             request_schema=AgentRequest,
             response_schema=AgentResponse,
+            subscription = subscriber,
+            consumer_name = consumer,
             timeout=timeout,
         )
 

trustgraph-flow/trustgraph/gateway/embeddings.py

@@ -1,20 +1,23 @@
 
 from .. schema import EmbeddingsRequest, EmbeddingsResponse
-from .. schema import embeddings_request_queue
-from .. schema import embeddings_response_queue
 
 from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 
 class EmbeddingsRequestor(ServiceRequestor):
-    def __init__(self, pulsar_client, timeout, auth):
+    def __init__(
+            self, pulsar_client, request_queue, response_queue, timeout, auth,
+            consumer, subscriber,
+    ):
 
         super(EmbeddingsRequestor, self).__init__(
             pulsar_client=pulsar_client,
-            request_queue=embeddings_request_queue,
+            request_queue=request_queue,
-            response_queue=embeddings_response_queue,
+            response_queue=response_queue,
             request_schema=EmbeddingsRequest,
             response_schema=EmbeddingsResponse,
+            subscription = subscriber,
+            consumer_name = consumer,
             timeout=timeout,
         )
 

trustgraph-flow/trustgraph/gateway/graph_embeddings_query.py

@@ -1,21 +1,24 @@
 
 from .. schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse
-from .. schema import graph_embeddings_request_queue
-from .. schema import graph_embeddings_response_queue
 
 from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 from . serialize import serialize_value
 
 class GraphEmbeddingsQueryRequestor(ServiceRequestor):
-    def __init__(self, pulsar_client, timeout, auth):
+    def __init__(
+            self, pulsar_client, request_queue, response_queue, timeout, auth,
+            consumer, subscriber,
+    ):
 
         super(GraphEmbeddingsQueryRequestor, self).__init__(
             pulsar_client=pulsar_client,
-            request_queue=graph_embeddings_request_queue,
+            request_queue=request_queue,
-            response_queue=graph_embeddings_response_queue,
+            response_queue=response_queue,
             request_schema=GraphEmbeddingsRequest,
             response_schema=GraphEmbeddingsResponse,
+            subscription = subscriber,
+            consumer_name = consumer,
             timeout=timeout,
         )
 

trustgraph-flow/trustgraph/gateway/graph_rag.py

@@ -1,20 +1,23 @@
 
 from .. schema import GraphRagQuery, GraphRagResponse
-from .. schema import graph_rag_request_queue
-from .. schema import graph_rag_response_queue
 
 from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 
 class GraphRagRequestor(ServiceRequestor):
-    def __init__(self, pulsar_client, timeout, auth):
+    def __init__(
+            self, pulsar_client, request_queue, response_queue, timeout, auth,
+            consumer, subscriber,
+    ):
 
         super(GraphRagRequestor, self).__init__(
             pulsar_client=pulsar_client,
-            request_queue=graph_rag_request_queue,
+            request_queue=request_queue,
-            response_queue=graph_rag_response_queue,
+            response_queue=response_queue,
             request_schema=GraphRagQuery,
             response_schema=GraphRagResponse,
+            subscription = subscriber,
+            consumer_name = consumer,
             timeout=timeout,
         )
 

trustgraph-flow/trustgraph/gateway/prompt.py

@@ -2,21 +2,24 @@
 import json
 
 from .. schema import PromptRequest, PromptResponse
-from .. schema import prompt_request_queue
-from .. schema import prompt_response_queue
 
 from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 
 class PromptRequestor(ServiceRequestor):
-    def __init__(self, pulsar_client, timeout, auth):
+    def __init__(
+            self, pulsar_client, request_queue, response_queue, timeout, auth,
+            consumer, subscriber,
+    ):
 
         super(PromptRequestor, self).__init__(
             pulsar_client=pulsar_client,
-            request_queue=prompt_request_queue,
+            request_queue=request_queue,
-            response_queue=prompt_response_queue,
+            response_queue=response_queue,
             request_schema=PromptRequest,
             response_schema=PromptResponse,
+            subscription = subscriber,
+            consumer_name = consumer,
             timeout=timeout,
         )
 

trustgraph-flow/trustgraph/gateway/requestor.py

@@ -34,14 +34,10 @@ class ServiceRequestor:
         self.timeout = timeout
 
     async def start(self):
-
         await self.pub.start()
         await self.sub.start()
 
     async def stop(self):
-
-        print("STOPPING")
-
         await self.pub.stop()
         await self.sub.stop()
 

trustgraph-flow/trustgraph/gateway/service.py

@@ -31,8 +31,8 @@ from . running import Running
 from .. schema import ConfigPush, config_push_queue
 
 from . text_completion import TextCompletionRequestor
-#from . prompt import PromptRequestor
+from . prompt import PromptRequestor
-#from . graph_rag import GraphRagRequestor
+from . graph_rag import GraphRagRequestor
 #from . document_rag import DocumentRagRequestor
 #from . triples_query import TriplesQueryRequestor
 #from . graph_embeddings_query import GraphEmbeddingsQueryRequestor
@@ -145,11 +145,11 @@ class Api:
             #     pulsar_client=self.pulsar_client, timeout=self.timeout,
             #     auth = self.auth,
             # ),
-            "config": ConfigRequestor(
+            (None, "config"): ConfigRequestor(
                 pulsar_client=self.pulsar_client, timeout=self.timeout,
                 auth = self.auth,
             ),
-            "flow": FlowRequestor(
+            (None, "flow"): FlowRequestor(
                 pulsar_client=self.pulsar_client, timeout=self.timeout,
                 auth = self.auth,
             ),
@@ -213,11 +213,11 @@ class Api:
             # ),
             ServiceEndpoint(
                 endpoint_path = "/api/v1/config", auth=self.auth,
-                requestor = self.services["config"],
+                requestor = self.services[(None, "config")],
             ),
             ServiceEndpoint(
                 endpoint_path = "/api/v1/flow", auth=self.auth,
-                requestor = self.services["flow"],
+                requestor = self.services[(None, "flow")],
             ),
             FlowEndpoint(
                 endpoint_path = "/api/v1/flow/{flow}/{kind}",
@@ -282,8 +282,6 @@ class Api:
 
         self.flows = {}
 
-#        self.services = {}
-
     async def on_config(self, msg, proc, flow):
 
         try:
@@ -313,32 +311,49 @@ class Api:
             print(f"Exception: {e}", flush=True)
 
     async def start_flow(self, id, flow):
+
         print("Start flow", id)
         intf = flow["interfaces"]
 
-        if "text-completion" in intf:
+        kinds = {
-            k = (id, "text-completion")
+            "agent": AgentRequestor,
-            if k in self.services:
+            "text-completion": TextCompletionRequestor,
-                await self.services[k].stop()
+            "prompt": PromptRequestor,
-                del self.services[k]
+            "graph-rag": GraphRagRequestor,
+            "embeddings": EmbeddingsRequestor,
+            "graph-embeddings": GraphEmbeddingsRequestor,
+            "triples-query": TriplesQueryRequestor,
+        }
 
-            self.services[k] = TextCompletionRequestor(
+        for api_kind, requestor in kinds.items():
-                pulsar_client=self.pulsar_client, timeout=self.timeout,
+
-                request_queue = intf["text-completion"]["request"],
+            if api_kind in intf:
-                response_queue = intf["text-completion"]["response"],
+                k = (id, api_kind)
-                consumer = f"api-gateway-{id}-text-completion-request",
+                if k in self.services:
-                subscriber = f"api-gateway-{id}-text-completion-request",
+                    await self.services[k].stop()
-                auth = self.auth,
+                    del self.services[k]
-            )
+
-            await self.services[k].start()
+                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",
+                    auth = self.auth,
+                )
+                await self.services[k].start()
 
     async def stop_flow(self, id, flow):
         print("Stop flow", id)
         intf = flow["interfaces"]
 
-        if "text-completion" in intf:
+        svc_list = list(self.services.keys())
-            k = (id, "text-completion")
+
-            if k in self.services:
+        for k in svc_list:
+
+            kid, kkind = k
+
+            if id == kid:
                 await self.services[k].stop()
                 del self.services[k]
 

trustgraph-flow/trustgraph/gateway/triples_query.py

@@ -1,21 +1,24 @@
 
 from .. schema import TriplesQueryRequest, TriplesQueryResponse, Triples
-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 TriplesQueryRequestor(ServiceRequestor):
-    def __init__(self, pulsar_client, timeout, auth):
+    def __init__(
+            self, pulsar_client, request_queue, response_queue, timeout, auth,
+            consumer, subscriber,
+    ):
 
         super(TriplesQueryRequestor, self).__init__(
             pulsar_client=pulsar_client,
-            request_queue=triples_request_queue,
+            request_queue=request_queue,
-            response_queue=triples_response_queue,
+            response_queue=response_queue,
             request_schema=TriplesQueryRequest,
             response_schema=TriplesQueryResponse,
+            subscription = subscriber,
+            consumer_name = consumer,
             timeout=timeout,
         )