From c03ba1e5e0144c616fab1959cf261493728aee82 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Tue, 29 Apr 2025 22:41:11 +0100 Subject: [PATCH] More API stuff --- trustgraph-flow/trustgraph/gateway/agent.py | 13 ++-- .../trustgraph/gateway/embeddings.py | 13 ++-- .../gateway/graph_embeddings_query.py | 13 ++-- .../trustgraph/gateway/graph_rag.py | 13 ++-- trustgraph-flow/trustgraph/gateway/prompt.py | 13 ++-- .../trustgraph/gateway/requestor.py | 4 -- trustgraph-flow/trustgraph/gateway/service.py | 65 ++++++++++++------- .../trustgraph/gateway/triples_query.py | 13 ++-- 8 files changed, 88 insertions(+), 59 deletions(-) diff --git a/trustgraph-flow/trustgraph/gateway/agent.py b/trustgraph-flow/trustgraph/gateway/agent.py index 5a54931b..c810b2dd 100644 --- a/trustgraph-flow/trustgraph/gateway/agent.py +++ b/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, - response_queue=agent_response_queue, + request_queue=request_queue, + response_queue=response_queue, request_schema=AgentRequest, response_schema=AgentResponse, + subscription = subscriber, + consumer_name = consumer, timeout=timeout, ) diff --git a/trustgraph-flow/trustgraph/gateway/embeddings.py b/trustgraph-flow/trustgraph/gateway/embeddings.py index 42ed91a1..3bce6010 100644 --- a/trustgraph-flow/trustgraph/gateway/embeddings.py +++ b/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, - response_queue=embeddings_response_queue, + request_queue=request_queue, + response_queue=response_queue, request_schema=EmbeddingsRequest, response_schema=EmbeddingsResponse, + subscription = subscriber, + consumer_name = consumer, timeout=timeout, ) diff --git a/trustgraph-flow/trustgraph/gateway/graph_embeddings_query.py b/trustgraph-flow/trustgraph/gateway/graph_embeddings_query.py index 8df38e97..8cba4f9a 100644 --- a/trustgraph-flow/trustgraph/gateway/graph_embeddings_query.py +++ b/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, - response_queue=graph_embeddings_response_queue, + request_queue=request_queue, + response_queue=response_queue, request_schema=GraphEmbeddingsRequest, response_schema=GraphEmbeddingsResponse, + subscription = subscriber, + consumer_name = consumer, timeout=timeout, ) diff --git a/trustgraph-flow/trustgraph/gateway/graph_rag.py b/trustgraph-flow/trustgraph/gateway/graph_rag.py index b2b69758..3dc9fcfe 100644 --- a/trustgraph-flow/trustgraph/gateway/graph_rag.py +++ b/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, - response_queue=graph_rag_response_queue, + request_queue=request_queue, + response_queue=response_queue, request_schema=GraphRagQuery, response_schema=GraphRagResponse, + subscription = subscriber, + consumer_name = consumer, timeout=timeout, ) diff --git a/trustgraph-flow/trustgraph/gateway/prompt.py b/trustgraph-flow/trustgraph/gateway/prompt.py index eb50ac73..86a9e788 100644 --- a/trustgraph-flow/trustgraph/gateway/prompt.py +++ b/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, - response_queue=prompt_response_queue, + request_queue=request_queue, + response_queue=response_queue, request_schema=PromptRequest, response_schema=PromptResponse, + subscription = subscriber, + consumer_name = consumer, timeout=timeout, ) diff --git a/trustgraph-flow/trustgraph/gateway/requestor.py b/trustgraph-flow/trustgraph/gateway/requestor.py index ec70812b..04837518 100644 --- a/trustgraph-flow/trustgraph/gateway/requestor.py +++ b/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() diff --git a/trustgraph-flow/trustgraph/gateway/service.py b/trustgraph-flow/trustgraph/gateway/service.py index 975b6968..1b91523d 100755 --- a/trustgraph-flow/trustgraph/gateway/service.py +++ b/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 . graph_rag import GraphRagRequestor +from . prompt import PromptRequestor +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: - k = (id, "text-completion") - if k in self.services: - await self.services[k].stop() - del self.services[k] + kinds = { + "agent": AgentRequestor, + "text-completion": TextCompletionRequestor, + "prompt": PromptRequestor, + "graph-rag": GraphRagRequestor, + "embeddings": EmbeddingsRequestor, + "graph-embeddings": GraphEmbeddingsRequestor, + "triples-query": TriplesQueryRequestor, + } - self.services[k] = TextCompletionRequestor( - pulsar_client=self.pulsar_client, timeout=self.timeout, - request_queue = intf["text-completion"]["request"], - response_queue = intf["text-completion"]["response"], - consumer = f"api-gateway-{id}-text-completion-request", - subscriber = f"api-gateway-{id}-text-completion-request", - auth = self.auth, - ) - await self.services[k].start() + 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() + del self.services[k] + + 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: - k = (id, "text-completion") - if k in self.services: + svc_list = list(self.services.keys()) + + for k in svc_list: + + kid, kkind = k + + if id == kid: await self.services[k].stop() del self.services[k] diff --git a/trustgraph-flow/trustgraph/gateway/triples_query.py b/trustgraph-flow/trustgraph/gateway/triples_query.py index 061bd4d8..3775b270 100644 --- a/trustgraph-flow/trustgraph/gateway/triples_query.py +++ b/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, - response_queue=triples_response_queue, + request_queue=request_queue, + response_queue=response_queue, request_schema=TriplesQueryRequest, response_schema=TriplesQueryResponse, + subscription = subscriber, + consumer_name = consumer, timeout=timeout, )