apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-05-26 07:45:13 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Feature/pulsar api key support (#308)

Browse source 
* Add pulsar API token check

* Added missing api_key references

---------

Co-authored-by: Tyler O <4535788+toliver38@users.noreply.github.com>
This commit is contained in:
cybermaggedon 2025-02-15 11:22:48 +00:00 • committed by GitHub
parent f7df2df266
commit 617eb7efd5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
40 changed files with 173 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

21
trustgraph-base/trustgraph/base/base_processor.py
View file

@ -12,6 +12,7 @@ from .. log_level import LogLevel
class BaseProcessor:
default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://pulsar:6650')
default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
def __init__(self, **params):
@ -30,15 +31,23 @@ class BaseProcessor:
pulsar_host = params.get("pulsar_host", self.default_pulsar_host)
pulsar_listener = params.get("pulsar_listener", None)
pulsar_api_key = params.get("pulsar_api_key", None)
log_level = params.get("log_level", LogLevel.INFO)
self.pulsar_host = pulsar_host
self.client = pulsar.Client(
if pulsar_api_key:
auth = pulsar.AuthenticationToken(pulsar_api_key)
self.client = pulsar.Client(
pulsar_host,
authentication=auth,
logger=pulsar.ConsoleLogger(log_level.to_pulsar())
)
else:
self.client = pulsar.Client(
pulsar_host,
listener_name=pulsar_listener,
logger=pulsar.ConsoleLogger(log_level.to_pulsar())
)
)
self.pulsar_listener = pulsar_listener
@ -56,6 +65,12 @@ class BaseProcessor:
default=__class__.default_pulsar_host,
help=f'Pulsar host (default: {__class__.default_pulsar_host})',
)
parser.add_argument(
'--pulsar-api-key',
default=__class__.default_pulsar_api_key,
help=f'Pulsar API key',
)
parser.add_argument(
'--pulsar-listener',

2
trustgraph-base/trustgraph/base/publisher.py
View file

@ -31,8 +31,6 @@ class Publisher:
while self.running:
try:
print(self.chunking_enabled)
producer = self.client.create_producer(
topic=self.topic,
schema=self.schema,

2
trustgraph-base/trustgraph/clients/agent_client.py
View file

@ -20,6 +20,7 @@ class AgentClient(BaseClient):
input_queue=None,
output_queue=None,
pulsar_host="pulsar://pulsar:6650",
pulsar_api_key=None,
):
if input_queue is None: input_queue = agent_request_queue
@ -33,6 +34,7 @@ class AgentClient(BaseClient):
pulsar_host=pulsar_host,
input_schema=AgentRequest,
output_schema=AgentResponse,
pulsar_api_key=pulsar_api_key
)
def request(

13
trustgraph-base/trustgraph/clients/base.py
View file

@ -27,6 +27,7 @@ class BaseClient:
input_schema=None,
output_schema=None,
pulsar_host="pulsar://pulsar:6650",
pulsar_api_key=None,
):
if input_queue == None: raise RuntimeError("Need input_queue")
@ -37,10 +38,18 @@ class BaseClient:
if subscriber == None:
subscriber = str(uuid.uuid4())
self.client = pulsar.Client(
if pulsar_api_key:
auth = pulsar.AuthenticationToken(pulsar_api_key)
self.client = pulsar.Client(
pulsar_host,
logger=pulsar.ConsoleLogger(log_level),
)
authentication=auth,
)
else:
self.client = pulsar.Client(
pulsar_host,
logger=pulsar.ConsoleLogger(log_level)
)
self.producer = self.client.create_producer(
topic=input_queue,

2
trustgraph-base/trustgraph/clients/document_embeddings_client.py
View file

@ -20,6 +20,7 @@ class DocumentEmbeddingsClient(BaseClient):
input_queue=None,
output_queue=None,
pulsar_host="pulsar://pulsar:6650",
pulsar_api_key=None,
):
if input_queue == None:
@ -34,6 +35,7 @@ class DocumentEmbeddingsClient(BaseClient):
input_queue=input_queue,
output_queue=output_queue,
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
input_schema=DocumentEmbeddingsRequest,
output_schema=DocumentEmbeddingsResponse,
)

