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

Migrate api-gateway to async

Browse source
This commit is contained in:
Cyber MacGeddon 2025-04-17 17:19:54 +01:00
parent 24d05cbcdf
commit 8b24fa4209
13 changed files with 37 additions and 56 deletions
Download patch file Download diff file Expand all files Collapse all files

7
trustgraph-base/trustgraph/base/publisher.py
View file

@ -1,4 +1,6 @@
from pulsar.schema import JsonSchema
import asyncio import asyncio
import time import time
import pulsar import pulsar
@ -16,7 +18,6 @@ class Publisher:
async def start(self): async def start(self):
self.task = asyncio.create_task(self.run()) self.task = asyncio.create_task(self.run())
await self.task.start()
async def stop(self): async def stop(self):
self.running = False self.running = False
@ -32,14 +33,14 @@ class Publisher:
try: try:
producer = self.client.create_producer( producer = self.client.create_producer(
topic=self.topic, topic=self.topic,
schema=self.schema, schema=JsonSchema(self.schema),
chunking_enabled=self.chunking_enabled, chunking_enabled=self.chunking_enabled,
) )
while self.running: while self.running:
try: try:
id, item = asyncio.wait_for( id, item = await asyncio.wait_for(
self.q.get(), self.q.get(),
timeout=0.5 timeout=0.5
) )

7
trustgraph-flow/trustgraph/gateway/document_embeddings_load.py
View file

@ -1,6 +1,5 @@
import asyncio import asyncio
from pulsar.schema import JsonSchema
import uuid import uuid
from aiohttp import WSMsgType from aiohttp import WSMsgType
@ -26,12 +25,12 @@ class DocumentEmbeddingsLoadEndpoint(SocketEndpoint):
self.publisher = Publisher( self.publisher = Publisher(
self.pulsar_client, document_embeddings_store_queue, self.pulsar_client, document_embeddings_store_queue,
schema=JsonSchema(DocumentEmbeddings) schema=DocumentEmbeddings
) )
async def start(self): async def start(self):
self.publisher.start() await self.publisher.start()
async def listener(self, ws, running): async def listener(self, ws, running):
@ -59,6 +58,6 @@ class DocumentEmbeddingsLoadEndpoint(SocketEndpoint):
], ],
) )
self.publisher.send(None, elt) await self.publisher.send(None, elt)
running.stop() running.stop()

11
trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py
View file

@ -1,7 +1,6 @@
import asyncio import asyncio
import queue import queue
from pulsar.schema import JsonSchema
import uuid import uuid
from .. schema import DocumentEmbeddings from .. schema import DocumentEmbeddings
@ -27,7 +26,7 @@ class DocumentEmbeddingsStreamEndpoint(SocketEndpoint):
self.subscriber = Subscriber( self.subscriber = Subscriber(
self.pulsar_client, document_embeddings_store_queue, self.pulsar_client, document_embeddings_store_queue,
"api-gateway", "api-gateway", "api-gateway", "api-gateway",
schema=JsonSchema(DocumentEmbeddings), schema=DocumentEmbeddings,
) )
async def listener(self, ws, running): async def listener(self, ws, running):
@ -44,17 +43,17 @@ class DocumentEmbeddingsStreamEndpoint(SocketEndpoint):
async def start(self): async def start(self):
self.subscriber.start() await self.subscriber.start()
async def async_thread(self, ws, running): async def async_thread(self, ws, running):
id = str(uuid.uuid4()) id = str(uuid.uuid4())
q = self.subscriber.subscribe_all(id) q = await self.subscriber.subscribe_all(id)
while running.get(): while running.get():
try: try:
resp = await asyncio.to_thread(q.get, timeout=0.5) resp = await asyncio.wait_for(q.get(), timeout=0.5)
await ws.send_json(serialize_document_embeddings(resp)) await ws.send_json(serialize_document_embeddings(resp))
except TimeoutError: except TimeoutError:
@ -67,7 +66,7 @@ class DocumentEmbeddingsStreamEndpoint(SocketEndpoint):
print(f"Exception: {str(e)}", flush=True) print(f"Exception: {str(e)}", flush=True)
break break
self.subscriber.unsubscribe_all(id) await self.subscriber.unsubscribe_all(id)
running.stop() running.stop()

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

