Restructure

This commit is contained in:
Cyber MacGeddon 2025-05-01 11:35:19 +01:00
parent 2b8ac5b9bc
commit a1afbbac05
25 changed files with 262 additions and 336 deletions

View file

@ -0,0 +1,121 @@
"""
API gateway. Offers HTTP services which are translated to interaction on the
Pulsar bus.
"""
module = "api-gateway"
# FIXME: Subscribes to Pulsar unnecessarily, should only do it when there
# are active listeners
# FIXME: Connection errors in publishers / subscribers cause those threads
# to fail and are not failed or retried
import asyncio
import argparse
from aiohttp import web
import logging
import os
import base64
import uuid
import json
import pulsar
from prometheus_client import start_http_server
from ... schema import ConfigPush, config_push_queue
from ... base import Consumer
logger = logging.getLogger("config.receiver")
logger.setLevel(logging.INFO)
class ConfigReceiver:
def __init__(self, pulsar_client):
self.pulsar_client = pulsar_client
self.flow_handlers = []
self.flows = {}
def add_handler(self, h):
self.flow_handlers.append(h)
async def on_config(self, msg, proc, flow):
try:
v = msg.value()
print(f"Config version", v.version)
if "flows" in v.config:
flows = v.config["flows"]
wanted = list(flows.keys())
current = list(self.flows.keys())
for k in wanted:
if k not in current:
self.flows[k] = json.loads(flows[k])
await self.start_flow(k, self.flows[k])
for k in current:
if k not in wanted:
await self.stop_flow(k, self.flows[k])
del self.flows[k]
except Exception as e:
print(f"Exception: {e}", flush=True)
async def start_flow(self, id, flow):
print("Start flow", id)
for handler in self.flow_handlers:
try:
await handler.start_flow(id, flow)
except Exception as e:
print(f"Exception: {e}", flush=True)
async def stop_flow(self, id, flow):
print("Stop flow", id)
for handler in self.flow_handlers:
try:
await handler.stop_flow(id, flow)
except Exception as e:
print(f"Exception: {e}", flush=True)
async def config_loader(self):
async with asyncio.TaskGroup() as tg:
id = str(uuid.uuid4())
self.config_cons = Consumer(
taskgroup = tg,
flow = None,
client = self.pulsar_client,
subscriber = f"gateway-{id}",
topic = config_push_queue,
schema = ConfigPush,
handler = self.on_config,
start_of_messages = True,
)
await self.config_cons.start()
print("Waiting...")
print("Config consumer done. :/")
async def start(self):
asyncio.create_task(self.config_loader())

View file