2
trustgraph-base/trustgraph/clients/document_rag_client.py
View file

@ -20,6 +20,7 @@ class DocumentRagClient(BaseClient):
input_queue=None,
output_queue=None,
pulsar_host="pulsar://pulsar:6650",
pulsar_api_key=None,
):
if input_queue == None:
@ -34,6 +35,7 @@ class DocumentRagClient(BaseClient):
input_queue=input_queue,
output_queue=output_queue,
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
input_schema=DocumentRagQuery,
output_schema=DocumentRagResponse,
)

2
trustgraph-base/trustgraph/clients/embeddings_client.py
View file

@ -20,6 +20,7 @@ class EmbeddingsClient(BaseClient):
output_queue=None,
subscriber=None,
pulsar_host="pulsar://pulsar:6650",
pulsar_api_key=None,
):
if input_queue == None:
@ -34,6 +35,7 @@ class EmbeddingsClient(BaseClient):
input_queue=input_queue,
output_queue=output_queue,
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
input_schema=EmbeddingsRequest,
output_schema=EmbeddingsResponse,
)

2
trustgraph-base/trustgraph/clients/graph_embeddings_client.py
View file

@ -20,6 +20,7 @@ class GraphEmbeddingsClient(BaseClient):
input_queue=None,
output_queue=None,
pulsar_host="pulsar://pulsar:6650",
pulsar_api_key=None,
):
if input_queue == None:
@ -34,6 +35,7 @@ class GraphEmbeddingsClient(BaseClient):
input_queue=input_queue,
output_queue=output_queue,
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
input_schema=GraphEmbeddingsRequest,
output_schema=GraphEmbeddingsResponse,
)

2
trustgraph-base/trustgraph/clients/graph_rag_client.py
View file

@ -20,6 +20,7 @@ class GraphRagClient(BaseClient):
input_queue=None,
output_queue=None,
pulsar_host="pulsar://pulsar:6650",
pulsar_api_key=None,
):
if input_queue == None:
@ -34,6 +35,7 @@ class GraphRagClient(BaseClient):
input_queue=input_queue,
output_queue=output_queue,
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
input_schema=GraphRagQuery,
output_schema=GraphRagResponse,
)

2
trustgraph-base/trustgraph/clients/llm_client.py
View file

@ -20,6 +20,7 @@ class LlmClient(BaseClient):
input_queue=None,
output_queue=None,
pulsar_host="pulsar://pulsar:6650",
pulsar_api_key=None,
):
if input_queue is None: input_queue = text_completion_request_queue
@ -31,6 +32,7 @@ class LlmClient(BaseClient):
input_queue=input_queue,
output_queue=output_queue,
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
input_schema=TextCompletionRequest,
output_schema=TextCompletionResponse,
)

2
trustgraph-base/trustgraph/clients/prompt_client.py
View file

@ -39,6 +39,7 @@ class PromptClient(BaseClient):
input_queue=None,
output_queue=None,
pulsar_host="pulsar://pulsar:6650",
pulsar_api_key=None,
):
if input_queue == None:
@ -53,6 +54,7 @@ class PromptClient(BaseClient):
input_queue=input_queue,
output_queue=output_queue,
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
input_schema=PromptRequest,
output_schema=PromptResponse,
)

2
trustgraph-base/trustgraph/clients/triples_query_client.py
View file

@ -21,6 +21,7 @@ class TriplesQueryClient(BaseClient):
input_queue=None,
output_queue=None,
pulsar_host="pulsar://pulsar:6650",
pulsar_api_key=None,
):
if input_queue == None:
@ -34,6 +35,7 @@ class TriplesQueryClient(BaseClient):
subscriber=subscriber,
input_queue=input_queue,
output_queue=output_queue,
pulsar_api_key=pulsar_api_key,
pulsar_host=pulsar_host,
input_schema=TriplesQueryRequest,
output_schema=TriplesQueryResponse,

