More dispatch code translated

This commit is contained in:
Cyber MacGeddon 2025-06-20 16:17:05 +01:00
parent 8bbf1cc6c0
commit 5ba7199895
14 changed files with 548 additions and 176 deletions

View file

@ -25,6 +25,14 @@ class MessageTranslator(Translator):
return self.from_pulsar(obj), True
class SendTranslator(Translator):
"""For fire-and-forget send operations (like ServiceSender)"""
def from_pulsar(self, obj: Record) -> Dict[str, Any]:
"""Usually not needed for send-only operations"""
raise NotImplementedError("Send translators typically don't need from_pulsar")
def handle_optional_fields(obj: Record, fields: list) -> Dict[str, Any]:
"""Helper to extract optional fields from Pulsar object"""
result = {}

View file

@ -0,0 +1,100 @@
from typing import Dict, Any, Tuple
from ...schema import ConfigRequest, ConfigResponse, ConfigKey, ConfigValue
from .base import MessageTranslator
class ConfigRequestTranslator(MessageTranslator):
"""Translator for ConfigRequest schema objects"""
def to_pulsar(self, data: Dict[str, Any]) -> ConfigRequest:
keys = None
if "keys" in data:
keys = [
ConfigKey(
type=k["type"],
key=k["key"]
)
for k in data["keys"]
]
values = None
if "values" in data:
values = [
ConfigValue(
type=v["type"],
key=v["key"],
value=v["value"]
)
for v in data["values"]
]
return ConfigRequest(
operation=data.get("operation"),
keys=keys,
type=data.get("type"),
values=values
)
def from_pulsar(self, obj: ConfigRequest) -> Dict[str, Any]:
result = {}
if obj.operation:
result["operation"] = obj.operation
if obj.type:
result["type"] = obj.type
if obj.keys:
result["keys"] = [
{
"type": k.type,
"key": k.key
}
for k in obj.keys
]
if obj.values:
result["values"] = [
{
"type": v.type,
"key": v.key,
"value": v.value
}
for v in obj.values
]
return result
class ConfigResponseTranslator(MessageTranslator):
"""Translator for ConfigResponse schema objects"""
def to_pulsar(self, data: Dict[str, Any]) -> ConfigResponse:
raise NotImplementedError("Response translation to Pulsar not typically needed")
def from_pulsar(self, obj: ConfigResponse) -> Dict[str, Any]:
result = {}
if obj.version is not None:
result["version"] = obj.version
if obj.values:
result["values"] = [
{
"type": v.type,
"key": v.key,
"value": v.value
}
for v in obj.values
]
if obj.directory:
result["directory"] = obj.directory
if obj.config:
result["config"] = obj.config
return result
def from_response_with_completion(self, obj: ConfigResponse) -> Tuple[Dict[str, Any], bool]:
"""Returns (response_dict, is_final)"""
return self.from_pulsar(obj), True

View file

