trustgraph/trustgraph-flow/trustgraph/gateway/graph_rag.py
Cyber MacGeddon 8830e99e57 Change document-rag and graph-rag processing so that the user can
specify parameters.  Changes in Pulsar services, Pulsar message
schemas, gateway and command-line tools.  User-visible changes in
new parameters on command-line tools.
2025-03-12 17:51:47 +00:00

33 lines
1.1 KiB
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_client, timeout, auth):
super(GraphRagRequestor, self).__init__(
pulsar_client=pulsar_client,
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"),
entity_limit=body.get("entity-limit", 50),
triple_limit=body.get("triple-limit", 30),
doc_limit=body.get("doc-limit", 1000),
)
def from_response(self, message):
return { "response": message.response }, True