diff --git a/trustgraph-flow/trustgraph/gateway/agent.py b/trustgraph-flow/trustgraph/gateway/agent.py index b6307a62..c7af947b 100644 --- a/trustgraph-flow/trustgraph/gateway/agent.py +++ b/trustgraph-flow/trustgraph/gateway/agent.py @@ -3,10 +3,10 @@ from .. schema import AgentRequest, AgentResponse from .. schema import agent_request_queue from .. schema import agent_response_queue -from . endpoint import MultiResponseServiceEndpoint -from . requestor import MultiResponseServiceRequestor +from . endpoint import ServiceEndpoint +from . requestor import ServiceRequestor -class AgentRequestor(MultiResponseServiceRequestor): +class AgentRequestor(ServiceRequestor): def __init__(self, pulsar_host, timeout, auth): super(AgentRequestor, self).__init__( @@ -39,14 +39,4 @@ class AgentRequestor(MultiResponseServiceRequestor): # The 2nd boolean expression indicates whether we're done responding return resp, (message.answer is not None) -class AgentEndpoint(MultiResponseServiceEndpoint): - def __init__(self, pulsar_host, timeout, auth): - - super(AgentEndpoint, self).__init__( - endpoint_path="/api/v1/agent", - 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 index 57460fbb..8e360ef7 100644 --- a/trustgraph-flow/trustgraph/gateway/command.py +++ b/trustgraph-flow/trustgraph/gateway/command.py @@ -44,9 +44,14 @@ class CommandEndpoint(SocketEndpoint): requestor = self.services[svc] - resp = await requestor.process(request) + async def responder(resp, fin): + await ws.send_json({ + "id": id, + "response": resp, + "complete": fin, + }) - await ws.send_json({"id": id, "response": resp }) + resp = await requestor.process(request, responder) except Exception as e: diff --git a/trustgraph-flow/trustgraph/gateway/dbpedia.py b/trustgraph-flow/trustgraph/gateway/dbpedia.py index 4991e81c..8ae4f695 100644 --- a/trustgraph-flow/trustgraph/gateway/dbpedia.py +++ b/trustgraph-flow/trustgraph/gateway/dbpedia.py @@ -25,16 +25,5 @@ class DbpediaRequestor(ServiceRequestor): ) 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 - ) - ) + return { "text": message.text }, True diff --git a/trustgraph-flow/trustgraph/gateway/embeddings.py b/trustgraph-flow/trustgraph/gateway/embeddings.py index d7fce0ff..74785563 100644 --- a/trustgraph-flow/trustgraph/gateway/embeddings.py +++ b/trustgraph-flow/trustgraph/gateway/embeddings.py @@ -24,7 +24,7 @@ class EmbeddingsRequestor(ServiceRequestor): ) def from_response(self, message): - return { "vectors": message.vectors } + return { "vectors": message.vectors }, True class EmbeddingsEndpoint(ServiceEndpoint): def __init__(self, pulsar_host, timeout, auth): diff --git a/trustgraph-flow/trustgraph/gateway/encyclopedia.py b/trustgraph-flow/trustgraph/gateway/encyclopedia.py index 7891dce3..b222f528 100644 --- a/trustgraph-flow/trustgraph/gateway/encyclopedia.py +++ b/trustgraph-flow/trustgraph/gateway/encyclopedia.py @@ -25,7 +25,7 @@ class EncyclopediaRequestor(ServiceRequestor): ) def from_response(self, message): - return { "text": message.text } + return { "text": message.text }, True class EncyclopediaEndpoint(ServiceEndpoint): diff --git a/trustgraph-flow/trustgraph/gateway/endpoint.py b/trustgraph-flow/trustgraph/gateway/endpoint.py index 8e844533..6d6ca8d5 100644 --- a/trustgraph-flow/trustgraph/gateway/endpoint.py +++ b/trustgraph-flow/trustgraph/gateway/endpoint.py @@ -53,30 +53,10 @@ class ServiceEndpoint: print(data) - resp = await self.requestor.process(data) - - return web.json_response(resp) - - except Exception as e: - logging.error(f"Exception: {e}") - - return web.json_response( - { "error": str(e) } - ) - -class MultiResponseServiceEndpoint(ServiceEndpoint): - - async def handle(self, request): - - try: - - data = await request.json() - print(data) - - def responder(x): + def responder(x, fin): print(x) - resp = await self.requestor.process(data, responder) + resp, fin = await self.requestor.process(data, responder) return web.json_response(resp) diff --git a/trustgraph-flow/trustgraph/gateway/graph_rag.py b/trustgraph-flow/trustgraph/gateway/graph_rag.py index 1438dec3..a9cb6a45 100644 --- a/trustgraph-flow/trustgraph/gateway/graph_rag.py +++ b/trustgraph-flow/trustgraph/gateway/graph_rag.py @@ -26,7 +26,7 @@ class GraphRagRequestor(ServiceRequestor): ) def from_response(self, message): - return { "response": message.response } + return { "response": message.response }, True class GraphRagEndpoint(ServiceEndpoint): def __init__(self, pulsar_host, timeout, auth): diff --git a/trustgraph-flow/trustgraph/gateway/internet_search.py b/trustgraph-flow/trustgraph/gateway/internet_search.py index ce36bd92..9843098b 100644 --- a/trustgraph-flow/trustgraph/gateway/internet_search.py +++ b/trustgraph-flow/trustgraph/gateway/internet_search.py @@ -25,7 +25,7 @@ class InternetSearchRequestor(ServiceRequestor): ) def from_response(self, message): - return { "text": message.text } + return { "text": message.text }, True class InternetSearchEndpoint(ServiceEndpoint): def __init__(self, pulsar_host, timeout, auth): diff --git a/trustgraph-flow/trustgraph/gateway/prompt.py b/trustgraph-flow/trustgraph/gateway/prompt.py index af121b02..31d8cbea 100644 --- a/trustgraph-flow/trustgraph/gateway/prompt.py +++ b/trustgraph-flow/trustgraph/gateway/prompt.py @@ -33,11 +33,11 @@ class PromptRequestor(ServiceRequestor): if message.object: return { "object": message.object - } + }, True else: return { "text": message.text - } + }, True class PromptEndpoint(ServiceEndpoint): def __init__(self, pulsar_host, timeout, auth): diff --git a/trustgraph-flow/trustgraph/gateway/requestor.py b/trustgraph-flow/trustgraph/gateway/requestor.py index ad6cf185..5f6e2692 100644 --- a/trustgraph-flow/trustgraph/gateway/requestor.py +++ b/trustgraph-flow/trustgraph/gateway/requestor.py @@ -45,7 +45,7 @@ class ServiceRequestor: def from_response(self, response): raise RuntimeError("Not defined") - async def process(self, request): + async def process(self, request, responder=None): id = str(uuid.uuid4()) @@ -57,41 +57,6 @@ class ServiceRequestor: self.pub.send, id, self.to_request(request) ) - 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 } - - 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: @@ -102,16 +67,16 @@ class MultiResponseServiceRequestor(ServiceRequestor): 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) + print(resp, fin) + + if responder: + await responder(resp, fin) if fin: return resp - # Not finished, so loop round and continue - except Exception as e: logging.error(f"Exception: {e}") diff --git a/trustgraph-flow/trustgraph/gateway/service.py b/trustgraph-flow/trustgraph/gateway/service.py index b0ad2fd1..3898242e 100755 --- a/trustgraph-flow/trustgraph/gateway/service.py +++ b/trustgraph-flow/trustgraph/gateway/service.py @@ -46,7 +46,7 @@ from . triples_load import TriplesLoadEndpoint from . graph_embeddings_load import GraphEmbeddingsLoadEndpoint from . command import CommandEndpoint -from . endpoint import ServiceEndpoint, MultiResponseServiceEndpoint +from . endpoint import ServiceEndpoint from . auth import Authenticator logger = logging.getLogger("api") @@ -138,7 +138,7 @@ class Api: endpoint_path = "/api/v1/embeddings", auth=self.auth, requestor = self.services["embeddings"], ), - MultiResponseServiceEndpoint( + ServiceEndpoint( endpoint_path = "/api/v1/agent", auth=self.auth, requestor = self.services["agent"], ), diff --git a/trustgraph-flow/trustgraph/gateway/text_completion.py b/trustgraph-flow/trustgraph/gateway/text_completion.py index 7819b3e6..80428c2f 100644 --- a/trustgraph-flow/trustgraph/gateway/text_completion.py +++ b/trustgraph-flow/trustgraph/gateway/text_completion.py @@ -25,7 +25,7 @@ class TextCompletionRequestor(ServiceRequestor): ) def from_response(self, message): - return { "response": message.response } + return { "response": message.response }, True class TextCompletionEndpoint(ServiceEndpoint): def __init__(self, pulsar_host, timeout, auth): diff --git a/trustgraph-flow/trustgraph/gateway/triples_query.py b/trustgraph-flow/trustgraph/gateway/triples_query.py index c95bc12a..d6ec5c54 100644 --- a/trustgraph-flow/trustgraph/gateway/triples_query.py +++ b/trustgraph-flow/trustgraph/gateway/triples_query.py @@ -49,7 +49,7 @@ class TriplesQueryRequestor(ServiceRequestor): print(message) return { "response": serialize_subgraph(message.triples) - } + }, True class TriplesQueryEndpoint(ServiceEndpoint): def __init__(self, pulsar_host, timeout, auth):