@ -0,0 +1,191 @@
import base64
from typing import Dict, Any
from ...schema import Document, TextDocument, Chunk, DocumentEmbeddings, ChunkEmbeddings
from .base import SendTranslator
from .metadata import DocumentMetadataTranslator
from .primitives import SubgraphTranslator
class DocumentTranslator(SendTranslator):
"""Translator for Document schema objects (PDF docs etc.)"""
def __init__(self):
self.subgraph_translator = SubgraphTranslator()
def to_pulsar(self, data: Dict[str, Any]) -> Document:
metadata = data.get("metadata", [])
# Handle base64 content validation
doc = base64.b64decode(data["data"])
from ...schema import Metadata
return Document(
metadata=Metadata(
id=data.get("id"),
metadata=self.subgraph_translator.to_pulsar(metadata) if metadata else [],
user=data.get("user", "trustgraph"),
collection=data.get("collection", "default"),
),
data=base64.b64encode(doc).decode("utf-8")
)
def from_pulsar(self, obj: Document) -> Dict[str, Any]:
result = {
"data": obj.data
}
if obj.metadata:
metadata_dict = {}
if obj.metadata.id:
metadata_dict["id"] = obj.metadata.id
if obj.metadata.user:
metadata_dict["user"] = obj.metadata.user
if obj.metadata.collection:
metadata_dict["collection"] = obj.metadata.collection
if obj.metadata.metadata:
metadata_dict["metadata"] = self.subgraph_translator.from_pulsar(obj.metadata.metadata)
result["metadata"] = metadata_dict
return result
class TextDocumentTranslator(SendTranslator):
"""Translator for TextDocument schema objects"""
def __init__(self):
self.subgraph_translator = SubgraphTranslator()
def to_pulsar(self, data: Dict[str, Any]) -> TextDocument:
metadata = data.get("metadata", [])
charset = data.get("charset", "utf-8")
# Text is base64 encoded in input
text = base64.b64decode(data["text"]).decode(charset)
from ...schema import Metadata
return TextDocument(
metadata=Metadata(
id=data.get("id"),
metadata=self.subgraph_translator.to_pulsar(metadata) if metadata else [],
user=data.get("user", "trustgraph"),
collection=data.get("collection", "default"),
),
text=text.encode("utf-8")
)
def from_pulsar(self, obj: TextDocument) -> Dict[str, Any]:
result = {
"text": obj.text.decode("utf-8") if isinstance(obj.text, bytes) else obj.text
}
if obj.metadata:
metadata_dict = {}
if obj.metadata.id:
metadata_dict["id"] = obj.metadata.id
if obj.metadata.user:
metadata_dict["user"] = obj.metadata.user
if obj.metadata.collection:
metadata_dict["collection"] = obj.metadata.collection
if obj.metadata.metadata:
metadata_dict["metadata"] = self.subgraph_translator.from_pulsar(obj.metadata.metadata)
result["metadata"] = metadata_dict
return result
class ChunkTranslator(SendTranslator):
"""Translator for Chunk schema objects"""
def __init__(self):
self.subgraph_translator = SubgraphTranslator()
def to_pulsar(self, data: Dict[str, Any]) -> Chunk:
metadata = data.get("metadata", [])
from ...schema import Metadata
return Chunk(
metadata=Metadata(
id=data.get("id"),
metadata=self.subgraph_translator.to_pulsar(metadata) if metadata else [],
user=data.get("user", "trustgraph"),
collection=data.get("collection", "default"),
),
chunk=data["chunk"].encode("utf-8") if isinstance(data["chunk"], str) else data["chunk"]
)
def from_pulsar(self, obj: Chunk) -> Dict[str, Any]:
result = {
"chunk": obj.chunk.decode("utf-8") if isinstance(obj.chunk, bytes) else obj.chunk
}
if obj.metadata:
metadata_dict = {}
if obj.metadata.id:
metadata_dict["id"] = obj.metadata.id
if obj.metadata.user:
metadata_dict["user"] = obj.metadata.user
if obj.metadata.collection:
metadata_dict["collection"] = obj.metadata.collection
if obj.metadata.metadata:
metadata_dict["metadata"] = self.subgraph_translator.from_pulsar(obj.metadata.metadata)
result["metadata"] = metadata_dict
return result
class DocumentEmbeddingsTranslator(SendTranslator):
"""Translator for DocumentEmbeddings schema objects"""
def __init__(self):
self.subgraph_translator = SubgraphTranslator()
def to_pulsar(self, data: Dict[str, Any]) -> DocumentEmbeddings:
metadata = data.get("metadata", {})
chunks = [
ChunkEmbeddings(
chunk=chunk["chunk"].encode("utf-8") if isinstance(chunk["chunk"], str) else chunk["chunk"],
vectors=chunk["vectors"]
)
for chunk in data.get("chunks", [])
]
from ...schema import Metadata
return DocumentEmbeddings(
metadata=Metadata(
id=metadata.get("id"),
metadata=self.subgraph_translator.to_pulsar(metadata.get("metadata", [])),
user=metadata.get("user", "trustgraph"),
collection=metadata.get("collection", "default"),
),
chunks=chunks
)
def from_pulsar(self, obj: DocumentEmbeddings) -> Dict[str, Any]:
result = {
"chunks": [
{
"chunk": chunk.chunk.decode("utf-8") if isinstance(chunk.chunk, bytes) else chunk.chunk,
"vectors": chunk.vectors
}
for chunk in obj.chunks
]
}
if obj.metadata:
metadata_dict = {}
if obj.metadata.id:
metadata_dict["id"] = obj.metadata.id
if obj.metadata.user:
metadata_dict["user"] = obj.metadata.user
if obj.metadata.collection:
metadata_dict["collection"] = obj.metadata.collection
if obj.metadata.metadata:
metadata_dict["metadata"] = self.subgraph_translator.from_pulsar(obj.metadata.metadata)
result["metadata"] = metadata_dict
return result

