Flow endpoints seem to be working

This commit is contained in:
Cyber MacGeddon 2025-05-01 12:06:16 +01:00
parent a1afbbac05
commit e473672297
10 changed files with 71 additions and 25 deletions

View file

@ -1,12 +1,11 @@
from .. schema import AgentRequest, AgentResponse
from ... schema import AgentRequest, AgentResponse
from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
class AgentRequestor(ServiceRequestor):
def __init__(
self, pulsar_client, request_queue, response_queue, timeout, auth,
self, pulsar_client, request_queue, response_queue, timeout,
consumer, subscriber,
):

View file

@ -1,12 +1,11 @@
from .. schema import DocumentRagQuery, DocumentRagResponse
from ... schema import DocumentRagQuery, DocumentRagResponse
from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
class DocumentRagRequestor(ServiceRequestor):
def __init__(
self, pulsar_client, request_queue, response_queue, timeout, auth,
self, pulsar_client, request_queue, response_queue, timeout,
consumer, subscriber,
):

View file

@ -1,12 +1,11 @@
from .. schema import EmbeddingsRequest, EmbeddingsResponse
from ... schema import EmbeddingsRequest, EmbeddingsResponse
from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
class EmbeddingsRequestor(ServiceRequestor):
def __init__(
self, pulsar_client, request_queue, response_queue, timeout, auth,
self, pulsar_client, request_queue, response_queue, timeout,
consumer, subscriber,
):

View file

@ -1,13 +1,12 @@
from .. schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse
from ... schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse
from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
from . serialize import serialize_value
class GraphEmbeddingsQueryRequestor(ServiceRequestor):
def __init__(
self, pulsar_client, request_queue, response_queue, timeout, auth,
self, pulsar_client, request_queue, response_queue, timeout,
consumer, subscriber,
):

View file

@ -1,12 +1,11 @@
from .. schema import GraphRagQuery, GraphRagResponse
from ... schema import GraphRagQuery, GraphRagResponse
from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
class GraphRagRequestor(ServiceRequestor):
def __init__(
self, pulsar_client, request_queue, response_queue, timeout, auth,
self, pulsar_client, request_queue, response_queue, timeout,
consumer, subscriber,
):

View file

@ -1,14 +1,13 @@
import json
from .. schema import PromptRequest, PromptResponse
from ... schema import PromptRequest, PromptResponse
from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
class PromptRequestor(ServiceRequestor):
def __init__(
self, pulsar_client, request_queue, response_queue, timeout, auth,
self, pulsar_client, request_queue, response_queue, timeout,
consumer, subscriber,
):

View file

@ -5,7 +5,7 @@ import asyncio
import uuid
import logging
from .. base import Publisher
from ... base import Publisher
logger = logging.getLogger("sender")
logger.setLevel(logging.INFO)

View file

@ -1,12 +1,11 @@
from .. schema import TextCompletionRequest, TextCompletionResponse
from ... schema import TextCompletionRequest, TextCompletionResponse
from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
class TextCompletionRequestor(ServiceRequestor):
def __init__(
self, pulsar_client, request_queue, response_queue, timeout, auth,
self, pulsar_client, request_queue, response_queue, timeout,
consumer, subscriber,
):

View file

@ -1,13 +1,12 @@
from .. schema import TriplesQueryRequest, TriplesQueryResponse, Triples
from ... schema import TriplesQueryRequest, TriplesQueryResponse, Triples
from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
from . serialize import to_value, serialize_subgraph
class TriplesQueryRequestor(ServiceRequestor):
def __init__(
self, pulsar_client, request_queue, response_queue, timeout, auth,
self, pulsar_client, request_queue, response_queue, timeout,
consumer, subscriber,
):

View file

@ -3,12 +3,23 @@ from . endpoint import ServiceEndpoint
from . flow_endpoint import FlowEndpoint
from .. dispatch.agent import AgentRequestor
from .. dispatch.text_completion import TextCompletionRequestor
from .. dispatch.prompt import PromptRequestor
from .. dispatch.graph_rag import GraphRagRequestor
from .. dispatch.document_rag import DocumentRagRequestor
from .. dispatch.triples_query import TriplesQueryRequestor
from .. dispatch.embeddings import EmbeddingsRequestor
from .. dispatch.graph_embeddings_query import GraphEmbeddingsQueryRequestor
from .. dispatch.prompt import PromptRequestor
class FlowEndpointManager:
def __init__(self, config_receiver, pulsar_client, auth, timeout=600):
self.config_receiver = config_receiver
self.pulsar_client = pulsar_client
self.timeout = timeout
self.services = {
}
@ -32,8 +43,51 @@ class FlowEndpointManager:
await ep.start()
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",
)
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]