6
trustgraph-cli/scripts/tg-invoke-agent
View file

@ -137,6 +137,12 @@ def main():
action="store_true",
help=f'Output thinking/observations'
)
# parser.add_argument(
# '--pulsar-api-key',
# default=default_pulsar_api_key,
# help=f'Pulsar API key',
# )
args = parser.parse_args()

8
trustgraph-cli/scripts/tg-invoke-document-rag
View file

@ -34,8 +34,14 @@ def main():
help=f'API URL (default: {default_url})',
)
# parser.add_argument(
# '--pulsar-api-key',
# default=default_pulsar_api_key,
# help=f'Pulsar API key',
# )
parser.add_argument(
'-q', '--question',
'-q', '--query',
required=True,
help=f'Question to answer',
)

6
trustgraph-cli/scripts/tg-invoke-llm
View file

@ -44,6 +44,12 @@ def main():
nargs=1,
help='LLM prompt e.g. What is 2 + 2?',
)
# parser.add_argument(
# '--pulsar-api-key',
# default=default_pulsar_api_key,
# help=f'Pulsar API key',
# )
args = parser.parse_args()

6
trustgraph-cli/scripts/tg-invoke-prompt
View file

@ -54,6 +54,12 @@ def main():
help='''Prompt template terms of the form variable=value, can be
specified multiple times''',
)
# parser.add_argument(
# '--pulsar-api-key',
# default=default_pulsar_api_key,
# help=f'Pulsar API key',
# )
args = parser.parse_args()

1
trustgraph-cli/scripts/tg-load-pdf
View file

@ -30,6 +30,7 @@ class Loader:
user,
collection,
metadata,
pulsar_api_key=None,
):
self.api = Api(url)

6
trustgraph-cli/scripts/tg-load-text
View file