View file

@ -0,0 +1,94 @@
from typing import Dict, Any, Tuple
from ...schema import (
DocumentEmbeddingsRequest, DocumentEmbeddingsResponse,
GraphEmbeddingsRequest, GraphEmbeddingsResponse
)
from .base import MessageTranslator
from .primitives import ValueTranslator
class DocumentEmbeddingsRequestTranslator(MessageTranslator):
"""Translator for DocumentEmbeddingsRequest schema objects"""
def to_pulsar(self, data: Dict[str, Any]) -> DocumentEmbeddingsRequest:
return DocumentEmbeddingsRequest(
vectors=data["vectors"],
limit=int(data.get("limit", 10)),
user=data.get("user", "trustgraph"),
collection=data.get("collection", "default")
)
def from_pulsar(self, obj: DocumentEmbeddingsRequest) -> Dict[str, Any]:
return {
"vectors": obj.vectors,
"limit": obj.limit,
"user": obj.user,
"collection": obj.collection
}
class DocumentEmbeddingsResponseTranslator(MessageTranslator):
"""Translator for DocumentEmbeddingsResponse schema objects"""
def to_pulsar(self, data: Dict[str, Any]) -> DocumentEmbeddingsResponse:
raise NotImplementedError("Response translation to Pulsar not typically needed")
def from_pulsar(self, obj: DocumentEmbeddingsResponse) -> Dict[str, Any]:
result = {}
if obj.documents:
result["documents"] = [
doc.decode("utf-8") if isinstance(doc, bytes) else doc
for doc in obj.documents
]
return result
def from_response_with_completion(self, obj: DocumentEmbeddingsResponse) -> Tuple[Dict[str, Any], bool]:
"""Returns (response_dict, is_final)"""
return self.from_pulsar(obj), True
class GraphEmbeddingsRequestTranslator(MessageTranslator):
"""Translator for GraphEmbeddingsRequest schema objects"""
def to_pulsar(self, data: Dict[str, Any]) -> GraphEmbeddingsRequest:
return GraphEmbeddingsRequest(
vectors=data["vectors"],
limit=int(data.get("limit", 10)),
user=data.get("user", "trustgraph"),
collection=data.get("collection", "default")
)
def from_pulsar(self, obj: GraphEmbeddingsRequest) -> Dict[str, Any]:
return {
"vectors": obj.vectors,
"limit": obj.limit,
"user": obj.user,
"collection": obj.collection
}
class GraphEmbeddingsResponseTranslator(MessageTranslator):
"""Translator for GraphEmbeddingsResponse schema objects"""
def __init__(self):
self.value_translator = ValueTranslator()
def to_pulsar(self, data: Dict[str, Any]) -> GraphEmbeddingsResponse:
raise NotImplementedError("Response translation to Pulsar not typically needed")
def from_pulsar(self, obj: GraphEmbeddingsResponse) -> Dict[str, Any]:
result = {}
if obj.entities:
result["entities"] = [
self.value_translator.from_pulsar(entity)
for entity in obj.entities
]
return result
def from_response_with_completion(self, obj: GraphEmbeddingsResponse) -> Tuple[Dict[str, Any], bool]:
"""Returns (response_dict, is_final)"""
return self.from_pulsar(obj), True

