apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-04-28 18:06:21 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Make embedding & text-completion queues configurable in graph-rag

Browse source
This commit is contained in:
Cyber MacGeddon 2024-08-05 22:37:11 +01:00
parent 963944f4aa
commit 103251dbf2
2 changed files with 70 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

22
trustgraph/graph_rag.py
View file

@ -4,6 +4,8 @@ from trustgraph.triple_vectors import TripleVectors
from trustgraph.trustgraph import TrustGraph
from trustgraph.llm_client import LlmClient
from trustgraph.embeddings_client import EmbeddingsClient
from . schema import text_completion_request_queue
from . schema import text_completion_response_queue
LABEL="http://www.w3.org/2000/01/rdf-schema#label"
DEFINITION="http://www.w3.org/2004/02/skos/core#definition"
@ -15,6 +17,10 @@ class GraphRag:
graph_hosts=None,
pulsar_host="pulsar://pulsar:6650",
vector_store="http://milvus:19530",
completion_request_queue=None,
completion_response_queue=None,
emb_request_queue=None,
emb_response_queue=None,
verbose=False,
entity_limit=50,
triple_limit=30,
@ -24,6 +30,18 @@ class GraphRag:
self.verbose=verbose
if completion_request_queue == None:
completion_request_queue = text_completion_request_queue
if completion_response_queue == None:
completion_response_queue = text_completion_response_queue
if emb_request_queue == None:
emb_request_queue = embeddings_request_queue
if emb_response_queue == None:
emb_response_queue = embeddings_response_queue
if graph_hosts == None:
graph_hosts = ["cassandra"]
@ -34,6 +52,8 @@ class GraphRag:
self.embeddings = EmbeddingsClient(
pulsar_host=pulsar_host,
input_queue=emb_request_queue,
output_queue=emb_response_queue,
subscriber=module + "-emb",
)
@ -47,6 +67,8 @@ class GraphRag:
self.llm = LlmClient(
pulsar_host=pulsar_host,
input_queue=completion_request_queue,
output_queue=completion_response_queue,
subscriber=module + "-llm",
)

48
trustgraph/retrieval/graph_rag/rag.py
View file

@ -6,6 +6,10 @@ Input is query, output is response.
from ... schema import GraphRagQuery, GraphRagResponse
from ... schema import graph_rag_request_queue, graph_rag_response_queue
from ... schema import text_completion_request_queue
from ... schema import text_completion_response_queue
from ... schema import embeddings_request_queue
from ... schema import embeddings_response_queue
from ... log_level import LogLevel
from ... graph_rag import GraphRag
from ... base import ConsumerProducer
@ -30,6 +34,18 @@ class Processor(ConsumerProducer):
entity_limit = params.get("entity_limit", 50)
triple_limit = params.get("triple_limit", 30)
max_subgraph_size = params.get("max_subgraph_size", 3000)
tc_request_queue = params.get(
"text_completion_request_queue", text_completion_request_queue
)
tc_response_queue = params.get(
"text_completion_response_queue", text_completion_response_queue
)
emb_request_queue = params.get(
"embeddings_request_queue", embeddings_request_queue
)
emb_response_queue = params.get(
"embeddings_response_queue", embeddings_response_queue
)
super(Processor, self).__init__(
**params | {
@ -41,12 +57,20 @@ class Processor(ConsumerProducer):
"entity_limit": entity_limit,
"triple_limit": triple_limit,
"max_subgraph_size": max_subgraph_size,
"text_completion_request_queue": tc_request_queue,
"text_completion_response_queue": tc_response_queue,
"embeddings_request_queue": emb_request_queue,
"embeddings_response_queue": emb_response_queue,
}
)
self.rag = GraphRag(
pulsar_host=self.pulsar_host,
graph_hosts=graph_hosts.split(","),
completion_request_queue=tc_request_queue,
completion_response_queue=tc_response_queue,
emb_request_queue=emb_request_queue,
emb_response_queue=emb_response_queue,
vector_store=vector_store,
verbose=True,
entity_limit=entity_limit,
@ -114,6 +138,30 @@ class Processor(ConsumerProducer):
help=f'Max subgraph size (default: 3000)'
)
parser.add_argument(
'--text-completion-request-queue',
default=text_completion_request_queue,
help=f'Text completion request queue (default: {text_completion_request_queue})',
)
parser.add_argument(
'--text-completion-response-queue',
default=text_completion_response_queue,
help=f'Text completion response queue (default: {text_completion_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 request queue (default: {embeddings_response_queue})',
)
def run():
Processor.start(module, __doc__)