diff --git a/trustgraph-flow/trustgraph/gateway/agent.py b/trustgraph-flow/trustgraph/gateway/agent.py index e8fd0e72..c86b5f02 100644 --- a/trustgraph-flow/trustgraph/gateway/agent.py +++ b/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 + ) + ) + diff --git a/trustgraph-flow/trustgraph/gateway/command.py b/trustgraph-flow/trustgraph/gateway/command.py new file mode 100644 index 00000000..026523d9 --- /dev/null +++ b/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() + diff --git a/trustgraph-flow/trustgraph/gateway/dbpedia.py b/trustgraph-flow/trustgraph/gateway/dbpedia.py index a61292a6..4991e81c 100644 --- a/trustgraph-flow/trustgraph/gateway/dbpedia.py +++ b/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 + ) + ) + diff --git a/trustgraph-flow/trustgraph/gateway/embeddings.py b/trustgraph-flow/trustgraph/gateway/embeddings.py index 6d3a9fe6..d7fce0ff 100644 --- a/trustgraph-flow/trustgraph/gateway/embeddings.py +++ b/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 + ) + ) + diff --git a/trustgraph-flow/trustgraph/gateway/encyclopedia.py b/trustgraph-flow/trustgraph/gateway/encyclopedia.py index 32eb5cd1..7891dce3 100644 --- a/trustgraph-flow/trustgraph/gateway/encyclopedia.py +++ b/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 + ) + ) + diff --git a/trustgraph-flow/trustgraph/gateway/endpoint.py b/trustgraph-flow/trustgraph/gateway/endpoint.py index 2b246361..40371cc9 100644 --- a/trustgraph-flow/trustgraph/gateway/endpoint.py +++ b/trustgraph-flow/trustgraph/gateway/endpoint.py @@ -13,38 +13,17 @@ logger.setLevel(logging.INFO) class ServiceEndpoint: - def __init__( - 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) - ) + def __init__(self, endpoint_path, auth, requestor): 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() - self.sub.start() + async def start(self): + 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( - 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) - ) + return web.json_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: - 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 + return web.json_response(resp) except Exception as e: logging.error(f"Exception: {e}") diff --git a/trustgraph-flow/trustgraph/gateway/graph_rag.py b/trustgraph-flow/trustgraph/gateway/graph_rag.py index 58679004..1438dec3 100644 --- a/trustgraph-flow/trustgraph/gateway/graph_rag.py +++ b/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 + ) + ) + diff --git a/trustgraph-flow/trustgraph/gateway/internet_search.py b/trustgraph-flow/trustgraph/gateway/internet_search.py index 5a5dc948..ce36bd92 100644 --- a/trustgraph-flow/trustgraph/gateway/internet_search.py +++ b/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 + ) + ) + diff --git a/trustgraph-flow/trustgraph/gateway/prompt.py b/trustgraph-flow/trustgraph/gateway/prompt.py index f09a0e0e..af121b02 100644 --- a/trustgraph-flow/trustgraph/gateway/prompt.py +++ b/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 + ) + ) + diff --git a/trustgraph-flow/trustgraph/gateway/requestor.py b/trustgraph-flow/trustgraph/gateway/requestor.py new file mode 100644 index 00000000..985ff2dd --- /dev/null +++ b/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) + diff --git a/trustgraph-flow/trustgraph/gateway/service.py b/trustgraph-flow/trustgraph/gateway/service.py index e927ecf6..81f02197 100755 --- a/trustgraph-flow/trustgraph/gateway/service.py +++ b/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( diff --git a/trustgraph-flow/trustgraph/gateway/text_completion.py b/trustgraph-flow/trustgraph/gateway/text_completion.py index d59737f0..7819b3e6 100644 --- a/trustgraph-flow/trustgraph/gateway/text_completion.py +++ b/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 + ) + ) + diff --git a/trustgraph-flow/trustgraph/gateway/triples_query.py b/trustgraph-flow/trustgraph/gateway/triples_query.py index 5a0cfff8..c95bc12a 100644 --- a/trustgraph-flow/trustgraph/gateway/triples_query.py +++ b/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 + ) + ) +