View file

@ -0,0 +1,59 @@
from typing import Dict, Any, Tuple
from ...schema import FlowRequest, FlowResponse
from .base import MessageTranslator
class FlowRequestTranslator(MessageTranslator):
"""Translator for FlowRequest schema objects"""
def to_pulsar(self, data: Dict[str, Any]) -> FlowRequest:
return FlowRequest(
operation=data.get("operation"),
class_name=data.get("class-name"),
class_definition=data.get("class-definition"),
description=data.get("description"),
flow_id=data.get("flow-id")
)
def from_pulsar(self, obj: FlowRequest) -> Dict[str, Any]:
result = {}
if obj.operation:
result["operation"] = obj.operation
if obj.class_name:
result["class-name"] = obj.class_name
if obj.class_definition:
result["class-definition"] = obj.class_definition
if obj.description:
result["description"] = obj.description
if obj.flow_id:
result["flow-id"] = obj.flow_id
return result
class FlowResponseTranslator(MessageTranslator):
"""Translator for FlowResponse schema objects"""
def to_pulsar(self, data: Dict[str, Any]) -> FlowResponse:
raise NotImplementedError("Response translation to Pulsar not typically needed")
def from_pulsar(self, obj: FlowResponse) -> Dict[str, Any]:
result = {}
if obj.class_names:
result["class-names"] = obj.class_names
if obj.flow_ids:
result["flow-ids"] = obj.flow_ids
if obj.class_definition:
result["class-definition"] = obj.class_definition
if obj.flow:
result["flow"] = obj.flow
if obj.description:
result["description"] = obj.description
return result
def from_response_with_completion(self, obj: FlowResponse) -> Tuple[Dict[str, Any], bool]:
"""Returns (response_dict, is_final)"""
return self.from_pulsar(obj), True

View file

@ -0,0 +1,54 @@
import json
from typing import Dict, Any, Tuple
from ...schema import PromptRequest, PromptResponse
from .base import MessageTranslator
class PromptRequestTranslator(MessageTranslator):
"""Translator for PromptRequest schema objects"""
def to_pulsar(self, data: Dict[str, Any]) -> PromptRequest:
# Handle both "terms" and "variables" input keys
terms = data.get("terms", {})
if "variables" in data:
# Convert variables to JSON strings as expected by the service
terms = {
k: json.dumps(v)
for k, v in data["variables"].items()
}
return PromptRequest(
id=data.get("id"),
terms=terms
)
def from_pulsar(self, obj: PromptRequest) -> Dict[str, Any]:
result = {}
if obj.id:
result["id"] = obj.id
if obj.terms:
result["terms"] = obj.terms
return result
class PromptResponseTranslator(MessageTranslator):
"""Translator for PromptResponse schema objects"""
def to_pulsar(self, data: Dict[str, Any]) -> PromptResponse:
raise NotImplementedError("Response translation to Pulsar not typically needed")
def from_pulsar(self, obj: PromptResponse) -> Dict[str, Any]:
result = {}
if obj.text:
result["text"] = obj.text
if obj.object:
result["object"] = obj.object
return result
def from_response_with_completion(self, obj: PromptResponse) -> Tuple[Dict[str, Any], bool]:
"""Returns (response_dict, is_final)"""
return self.from_pulsar(obj), True

View file

