mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 00:16:23 +02:00
a4e2f67cb1
Pull the JSON/Pulsar message translation into a separate module, will be useful for other comms channels
51 lines
1.1 KiB
Python
51 lines
1.1 KiB
Python
|
|
import asyncio
|
|
import uuid
|
|
from aiohttp import WSMsgType
|
|
|
|
from ... schema import Metadata
|
|
from ... schema import DocumentEmbeddings, ChunkEmbeddings
|
|
from ... base import Publisher
|
|
from ... messaging.translators.document_loading import DocumentEmbeddingsTranslator
|
|
|
|
class DocumentEmbeddingsImport:
|
|
|
|
def __init__(
|
|
self, ws, running, pulsar_client, queue
|
|
):
|
|
|
|
self.ws = ws
|
|
self.running = running
|
|
self.translator = DocumentEmbeddingsTranslator()
|
|
|
|
self.publisher = Publisher(
|
|
pulsar_client, topic = queue, schema = DocumentEmbeddings
|
|
)
|
|
|
|
async def start(self):
|
|
await self.publisher.start()
|
|
|
|
async def destroy(self):
|
|
self.running.stop()
|
|
|
|
if self.ws:
|
|
await self.ws.close()
|
|
|
|
await self.publisher.stop()
|
|
|
|
async def receive(self, msg):
|
|
|
|
data = msg.json()
|
|
elt = self.translator.to_pulsar(data)
|
|
await self.publisher.send(None, elt)
|
|
|
|
async def run(self):
|
|
|
|
while self.running.get():
|
|
await asyncio.sleep(0.5)
|
|
|
|
if self.ws:
|
|
await self.ws.close()
|
|
|
|
self.ws = None
|
|
|