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

Socket structure working

Browse source
This commit is contained in:
Cyber MacGeddon 2025-05-01 16:58:38 +01:00
parent e473672297
commit f5c398cd41
6 changed files with 108 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

4
trustgraph-flow/trustgraph/gateway/dispatch/librarian.py
View file

@ -21,20 +21,16 @@ class LibrarianRequestor(ServiceRequestor):
def to_request(self, body):
print("TRR")
if "document" in body:
dp = to_document_package(body["document"])
else:
dp = None
print("GOT")
if "criteria" in body:
criteria = to_criteria(body["criteria"])
else:
criteria = None
print("ASLDKJ")
return LibrarianRequest(
operation = body.get("operation", None),
id = body.get("id", None),

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

@ -3,9 +3,9 @@ import asyncio
import queue
import uuid
from .. schema import DocumentEmbeddings
from .. schema import document_embeddings_store_queue
from .. base import Subscriber
from ... schema import DocumentEmbeddings
from ... schema import document_embeddings_store_queue
from ... base import Subscriber
from . socket import SocketEndpoint
from . serialize import serialize_document_embeddings
@ -14,7 +14,8 @@ class DocumentEmbeddingsStreamEndpoint(SocketEndpoint):
def __init__(
self, pulsar_client, auth,
path="/api/v1/stream/document-embeddings"
queue,
path="/api/v1/stream/document-embeddings",
):
super(DocumentEmbeddingsStreamEndpoint, self).__init__(
@ -24,7 +25,7 @@ class DocumentEmbeddingsStreamEndpoint(SocketEndpoint):
self.pulsar_client=pulsar_client
self.subscriber = Subscriber(
self.pulsar_client, document_embeddings_store_queue,
self.pulsar_client, queue,
"api-gateway", "api-gateway",
schema=DocumentEmbeddings,
)

1
trustgraph-flow/trustgraph/gateway/endpoint/flow_endpoint.py
View file

@ -23,7 +23,6 @@ class FlowEndpoint:
def add_routes(self, app):
pass
app.add_routes([
web.post(self.path, self.handle),
])

34
trustgraph-flow/trustgraph/gateway/endpoint/flows.py
View file

@ -12,6 +12,7 @@ from .. dispatch.triples_query import TriplesQueryRequestor
from .. dispatch.embeddings import EmbeddingsRequestor
from .. dispatch.graph_embeddings_query import GraphEmbeddingsQueryRequestor
from .. dispatch.prompt import PromptRequestor
from . stream import StreamEndpoint
class FlowEndpointManager:
@ -30,6 +31,11 @@ class FlowEndpointManager:
auth = auth,
requestors = self.services,
),
StreamEndpoint(
endpoint_path = "/api/v1/flow/{flow}/stream/{kind}",
auth = auth,
dispatchers = self.services,
),
]
self.config_receiver.add_handler(self)
@ -62,6 +68,7 @@ class FlowEndpointManager:
for api_kind, requestor in kinds.items():
if api_kind in intf:
k = (id, api_kind)
if k in self.services:
await self.services[k].stop()
@ -76,12 +83,35 @@ class FlowEndpointManager:
)
await self.services[k].start()
kinds = {
# "document-embeddings-stream": DocumentEmbeddingsStreamEndpoint,
# "triples-store"
# "bunch":
}
for api_kind, streamer in kinds.items():
# if api_kind in intf:
if True:
k = (id, api_kind)
if k in self.services:
await self.services[k].stop()
del self.services[k]
self.services[k] = steamer(
# pulsar_client=self.pulsar_client,
# timeout = self.timeout,
# input_queue = intf[api_kind],
# consumer = f"api-gateway-{id}-{api_kind}-stream",
# subscriber = f"api-gateway-{id}-{api_kind}-stream",
)
await self.services[k].start()
async def stop_flow(self, id, flow):
print("STOP FLOW", id)
intf = flow["interfaces"]
svc_list = list(self.services.keys())
for k in svc_list:

33
trustgraph-flow/trustgraph/gateway/socket.py → trustgraph-flow/trustgraph/gateway/endpoint/socket.py
View file

@ -3,7 +3,7 @@ import asyncio
from aiohttp import web, WSMsgType
import logging
from . running import Running
from .. running import Running
logger = logging.getLogger("socket")
logger.setLevel(logging.INFO)
@ -18,7 +18,12 @@ class SocketEndpoint:
self.auth = auth
self.operation = "socket"
async def listener(self, ws, running):
self.running = Running()
async def dispatcher(self, ws):
pass
async def listener(self, ws):
async for msg in ws:
# On error, finish
@ -31,7 +36,8 @@ class SocketEndpoint:
else:
break
running.stop()
self.running.stop()
await ws.close()
async def handle(self, request):
@ -43,20 +49,28 @@ class SocketEndpoint:
if not self.auth.permitted(token, self.operation):
return web.HTTPUnauthorized()
running = Running()
# 50MB max message size
ws = web.WebSocketResponse(max_msg_size=52428800)
await ws.prepare(request)
try:
await self.listener(ws, running)
async with asyncio.TaskGroup() as tg:
worker = tg.create_task(
self.dispatcher(ws)
)
worker = tg.create_task(
self.listener(ws)
)
# Wait for threads to complete
except Exception as e:
print("Socket exception:", e, flush=True)
running.stop()
await ws.close()
return ws
@ -64,6 +78,9 @@ class SocketEndpoint:
async def start(self):
pass
async def stop(self):
self.running.stop()
def add_routes(self, app):
app.add_routes([

45
trustgraph-flow/trustgraph/gateway/endpoint/stream.py Normal file
View file

@ -0,0 +1,45 @@
import asyncio
from aiohttp import web, WSMsgType
import logging
from .. running import Running
from . socket import SocketEndpoint
logger = logging.getLogger("socket")
logger.setLevel(logging.INFO)
class StreamEndpoint(SocketEndpoint):
def __init__(
self, endpoint_path, auth, dispatchers,
):
super(StreamEndpoint, self).__init__(endpoint_path, auth)
async def listener(self, ws):
async for msg in ws:
# On error, finish
if msg.type == WSMsgType.TEXT:
print("Received:", msg)
continue
elif msg.type == WSMsgType.BINARY:
# Ignore incoming message
continue
else:
break
self.running.stop()
async def dispatcher(self, ws):
for i in range(0, 10):
await asyncio.sleep(1)
await ws.send_json({"number": i})
await ws.close()
self.running.stop()