apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-17 09:11:03 +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 time
import pulsar
@ -16,7 +18,6 @@ class Publisher:
async def start(self):
self.task = asyncio.create_task(self.run())
await self.task.start()
async def stop(self):
self.running = False
@ -32,14 +33,14 @@ class Publisher:
try:
producer = self.client.create_producer(
topic=self.topic,
schema=self.schema,
schema=JsonSchema(self.schema),
chunking_enabled=self.chunking_enabled,
)
while self.running:
try:
id, item = asyncio.wait_for(
id, item = await asyncio.wait_for(
self.q.get(),
timeout=0.5
)

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

@ -1,6 +1,5 @@
import asyncio
from pulsar.schema import JsonSchema
import uuid
from aiohttp import WSMsgType
@ -26,12 +25,12 @@ class DocumentEmbeddingsLoadEndpoint(SocketEndpoint):
self.publisher = Publisher(
self.pulsar_client, document_embeddings_store_queue,
schema=JsonSchema(DocumentEmbeddings)
schema=DocumentEmbeddings
)
async def start(self):
self.publisher.start()
await self.publisher.start()
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()

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

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

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

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

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

@ -1,6 +1,5 @@
import asyncio
from pulsar.schema import JsonSchema
import uuid
from aiohttp import WSMsgType
@ -26,12 +25,12 @@ class GraphEmbeddingsLoadEndpoint(SocketEndpoint):
self.publisher = Publisher(
self.pulsar_client, graph_embeddings_store_queue,
schema=JsonSchema(GraphEmbeddings)
schema=GraphEmbeddings
)
async def start(self):
self.publisher.start()
await self.publisher.start()
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()

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

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

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

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

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

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

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

@ -1,6 +1,5 @@
import asyncio
from pulsar.schema import JsonSchema
import uuid
import logging
@ -23,21 +22,21 @@ class ServiceRequestor:
self.pub = Publisher(
pulsar_client, request_queue,
schema=JsonSchema(request_schema),
schema=request_schema,
)
self.sub = Subscriber(
pulsar_client, response_queue,
subscription, consumer_name,
JsonSchema(response_schema)
response_schema
)
self.timeout = timeout
async def start(self):
self.pub.start()
self.sub.start()
await self.pub.start()
await self.sub.start()
def to_request(self, request):
raise RuntimeError("Not defined")
@ -51,18 +50,15 @@ class ServiceRequestor:
try:
q = self.sub.subscribe(id)
q = await self.sub.subscribe(id)
await asyncio.to_thread(
self.pub.send, id, self.to_request(request)
)
await self.pub.send(id, self.to_request(request))
while True:
try:
resp = await asyncio.to_thread(
q.get,
timeout=self.timeout
resp = await asyncio.wait_for(
q.get(), timeout=self.timeout
)
except Exception as e:
raise RuntimeError("Timeout")
@ -99,5 +95,5 @@ class ServiceRequestor:
return err
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
import asyncio
from pulsar.schema import JsonSchema
import uuid
import logging
@ -21,12 +20,12 @@ class ServiceSender:
self.pub = Publisher(
pulsar_client, request_queue,
schema=JsonSchema(request_schema),
schema=request_schema,
)
async def start(self):
self.pub.start()
await self.pub.start()
def to_request(self, request):
raise RuntimeError("Not defined")
@ -35,9 +34,7 @@ class ServiceSender:
try:
await asyncio.to_thread(
self.pub.send, None, self.to_request(request)
)
await self.pub.send(None, self.to_request(request))
if responder:
await responder({}, True)

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

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

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

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

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

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