mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-22 19:51:02 +02:00
31 lines
1,014 B
Python
31 lines
1,014 B
Python
|
|
from .. schema import GraphRagQuery, GraphRagResponse
|
|
from .. schema import graph_rag_request_queue
|
|
from .. schema import graph_rag_response_queue
|
|
|
|
from . endpoint import ServiceEndpoint
|
|
from . requestor import ServiceRequestor
|
|
|
|
class GraphRagRequestor(ServiceRequestor):
|
|
def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
|
|
|
|
super(GraphRagRequestor, self).__init__(
|
|
pulsar_host=pulsar_host,
|
|
pulsar_api_key=pulsar_api_key,
|
|
request_queue=graph_rag_request_queue,
|
|
response_queue=graph_rag_response_queue,
|
|
request_schema=GraphRagQuery,
|
|
response_schema=GraphRagResponse,
|
|
timeout=timeout,
|
|
)
|
|
|
|
def to_request(self, body):
|
|
return GraphRagQuery(
|
|
query=body["query"],
|
|
user=body.get("user", "trustgraph"),
|
|
collection=body.get("collection", "default"),
|
|
)
|
|
|
|
def from_response(self, message):
|
|
return { "response": message.response }, True
|
|
|