Embeddings round

This commit is contained in:
Cyber MacGeddon 2025-05-02 15:17:43 +01:00
parent e9ae0b27f3
commit 54bd509d86
3 changed files with 69 additions and 15 deletions

View file

@ -1,6 +1,8 @@
import asyncio
from . embeddings import EmbeddingsRequestor
class TestDispatcher:
def __init__(self, pulsar_client, timeout=120):
self.pulsar_client = pulsar_client
@ -66,11 +68,63 @@ class TestDispatcher3:
return Runner(ws, running)
class DispatcherManager:
def __init__(self, pulsar_client):
def __init__(self, pulsar_client, config_receiver):
self.pulsar_client = pulsar_client
self.config_receiver = config_receiver
self.config_receiver.add_handler(self)
self.flows = {}
self.dispatchers = {}
async def start_flow(self, id, flow):
print("Start flow", id)
self.flows[id] = flow
return
async def stop_flow(self, id, flow):
print("Stop flow", id)
del self.flows[id]
return
def dispatch_test_service(self):
return TestDispatcher(pulsar_client = self.pulsar_client)
def dispatch_flow_service(self):
return TestDispatcher2(pulsar_client = self.pulsar_client)
# return TestDispatcher2(pulsar_client = self.pulsar_client)
return self
def dispatch_socket_service(self):
return TestDispatcher3(pulsar_client = self.pulsar_client).dispatch
async def process(self, data, responder, params):
flow = params.get("flow")
kind = params.get("kind")
key = (flow, kind)
if flow not in self.flows:
raise RuntimeError("Invalid flow")
if key in self.dispatchers:
return await self.dispatchers[key].process(data, responder)
qconfig = self.flows[flow]["interfaces"]["embeddings"]
dispatcher = EmbeddingsRequestor(
pulsar_client = self.pulsar_client,
request_queue = qconfig["request"],
response_queue = qconfig["response"],
timeout = 120,
consumer = f"api-gateway-{flow}-{kind}-request",
subscriber = f"api-gateway-{flow}-{kind}-request",
)
await dispatcher.start()
self.dispatchers[key] = dispatcher
print("CREATE")
return await dispatcher.process(data, responder)

View file

@ -13,37 +13,32 @@ from .. dispatch.manager import DispatcherManager
class FlowEndpointManager:
def __init__(self, config_receiver, pulsar_client, auth, timeout=600):
def __init__(self, dispatcher_manager, auth, timeout=600):
self.config_receiver = config_receiver
self.pulsar_client = pulsar_client
self.dispatcher_manager = dispatcher_manager
self.timeout = timeout
self.services = {
}
dm = DispatcherManager(pulsar_client)
self.endpoints = [
ConstantEndpoint(
endpoint_path = "/api/v1/test",
auth = auth,
dispatcher = dm.dispatch_test_service(),
dispatcher = dispatcher_manager.dispatch_test_service(),
),
VariableEndpoint(
endpoint_path = "/api/v1/test/{thing}",
endpoint_path = "/api/v1/flow/{flow}/{kind}",
auth = auth,
dispatcher = dm.dispatch_flow_service(),
dispatcher = dispatcher_manager.dispatch_flow_service(),
),
SocketEndpoint(
endpoint_path = "/api/v1/test2",
auth = auth,
dispatcher = dm.dispatch_socket_service()
dispatcher = dispatcher_manager.dispatch_socket_service()
),
]
self.config_receiver.add_handler(self)
def add_routes(self, app):
for ep in self.endpoints:
ep.add_routes(app)

View file

@ -13,6 +13,7 @@ from .. log_level import LogLevel
from . auth import Authenticator
from . config.receiver import ConfigReceiver
from . dispatch.manager import DispatcherManager
from . endpoint.globals import GlobalEndpointManager
from . endpoint.flows import FlowEndpointManager
@ -69,6 +70,11 @@ class Api:
self.config_receiver = ConfigReceiver(self.pulsar_client)
self.dispatcher_manager = DispatcherManager(
pulsar_client = self.pulsar_client,
config_receiver = self.config_receiver,
)
self.global_manager = GlobalEndpointManager(
pulsar_client = self.pulsar_client,
auth = self.auth,
@ -77,8 +83,7 @@ class Api:
)
self.flow_manager = FlowEndpointManager(
config_receiver = self.config_receiver,
pulsar_client = self.pulsar_client,
dispatcher_manager = self.dispatcher_manager,
auth = self.auth,
timeout = self.timeout,