apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-16 16:51:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Split API endpoint into endpoint and requestor

Browse source
This commit is contained in:
Cyber MacGeddon 2024-12-06 17:41:32 +00:00
parent fd3db3c925
commit f82d96c3e1
13 changed files with 333 additions and 115 deletions
Download patch file Download diff file Expand all files Collapse all files

19
trustgraph-flow/trustgraph/gateway/agent.py
View file

@ -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
)
)

58
trustgraph-flow/trustgraph/gateway/command.py Normal file
View file

@ -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()

18
trustgraph-flow/trustgraph/gateway/dbpedia.py
View file

@ -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
)
)

19
trustgraph-flow/trustgraph/gateway/embeddings.py
View file

@ -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
)
)

19
trustgraph-flow/trustgraph/gateway/encyclopedia.py
View file

@ -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
)
)

90
trustgraph-flow/trustgraph/gateway/endpoint.py
View file

@ -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}")

18
trustgraph-flow/trustgraph/gateway/graph_rag.py
View file

@ -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
)
)

18
trustgraph-flow/trustgraph/gateway/internet_search.py
View file

@ -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
)
)

18
trustgraph-flow/trustgraph/gateway/prompt.py
View file

@ -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
)
)

128
trustgraph-flow/trustgraph/gateway/requestor.py Normal file
View file

@ -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)

6
trustgraph-flow/trustgraph/gateway/service.py
View file

@ -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(

19
trustgraph-flow/trustgraph/gateway/text_completion.py
View file

@ -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
)
)

18
trustgraph-flow/trustgraph/gateway/triples_query.py
View file

@ -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
)
)