Tryna get knowledge core to load

This commit is contained in:
Cyber MacGeddon 2025-05-02 20:44:17 +01:00
parent 03d59b62af
commit dddf47d161
5 changed files with 33 additions and 24 deletions

View file

@ -28,7 +28,7 @@ async def load_ge(running, queue, url):
async with aiohttp.ClientSession() as session: async with aiohttp.ClientSession() as session:
async with session.ws_connect(f"{url}load/graph-embeddings") as ws: async with session.ws_connect(url) as ws:
while running.get(): while running.get():
@ -73,7 +73,7 @@ async def load_triples(running, queue, url):
async with aiohttp.ClientSession() as session: async with aiohttp.ClientSession() as session:
async with session.ws_connect(f"{url}load/triples") as ws: async with session.ws_connect(url) as ws:
while running.get(): while running.get():
@ -200,6 +200,9 @@ async def run(running, **args):
ge_q = asyncio.Queue(maxsize=10) ge_q = asyncio.Queue(maxsize=10)
t_q = asyncio.Queue(maxsize=10) t_q = asyncio.Queue(maxsize=10)
flow_id = args["flow"]
url = args["url"]
load_task = asyncio.create_task( load_task = asyncio.create_task(
loader( loader(
running=running, running=running,
@ -212,15 +215,17 @@ async def run(running, **args):
ge_task = asyncio.create_task( ge_task = asyncio.create_task(
load_ge( load_ge(
running=running, running = running,
queue=ge_q, url=args["url"] + "api/v1/" queue = ge_q,
url = f"{url}api/v1/{flow_id}/import/graph-embeddings"
) )
) )
triples_task = asyncio.create_task( triples_task = asyncio.create_task(
load_triples( load_triples(
running=running, running = running,
queue=t_q, url=args["url"] + "api/v1/" queue = t_q,
url = f"{url}api/v1/{flow_id}/import/triples"
) )
) )
@ -258,6 +263,13 @@ async def main(running):
help=f'Output file' help=f'Output file'
) )
parser.add_argument(
'-f', '--flow-id',
# Make it mandatory, difficult to over-write an existing file
default="0000",
help=f'Flow ID'
)
parser.add_argument( parser.add_argument(
'--format', '--format',
default="msgpack", default="msgpack",

View file

@ -3,11 +3,10 @@ import asyncio
import uuid import uuid
from aiohttp import WSMsgType from aiohttp import WSMsgType
from .. schema import Metadata from ... schema import Metadata
from .. schema import DocumentEmbeddings, ChunkEmbeddings from ... schema import DocumentEmbeddings, ChunkEmbeddings
from .. base import Publisher from ... base import Publisher
from . socket import SocketEndpoint
from . serialize import to_subgraph from . serialize import to_subgraph
class DocumentEmbeddingsImport: class DocumentEmbeddingsImport:
@ -20,7 +19,7 @@ class DocumentEmbeddingsImport:
self.running = running self.running = running
self.publisher = Publisher( self.publisher = Publisher(
pulsar_client, queue = queue, schema = DocumentEmbeddings pulsar_client, topic = queue, schema = DocumentEmbeddings
) )
async def destroy(self): async def destroy(self):

View file

@ -3,11 +3,10 @@ import asyncio
import uuid import uuid
from aiohttp import WSMsgType from aiohttp import WSMsgType
from .. schema import Metadata from ... schema import Metadata
from .. schema import GraphEmbeddings, EntityEmbeddings from ... schema import GraphEmbeddings, EntityEmbeddings
from .. base import Publisher from ... base import Publisher
from . socket import SocketEndpoint
from . serialize import to_subgraph, to_value from . serialize import to_subgraph, to_value
class GraphEmbeddingsImport: class GraphEmbeddingsImport:
@ -20,7 +19,7 @@ class GraphEmbeddingsImport:
self.running = running self.running = running
self.publisher = Publisher( self.publisher = Publisher(
pulsar_client, queue = queue, schema = GraphEmbeddings pulsar_client, topic = queue, schema = GraphEmbeddings
) )
async def destroy(self): async def destroy(self):

View file

@ -44,9 +44,9 @@ sender_dispatchers = {
} }
export_dispatchers = { export_dispatchers = {
"triples": TriplesStream, "triples": TriplesExport,
"graph-embeddings": GraphEmbeddingsStream, "graph-embeddings": GraphEmbeddingsExport,
"document-embeddings": DocumentEmbeddingsStream, "document-embeddings": DocumentEmbeddingsExport,
} }
import_dispatchers = { import_dispatchers = {

View file

@ -3,11 +3,10 @@ import asyncio
import uuid import uuid
from aiohttp import WSMsgType from aiohttp import WSMsgType
from .. schema import Metadata from ... schema import Metadata
from .. schema import Triples from ... schema import Triples
from .. base import Publisher from ... base import Publisher
from . socket import SocketEndpoint
from . serialize import to_subgraph from . serialize import to_subgraph
class TriplesImport: class TriplesImport:
@ -20,7 +19,7 @@ class TriplesImport:
self.running = running self.running = running
self.publisher = Publisher( self.publisher = Publisher(
pulsar_client, queue = queue, schema = Triples pulsar_client, topic = queue, schema = Triples
) )
async def destroy(self): async def destroy(self):