@ -2,6 +2,7 @@
from ... schema import ConfigRequest, ConfigResponse, ConfigKey, ConfigValue
from ... schema import config_request_queue
from ... schema import config_response_queue
from .... base.messaging import TranslatorRegistry
from . requestor import ServiceRequestor
@ -19,60 +20,12 @@ class ConfigRequestor(ServiceRequestor):
timeout=timeout,
)
self.request_translator = TranslatorRegistry.get_request_translator("config")
self.response_translator = TranslatorRegistry.get_response_translator("config")
def to_request(self, body):
if "keys" in body:
keys = [
ConfigKey(
type = k["type"],
key = k["key"],
)
for k in body["keys"]
]
else:
keys = None
if "values" in body:
values = [
ConfigValue(
type = v["type"],
key = v["key"],
value = v["value"],
)
for v in body["values"]
]
else:
values = None
return ConfigRequest(
operation = body.get("operation", None),
keys = keys,
type = body.get("type", None),
values = values
)
return self.request_translator.to_pulsar(body)
def from_response(self, message):
response = { }
if message.version is not None:
response["version"] = message.version
if message.values is not None:
response["values"] = [
{
"type": v.type,
"key": v.key,
"value": v.value,
}
for v in message.values
]
if message.directory is not None:
response["directory"] = message.directory
if message.config is not None:
response["config"] = message.config
return response, True
return self.response_translator.from_response_with_completion(message)

View file

@ -6,8 +6,7 @@ from aiohttp import WSMsgType
from ... schema import Metadata
from ... schema import DocumentEmbeddings, ChunkEmbeddings
from ... base import Publisher
from . serialize import to_subgraph
from .... base.messaging.translators.document_loading import DocumentEmbeddingsTranslator
class DocumentEmbeddingsImport:
@ -17,6 +16,7 @@ class DocumentEmbeddingsImport:
self.ws = ws
self.running = running
self.translator = DocumentEmbeddingsTranslator()
self.publisher = Publisher(
pulsar_client, topic = queue, schema = DocumentEmbeddings
@ -36,23 +36,7 @@ class DocumentEmbeddingsImport:
async def receive(self, msg):
data = msg.json()
elt = DocumentEmbeddings(
metadata=Metadata(
id=data["metadata"]["id"],
metadata=to_subgraph(data["metadata"]["metadata"]),
user=data["metadata"]["user"],
collection=data["metadata"]["collection"],
),
chunks=[
ChunkEmbeddings(
chunk=de["chunk"].encode("utf-8"),
vectors=de["vectors"],
)
for de in data["chunks"]
],
)
elt = self.translator.to_pulsar(data)
await self.publisher.send(None, elt)
async def run(self):

View file

@ -2,9 +2,9 @@
import base64
from ... schema import Document, Metadata
from .... base.messaging import TranslatorRegistry
from . sender import ServiceSender
from . serialize import to_subgraph
class DocumentLoad(ServiceSender):
def __init__(self, pulsar_client, queue):
@ -15,26 +15,9 @@ class DocumentLoad(ServiceSender):
schema = Document,
)
self.translator = TranslatorRegistry.get_request_translator("document")
def to_request(self, body):
if "metadata" in body:
metadata = to_subgraph(body["metadata"])
else:
metadata = []
# Doing a base64 decoe/encode here to make sure the
# content is valid base64
doc = base64.b64decode(body["data"])
print("Document received")
return Document(
metadata=Metadata(
id=body.get("id"),
metadata=metadata,
user=body.get("user", "trustgraph"),
collection=body.get("collection", "default"),
),
data=base64.b64encode(doc).decode("utf-8")
)
return self.translator.to_pulsar(body)

View file

