apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-22 11:41:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Refactoring dispatch to breakout messaging

Browse source
This commit is contained in:
Cyber MacGeddon 2025-06-20 16:07:23 +01:00
parent 3fa004d628
commit 8bbf1cc6c0
21 changed files with 958 additions and 201 deletions
Download patch file Download diff file Expand all files Collapse all files

23
trustgraph-flow/trustgraph/gateway/dispatch/agent.py
View file

@ -1,5 +1,6 @@
from ... schema import AgentRequest, AgentResponse
from .... base.messaging import TranslatorRegistry
from . requestor import ServiceRequestor
@ -20,24 +21,12 @@ class AgentRequestor(ServiceRequestor):
timeout=timeout,
)
self.request_translator = TranslatorRegistry.get_request_translator("agent")
self.response_translator = TranslatorRegistry.get_response_translator("agent")
def to_request(self, body):
return AgentRequest(
question=body["question"]
)
return self.request_translator.to_pulsar(body)
def from_response(self, message):
resp = {
}
if message.answer:
resp["answer"] = message.answer
if message.thought:
resp["thought"] = message.thought
if message.observation:
resp["observation"] = message.observation
# The 2nd boolean expression indicates whether we're done responding
return resp, (message.answer is not None)
return self.response_translator.from_response_with_completion(message)

13
trustgraph-flow/trustgraph/gateway/dispatch/document_rag.py
View file

@ -1,5 +1,6 @@
from ... schema import DocumentRagQuery, DocumentRagResponse
from .... base.messaging import TranslatorRegistry
from . requestor import ServiceRequestor
@ -20,14 +21,12 @@ class DocumentRagRequestor(ServiceRequestor):
timeout=timeout,
)
self.request_translator = TranslatorRegistry.get_request_translator("document-rag")
self.response_translator = TranslatorRegistry.get_response_translator("document-rag")
def to_request(self, body):
return DocumentRagQuery(
query=body["query"],
user=body.get("user", "trustgraph"),
collection=body.get("collection", "default"),
doc_limit=int(body.get("doc-limit", 20)),
)
return self.request_translator.to_pulsar(body)
def from_response(self, message):
return { "response": message.response }, True
return self.response_translator.from_response_with_completion(message)

10
trustgraph-flow/trustgraph/gateway/dispatch/embeddings.py
View file

@ -1,5 +1,6 @@
from ... schema import EmbeddingsRequest, EmbeddingsResponse
from .... base.messaging import TranslatorRegistry
from . requestor import ServiceRequestor
@ -20,11 +21,12 @@ class EmbeddingsRequestor(ServiceRequestor):
timeout=timeout,
)
self.request_translator = TranslatorRegistry.get_request_translator("embeddings")
self.response_translator = TranslatorRegistry.get_response_translator("embeddings")
def to_request(self, body):
return EmbeddingsRequest(
text=body["text"]
)
return self.request_translator.to_pulsar(body)
def from_response(self, message):
return { "vectors": message.vectors }, True
return self.response_translator.from_response_with_completion(message)

16
trustgraph-flow/trustgraph/gateway/dispatch/graph_rag.py
View file

@ -1,5 +1,6 @@
from ... schema import GraphRagQuery, GraphRagResponse
from .... base.messaging import TranslatorRegistry
from . requestor import ServiceRequestor
@ -20,17 +21,12 @@ class GraphRagRequestor(ServiceRequestor):
timeout=timeout,
)
self.request_translator = TranslatorRegistry.get_request_translator("graph-rag")
self.response_translator = TranslatorRegistry.get_response_translator("graph-rag")
def to_request(self, body):
return GraphRagQuery(
query=body["query"],
user=body.get("user", "trustgraph"),
collection=body.get("collection", "default"),
entity_limit=int(body.get("entity-limit", 50)),
triple_limit=int(body.get("triple-limit", 30)),
max_subgraph_size=int(body.get("max-subgraph-size", 1000)),
max_path_length=int(body.get("max-path-length", 2)),
)
return self.request_translator.to_pulsar(body)
def from_response(self, message):
return { "response": message.response }, True
return self.response_translator.from_response_with_completion(message)

75
trustgraph-flow/trustgraph/gateway/dispatch/knowledge.py
View file

@ -5,11 +5,9 @@ from ... schema import KnowledgeRequest, KnowledgeResponse, Triples
from ... schema import GraphEmbeddings, Metadata, EntityEmbeddings
from ... schema import knowledge_request_queue
from ... schema import knowledge_response_queue
from .... base.messaging import TranslatorRegistry
from . requestor import ServiceRequestor
from . serialize import serialize_graph_embeddings
from . serialize import serialize_triples, to_subgraph, to_value
from . serialize import to_document_metadata, to_processing_metadata
class KnowledgeRequestor(ServiceRequestor):
def __init__(self, pulsar_client, consumer, subscriber, timeout=120):
@ -25,73 +23,12 @@ class KnowledgeRequestor(ServiceRequestor):
timeout=timeout,
)
self.request_translator = TranslatorRegistry.get_request_translator("knowledge")
self.response_translator = TranslatorRegistry.get_response_translator("knowledge")
def to_request(self, body):
if "triples" in body:
triples = Triples(
metadata=Metadata(
id = body["triples"]["metadata"]["id"],
metadata = to_subgraph(body["triples"]["metadata"]["metadata"]),
user = body["triples"]["metadata"]["user"],
),
triples = to_subgraph(body["triples"]["triples"]),
)
else:
triples = None
if "graph-embeddings" in body:
ge = GraphEmbeddings(
metadata = Metadata(
id = body["graph-embeddings"]["metadata"]["id"],
metadata = to_subgraph(body["graph-embeddings"]["metadata"]["metadata"]),
user = body["graph-embeddings"]["metadata"]["user"],
),
entities=[
EntityEmbeddings(
entity = to_value(ent["entity"]),
vectors = ent["vectors"],
)
for ent in body["graph-embeddings"]["entities"]
]
)
else:
ge = None
return KnowledgeRequest(
operation = body.get("operation", None),
user = body.get("user", None),
id = body.get("id", None),
flow = body.get("flow", None),
collection = body.get("collection", None),
triples = triples,
graph_embeddings = ge,
)
return self.request_translator.to_pulsar(body)
def from_response(self, message):
# Response to list,
if message.ids is not None:
return {
"ids": message.ids
}, True
if message.triples:
return {
"triples": serialize_triples(message.triples)
}, False
if message.graph_embeddings:
return {
"graph-embeddings": serialize_graph_embeddings(
message.graph_embeddings
)
}, False
if message.eos is True:
return {
"eos": True
}, True
# Empty case, return from successful delete.
return {}, True
return self.response_translator.from_response_with_completion(message)