@ -1,13 +1,9 @@
import asyncio import asyncio
from pulsar.schema import JsonSchema
from aiohttp import web from aiohttp import web
import uuid import uuid
import logging import logging
from .. base import Publisher
from .. base import Subscriber
logger = logging.getLogger("endpoint") logger = logging.getLogger("endpoint")
logger.setLevel(logging.INFO) logger.setLevel(logging.INFO)

7
trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py
View file

@ -1,6 +1,5 @@
import asyncio import asyncio
from pulsar.schema import JsonSchema
import uuid import uuid
from aiohttp import WSMsgType from aiohttp import WSMsgType
@ -26,12 +25,12 @@ class GraphEmbeddingsLoadEndpoint(SocketEndpoint):
self.publisher = Publisher( self.publisher = Publisher(
self.pulsar_client, graph_embeddings_store_queue, self.pulsar_client, graph_embeddings_store_queue,
schema=JsonSchema(GraphEmbeddings) schema=GraphEmbeddings
) )
async def start(self): async def start(self):
self.publisher.start() await self.publisher.start()
async def listener(self, ws, running): async def listener(self, ws, running):
@ -60,6 +59,6 @@ class GraphEmbeddingsLoadEndpoint(SocketEndpoint):
] ]
) )
self.publisher.send(None, elt) await self.publisher.send(None, elt)
running.stop() running.stop()

11
trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py
View file

@ -1,7 +1,6 @@
import asyncio import asyncio
import queue import queue
from pulsar.schema import JsonSchema
import uuid import uuid
from .. schema import GraphEmbeddings from .. schema import GraphEmbeddings
@ -26,7 +25,7 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint):
self.subscriber = Subscriber( self.subscriber = Subscriber(
self.pulsar_client, graph_embeddings_store_queue, self.pulsar_client, graph_embeddings_store_queue,
"api-gateway", "api-gateway", "api-gateway", "api-gateway",
schema=JsonSchema(GraphEmbeddings) schema=GraphEmbeddings
) )
async def listener(self, ws, running): async def listener(self, ws, running):
@ -41,17 +40,17 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint):
async def start(self): async def start(self):
self.subscriber.start() await self.subscriber.start()
async def async_thread(self, ws, running): async def async_thread(self, ws, running):
id = str(uuid.uuid4()) id = str(uuid.uuid4())
q = self.subscriber.subscribe_all(id) q = await self.subscriber.subscribe_all(id)
while running.get(): while running.get():
try: try:
resp = await asyncio.to_thread(q.get, timeout=0.5) resp = await asyncio.wait_for(q.get, timeout=0.5)
await ws.send_json(serialize_graph_embeddings(resp)) await ws.send_json(serialize_graph_embeddings(resp))
except TimeoutError: except TimeoutError:
@ -64,7 +63,7 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint):
print(f"Exception: {str(e)}", flush=True) print(f"Exception: {str(e)}", flush=True)
break break
self.subscriber.unsubscribe_all(id) await self.subscriber.unsubscribe_all(id)
running.stop() running.stop()

1
trustgraph-flow/trustgraph/gateway/metrics.py
View file

@ -7,7 +7,6 @@
import aiohttp import aiohttp
from aiohttp import web from aiohttp import web
import asyncio import asyncio
from pulsar.schema import JsonSchema
import uuid import uuid
import logging import logging

1
trustgraph-flow/trustgraph/gateway/mux.py
View file

@ -1,7 +1,6 @@
import asyncio import asyncio
import queue import queue
from pulsar.schema import JsonSchema
import uuid import uuid
from aiohttp import web, WSMsgType from aiohttp import web, WSMsgType

22
trustgraph-flow/trustgraph/gateway/requestor.py
View file

