diff --git a/trustgraph-base/trustgraph/messaging/translators/base.py b/trustgraph-base/trustgraph/messaging/translators/base.py index 284b5bc3..64e2b635 100644 --- a/trustgraph-base/trustgraph/messaging/translators/base.py +++ b/trustgraph-base/trustgraph/messaging/translators/base.py @@ -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 = {} diff --git a/trustgraph-base/trustgraph/messaging/translators/config.py b/trustgraph-base/trustgraph/messaging/translators/config.py new file mode 100644 index 00000000..10e023f6 --- /dev/null +++ b/trustgraph-base/trustgraph/messaging/translators/config.py @@ -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 \ No newline at end of file diff --git a/trustgraph-base/trustgraph/messaging/translators/document_loading.py b/trustgraph-base/trustgraph/messaging/translators/document_loading.py new file mode 100644 index 00000000..3dfef718 --- /dev/null +++ b/trustgraph-base/trustgraph/messaging/translators/document_loading.py @@ -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 \ No newline at end of file diff --git a/trustgraph-base/trustgraph/messaging/translators/embeddings_query.py b/trustgraph-base/trustgraph/messaging/translators/embeddings_query.py new file mode 100644 index 00000000..d69e7bef --- /dev/null +++ b/trustgraph-base/trustgraph/messaging/translators/embeddings_query.py @@ -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 \ No newline at end of file diff --git a/trustgraph-base/trustgraph/messaging/translators/flow.py b/trustgraph-base/trustgraph/messaging/translators/flow.py new file mode 100644 index 00000000..212a9992 --- /dev/null +++ b/trustgraph-base/trustgraph/messaging/translators/flow.py @@ -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 \ No newline at end of file diff --git a/trustgraph-base/trustgraph/messaging/translators/prompt.py b/trustgraph-base/trustgraph/messaging/translators/prompt.py new file mode 100644 index 00000000..b0e7351f --- /dev/null +++ b/trustgraph-base/trustgraph/messaging/translators/prompt.py @@ -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 \ No newline at end of file diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/config.py b/trustgraph-flow/trustgraph/gateway/dispatch/config.py index 3aeedb6f..c44e56f9 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/config.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/config.py @@ -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) diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_import.py b/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_import.py index e486f613..1c6177bf 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_import.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_import.py @@ -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): diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/document_load.py b/trustgraph-flow/trustgraph/gateway/dispatch/document_load.py index f92fc34f..4991de83 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/document_load.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/document_load.py @@ -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) diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/flow.py b/trustgraph-flow/trustgraph/gateway/dispatch/flow.py index 0b38e9be..c45f5066 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/flow.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/flow.py @@ -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) diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_query.py b/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_query.py index 27ceb702..683b10c8 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_query.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_query.py @@ -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) diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/prompt.py b/trustgraph-flow/trustgraph/gateway/dispatch/prompt.py index 496d01e5..ab73dad5 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/prompt.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/prompt.py @@ -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) diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/serialize.py b/trustgraph-flow/trustgraph/gateway/dispatch/serialize.py index bde3553a..1ee0f772 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/serialize.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/serialize.py @@ -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"]) diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/text_load.py b/trustgraph-flow/trustgraph/gateway/dispatch/text_load.py index 53ea7452..018e1324 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/text_load.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/text_load.py @@ -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)