72
trustgraph-flow/trustgraph/gateway/dispatch/librarian.py
View file

@ -4,12 +4,9 @@ import base64
from ... schema import LibrarianRequest, LibrarianResponse
from ... schema import librarian_request_queue
from ... schema import librarian_response_queue
from .... base.messaging import TranslatorRegistry
from . requestor import ServiceRequestor
from . serialize import serialize_document_metadata
from . serialize import serialize_processing_metadata
from . serialize import to_document_metadata, to_processing_metadata
from . serialize import to_criteria
class LibrarianRequestor(ServiceRequestor):
def __init__(self, pulsar_client, consumer, subscriber, timeout=120):
@ -25,67 +22,20 @@ class LibrarianRequestor(ServiceRequestor):
timeout=timeout,
)
self.request_translator = TranslatorRegistry.get_request_translator("librarian")
self.response_translator = TranslatorRegistry.get_response_translator("librarian")
def to_request(self, body):
# Content gets base64 decoded & encoded again. It at least makes
# sure payload is valid base64.
if "document-metadata" in body:
dm = to_document_metadata(body["document-metadata"])
else:
dm = None
if "processing-metadata" in body:
pm = to_processing_metadata(body["processing-metadata"])
else:
pm = None
if "criteria" in body:
criteria = to_criteria(body["criteria"])
else:
criteria = None
# Handle base64 content processing
if "content" in body:
# Content gets base64 decoded & encoded again to ensure valid base64
content = base64.b64decode(body["content"].encode("utf-8"))
content = base64.b64encode(content).decode("utf-8")
else:
content = None
return LibrarianRequest(
operation = body.get("operation", None),
document_id = body.get("document-id", None),
processing_id = body.get("processing-id", None),
document_metadata = dm,
processing_metadata = pm,
content = content,
user = body.get("user", None),
collection = body.get("collection", None),
criteria = criteria,
)
body = body.copy()
body["content"] = content
return self.request_translator.to_pulsar(body)
def from_response(self, message):
response = {}
if message.document_metadata:
response["document-metadata"] = serialize_document_metadata(
message.document_metadata
)
if message.content:
response["content"] = message.content.decode("utf-8")
if message.document_metadatas != None:
response["document-metadatas"] = [
serialize_document_metadata(v)
for v in message.document_metadatas
]
if message.processing_metadatas != None:
response["processing-metadatas"] = [
serialize_processing_metadata(v)
for v in message.processing_metadatas
]
return response, True
return self.response_translator.from_response_with_completion(message)

11
trustgraph-flow/trustgraph/gateway/dispatch/text_completion.py
View file

@ -1,5 +1,6 @@
from ... schema import TextCompletionRequest, TextCompletionResponse
from .... base.messaging import TranslatorRegistry
from . requestor import ServiceRequestor
@ -20,12 +21,12 @@ class TextCompletionRequestor(ServiceRequestor):
timeout=timeout,
)
self.request_translator = TranslatorRegistry.get_request_translator("text-completion")
self.response_translator = TranslatorRegistry.get_response_translator("text-completion")
def to_request(self, body):
return TextCompletionRequest(
system=body["system"],
prompt=body["prompt"]
)
return self.request_translator.to_pulsar(body)
def from_response(self, message):
return { "response": message.response }, True
return self.response_translator.from_response_with_completion(message)

34
trustgraph-flow/trustgraph/gateway/dispatch/triples_query.py
View file

@ -1,8 +1,8 @@
from ... schema import TriplesQueryRequest, TriplesQueryResponse, Triples
from .... base.messaging import TranslatorRegistry
from . requestor import ServiceRequestor
from . serialize import to_value, serialize_subgraph
class TriplesQueryRequestor(ServiceRequestor):
def __init__(
@ -21,34 +21,12 @@ class TriplesQueryRequestor(ServiceRequestor):
timeout=timeout,
)
self.request_translator = TranslatorRegistry.get_request_translator("triples-query")
self.response_translator = TranslatorRegistry.get_response_translator("triples-query")
def to_request(self, body):
if "s" in body:
s = to_value(body["s"])
else:
s = None
if "p" in body:
p = to_value(body["p"])
else:
p = None
if "o" in body:
o = to_value(body["o"])
else:
o = None
limit = int(body.get("limit", 10000))
return TriplesQueryRequest(
s = s, p = p, o = o,
limit = limit,
user = body.get("user", "trustgraph"),
collection = body.get("collection", "default"),
)
return self.request_translator.to_pulsar(body)
def from_response(self, message):
return {
"response": serialize_subgraph(message.triples)
}, True
return self.response_translator.from_response_with_completion(message)