@ -80,6 +80,12 @@ def main():
default=default_url,
help=f'API URL (default: {default_url})',
)
# parser.add_argument(
# '--pulsar-api-key',
# default=default_pulsar_api_key,
# help=f'Pulsar API key',
# )
parser.add_argument(
'-U', '--user',

26
trustgraph-cli/scripts/tg-load-turtle
View file

@ -19,6 +19,8 @@ from trustgraph.log_level import LogLevel
default_user = 'trustgraph'
default_collection = 'default'
default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
default_output_queue = triples_store_queue
class Loader:
@ -31,12 +33,21 @@ class Loader:
files,
user,
collection,
pulsar_api_key=None,
):
self.client = pulsar.Client(
pulsar_host,
logger=pulsar.ConsoleLogger(log_level.to_pulsar())
)
if pulsar_api_key:
auth = pulsar.AuthenticationToken(pulsar_api_key)
self.client = pulsar.Client(
pulsar_host,
authentication=auth,
logger=pulsar.ConsoleLogger(log_level.to_pulsar())
)
else:
self.client = pulsar.Client(
pulsar_host,
logger=pulsar.ConsoleLogger(log_level.to_pulsar())
)
self.producer = self.client.create_producer(
topic=output_queue,
@ -98,6 +109,12 @@ def main():
default=default_pulsar_host,
help=f'Pulsar host (default: {default_pulsar_host})',
)
parser.add_argument(
'--pulsar-api-key',
default=default_pulsar_api_key,
help=f'Pulsar API key',
)
parser.add_argument(
'-o', '--output-queue',
@ -137,6 +154,7 @@ def main():
try:
p = Loader(
pulsar_host=args.pulsar_host,
pulsar_api_key=args.pulsar_api_key,
output_queue=args.output_queue,
log_level=args.log_level,
files=args.files,

6
trustgraph-flow/trustgraph/agent/react/service.py
View file

@ -156,14 +156,16 @@ class Processor(ConsumerProducer):
subscriber=subscriber,
input_queue=prompt_request_queue,
output_queue=prompt_response_queue,
pulsar_host = self.pulsar_host
pulsar_host = self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
)
self.graph_rag = GraphRagClient(
subscriber=subscriber,
input_queue=graph_rag_request_queue,
output_queue=graph_rag_response_queue,
pulsar_host = self.pulsar_host
pulsar_host = self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
)
# Need to be able to feed requests to myself

4
trustgraph-flow/trustgraph/document_rag.py
View file

@ -59,6 +59,7 @@ class DocumentRag:
def __init__(
self,
pulsar_host="pulsar://pulsar:6650",
pulsar_api_key=None,
pr_request_queue=None,
pr_response_queue=None,
emb_request_queue=None,
@ -100,6 +101,7 @@ class DocumentRag:
subscriber=module + "-de",
input_queue=de_request_queue,
output_queue=de_response_queue,
pulsar_api_key=pulsar_api_key,
)
self.embeddings = EmbeddingsClient(
@ -107,6 +109,7 @@ class DocumentRag:
input_queue=emb_request_queue,
output_queue=emb_response_queue,
subscriber=module + "-emb",
pulsar_api_key=pulsar_api_key,
)
self.lang = PromptClient(
@ -114,6 +117,7 @@ class DocumentRag:
input_queue=pr_request_queue,
output_queue=pr_response_queue,
subscriber=module + "-de-prompt",
pulsar_api_key=pulsar_api_key,
)
if self.verbose:

1
trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py
View file

@ -47,6 +47,7 @@ class Processor(ConsumerProducer):
self.embeddings = EmbeddingsClient(
pulsar_host=self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
input_queue=emb_request_queue,
output_queue=emb_response_queue,
subscriber=module + "-emb",

1
trustgraph-flow/trustgraph/extract/kg/definitions/extract.py
View file

@ -79,6 +79,7 @@ class Processor(ConsumerProducer):
self.prompt = PromptClient(
pulsar_host=self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
input_queue=pr_request_queue,
output_queue=pr_response_queue,
subscriber = module + "-prompt",

1
trustgraph-flow/trustgraph/extract/kg/relationships/extract.py
View file

@ -54,6 +54,7 @@ class Processor(ConsumerProducer):
self.prompt = PromptClient(
pulsar_host=self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
input_queue=pr_request_queue,
output_queue=pr_response_queue,
subscriber = module + "-prompt",

1
trustgraph-flow/trustgraph/extract/kg/topics/extract.py
View file

@ -52,6 +52,7 @@ class Processor(ConsumerProducer):
self.prompt = PromptClient(
pulsar_host=self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
input_queue=pr_request_queue,
output_queue=pr_response_queue,
subscriber = module + "-prompt",

1
trustgraph-flow/trustgraph/extract/object/row/extract.py
View file

@ -112,6 +112,7 @@ class Processor(ConsumerProducer):
self.prompt = PromptClient(
pulsar_host=self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
input_queue=pr_request_queue,
output_queue=pr_response_queue,
subscriber = module + "-prompt",

8
trustgraph-flow/trustgraph/gateway/agent.py
View file

@ -7,10 +7,18 @@ from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
class AgentRequestor(ServiceRequestor):
<<<<<<< HEAD
def __init__(self, pulsar_client, timeout, auth):
super(AgentRequestor, self).__init__(
pulsar_client=pulsar_client,
=======
def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
super(AgentRequestor, self).__init__(
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
>>>>>>> a5d5b4c (Add pulsar API token check)
request_queue=agent_request_queue,
response_queue=agent_response_queue,
request_schema=AgentRequest,

3
trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py
View file

@ -27,7 +27,8 @@ class DocumentEmbeddingsStreamEndpoint(SocketEndpoint):
self.subscriber = Subscriber(
self.pulsar_client, document_embeddings_store_queue,
"api-gateway", "api-gateway",
schema=JsonSchema(DocumentEmbeddings)
schema=JsonSchema(DocumentEmbeddings),
pulsar_api_key=self.pulsar_api_key
)
async def listener(self, ws, running):

1
trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py
View file

@ -26,6 +26,7 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint):
self.subscriber = Subscriber(
self.pulsar_client, graph_embeddings_store_queue,
"api-gateway", "api-gateway",
pulsar_api_key=self.pulsar_api_key,
schema=JsonSchema(GraphEmbeddings)
)

1
trustgraph-flow/trustgraph/gateway/mux.py
View file

@ -21,6 +21,7 @@ class MuxEndpoint(SocketEndpoint):
self, pulsar_client, auth,
services,
path="/api/v1/socket",
pulsar_api_key=None
):
super(MuxEndpoint, self).__init__(

2
trustgraph-flow/trustgraph/gateway/requestor.py
View file

@ -19,6 +19,7 @@ class ServiceRequestor:
response_queue, response_schema,
subscription="api-gateway", consumer_name="api-gateway",
timeout=600,
pulsar_api_key=None,
):
self.pub = Publisher(
@ -29,6 +30,7 @@ class ServiceRequestor:
self.sub = Subscriber(
pulsar_client, response_queue,
subscription, consumer_name,
pulsar_api_key,
JsonSchema(response_schema)
)

1
trustgraph-flow/trustgraph/gateway/sender.py
View file

@ -17,6 +17,7 @@ class ServiceSender:
self,
pulsar_client,
request_queue, request_schema,
pulsar_api_key=None,
):
self.pub = Publisher(

24
trustgraph-flow/trustgraph/gateway/service.py
View file

@ -57,6 +57,7 @@ logger.setLevel(logging.INFO)
default_pulsar_host = os.getenv("PULSAR_HOST", "pulsar://pulsar:6650")
default_prometheus_url = os.getenv("PROMETHEUS_URL", "http://prometheus:9090")
default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
default_timeout = 600
default_port = 8088
default_api_token = os.getenv("GATEWAY_SECRET", "")
@ -73,11 +74,20 @@ class Api:
self.port = int(config.get("port", default_port))
self.timeout = int(config.get("timeout", default_timeout))
self.pulsar_host = config.get("pulsar_host", default_pulsar_host)
self.pulsar_api_key = config.get(
"pulsar_api_key", default_pulsar_api_key
)
self.pulsar_listener = config.get("pulsar_listener", None)
self.pulsar_client = pulsar.Client(
self.pulsar_host, listener_name=self.pulsar_listener
)
if self.pulsar_api_key:
self.pulsar_client = pulsar.Client(
self.pulsar_host, listener_name=self.pulsar_listener,
authentication=pulsar.AuthenticationToken(self.pulsar_api_key)
)
else:
self.pulsar_client = pulsar.Client(
self.pulsar_host, listener_name=self.pulsar_listener,
)
self.prometheus_url = config.get(
"prometheus_url", default_prometheus_url,
@ -224,6 +234,7 @@ class Api:
TriplesLoadEndpoint(
pulsar_client=self.pulsar_client,
auth = self.auth,
pulsar_api_key=self.pulsar_api_key,
),
GraphEmbeddingsLoadEndpoint(
pulsar_client=self.pulsar_client,
@ -237,6 +248,7 @@ class Api:
pulsar_client=self.pulsar_client,
auth = self.auth,
services = self.services,
pulsar_api_key=self.pulsar_api_key,
),
MetricsEndpoint(
endpoint_path = "/api/v1/metrics",
@ -270,6 +282,12 @@ def run():
default=default_pulsar_host,
help=f'Pulsar host (default: {default_pulsar_host})',
)
parser.add_argument(
'--pulsar-api-key',
default=default_pulsar_api_key,
help=f'Pulsar API key',
)
parser.add_argument(
'--pulsar-listener',

1
trustgraph-flow/trustgraph/gateway/triples_stream.py
View file

@ -24,6 +24,7 @@ class TriplesStreamEndpoint(SocketEndpoint):
self.subscriber = Subscriber(
self.pulsar_client, triples_store_queue,
"api-gateway", "api-gateway",
pulsar_api_key=self.pulsar_api_key,
schema=JsonSchema(Triples)
)

5
trustgraph-flow/trustgraph/graph_rag.py
View file

@ -161,6 +161,7 @@ class GraphRag:
def __init__(
self,
pulsar_host="pulsar://pulsar:6650",
pulsar_api_key=None,
pr_request_queue=None,
pr_response_queue=None,
emb_request_queue=None,
@ -207,6 +208,7 @@ class GraphRag:
self.ge_client = GraphEmbeddingsClient(
pulsar_host=pulsar_host,
pulsar_api_key=-pulsar_api_key,
subscriber=module + "-ge",
input_queue=ge_request_queue,
output_queue=ge_response_queue,
@ -214,6 +216,7 @@ class GraphRag:
self.triples_client = TriplesQueryClient(
pulsar_host=pulsar_host,
pulsar_api_key=-pulsar_api_key,
subscriber=module + "-tpl",
input_queue=tpl_request_queue,
output_queue=tpl_response_queue
@ -221,6 +224,7 @@ class GraphRag:
self.embeddings = EmbeddingsClient(
pulsar_host=pulsar_host,
pulsar_api_key=-pulsar_api_key,
input_queue=emb_request_queue,
output_queue=emb_response_queue,
subscriber=module + "-emb",
@ -234,6 +238,7 @@ class GraphRag:
self.prompt = PromptClient(
pulsar_host=pulsar_host,
pulsar_api_key=-pulsar_api_key,
input_queue=pr_request_queue,
output_queue=pr_response_queue,
subscriber=module + "-prompt",

3
trustgraph-flow/trustgraph/model/prompt/generic/service.py
View file

@ -63,7 +63,8 @@ class Processor(ConsumerProducer):
subscriber=subscriber,
input_queue=tc_request_queue,
output_queue=tc_response_queue,
pulsar_host = self.pulsar_host
pulsar_host = self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
)
def parse_json(self, text):

3
trustgraph-flow/trustgraph/model/prompt/template/service.py
View file

@ -136,7 +136,8 @@ class Processor(ConsumerProducer):
subscriber=subscriber,
input_queue=tc_request_queue,
output_queue=tc_response_queue,
pulsar_host = self.pulsar_host
pulsar_host = self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
)
# System prompt hack

12
trustgraph-flow/trustgraph/processing/processing.py
View file

@ -49,11 +49,12 @@ class Processing:
pulsar_host,
log_level,
file,
pulsar_api_key=None,
):
self.pulsar_host = pulsar_host
self.log_level = log_level
self.file = file
self.pulsar_api_key = pulsar_api_key
self.defs = load(open(file, "r"), Loader=Loader)
def run(self):
@ -68,6 +69,7 @@ class Processing:
params = {
"pulsar_host": self.pulsar_host,
"pulsar_api_key": self.pulsar_api_key,
"log_level": str(self.log_level),
}
@ -125,12 +127,19 @@ def run():
)
default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://pulsar:6650')
default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
parser.add_argument(
'-p', '--pulsar-host',
default=default_pulsar_host,
help=f'Pulsar host (default: {default_pulsar_host})',
)
parser.add_argument(
'--pulsar-api-key',
default=default_pulsar_api_key,
help=f'Pulsar API key',
)
parser.add_argument(
'-l', '--log-level',
@ -153,6 +162,7 @@ def run():
try:
p = Processing(
pulsar_host=args.pulsar_host,
pulsar_api_key=args.pulsar_api_key,
file=args.file,
log_level=args.log_level,
)

1
trustgraph-flow/trustgraph/retrieval/document_rag/rag.py
View file

@ -68,6 +68,7 @@ class Processor(ConsumerProducer):
self.rag = DocumentRag(
pulsar_host=self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
pr_request_queue=pr_request_queue,
pr_response_queue=pr_response_queue,
emb_request_queue=emb_request_queue,

1
trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py
View file

@ -82,6 +82,7 @@ class Processor(ConsumerProducer):
self.rag = GraphRag(
pulsar_host=self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
pr_request_queue=pr_request_queue,
pr_response_queue=pr_response_queue,
emb_request_queue=emb_request_queue,