mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-20 18:51:03 +02:00
Refactor to simplify config etc.
This commit is contained in:
Cyber MacGeddon
parent
5f08220ef1
commit
67e5a968e6
8 changed files with 52 additions and 213 deletions
|
|
@ -45,6 +45,8 @@ class Subscriber:
|
|||
|
||||
try:
|
||||
|
||||
# FIXME: Create consumer in start method so we know
|
||||
# it is definitely running when start completes
|
||||
consumer = self.client.subscribe(
|
||||
topic = self.topic,
|
||||
subscription_name = self.subscription,
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from ... schema import config_response_queue
|
|||
from . requestor import ServiceRequestor
|
||||
|
||||
class ConfigRequestor(ServiceRequestor):
|
||||
def __init__(self, pulsar_client, timeout):
|
||||
def __init__(self, pulsar_client, timeout=120):
|
||||
|
||||
super(ConfigRequestor, self).__init__(
|
||||
pulsar_client=pulsar_client,
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from ... schema import flow_response_queue
|
|||
from . requestor import ServiceRequestor
|
||||
|
||||
class FlowRequestor(ServiceRequestor):
|
||||
def __init__(self, pulsar_client, timeout):
|
||||
def __init__(self, pulsar_client, timeout=120):
|
||||
|
||||
super(FlowRequestor, self).__init__(
|
||||
pulsar_client=pulsar_client,
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from . serialize import serialize_document_package, serialize_document_info
|
|||
from . serialize import to_document_package, to_document_info, to_criteria
|
||||
|
||||
class LibrarianRequestor(ServiceRequestor):
|
||||
def __init__(self, pulsar_client, timeout):
|
||||
def __init__(self, pulsar_client, timeout=120):
|
||||
|
||||
super(LibrarianRequestor, self).__init__(
|
||||
pulsar_client=pulsar_client,
|
||||
|
|
|
|||
|
|
@ -2,6 +2,10 @@
|
|||
import asyncio
|
||||
import uuid
|
||||
|
||||
from . config import ConfigRequestor
|
||||
from . flow import FlowRequestor
|
||||
from . librarian import LibrarianRequestor
|
||||
|
||||
from . embeddings import EmbeddingsRequestor
|
||||
from . agent import AgentRequestor
|
||||
from . text_completion import TextCompletionRequestor
|
||||
|
|
@ -12,6 +16,7 @@ from . triples_query import TriplesQueryRequestor
|
|||
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
|
||||
|
|
@ -33,69 +38,15 @@ receive_dispatchers = {
|
|||
"document-embeddings": DocumentEmbeddingsStream,
|
||||
}
|
||||
|
||||
class TestDispatcher:
|
||||
def __init__(self, pulsar_client, timeout=120):
|
||||
self.pulsar_client = pulsar_client
|
||||
timeout = timeout
|
||||
class DispatcherWrapper:
|
||||
def __init__(self, mgr, name, impl):
|
||||
self.mgr = mgr
|
||||
self.name = name
|
||||
self.impl = impl
|
||||
async def process(self, data, responder):
|
||||
result = { "result": "Hello world!" }
|
||||
|
||||
if responder:
|
||||
await responder(result, True)
|
||||
|
||||
return result
|
||||
|
||||
class TestDispatcher2:
|
||||
def __init__(self, pulsar_client, timeout=120):
|
||||
self.pulsar_client = pulsar_client
|
||||
timeout = timeout
|
||||
async def process(self, data, responder, params):
|
||||
|
||||
thing = params['thing']
|
||||
|
||||
result = { "result": "Hello world!!", "thing": thing }
|
||||
|
||||
if responder:
|
||||
await responder(result, True)
|
||||
|
||||
return result
|
||||
|
||||
class TestDispatcher3:
|
||||
def __init__(self, pulsar_client, timeout=120):
|
||||
self.pulsar_client = pulsar_client
|
||||
self.timeout = timeout
|
||||
|
||||
async def dispatch(self, ws, running, request):
|
||||
|
||||
class Runner:
|
||||
def __init__(self, ws, running):
|
||||
self.ws = ws
|
||||
self.running = running
|
||||
|
||||
async def destroy(self):
|
||||
|
||||
if self.ws:
|
||||
await self.ws.close()
|
||||
self.ws = None
|
||||
|
||||
self.running.stop()
|
||||
|
||||
async def run(self):
|
||||
|
||||
i = 0
|
||||
|
||||
while self.running.get():
|
||||
await self.ws.send_json({"i": i})
|
||||
i += 1
|
||||
await asyncio.sleep(1)
|
||||
|
||||
await self.ws.close()
|
||||
self.ws = None
|
||||
|
||||
async def receive(self, msg):
|
||||
print("Receive:", msg.data)
|
||||
|
||||
return Runner(ws, running)
|
||||
return await self.mgr.process_impl(
|
||||
data, responder, self.name, self.impl
|
||||
)
|
||||
|
||||
class DispatcherManager:
|
||||
|
||||
|
|
@ -117,8 +68,31 @@ class DispatcherManager:
|
|||
del self.flows[id]
|
||||
return
|
||||
|
||||
def dispatch_test_service(self):
|
||||
return TestDispatcher(pulsar_client = self.pulsar_client)
|
||||
def dispatch_config(self):
|
||||
return DispatcherWrapper(self, "config", ConfigRequestor)
|
||||
|
||||
def dispatch_flow(self):
|
||||
return DispatcherWrapper(self, "flow", FlowRequestor)
|
||||
|
||||
def dispatch_librarian(self):
|
||||
return DispatcherWrapper(self, "librarian", LibrarianRequestor)
|
||||
|
||||
async def process_impl(self, data, responder, name, impl):
|
||||
|
||||
key = (None, name)
|
||||
|
||||
if key in self.dispatchers:
|
||||
return await self.dispatchers[key].process(data, responder)
|
||||
|
||||
dispatcher = impl(
|
||||
pulsar_client = self.pulsar_client
|
||||
)
|
||||
|
||||
await dispatcher.start()
|
||||
|
||||
self.dispatchers[key] = dispatcher
|
||||
|
||||
return await dispatcher.process(data, responder)
|
||||
|
||||
def dispatch_flow_service(self):
|
||||
return self
|
||||
|
|
@ -152,7 +126,6 @@ class DispatcherManager:
|
|||
pulsar_client = self.pulsar_client,
|
||||
ws = ws,
|
||||
running = running,
|
||||
# FIXME!
|
||||
queue = qconfig,
|
||||
consumer = f"api-gateway-{id}",
|
||||
subscriber = f"api-gateway-{id}",
|
||||
|
|
|
|||
|
|
@ -36,9 +36,9 @@ class ServiceRequestor:
|
|||
self.running = True
|
||||
|
||||
async def start(self):
|
||||
await self.pub.start()
|
||||
await self.sub.start()
|
||||
self.running = True
|
||||
await self.sub.start()
|
||||
await self.pub.start()
|
||||
|
||||
async def stop(self):
|
||||
await self.pub.stop()
|
||||
|
|
@ -68,6 +68,7 @@ class ServiceRequestor:
|
|||
q.get(), timeout=self.timeout
|
||||
)
|
||||
except Exception as e:
|
||||
print("Exception", e)
|
||||
raise RuntimeError("Timeout")
|
||||
|
||||
if resp.error:
|
||||
|
|
|
|||
|
|
@ -1,128 +0,0 @@
|
|||
|
||||
import asyncio
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
#from . endpoint import ServiceEndpoint
|
||||
|
||||
from . constant_endpoint import ConstantEndpoint
|
||||
from . variable_endpoint import VariableEndpoint
|
||||
from . socket import SocketEndpoint
|
||||
|
||||
from .. dispatch.manager import DispatcherManager
|
||||
|
||||
class FlowEndpointManager:
|
||||
|
||||
def __init__(self, dispatcher_manager, auth, timeout=600):
|
||||
|
||||
self.dispatcher_manager = dispatcher_manager
|
||||
self.timeout = timeout
|
||||
|
||||
self.services = {
|
||||
}
|
||||
|
||||
self.endpoints = [
|
||||
ConstantEndpoint(
|
||||
endpoint_path = "/api/v1/test",
|
||||
auth = auth,
|
||||
dispatcher = dispatcher_manager.dispatch_test_service(),
|
||||
),
|
||||
VariableEndpoint(
|
||||
endpoint_path = "/api/v1/flow/{flow}/{kind}",
|
||||
auth = auth,
|
||||
dispatcher = dispatcher_manager.dispatch_flow_service(),
|
||||
),
|
||||
SocketEndpoint(
|
||||
endpoint_path = "/api/v1/flow/{flow}/receive/{kind}",
|
||||
auth = auth,
|
||||
dispatcher = dispatcher_manager.dispatch_flow_receive()
|
||||
),
|
||||
]
|
||||
|
||||
def add_routes(self, app):
|
||||
for ep in self.endpoints:
|
||||
ep.add_routes(app)
|
||||
|
||||
async def start(self):
|
||||
for ep in self.endpoints:
|
||||
await ep.start()
|
||||
|
||||
async def start_flow(self, id, flow):
|
||||
|
||||
print("START FLOW", id)
|
||||
|
||||
return
|
||||
|
||||
intf = flow["interfaces"]
|
||||
|
||||
kinds = {
|
||||
"agent": AgentRequestor,
|
||||
"text-completion": TextCompletionRequestor,
|
||||
"prompt": PromptRequestor,
|
||||
"graph-rag": GraphRagRequestor,
|
||||
"document-rag": DocumentRagRequestor,
|
||||
"embeddings": EmbeddingsRequestor,
|
||||
"graph-embeddings": GraphEmbeddingsQueryRequestor,
|
||||
"triples-query": TriplesQueryRequestor,
|
||||
}
|
||||
|
||||
for api_kind, requestor in kinds.items():
|
||||
|
||||
if api_kind in intf:
|
||||
|
||||
k = (id, api_kind)
|
||||
if k in self.services:
|
||||
await self.services[k].stop()
|
||||
del self.services[k]
|
||||
|
||||
|
||||
self.services[k] = requestor(
|
||||
pulsar_client=self.pulsar_client, timeout = self.timeout,
|
||||
request_queue = intf[api_kind]["request"],
|
||||
response_queue = intf[api_kind]["response"],
|
||||
consumer = f"api-gateway-{id}-{api_kind}-request",
|
||||
subscriber = f"api-gateway-{id}-{api_kind}-request",
|
||||
)
|
||||
await self.services[k].start()
|
||||
|
||||
kinds = {
|
||||
# "document-embeddings-stream": DocumentEmbeddingsStreamEndpoint,
|
||||
"triples-stream": TriplesStream,
|
||||
# "bunch":
|
||||
}
|
||||
|
||||
for api_kind, streamer in kinds.items():
|
||||
|
||||
# if api_kind in intf:
|
||||
if True:
|
||||
|
||||
k = (id, api_kind)
|
||||
if k in self.services:
|
||||
await self.services[k].stop()
|
||||
del self.services[k]
|
||||
|
||||
self.services[k] = Dispatcher(
|
||||
pulsar_client=self.pulsar_client,
|
||||
timeout = self.timeout,
|
||||
input_queue = intf[api_kind],
|
||||
consumer = f"api-gateway-{id}-{api_kind}-stream",
|
||||
subscriber = f"api-gateway-{id}-{api_kind}-stream",
|
||||
impl=streamer,
|
||||
)
|
||||
await self.services[k].start()
|
||||
|
||||
async def stop_flow(self, id, flow):
|
||||
|
||||
print("STOP FLOW", id)
|
||||
|
||||
return
|
||||
|
||||
svc_list = list(self.services.keys())
|
||||
|
||||
for k in svc_list:
|
||||
|
||||
kid, kkind = k
|
||||
|
||||
if id == kid:
|
||||
await self.services[k].stop()
|
||||
del self.services[k]
|
||||
|
|
@ -14,8 +14,8 @@ 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
|
||||
|
||||
from . endpoint.manager import EndpointManager
|
||||
|
||||
import pulsar
|
||||
from prometheus_client import start_http_server
|
||||
|
|
@ -75,16 +75,10 @@ class Api:
|
|||
config_receiver = self.config_receiver,
|
||||
)
|
||||
|
||||
self.global_manager = GlobalEndpointManager(
|
||||
pulsar_client = self.pulsar_client,
|
||||
auth = self.auth,
|
||||
prometheus_url = self.prometheus_url,
|
||||
timeout = self.timeout,
|
||||
)
|
||||
|
||||
self.flow_manager = FlowEndpointManager(
|
||||
self.endpoint_manager = EndpointManager(
|
||||
dispatcher_manager = self.dispatcher_manager,
|
||||
auth = self.auth,
|
||||
prometheus_url = self.prometheus_url,
|
||||
timeout = self.timeout,
|
||||
|
||||
)
|
||||
|
|
@ -108,11 +102,8 @@ class Api:
|
|||
for ep in self.endpoints:
|
||||
await ep.start()
|
||||
|
||||
self.global_manager.add_routes(self.app)
|
||||
await self.global_manager.start()
|
||||
|
||||
self.flow_manager.add_routes(self.app)
|
||||
await self.flow_manager.start()
|
||||
self.endpoint_manager.add_routes(self.app)
|
||||
await self.endpoint_manager.start()
|
||||
|
||||
return self.app
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue