mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-20 18:51:03 +02:00
All streamers working
This commit is contained in:
Cyber MacGeddon
parent
873f7a84e4
commit
e606753f51
6 changed files with 140 additions and 148 deletions
|
|
@ -0,0 +1,67 @@
|
|||
|
||||
import asyncio
|
||||
import queue
|
||||
import uuid
|
||||
|
||||
from ... schema import DocumentEmbeddings
|
||||
from ... base import Subscriber
|
||||
|
||||
from . serialize import serialize_document_embeddings
|
||||
|
||||
class DocumentEmbeddingsStream:
|
||||
|
||||
def __init__(
|
||||
self, ws, running, pulsar_client, queue, consumer, subscriber
|
||||
):
|
||||
|
||||
self.ws = ws
|
||||
self.running = running
|
||||
self.pulsar_client = pulsar_client
|
||||
self.queue = queue
|
||||
self.consumer = consumer
|
||||
self.subscriber = subscriber
|
||||
|
||||
async def destroy(self):
|
||||
self.running.stop()
|
||||
await self.ws.close()
|
||||
|
||||
async def receive(self, msg):
|
||||
# Ignore incoming info from websocket
|
||||
pass
|
||||
|
||||
async def run(self):
|
||||
|
||||
subs = Subscriber(
|
||||
client = self.pulsar_client, topic = self.queue,
|
||||
consumer_name = self.consumer, subscription = self.subscriber,
|
||||
schema = DocumentEmbeddings
|
||||
)
|
||||
|
||||
await subs.start()
|
||||
|
||||
id = str(uuid.uuid4())
|
||||
q = await subs.subscribe_all(id)
|
||||
|
||||
while self.running.get():
|
||||
try:
|
||||
|
||||
resp = await asyncio.wait_for(q.get(), timeout=0.5)
|
||||
await self.ws.send_json(serialize_document_embeddings(resp))
|
||||
|
||||
except TimeoutError:
|
||||
continue
|
||||
|
||||
except queue.Empty:
|
||||
continue
|
||||
|
||||
except Exception as e:
|
||||
print(f"Exception: {str(e)}", flush=True)
|
||||
break
|
||||
|
||||
await subs.unsubscribe_all(id)
|
||||
|
||||
await subs.stop()
|
||||
|
||||
await self.ws.close()
|
||||
self.running.stop()
|
||||
|
||||
|
|
@ -0,0 +1,67 @@
|
|||
|
||||
import asyncio
|
||||
import queue
|
||||
import uuid
|
||||
|
||||
from ... schema import GraphEmbeddings
|
||||
from ... base import Subscriber
|
||||
|
||||
from . serialize import serialize_graph_embeddings
|
||||
|
||||
class GraphEmbeddingsStream:
|
||||
|
||||
def __init__(
|
||||
self, ws, running, pulsar_client, queue, consumer, subscriber
|
||||
):
|
||||
|
||||
self.ws = ws
|
||||
self.running = running
|
||||
self.pulsar_client = pulsar_client
|
||||
self.queue = queue
|
||||
self.consumer = consumer
|
||||
self.subscriber = subscriber
|
||||
|
||||
async def destroy(self):
|
||||
self.running.stop()
|
||||
await self.ws.close()
|
||||
|
||||
async def receive(self, msg):
|
||||
# Ignore incoming info from websocket
|
||||
pass
|
||||
|
||||
async def run(self):
|
||||
|
||||
subs = Subscriber(
|
||||
client = self.pulsar_client, topic = self.queue,
|
||||
consumer_name = self.consumer, subscription = self.subscriber,
|
||||
schema = GraphEmbeddings
|
||||
)
|
||||
|
||||
await subs.start()
|
||||
|
||||
id = str(uuid.uuid4())
|
||||
q = await subs.subscribe_all(id)
|
||||
|
||||
while self.running.get():
|
||||
try:
|
||||
|
||||
resp = await asyncio.wait_for(q.get(), timeout=0.5)
|
||||
await self.ws.send_json(serialize_graph_embeddings(resp))
|
||||
|
||||
except TimeoutError:
|
||||
continue
|
||||
|
||||
except queue.Empty:
|
||||
continue
|
||||
|
||||
except Exception as e:
|
||||
print(f"Exception: {str(e)}", flush=True)
|
||||
break
|
||||
|
||||
await subs.unsubscribe_all(id)
|
||||
|
||||
await subs.stop()
|
||||
|
||||
await self.ws.close()
|
||||
self.running.stop()
|
||||
|
||||
|
|
@ -12,6 +12,8 @@ from . embeddings import EmbeddingsRequestor
|
|||
from . graph_embeddings_query import GraphEmbeddingsQueryRequestor
|
||||
from . prompt import PromptRequestor
|
||||
from . triples_stream import TriplesStream
|
||||
from . graph_embeddings_stream import GraphEmbeddingsStream
|
||||
from . document_embeddings_stream import DocumentEmbeddingsStream
|
||||
|
||||
request_response_dispatchers = {
|
||||
"agent": AgentRequestor,
|
||||
|
|
@ -26,6 +28,8 @@ request_response_dispatchers = {
|
|||
|
||||
receive_dispatchers = {
|
||||
"triples-store": TriplesStream,
|
||||
"graph-embeddings-store": GraphEmbeddingsStream,
|
||||
"document-embeddings-store": DocumentEmbeddingsStream,
|
||||
}
|
||||
|
||||
class TestDispatcher:
|
||||
|
|
@ -123,7 +127,6 @@ class DispatcherManager:
|
|||
|
||||
async def dispatch_receive(self, ws, running, params):
|
||||
|
||||
print("HERE")
|
||||
flow = params.get("flow")
|
||||
kind = params.get("kind")
|
||||
|
||||
|
|
|
|||
|
|
@ -26,14 +26,11 @@ class TriplesStream:
|
|||
await self.ws.close()
|
||||
|
||||
async def receive(self, msg):
|
||||
print(msg.data)
|
||||
# Ignore incoming info from websocket
|
||||
pass
|
||||
|
||||
async def run(self):
|
||||
|
||||
print("c", self.consumer)
|
||||
print("s", self.subscriber)
|
||||
print("q", self.queue)
|
||||
|
||||
subs = Subscriber(
|
||||
client = self.pulsar_client, topic = self.queue,
|
||||
consumer_name = self.consumer, subscription = self.subscriber,
|
||||
|
|
|
|||
|
|
@ -1,73 +0,0 @@
|
|||
|
||||
import asyncio
|
||||
import queue
|
||||
import uuid
|
||||
|
||||
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
|
||||
|
||||
class DocumentEmbeddingsStreamEndpoint(SocketEndpoint):
|
||||
|
||||
def __init__(
|
||||
self, pulsar_client, auth,
|
||||
queue,
|
||||
path="/api/v1/stream/document-embeddings",
|
||||
):
|
||||
|
||||
super(DocumentEmbeddingsStreamEndpoint, self).__init__(
|
||||
endpoint_path=path, auth=auth,
|
||||
)
|
||||
|
||||
self.pulsar_client=pulsar_client
|
||||
|
||||
self.subscriber = Subscriber(
|
||||
self.pulsar_client, queue,
|
||||
"api-gateway", "api-gateway",
|
||||
schema=DocumentEmbeddings,
|
||||
)
|
||||
|
||||
async def listener(self, ws, running):
|
||||
|
||||
worker = asyncio.create_task(
|
||||
self.async_thread(ws, running)
|
||||
)
|
||||
|
||||
await super(DocumentEmbeddingsStreamEndpoint, self).listener(
|
||||
ws, running
|
||||
)
|
||||
|
||||
await worker
|
||||
|
||||
async def start(self):
|
||||
|
||||
await self.subscriber.start()
|
||||
|
||||
async def async_thread(self, ws, running):
|
||||
|
||||
id = str(uuid.uuid4())
|
||||
|
||||
q = await self.subscriber.subscribe_all(id)
|
||||
|
||||
while running.get():
|
||||
try:
|
||||
resp = await asyncio.wait_for(q.get(), timeout=0.5)
|
||||
await ws.send_json(serialize_document_embeddings(resp))
|
||||
|
||||
except TimeoutError:
|
||||
continue
|
||||
|
||||
except queue.Empty:
|
||||
continue
|
||||
|
||||
except Exception as e:
|
||||
print(f"Exception: {str(e)}", flush=True)
|
||||
break
|
||||
|
||||
await self.subscriber.unsubscribe_all(id)
|
||||
|
||||
running.stop()
|
||||
|
||||
|
|
@ -1,69 +0,0 @@
|
|||
|
||||
import asyncio
|
||||
import queue
|
||||
import uuid
|
||||
|
||||
from .. schema import GraphEmbeddings
|
||||
from .. schema import graph_embeddings_store_queue
|
||||
from .. base import Subscriber
|
||||
|
||||
from . socket import SocketEndpoint
|
||||
from . serialize import serialize_graph_embeddings
|
||||
|
||||
class GraphEmbeddingsStreamEndpoint(SocketEndpoint):
|
||||
|
||||
def __init__(
|
||||
self, pulsar_client, auth, path="/api/v1/stream/graph-embeddings"
|
||||
):
|
||||
|
||||
super(GraphEmbeddingsStreamEndpoint, self).__init__(
|
||||
endpoint_path=path, auth=auth,
|
||||
)
|
||||
|
||||
self.pulsar_client=pulsar_client
|
||||
|
||||
self.subscriber = Subscriber(
|
||||
self.pulsar_client, graph_embeddings_store_queue,
|
||||
"api-gateway", "api-gateway",
|
||||
schema=GraphEmbeddings
|
||||
)
|
||||
|
||||
async def listener(self, ws, running):
|
||||
|
||||
worker = asyncio.create_task(
|
||||
self.async_thread(ws, running)
|
||||
)
|
||||
|
||||
await super(GraphEmbeddingsStreamEndpoint, self).listener(ws, running)
|
||||
|
||||
await worker
|
||||
|
||||
async def start(self):
|
||||
|
||||
await self.subscriber.start()
|
||||
|
||||
async def async_thread(self, ws, running):
|
||||
|
||||
id = str(uuid.uuid4())
|
||||
|
||||
q = await self.subscriber.subscribe_all(id)
|
||||
|
||||
while running.get():
|
||||
try:
|
||||
resp = await asyncio.wait_for(q.get, timeout=0.5)
|
||||
await ws.send_json(serialize_graph_embeddings(resp))
|
||||
|
||||
except TimeoutError:
|
||||
continue
|
||||
|
||||
except queue.Empty:
|
||||
continue
|
||||
|
||||
except Exception as e:
|
||||
print(f"Exception: {str(e)}", flush=True)
|
||||
break
|
||||
|
||||
await self.subscriber.unsubscribe_all(id)
|
||||
|
||||
running.stop()
|
||||
|
||||
Loading…
Add table
Add a link
Reference in a new issue