From ae00831e395511dd028c0a1ad9ce4ce4fd540fe0 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Sat, 4 Jan 2025 22:08:24 +0000 Subject: [PATCH] Added missing file (#249) --- .../trustgraph/gateway/document_rag.py | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) create mode 100644 trustgraph-flow/trustgraph/gateway/document_rag.py diff --git a/trustgraph-flow/trustgraph/gateway/document_rag.py b/trustgraph-flow/trustgraph/gateway/document_rag.py new file mode 100644 index 00000000..6665a7dc --- /dev/null +++ b/trustgraph-flow/trustgraph/gateway/document_rag.py @@ -0,0 +1,30 @@ + +from .. schema import DocumentRagQuery, DocumentRagResponse +from .. schema import document_rag_request_queue +from .. schema import document_rag_response_queue + +from . endpoint import ServiceEndpoint +from . requestor import ServiceRequestor + +class DocumentRagRequestor(ServiceRequestor): + def __init__(self, pulsar_host, timeout, auth): + + super(DocumentRagRequestor, self).__init__( + pulsar_host=pulsar_host, + request_queue=document_rag_request_queue, + response_queue=document_rag_response_queue, + request_schema=DocumentRagQuery, + response_schema=DocumentRagResponse, + timeout=timeout, + ) + + def to_request(self, body): + return DocumentRagQuery( + query=body["query"], + user=body.get("user", "trustgraph"), + collection=body.get("collection", "default"), + ) + + def from_response(self, message): + return { "response": message.response }, True +