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

Rework loaders, refactored stream/load to be export/import

Browse source
This commit is contained in:
Cyber MacGeddon 2025-05-02 20:03:17 +01:00
parent 8e9c72fda5
commit 03d59b62af
11 changed files with 240 additions and 200 deletions
Download patch file Download diff file Expand all files Collapse all files

2
trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_stream.py → trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_export.py
View file

@ -8,7 +8,7 @@ from ... base import Subscriber
from . serialize import serialize_document_embeddings
class DocumentEmbeddingsStream:
class DocumentEmbeddingsExport:
def __init__(
self, ws, running, pulsar_client, queue, consumer, subscriber

60
trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_import.py Normal file
View file

@ -0,0 +1,60 @@
import asyncio
import uuid
from aiohttp import WSMsgType
from .. schema import Metadata
from .. schema import DocumentEmbeddings, ChunkEmbeddings
from .. base import Publisher
from . socket import SocketEndpoint
from . serialize import to_subgraph
class DocumentEmbeddingsImport:
def __init__(
self, ws, running, pulsar_client, queue
):
self.ws = ws
self.running = running
self.publisher = Publisher(
pulsar_client, queue = queue, schema = DocumentEmbeddings
)
async def destroy(self):
self.running.stop()
await self.ws.close()
await self.publisher.stop()
async def receive(self, msg):
data = msg.json()
elt = DocumentEmbeddings(
metadata=Metadata(
id=data["metadata"]["id"],
metadata=to_subgraph(data["metadata"]["metadata"]),
user=data["metadata"]["user"],
collection=data["metadata"]["collection"],
),
chunks=[
ChunkEmbeddings(
chunk=de["chunk"].encode("utf-8"),
vectors=de["vectors"],
)
for de in data["chunks"]
],
)
await self.publisher.send(None, elt)
async def run(self):
while self.running.get():
await asyncio.sleep(0.5)
await self.ws.close()
self.ws = None

63
trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_load.py
View file

@ -1,63 +0,0 @@
import asyncio
import uuid
from aiohttp import WSMsgType
from .. schema import Metadata
from .. schema import DocumentEmbeddings, ChunkEmbeddings
from .. schema import document_embeddings_store_queue
from .. base import Publisher
from . socket import SocketEndpoint
from . serialize import to_subgraph
class DocumentEmbeddingsLoadEndpoint(SocketEndpoint):
def __init__(
self, pulsar_client, auth, path="/api/v1/load/document-embeddings",
):
super(DocumentEmbeddingsLoadEndpoint, self).__init__(
endpoint_path=path, auth=auth,
)
self.pulsar_client=pulsar_client
self.publisher = Publisher(
self.pulsar_client, document_embeddings_store_queue,
schema=DocumentEmbeddings
)
async def start(self):
await self.publisher.start()
async def listener(self, ws, running):
async for msg in ws:
# On error, finish
if msg.type == WSMsgType.ERROR:
break
else:
data = msg.json()
elt = DocumentEmbeddings(
metadata=Metadata(
id=data["metadata"]["id"],
metadata=to_subgraph(data["metadata"]["metadata"]),
user=data["metadata"]["user"],
collection=data["metadata"]["collection"],
),
chunks=[
ChunkEmbeddings(
chunk=de["chunk"].encode("utf-8"),
vectors=de["vectors"],
)
for de in data["chunks"]
],
)
await self.publisher.send(None, elt)
running.stop()

2
trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_stream.py → trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_export.py
View file

@ -8,7 +8,7 @@ from ... base import Subscriber
from . serialize import serialize_graph_embeddings
class GraphEmbeddingsStream:
class GraphEmbeddingsExport:
def __init__(
self, ws, running, pulsar_client, queue, consumer, subscriber

60
trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_import.py Normal file
View file

@ -0,0 +1,60 @@
import asyncio
import uuid
from aiohttp import WSMsgType
from .. schema import Metadata
from .. schema import GraphEmbeddings, EntityEmbeddings
from .. base import Publisher
from . socket import SocketEndpoint
from . serialize import to_subgraph, to_value
class GraphEmbeddingsImport:
def __init__(
self, ws, running, pulsar_client, queue
):
self.ws = ws
self.running = running
self.publisher = Publisher(
pulsar_client, queue = queue, schema = GraphEmbeddings
)
async def destroy(self):
self.running.stop()
await self.ws.close()
await self.publisher.stop()
async def receive(self, msg):
data = msg.json()
elt = GraphEmbeddings(
metadata=Metadata(
id=data["metadata"]["id"],
metadata=to_subgraph(data["metadata"]["metadata"]),
user=data["metadata"]["user"],
collection=data["metadata"]["collection"],
),
entities=[
EntityEmbeddings(
entity=to_value(ent["entity"]),
vectors=ent["vectors"],
)
for ent in data["entities"]
]
)
await self.publisher.send(None, elt)
async def run(self):
while self.running.get():
await asyncio.sleep(0.5)
await self.ws.close()
self.ws = None

64
trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_load.py
View file

@ -1,64 +0,0 @@
import asyncio
import uuid
from aiohttp import WSMsgType
from .. schema import Metadata
from .. schema import GraphEmbeddings, EntityEmbeddings
from .. schema import graph_embeddings_store_queue
from .. base import Publisher
from . socket import SocketEndpoint
from . serialize import to_subgraph, to_value
class GraphEmbeddingsLoadEndpoint(SocketEndpoint):
def __init__(
self, pulsar_client, auth, path="/api/v1/load/graph-embeddings",
):
super(GraphEmbeddingsLoadEndpoint, self).__init__(
endpoint_path=path, auth=auth,
)
self.pulsar_client=pulsar_client
self.publisher = Publisher(
self.pulsar_client, graph_embeddings_store_queue,
schema=GraphEmbeddings
)
async def start(self):
await self.publisher.start()
async def listener(self, ws, running):
async for msg in ws:
# On error, finish
if msg.type == WSMsgType.ERROR:
break
else:
data = msg.json()
elt = GraphEmbeddings(
metadata=Metadata(
id=data["metadata"]["id"],
metadata=to_subgraph(data["metadata"]["metadata"]),
user=data["metadata"]["user"],
collection=data["metadata"]["collection"],
),
entities=[
EntityEmbeddings(
entity=to_value(ent["entity"]),
vectors=ent["vectors"],
)
for ent in data["entities"]
]
)
await self.publisher.send(None, elt)
running.stop()

64
trustgraph-flow/trustgraph/gateway/dispatch/manager.py
View file

@ -19,9 +19,13 @@ from . prompt import PromptRequestor
from . text_load import TextLoad
from . document_load import DocumentLoad
from . triples_stream import TriplesStream
from . graph_embeddings_stream import GraphEmbeddingsStream
from . document_embeddings_stream import DocumentEmbeddingsStream
from . triples_export import TriplesExport
from . graph_embeddings_export import GraphEmbeddingsExport
from . document_embeddings_export import DocumentEmbeddingsExport
from . triples_import import TriplesImport
from . graph_embeddings_import import GraphEmbeddingsImport
from . document_embeddings_import import DocumentEmbeddingsImport
request_response_dispatchers = {
"agent": AgentRequestor,
@ -39,12 +43,18 @@ sender_dispatchers = {
"document-load": DocumentLoad,
}
receive_dispatchers = {
export_dispatchers = {
"triples": TriplesStream,
"graph-embeddings": GraphEmbeddingsStream,
"document-embeddings": DocumentEmbeddingsStream,
}
import_dispatchers = {
"triples": TriplesImport,
"graph-embeddings": GraphEmbeddingsImport,
"document-embeddings": DocumentEmbeddingsImport,
}
class DispatcherWrapper:
def __init__(self, mgr, name, impl):
self.mgr = mgr
@ -101,13 +111,16 @@ class DispatcherManager:
return await dispatcher.process(data, responder)
def dispatch_flow_service(self):
def dispatch_service(self):
return self
def dispatch_flow_receive(self):
return self.dispatch_receive
def dispatch_import(self):
return self.invoke_import
async def dispatch_receive(self, ws, running, params):
def dispatch_export(self):
return self.invoke_export
async def invoke_import(self, ws, running, params):
flow = params.get("flow")
kind = params.get("kind")
@ -115,7 +128,7 @@ class DispatcherManager:
if flow not in self.flows:
raise RuntimeError("Invalid flow")
if kind not in receive_dispatchers:
if kind not in import_dispatchers:
raise RuntimeError("Invalid kind")
key = (flow, kind)
@ -129,7 +142,38 @@ class DispatcherManager:
qconfig = intf_defs[kind + "-store"]
id = str(uuid.uuid4())
dispatcher = receive_dispatchers[kind](
dispatcher = import_dispatchers[kind](
pulsar_client = self.pulsar_client,
ws = ws,
running = running,
queue = qconfig,
)
return dispatcher
async def invoke_export(self, ws, running, params):
flow = params.get("flow")
kind = params.get("kind")
if flow not in self.flows:
raise RuntimeError("Invalid flow")
if kind not in export_dispatchers:
raise RuntimeError("Invalid kind")
key = (flow, kind)
intf_defs = self.flows[flow]["interfaces"]
if kind not in intf_defs:
raise RuntimeError("This kind not supported by flow")
# FIXME: The -store bit, does it make sense?
qconfig = intf_defs[kind + "-store"]
id = str(uuid.uuid4())
dispatcher = export_dispatchers[kind](
pulsar_client = self.pulsar_client,
ws = ws,
running = running,

2
trustgraph-flow/trustgraph/gateway/dispatch/triples_stream.py → trustgraph-flow/trustgraph/gateway/dispatch/triples_export.py
View file

@ -8,7 +8,7 @@ from ... base import Subscriber
from . serialize import serialize_triples
class TriplesStream:
class TriplesExport:
def __init__(
self, ws, running, pulsar_client, queue, consumer, subscriber

54
trustgraph-flow/trustgraph/gateway/dispatch/triples_import.py Normal file
View file

@ -0,0 +1,54 @@
import asyncio
import uuid
from aiohttp import WSMsgType
from .. schema import Metadata
from .. schema import Triples
from .. base import Publisher
from . socket import SocketEndpoint
from . serialize import to_subgraph
class TriplesImport:
def __init__(
self, ws, running, pulsar_client, queue
):
self.ws = ws
self.running = running
self.publisher = Publisher(
pulsar_client, queue = queue, schema = Triples
)
async def destroy(self):
self.running.stop()
await self.ws.close()
await self.publisher.stop()
async def receive(self, msg):
data = msg.json()
elt = Triples(
metadata=Metadata(
id=data["metadata"]["id"],
metadata=to_subgraph(data["metadata"]["metadata"]),
user=data["metadata"]["user"],
collection=data["metadata"]["collection"],
),
triples=to_subgraph(data["triples"]),
)
await self.publisher.send(None, elt)
async def run(self):
while self.running.get():
await asyncio.sleep(0.5)
await self.ws.close()
self.ws = None

56
trustgraph-flow/trustgraph/gateway/dispatch/triples_load.py
View file

@ -1,56 +0,0 @@
import asyncio
import uuid
from aiohttp import WSMsgType
from .. schema import Metadata
from .. schema import Triples
from .. schema import triples_store_queue
from .. base import Publisher
from . socket import SocketEndpoint
from . serialize import to_subgraph
class TriplesLoadEndpoint(SocketEndpoint):
def __init__(self, pulsar_client, auth, path="/api/v1/load/triples"):
super(TriplesLoadEndpoint, self).__init__(
endpoint_path=path, auth=auth,
)
self.pulsar_client=pulsar_client
self.publisher = Publisher(
self.pulsar_client, triples_store_queue,
schema=Triples
)
async def start(self):
await self.publisher.start()
async def listener(self, ws, running):
async for msg in ws:
# On error, finish
if msg.type == WSMsgType.ERROR:
break
else:
data = msg.json()
elt = Triples(
metadata=Metadata(
id=data["metadata"]["id"],
metadata=to_subgraph(data["metadata"]["metadata"]),
user=data["metadata"]["user"],
collection=data["metadata"]["collection"],
),
triples=to_subgraph(data["triples"]),
)
await self.publisher.send(None, elt)
running.stop()

13
trustgraph-flow/trustgraph/gateway/endpoint/manager.py
View file

@ -41,14 +41,19 @@ class EndpointManager:
auth = auth,
),
VariableEndpoint(
endpoint_path = "/api/v1/flow/{flow}/{kind}",
endpoint_path = "/api/v1/flow/{flow}/service/{kind}",
auth = auth,
dispatcher = dispatcher_manager.dispatch_flow_service(),
dispatcher = dispatcher_manager.dispatch_service(),
),
SocketEndpoint(
endpoint_path = "/api/v1/flow/{flow}/receive/{kind}",
endpoint_path = "/api/v1/flow/{flow}/import/{kind}",
auth = auth,
dispatcher = dispatcher_manager.dispatch_flow_receive()
dispatcher = dispatcher_manager.dispatch_import()
),
SocketEndpoint(
endpoint_path = "/api/v1/flow/{flow}/export/{kind}",
auth = auth,
dispatcher = dispatcher_manager.dispatch_export()
),
]