@ -1,13 +1,12 @@
from .. schema import ConfigRequest, ConfigResponse, ConfigKey, ConfigValue
from .. schema import config_request_queue
from .. schema import config_response_queue
from ... schema import ConfigRequest, ConfigResponse, ConfigKey, ConfigValue
from ... schema import config_request_queue
from ... schema import config_response_queue
from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
class ConfigRequestor(ServiceRequestor):
def __init__(self, pulsar_client, timeout, auth):
def __init__(self, pulsar_client, timeout):
super(ConfigRequestor, self).__init__(
pulsar_client=pulsar_client,

View file

@ -1,13 +1,12 @@
from .. schema import FlowRequest, FlowResponse, ConfigKey, ConfigValue
from .. schema import flow_request_queue
from .. schema import flow_response_queue
from ... schema import FlowRequest, FlowResponse
from ... schema import flow_request_queue
from ... schema import flow_response_queue
from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
class FlowRequestor(ServiceRequestor):
def __init__(self, pulsar_client, timeout, auth):
def __init__(self, pulsar_client, timeout):
super(FlowRequestor, self).__init__(
pulsar_client=pulsar_client,

View file

@ -1,15 +1,14 @@
from .. schema import LibrarianRequest, LibrarianResponse, Triples
from .. schema import librarian_request_queue
from .. schema import librarian_response_queue
from ... schema import LibrarianRequest, LibrarianResponse
from ... schema import librarian_request_queue
from ... schema import librarian_response_queue
from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
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, auth):
def __init__(self, pulsar_client, timeout):
super(LibrarianRequestor, self).__init__(
pulsar_client=pulsar_client,

View file

@ -3,8 +3,8 @@ import asyncio
import uuid
import logging
from .. base import Publisher
from .. base import Subscriber
from ... base import Publisher
from ... base import Subscriber
logger = logging.getLogger("requestor")
logger.setLevel(logging.INFO)

View file

@ -1,7 +1,7 @@
import base64
from .. schema import Value, Triple, DocumentPackage, DocumentInfo
from ... schema import Value, Triple, DocumentPackage, DocumentInfo
def to_value(x):
return Value(value=x["v"], is_uri=x["e"])

View file

@ -0,0 +1,39 @@
from . endpoint import ServiceEndpoint
from . flow_endpoint import FlowEndpoint
class FlowEndpointManager:
def __init__(self, config_receiver, pulsar_client, auth, timeout=600):
self.config_receiver = config_receiver
self.pulsar_client = pulsar_client
self.services = {
}
self.endpoints = [
FlowEndpoint(
endpoint_path = "/api/v1/flow/{flow}/{kind}",
auth = auth,
requestors = self.services,
),
]
self.config_receiver.add_handler(self)
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)
async def stop_flow(self, id, flow):
print("STOP FLOW", id)

View file

@ -0,0 +1,56 @@
from . endpoint import ServiceEndpoint
from .. dispatch.librarian import LibrarianRequestor
from .. dispatch.config import ConfigRequestor
from .. dispatch.flow import FlowRequestor
from . metrics import MetricsEndpoint
class GlobalEndpointManager:
def __init__(self, pulsar_client, auth, prometheus_url, timeout=600):
self.pulsar_client = pulsar_client
self.services = {
"librarian": LibrarianRequestor(
pulsar_client=self.pulsar_client, timeout=timeout,
),
"config": ConfigRequestor(
pulsar_client=self.pulsar_client, timeout=timeout,
),
"flow": FlowRequestor(
pulsar_client=self.pulsar_client, timeout=timeout,
)
}
self.endpoints = [
ServiceEndpoint(
endpoint_path = "/api/v1/librarian", auth = auth,
requestor = self.services["librarian"],
),
ServiceEndpoint(
endpoint_path = "/api/v1/config", auth = auth,
requestor = self.services["config"],
),
ServiceEndpoint(
endpoint_path = "/api/v1/flow", auth = auth,
requestor = self.services["flow"],
),
MetricsEndpoint(
endpoint_path = "/api/v1/metrics",
prometheus_url = prometheus_url,
auth = auth,
),
]
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()

View file

@ -3,63 +3,21 @@ API gateway. Offers HTTP services which are translated to interaction on the
Pulsar bus.
"""
module = "api-gateway"
# FIXME: Subscribes to Pulsar unnecessarily, should only do it when there
# are active listeners
# FIXME: Connection errors in publishers / subscribers cause those threads
# to fail and are not failed or retried
import asyncio
import argparse
from aiohttp import web
import logging
import os
import base64
import uuid
import json
import pulsar
from prometheus_client import start_http_server
from .. log_level import LogLevel
from . serialize import to_subgraph
from . running import Running
from .. schema import ConfigPush, config_push_queue
from . text_completion import TextCompletionRequestor
from . prompt import PromptRequestor
from . graph_rag import GraphRagRequestor
from . document_rag import DocumentRagRequestor
from . triples_query import TriplesQueryRequestor
from . graph_embeddings_query import GraphEmbeddingsQueryRequestor
from . embeddings import EmbeddingsRequestor
#from . encyclopedia import EncyclopediaRequestor
from . agent import AgentRequestor
#from . dbpedia import DbpediaRequestor
#from . internet_search import InternetSearchRequestor
from . librarian import LibrarianRequestor
from . config import ConfigRequestor
from . flow import FlowRequestor
#from . triples_stream import TriplesStreamEndpoint
#from . graph_embeddings_stream import GraphEmbeddingsStreamEndpoint
#from . document_embeddings_stream import DocumentEmbeddingsStreamEndpoint
#from . triples_load import TriplesLoadEndpoint
#from . graph_embeddings_load import GraphEmbeddingsLoadEndpoint
#from . document_embeddings_load import DocumentEmbeddingsLoadEndpoint
from . mux import MuxEndpoint
#from . document_load import DocumentLoadSender
#from . text_load import TextLoadSender
from . metrics import MetricsEndpoint
from . endpoint import ServiceEndpoint
from . flow_endpoint import FlowEndpoint
from . auth import Authenticator
from .. base import Subscriber
from .. base import Consumer
from . config.receiver import ConfigReceiver
from . endpoint.globals import GlobalEndpointManager
from . endpoint.flows import FlowEndpointManager
import pulsar
from prometheus_client import start_http_server
logger = logging.getLogger("api")
logger.setLevel(logging.INFO)
@ -81,6 +39,7 @@ class Api:
self.pulsar_api_key = config.get(
"pulsar_api_key", default_pulsar_api_key
)
self.pulsar_listener = config.get("pulsar_listener", None)
if self.pulsar_api_key:
@ -108,279 +67,26 @@ class Api:
else:
self.auth = Authenticator(allow_all=True)
self.services = {
# "text-completion": TextCompletionRequestor(
# pulsar_client=self.pulsar_client, timeout=self.timeout,
# auth = self.auth,
# ),
# "prompt": PromptRequestor(
# pulsar_client=self.pulsar_client, timeout=self.timeout,
# auth = self.auth,
# ),
# "graph-rag": GraphRagRequestor(
# pulsar_client=self.pulsar_client, timeout=self.timeout,
# auth = self.auth,
# ),
# "document-rag": DocumentRagRequestor(
# pulsar_client=self.pulsar_client, timeout=self.timeout,
# auth = self.auth,
# ),
# "triples-query": TriplesQueryRequestor(
# pulsar_client=self.pulsar_client, timeout=self.timeout,
# auth = self.auth,
# ),
# "graph-embeddings-query": GraphEmbeddingsQueryRequestor(
# pulsar_client=self.pulsar_client, timeout=self.timeout,
# auth = self.auth,
# ),
# "embeddings": EmbeddingsRequestor(
# pulsar_client=self.pulsar_client, timeout=self.timeout,
# auth = self.auth,
# ),
# "agent": AgentRequestor(
# pulsar_client=self.pulsar_client, timeout=self.timeout,
# auth = self.auth,
# ),
(None, "librarian"): LibrarianRequestor(
pulsar_client=self.pulsar_client, timeout=self.timeout,
auth = self.auth,
),
(None, "config"): ConfigRequestor(
pulsar_client=self.pulsar_client, timeout=self.timeout,
auth = self.auth,
),
(None, "flow"): FlowRequestor(
pulsar_client=self.pulsar_client, timeout=self.timeout,
auth = self.auth,
),
# "encyclopedia": EncyclopediaRequestor(
# pulsar_client=self.pulsar_client, timeout=self.timeout,
# auth = self.auth,
# ),
# "dbpedia": DbpediaRequestor(
# pulsar_client=self.pulsar_client, timeout=self.timeout,
# auth = self.auth,
# ),
# "internet-search": InternetSearchRequestor(
# pulsar_client=self.pulsar_client, timeout=self.timeout,
# auth = self.auth,
# ),
# "document-load": DocumentLoadSender(
# pulsar_client=self.pulsar_client,
# ),
# "text-load": TextLoadSender(
# pulsar_client=self.pulsar_client,
# ),
}
self.config_receiver = ConfigReceiver(self.pulsar_client)
self.global_manager = GlobalEndpointManager(
pulsar_client = self.pulsar_client,
auth = self.auth,
prometheus_url = self.prometheus_url,
timeout = self.timeout,
)
self.flow_manager = FlowEndpointManager(
config_receiver = self.config_receiver,
pulsar_client = self.pulsar_client,
auth = self.auth,
timeout = self.timeout,
)
self.endpoints = [
# ServiceEndpoint(
# endpoint_path = "/api/v1/text-completion", auth=self.auth,
# requestor = self.services["text-completion"],
# ),
# ServiceEndpoint(
# endpoint_path = "/api/v1/prompt", auth=self.auth,
# requestor = self.services["prompt"],
# ),
# ServiceEndpoint(
# endpoint_path = "/api/v1/graph-rag", auth=self.auth,
# requestor = self.services["graph-rag"],
# ),
# ServiceEndpoint(
# endpoint_path = "/api/v1/document-rag", auth=self.auth,
# requestor = self.services["document-rag"],
# ),
# ServiceEndpoint(
# endpoint_path = "/api/v1/triples-query", auth=self.auth,
# requestor = self.services["triples-query"],
# ),
# ServiceEndpoint(
# endpoint_path = "/api/v1/graph-embeddings-query",
# auth=self.auth,
# requestor = self.services["graph-embeddings-query"],
# ),
# ServiceEndpoint(
# endpoint_path = "/api/v1/embeddings", auth=self.auth,
# requestor = self.services["embeddings"],
# ),
# ServiceEndpoint(
# endpoint_path = "/api/v1/agent", auth=self.auth,
# requestor = self.services["agent"],
# ),
ServiceEndpoint(
endpoint_path = "/api/v1/librarian", auth=self.auth,
requestor = self.services[(None, "librarian")],
),
ServiceEndpoint(
endpoint_path = "/api/v1/config", auth=self.auth,
requestor = self.services[(None, "config")],
),
ServiceEndpoint(
endpoint_path = "/api/v1/flow", auth=self.auth,
requestor = self.services[(None, "flow")],
),
FlowEndpoint(
endpoint_path = "/api/v1/flow/{flow}/{kind}",
auth=self.auth,
requestors = self.services,
),
# ServiceEndpoint(
# endpoint_path = "/api/v1/encyclopedia", auth=self.auth,
# requestor = self.services["encyclopedia"],
# ),
# ServiceEndpoint(
# endpoint_path = "/api/v1/dbpedia", auth=self.auth,
# requestor = self.services["dbpedia"],
# ),
# ServiceEndpoint(
# endpoint_path = "/api/v1/internet-search", auth=self.auth,
# requestor = self.services["internet-search"],
# ),
# ServiceEndpoint(
# endpoint_path = "/api/v1/load/document", auth=self.auth,
# requestor = self.services["document-load"],
# ),
# ServiceEndpoint(
# endpoint_path = "/api/v1/load/text", auth=self.auth,
# requestor = self.services["text-load"],
# ),
# TriplesStreamEndpoint(
# pulsar_client=self.pulsar_client,
# auth = self.auth,
# ),
# GraphEmbeddingsStreamEndpoint(
# pulsar_client=self.pulsar_client,
# auth = self.auth,
# ),
# DocumentEmbeddingsStreamEndpoint(
# pulsar_client=self.pulsar_client,
# auth = self.auth,
# ),
# TriplesLoadEndpoint(
# pulsar_client=self.pulsar_client,
# auth = self.auth,
# ),
# GraphEmbeddingsLoadEndpoint(
# pulsar_client=self.pulsar_client,
# auth = self.auth,
# ),
# DocumentEmbeddingsLoadEndpoint(
# pulsar_client=self.pulsar_client,
# auth = self.auth,
# ),
# MuxEndpoint(
# pulsar_client=self.pulsar_client,
# auth = self.auth,
# services = self.services,
# ),
MetricsEndpoint(
endpoint_path = "/api/v1/metrics",
prometheus_url = self.prometheus_url,
auth = self.auth,
),
]
self.flows = {}
async def on_config(self, msg, proc, flow):
try:
v = msg.value()
print(f"Config version", v.version)
if "flows" in v.config:
flows = v.config["flows"]
wanted = list(flows.keys())
current = list(self.flows.keys())
for k in wanted:
if k not in current:
self.flows[k] = json.loads(flows[k])
await self.start_flow(k, self.flows[k])
for k in current:
if k not in wanted:
await self.stop_flow(k, self.flows[k])
del self.flows[k]
except Exception as e:
print(f"Exception: {e}", flush=True)
async def start_flow(self, id, flow):
print("Start flow", id)
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",
auth = self.auth,
)
await self.services[k].start()
async def stop_flow(self, id, flow):
print("Stop flow", id)
intf = flow["interfaces"]
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]
async def config_loader(self):
async with asyncio.TaskGroup() as tg:
id = str(uuid.uuid4())
self.config_cons = Consumer(
taskgroup = tg,
flow = None,
client = self.pulsar_client,
subscriber = f"gateway-{id}",
topic = config_push_queue,
schema = ConfigPush,
handler = self.on_config,
start_of_messages = True,
)
await self.config_cons.start()
print("Waiting...")
print("Config consumer done. :/")
async def app_factory(self):
self.app = web.Application(
@ -388,7 +94,8 @@ class Api:
client_max_size=256 * 1024 * 1024
)
asyncio.create_task(self.config_loader())
await self.config_receiver.start()
for ep in self.endpoints:
ep.add_routes(self.app)
@ -396,6 +103,12 @@ 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()
return self.app
def run(self):