@ -1,6 +1,5 @@
import asyncio import asyncio
from pulsar.schema import JsonSchema
import uuid import uuid
import logging import logging
@ -23,21 +22,21 @@ class ServiceRequestor:
self.pub = Publisher( self.pub = Publisher(
pulsar_client, request_queue, pulsar_client, request_queue,
schema=JsonSchema(request_schema), schema=request_schema,
) )
self.sub = Subscriber( self.sub = Subscriber(
pulsar_client, response_queue, pulsar_client, response_queue,
subscription, consumer_name, subscription, consumer_name,
JsonSchema(response_schema) response_schema
) )
self.timeout = timeout self.timeout = timeout
async def start(self): async def start(self):
self.pub.start() await self.pub.start()
self.sub.start() await self.sub.start()
def to_request(self, request): def to_request(self, request):
raise RuntimeError("Not defined") raise RuntimeError("Not defined")
@ -51,18 +50,15 @@ class ServiceRequestor:
try: try:
q = self.sub.subscribe(id) q = await self.sub.subscribe(id)
await asyncio.to_thread( await self.pub.send(id, self.to_request(request))
self.pub.send, id, self.to_request(request)
)
while True: while True:
try: try:
resp = await asyncio.to_thread( resp = await asyncio.wait_for(
q.get, q.get(), timeout=self.timeout
timeout=self.timeout
) )
except Exception as e: except Exception as e:
raise RuntimeError("Timeout") raise RuntimeError("Timeout")
@ -99,5 +95,5 @@ class ServiceRequestor:
return err return err
finally: finally:
self.sub.unsubscribe(id) await self.sub.unsubscribe(id)

9
trustgraph-flow/trustgraph/gateway/sender.py
View file

@ -2,7 +2,6 @@
# Like ServiceRequestor, but just fire-and-forget instead of request/response # Like ServiceRequestor, but just fire-and-forget instead of request/response
import asyncio import asyncio
from pulsar.schema import JsonSchema
import uuid import uuid
import logging import logging
@ -21,12 +20,12 @@ class ServiceSender:
self.pub = Publisher( self.pub = Publisher(
pulsar_client, request_queue, pulsar_client, request_queue,
schema=JsonSchema(request_schema), schema=request_schema,
) )
async def start(self): async def start(self):
self.pub.start() await self.pub.start()
def to_request(self, request): def to_request(self, request):
raise RuntimeError("Not defined") raise RuntimeError("Not defined")
@ -35,9 +34,7 @@ class ServiceSender:
try: try:
await asyncio.to_thread( await self.pub.send(None, self.to_request(request))
self.pub.send, None, self.to_request(request)
)
if responder: if responder:
await responder({}, True) await responder({}, True)

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

@ -19,7 +19,6 @@ import os
import base64 import base64
import pulsar import pulsar
from pulsar.schema import JsonSchema
from prometheus_client import start_http_server from prometheus_client import start_http_server
from .. log_level import LogLevel from .. log_level import LogLevel

7
trustgraph-flow/trustgraph/gateway/triples_load.py
View file

@ -1,6 +1,5 @@
import asyncio import asyncio
from pulsar.schema import JsonSchema
import uuid import uuid
from aiohttp import WSMsgType from aiohttp import WSMsgType
@ -24,12 +23,12 @@ class TriplesLoadEndpoint(SocketEndpoint):
self.publisher = Publisher( self.publisher = Publisher(
self.pulsar_client, triples_store_queue, self.pulsar_client, triples_store_queue,
schema=JsonSchema(Triples) schema=Triples
) )
async def start(self): async def start(self):
self.publisher.start() await self.publisher.start()
async def listener(self, ws, running): async def listener(self, ws, running):
@ -51,7 +50,7 @@ class TriplesLoadEndpoint(SocketEndpoint):
triples=to_subgraph(data["triples"]), triples=to_subgraph(data["triples"]),
) )
self.publisher.send(None, elt) await self.publisher.send(None, elt)
running.stop() running.stop()

5
trustgraph-flow/trustgraph/gateway/triples_stream.py
View file

@ -1,7 +1,6 @@
import asyncio import asyncio
import queue import queue
from pulsar.schema import JsonSchema
import uuid import uuid
from .. schema import Triples from .. schema import Triples
@ -24,7 +23,7 @@ class TriplesStreamEndpoint(SocketEndpoint):
self.subscriber = Subscriber( self.subscriber = Subscriber(
self.pulsar_client, triples_store_queue, self.pulsar_client, triples_store_queue,
"api-gateway", "api-gateway", "api-gateway", "api-gateway",
schema=JsonSchema(Triples) schema=Triples
) )
async def listener(self, ws, running): async def listener(self, ws, running):
@ -39,7 +38,7 @@ class TriplesStreamEndpoint(SocketEndpoint):
async def start(self): async def start(self):
self.subscriber.start() await self.subscriber.start()
async def async_thread(self, ws, running): async def async_thread(self, ws, running):