mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-26 00:46:22 +02:00
Feature/pkgsplit (#83)
* Starting to spawn base package * More package hacking * Bedrock and VertexAI * Parquet split * Updated templates * Utils
This commit is contained in:
cybermaggedon
GitHub
parent
3fb75c617b
commit
9b91d5eee3
262 changed files with 630 additions and 420 deletions
217
trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py
Executable file
217
trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py
Executable file
|
|
@ -0,0 +1,217 @@
|
|||
|
||||
"""
|
||||
Simple RAG service, performs query using graph RAG an LLM.
|
||||
Input is query, output is response.
|
||||
"""
|
||||
|
||||
from ... schema import GraphRagQuery, GraphRagResponse, Error
|
||||
from ... schema import graph_rag_request_queue, graph_rag_response_queue
|
||||
from ... schema import prompt_request_queue
|
||||
from ... schema import prompt_response_queue
|
||||
from ... schema import embeddings_request_queue
|
||||
from ... schema import embeddings_response_queue
|
||||
from ... schema import graph_embeddings_request_queue
|
||||
from ... schema import graph_embeddings_response_queue
|
||||
from ... schema import triples_request_queue
|
||||
from ... schema import triples_response_queue
|
||||
from ... log_level import LogLevel
|
||||
from ... graph_rag import GraphRag
|
||||
from ... base import ConsumerProducer
|
||||
|
||||
module = ".".join(__name__.split(".")[1:-1])
|
||||
|
||||
default_input_queue = graph_rag_request_queue
|
||||
default_output_queue = graph_rag_response_queue
|
||||
default_subscriber = module
|
||||
|
||||
class Processor(ConsumerProducer):
|
||||
|
||||
def __init__(self, **params):
|
||||
|
||||
input_queue = params.get("input_queue", default_input_queue)
|
||||
output_queue = params.get("output_queue", default_output_queue)
|
||||
subscriber = params.get("subscriber", default_subscriber)
|
||||
entity_limit = params.get("entity_limit", 50)
|
||||
triple_limit = params.get("triple_limit", 30)
|
||||
max_subgraph_size = params.get("max_subgraph_size", 3000)
|
||||
pr_request_queue = params.get(
|
||||
"prompt_request_queue", prompt_request_queue
|
||||
)
|
||||
pr_response_queue = params.get(
|
||||
"prompt_response_queue", prompt_response_queue
|
||||
)
|
||||
emb_request_queue = params.get(
|
||||
"embeddings_request_queue", embeddings_request_queue
|
||||
)
|
||||
emb_response_queue = params.get(
|
||||
"embeddings_response_queue", embeddings_response_queue
|
||||
)
|
||||
ge_request_queue = params.get(
|
||||
"graph_embeddings_request_queue", graph_embeddings_request_queue
|
||||
)
|
||||
ge_response_queue = params.get(
|
||||
"graph_embeddings_response_queue", graph_embeddings_response_queue
|
||||
)
|
||||
tpl_request_queue = params.get(
|
||||
"triples_request_queue", triples_request_queue
|
||||
)
|
||||
tpl_response_queue = params.get(
|
||||
"triples_response_queue", triples_response_queue
|
||||
)
|
||||
|
||||
super(Processor, self).__init__(
|
||||
**params | {
|
||||
"input_queue": input_queue,
|
||||
"output_queue": output_queue,
|
||||
"subscriber": subscriber,
|
||||
"input_schema": GraphRagQuery,
|
||||
"output_schema": GraphRagResponse,
|
||||
"entity_limit": entity_limit,
|
||||
"triple_limit": triple_limit,
|
||||
"max_subgraph_size": max_subgraph_size,
|
||||
"prompt_request_queue": pr_request_queue,
|
||||
"prompt_response_queue": pr_response_queue,
|
||||
"embeddings_request_queue": emb_request_queue,
|
||||
"embeddings_response_queue": emb_response_queue,
|
||||
"graph_embeddings_request_queue": ge_request_queue,
|
||||
"graph_embeddings_response_queue": ge_response_queue,
|
||||
"triples_request_queue": triples_request_queue,
|
||||
"triples_response_queue": triples_response_queue,
|
||||
}
|
||||
)
|
||||
|
||||
self.rag = GraphRag(
|
||||
pulsar_host=self.pulsar_host,
|
||||
pr_request_queue=pr_request_queue,
|
||||
pr_response_queue=pr_response_queue,
|
||||
emb_request_queue=emb_request_queue,
|
||||
emb_response_queue=emb_response_queue,
|
||||
ge_request_queue=ge_request_queue,
|
||||
ge_response_queue=ge_response_queue,
|
||||
tpl_request_queue=triples_request_queue,
|
||||
tpl_response_queue=triples_response_queue,
|
||||
verbose=True,
|
||||
entity_limit=entity_limit,
|
||||
triple_limit=triple_limit,
|
||||
max_subgraph_size=max_subgraph_size,
|
||||
module=module,
|
||||
)
|
||||
|
||||
def handle(self, msg):
|
||||
|
||||
try:
|
||||
|
||||
v = msg.value()
|
||||
|
||||
# Sender-produced ID
|
||||
id = msg.properties()["id"]
|
||||
|
||||
print(f"Handling input {id}...", flush=True)
|
||||
|
||||
response = self.rag.query(v.query)
|
||||
|
||||
print("Send response...", flush=True)
|
||||
r = GraphRagResponse(response = response, error=None)
|
||||
self.producer.send(r, properties={"id": id})
|
||||
|
||||
print("Done.", flush=True)
|
||||
|
||||
except Exception as e:
|
||||
|
||||
print(f"Exception: {e}")
|
||||
|
||||
print("Send error response...", flush=True)
|
||||
|
||||
r = GraphRagResponse(
|
||||
error=Error(
|
||||
type = "llm-error",
|
||||
message = str(e),
|
||||
),
|
||||
response=None,
|
||||
)
|
||||
|
||||
self.producer.send(r, properties={"id": id})
|
||||
|
||||
self.consumer.acknowledge(msg)
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
ConsumerProducer.add_args(
|
||||
parser, default_input_queue, default_subscriber,
|
||||
default_output_queue,
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-e', '--entity-limit',
|
||||
type=int,
|
||||
default=50,
|
||||
help=f'Entity vector fetch limit (default: 50)'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-t', '--triple-limit',
|
||||
type=int,
|
||||
default=30,
|
||||
help=f'Triple query limit, per query (default: 30)'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-u', '--max-subgraph-size',
|
||||
type=int,
|
||||
default=3000,
|
||||
help=f'Max subgraph size (default: 3000)'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--prompt-request-queue',
|
||||
default=prompt_request_queue,
|
||||
help=f'Prompt request queue (default: {prompt_request_queue})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--prompt-response-queue',
|
||||
default=prompt_response_queue,
|
||||
help=f'Prompt response queue (default: {prompt_response_queue})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--embeddings-request-queue',
|
||||
default=embeddings_request_queue,
|
||||
help=f'Embeddings request queue (default: {embeddings_request_queue})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--embeddings-response-queue',
|
||||
default=embeddings_response_queue,
|
||||
help=f'Embeddings response queue (default: {embeddings_response_queue})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--graph-embeddings-request-queue',
|
||||
default=graph_embeddings_request_queue,
|
||||
help=f'Graph embeddings request queue (default: {graph_embeddings_request_queue})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--graph-embeddings-response-queue',
|
||||
default=graph_embeddings_response_queue,
|
||||
help=f'Graph embeddings response queue (default: {graph_embeddings_response_queue})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--triples-request-queue',
|
||||
default=triples_request_queue,
|
||||
help=f'Triples request queue (default: {triples_request_queue})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--triples-response-queue',
|
||||
default=triples_response_queue,
|
||||
help=f'Triples response queue (default: {triples_response_queue})',
|
||||
)
|
||||
|
||||
def run():
|
||||
|
||||
Processor.start(module, __doc__)
|
||||
|
||||
Loading…
Add table
Add a link
Reference in a new issue