@ -2,6 +2,7 @@
from ... schema import FlowRequest, FlowResponse
from ... schema import flow_request_queue
from ... schema import flow_response_queue
from .... base.messaging import TranslatorRegistry
from . requestor import ServiceRequestor
@ -19,34 +20,12 @@ class FlowRequestor(ServiceRequestor):
timeout=timeout,
)
def to_request(self, body):
self.request_translator = TranslatorRegistry.get_request_translator("flow")
self.response_translator = TranslatorRegistry.get_response_translator("flow")
return FlowRequest(
operation = body.get("operation", None),
class_name = body.get("class-name", None),
class_definition = body.get("class-definition", None),
description = body.get("description", None),
flow_id = body.get("flow-id", None),
)
def to_request(self, body):
return self.request_translator.to_pulsar(body)
def from_response(self, message):
response = { }
if message.class_names is not None:
response["class-names"] = message.class_names
if message.flow_ids is not None:
response["flow-ids"] = message.flow_ids
if message.class_definition is not None:
response["class-definition"] = message.class_definition
if message.flow is not None:
response["flow"] = message.flow
if message.description is not None:
response["description"] = message.description
return response, True
return self.response_translator.from_response_with_completion(message)

View file

@ -1,8 +1,8 @@
from ... schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse
from .... base.messaging import TranslatorRegistry
from . requestor import ServiceRequestor
from . serialize import serialize_value
class GraphEmbeddingsQueryRequestor(ServiceRequestor):
def __init__(
@ -21,22 +21,12 @@ class GraphEmbeddingsQueryRequestor(ServiceRequestor):
timeout=timeout,
)
self.request_translator = TranslatorRegistry.get_request_translator("graph-embeddings-query")
self.response_translator = TranslatorRegistry.get_response_translator("graph-embeddings-query")
def to_request(self, body):
limit = int(body.get("limit", 20))
return GraphEmbeddingsRequest(
vectors = body["vectors"],
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 {
"entities": [
serialize_value(ent) for ent in message.entities
]
}, True
return self.response_translator.from_response_with_completion(message)

View file

@ -2,6 +2,7 @@
import json
from ... schema import PromptRequest, PromptResponse
from .... base.messaging import TranslatorRegistry
from . requestor import ServiceRequestor
@ -22,22 +23,12 @@ class PromptRequestor(ServiceRequestor):
timeout=timeout,
)
self.request_translator = TranslatorRegistry.get_request_translator("prompt")
self.response_translator = TranslatorRegistry.get_response_translator("prompt")
def to_request(self, body):
return PromptRequest(
id=body["id"],
terms={
k: json.dumps(v)
for k, v in body["variables"].items()
}
)
return self.request_translator.to_pulsar(body)
def from_response(self, message):
if message.object:
return {
"object": message.object
}, True
else:
return {
"text": message.text
}, True
return self.response_translator.from_response_with_completion(message)

View file

@ -3,6 +3,13 @@ import base64
from ... schema import Value, Triple, DocumentMetadata, ProcessingMetadata
# DEPRECATED: These functions have been moved to trustgraph.base.messaging.translators
# Use the new messaging translation system instead for consistency and reusability.
# Examples:
# from trustgraph.base.messaging.translators.primitives import ValueTranslator
# value_translator = ValueTranslator()
# pulsar_value = value_translator.to_pulsar({"v": "example", "e": True})
def to_value(x):
return Value(value=x["v"], is_uri=x["e"])

View file

@ -2,9 +2,9 @@
import base64
from ... schema import TextDocument, Metadata
from .... base.messaging import TranslatorRegistry
from . sender import ServiceSender
from . serialize import to_subgraph
class TextLoad(ServiceSender):
def __init__(self, pulsar_client, queue):
@ -15,30 +15,9 @@ class TextLoad(ServiceSender):
schema = TextDocument,
)
self.translator = TranslatorRegistry.get_request_translator("text-document")
def to_request(self, body):
if "metadata" in body:
metadata = to_subgraph(body["metadata"])
else:
metadata = []
if "charset" in body:
charset = body["charset"]
else:
charset = "utf-8"
# Text is base64 encoded
text = base64.b64decode(body["text"]).decode(charset)
print("Text document received")
return TextDocument(
metadata=Metadata(
id=body.get("id"),
metadata=metadata,
user=body.get("user", "trustgraph"),
collection=body.get("collection", "default"),
),
text=text,
)
return self.translator.to_pulsar(body)