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 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]: def handle_optional_fields(obj: Record, fields: list) -> Dict[str, Any]:
"""Helper to extract optional fields from Pulsar object""" """Helper to extract optional fields from Pulsar object"""
result = {} 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 ConfigRequest, ConfigResponse, ConfigKey, ConfigValue
from ... schema import config_request_queue from ... schema import config_request_queue
from ... schema import config_response_queue from ... schema import config_response_queue
from .... base.messaging import TranslatorRegistry
from . requestor import ServiceRequestor from . requestor import ServiceRequestor
@ -19,60 +20,12 @@ class ConfigRequestor(ServiceRequestor):
timeout=timeout, timeout=timeout,
) )
self.request_translator = TranslatorRegistry.get_request_translator("config")
self.response_translator = TranslatorRegistry.get_response_translator("config")
def to_request(self, body): def to_request(self, body):
return self.request_translator.to_pulsar(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
)
def from_response(self, message): def from_response(self, message):
return self.response_translator.from_response_with_completion(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

View file

@ -6,8 +6,7 @@ from aiohttp import WSMsgType
from ... schema import Metadata from ... schema import Metadata
from ... schema import DocumentEmbeddings, ChunkEmbeddings from ... schema import DocumentEmbeddings, ChunkEmbeddings
from ... base import Publisher from ... base import Publisher
from .... base.messaging.translators.document_loading import DocumentEmbeddingsTranslator
from . serialize import to_subgraph
class DocumentEmbeddingsImport: class DocumentEmbeddingsImport:
@ -17,6 +16,7 @@ class DocumentEmbeddingsImport:
self.ws = ws self.ws = ws
self.running = running self.running = running
self.translator = DocumentEmbeddingsTranslator()
self.publisher = Publisher( self.publisher = Publisher(
pulsar_client, topic = queue, schema = DocumentEmbeddings pulsar_client, topic = queue, schema = DocumentEmbeddings
@ -36,23 +36,7 @@ class DocumentEmbeddingsImport:
async def receive(self, msg): async def receive(self, msg):
data = msg.json() data = msg.json()
elt = self.translator.to_pulsar(data)
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"]
],
)
await self.publisher.send(None, elt) await self.publisher.send(None, elt)
async def run(self): async def run(self):

View file

@ -2,9 +2,9 @@
import base64 import base64
from ... schema import Document, Metadata from ... schema import Document, Metadata
from .... base.messaging import TranslatorRegistry
from . sender import ServiceSender from . sender import ServiceSender
from . serialize import to_subgraph
class DocumentLoad(ServiceSender): class DocumentLoad(ServiceSender):
def __init__(self, pulsar_client, queue): def __init__(self, pulsar_client, queue):
@ -15,26 +15,9 @@ class DocumentLoad(ServiceSender):
schema = Document, schema = Document,
) )
self.translator = TranslatorRegistry.get_request_translator("document")
def to_request(self, body): 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") print("Document received")
return self.translator.to_pulsar(body)
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")
)

View file

@ -2,6 +2,7 @@
from ... schema import FlowRequest, FlowResponse from ... schema import FlowRequest, FlowResponse
from ... schema import flow_request_queue from ... schema import flow_request_queue
from ... schema import flow_response_queue from ... schema import flow_response_queue
from .... base.messaging import TranslatorRegistry
from . requestor import ServiceRequestor from . requestor import ServiceRequestor
@ -19,34 +20,12 @@ class FlowRequestor(ServiceRequestor):
timeout=timeout, 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( def to_request(self, body):
operation = body.get("operation", None), return self.request_translator.to_pulsar(body)
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 from_response(self, message): def from_response(self, message):
return self.response_translator.from_response_with_completion(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

View file

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

View file

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

View file

@ -3,6 +3,13 @@ import base64
from ... schema import Value, Triple, DocumentMetadata, ProcessingMetadata 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): def to_value(x):
return Value(value=x["v"], is_uri=x["e"]) return Value(value=x["v"], is_uri=x["e"])

View file

@ -2,9 +2,9 @@
import base64 import base64
from ... schema import TextDocument, Metadata from ... schema import TextDocument, Metadata
from .... base.messaging import TranslatorRegistry
from . sender import ServiceSender from . sender import ServiceSender
from . serialize import to_subgraph
class TextLoad(ServiceSender): class TextLoad(ServiceSender):
def __init__(self, pulsar_client, queue): def __init__(self, pulsar_client, queue):
@ -15,30 +15,9 @@ class TextLoad(ServiceSender):
schema = TextDocument, schema = TextDocument,
) )
self.translator = TranslatorRegistry.get_request_translator("text-document")
def to_request(self, body): 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") print("Text document received")
return self.translator.to_pulsar(body)
return TextDocument(
metadata=Metadata(
id=body.get("id"),
metadata=metadata,
user=body.get("user", "trustgraph"),
collection=body.get("collection", "default"),
),
text=text,
)