mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-21 19:21:03 +02:00
More dispatch code translated
This commit is contained in:
parent
8bbf1cc6c0
commit
5ba7199895
14 changed files with 548 additions and 176 deletions
|
|
@ -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 = {}
|
||||
|
|
|
|||
100
trustgraph-base/trustgraph/messaging/translators/config.py
Normal file
100
trustgraph-base/trustgraph/messaging/translators/config.py
Normal 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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
59
trustgraph-base/trustgraph/messaging/translators/flow.py
Normal file
59
trustgraph-base/trustgraph/messaging/translators/flow.py
Normal 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
|
||||
54
trustgraph-base/trustgraph/messaging/translators/prompt.py
Normal file
54
trustgraph-base/trustgraph/messaging/translators/prompt.py
Normal 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
|
||||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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"])
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue