apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-22 03:31:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

Add pulsar API token check

Browse source
This commit is contained in:
Tyler O 2025-02-10 16:30:53 +00:00
parent d0ae772fd6
commit a5d5b4ca4a
56 changed files with 319 additions and 82 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -166,21 +166,24 @@ 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.llm = LlmClient(
subscriber=subscriber,
input_queue=text_completion_request_queue,
output_queue=text_completion_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

@ -21,6 +21,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,
@ -62,6 +63,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(
@ -69,6 +71,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(
@ -76,6 +79,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/vectorize/vectorize.py
View file

@ -45,6 +45,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

@ -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/relationships/extract.py
View file

@ -76,6 +76,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",

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

@ -7,10 +7,11 @@ from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
class AgentRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth):
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,
request_queue=agent_request_queue,
response_queue=agent_response_queue,
request_schema=AgentRequest,

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

@ -7,10 +7,11 @@ from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
class DbpediaRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth):
def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
super(DbpediaRequestor, self).__init__(
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
request_queue=dbpedia_lookup_request_queue,
response_queue=dbpedia_lookup_response_queue,
request_schema=LookupRequest,

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

@ -8,10 +8,11 @@ from . sender import ServiceSender
from . serialize import to_subgraph
class DocumentLoadSender(ServiceSender):
def __init__(self, pulsar_host):
def __init__(self, pulsar_host, pulsar_api_key=None):
super(DocumentLoadSender, self).__init__(
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
request_queue=document_ingest_queue,
request_schema=Document,
)

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

@ -7,10 +7,11 @@ from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
class EmbeddingsRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth):
def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
super(EmbeddingsRequestor, self).__init__(
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
request_queue=embeddings_request_queue,
response_queue=embeddings_response_queue,
request_schema=EmbeddingsRequest,

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

@ -7,10 +7,11 @@ from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
class EncyclopediaRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth):
def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
super(EncyclopediaRequestor, self).__init__(
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
request_queue=encyclopedia_lookup_request_queue,
response_queue=encyclopedia_lookup_response_queue,
request_schema=LookupRequest,

4
trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py
View file

@ -15,7 +15,7 @@ from . serialize import to_subgraph, to_value
class GraphEmbeddingsLoadEndpoint(SocketEndpoint):
def __init__(
self, pulsar_host, auth, path="/api/v1/load/graph-embeddings",
self, pulsar_host, auth, pulsar_api_key=None, path="/api/v1/load/graph-embeddings",
):
super(GraphEmbeddingsLoadEndpoint, self).__init__(
@ -23,9 +23,11 @@ class GraphEmbeddingsLoadEndpoint(SocketEndpoint):
)
self.pulsar_host=pulsar_host
self.pulsar_api_key=pulsar_api_key
self.publisher = Publisher(
self.pulsar_host, graph_embeddings_store_queue,
self.pulsar_api_key,
schema=JsonSchema(GraphEmbeddings)
)

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

@ -8,10 +8,11 @@ from . requestor import ServiceRequestor
from . serialize import serialize_value
class GraphEmbeddingsQueryRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth):
def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
super(GraphEmbeddingsQueryRequestor, self).__init__(
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
request_queue=graph_embeddings_request_queue,
response_queue=graph_embeddings_response_queue,
request_schema=GraphEmbeddingsRequest,

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

@ -14,7 +14,7 @@ from . serialize import serialize_graph_embeddings
class GraphEmbeddingsStreamEndpoint(SocketEndpoint):
def __init__(
self, pulsar_host, auth, path="/api/v1/stream/graph-embeddings"
self, pulsar_host, auth, path="/api/v1/stream/graph-embeddings", pulsar_api_key=None
):
super(GraphEmbeddingsStreamEndpoint, self).__init__(
@ -22,10 +22,12 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint):
)
self.pulsar_host=pulsar_host
self.pulsar_api_key=pulsar_api_key
self.subscriber = Subscriber(
self.pulsar_host, graph_embeddings_store_queue,
"api-gateway", "api-gateway",
pulsar_api_key=self.pulsar_api_key,
schema=JsonSchema(GraphEmbeddings)
)

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

@ -7,10 +7,11 @@ from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
class GraphRagRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth):
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,

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

@ -7,10 +7,11 @@ from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
class InternetSearchRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth):
def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
super(InternetSearchRequestor, self).__init__(
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
request_queue=internet_search_request_queue,
response_queue=internet_search_response_queue,
request_schema=LookupRequest,

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

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

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

@ -9,10 +9,11 @@ from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
class PromptRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth):
def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
super(PromptRequestor, self).__init__(
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
request_queue=prompt_request_queue,
response_queue=prompt_response_queue,
request_schema=PromptRequest,

17
trustgraph-flow/trustgraph/gateway/publisher.py
View file

@ -7,8 +7,9 @@ import threading
class Publisher:
def __init__(self, pulsar_host, topic, schema=None, max_size=10,
chunking_enabled=False):
chunking_enabled=False, pulsar_api_key=None):
self.pulsar_host = pulsar_host
self.pulsar_api_key = pulsar_api_key,
self.topic = topic
self.schema = schema
self.q = queue.Queue(maxsize=max_size)
@ -23,10 +24,16 @@ class Publisher:
while True:
try:
client = pulsar.Client(
self.pulsar_host,
)
if self.pulsar_api_key:
client = pulsar.Client(
self.pulsar_host,
authentication=pulsar.AuthenticationToken(self.pulsar_api_key)
)
else:
client = pulsar.Client(
self.pulsar_host,
)
producer = client.create_producer(
topic=self.topic,

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

@ -19,16 +19,19 @@ class ServiceRequestor:
response_queue, response_schema,
subscription="api-gateway", consumer_name="api-gateway",
timeout=600,
pulsar_api_key=None,
):
self.pub = Publisher(
pulsar_host, request_queue,
schema=JsonSchema(request_schema)
pulsar_api_key,
schema=JsonSchema(request_schema),
)
self.sub = Subscriber(
pulsar_host, response_queue,
subscription, consumer_name,
pulsar_api_key,
JsonSchema(response_schema)
)

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

@ -17,10 +17,12 @@ class ServiceSender:
self,
pulsar_host,
request_queue, request_schema,
pulsar_api_key=None,
):
self.pub = Publisher(
pulsar_host, request_queue,
pulsar_api_key,
schema=JsonSchema(request_schema)
)

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

@ -53,6 +53,7 @@ logger = logging.getLogger("api")
logger.setLevel(logging.INFO)
default_pulsar_host = os.getenv("PULSAR_HOST", "pulsar://pulsar:6650")
default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
default_timeout = 600
default_port = 8088
default_api_token = os.getenv("GATEWAY_SECRET", "")
@ -69,6 +70,7 @@ 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)
api_token = config.get("api_token", default_api_token)
@ -81,49 +83,49 @@ class Api:
self.services = {
"text-completion": TextCompletionRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout,
auth = self.auth,
auth = self.auth, pulsar_api_key=self.pulsar_api_key,
),
"prompt": PromptRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout,
auth = self.auth,
auth = self.auth, pulsar_api_key=self.pulsar_api_key,
),
"graph-rag": GraphRagRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout,
auth = self.auth,
auth = self.auth, pulsar_api_key=self.pulsar_api_key,
),
"triples-query": TriplesQueryRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout,
auth = self.auth,
auth = self.auth, pulsar_api_key=self.pulsar_api_key,
),
"graph-embeddings-query": GraphEmbeddingsQueryRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout,
auth = self.auth,
auth = self.auth, pulsar_api_key=self.pulsar_api_key,
),
"embeddings": EmbeddingsRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout,
auth = self.auth,
auth = self.auth, pulsar_api_key=self.pulsar_api_key,
),
"agent": AgentRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout,
auth = self.auth,
auth = self.auth, pulsar_api_key=self.pulsar_api_key,
),
"encyclopedia": EncyclopediaRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout,
auth = self.auth,
auth = self.auth, pulsar_api_key=self.pulsar_api_key,
),
"dbpedia": DbpediaRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout,
auth = self.auth,
auth = self.auth, pulsar_api_key=self.pulsar_api_key,
),
"internet-search": InternetSearchRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout,
auth = self.auth,
auth = self.auth, pulsar_api_key=self.pulsar_api_key,
),
"document-load": DocumentLoadSender(
pulsar_host=self.pulsar_host,
pulsar_host=self.pulsar_host, pulsar_api_key=self.pulsar_api_key,
),
"text-load": TextLoadSender(
pulsar_host=self.pulsar_host,
pulsar_host=self.pulsar_host, pulsar_api_key=self.pulsar_api_key,
),
}
@ -179,24 +181,29 @@ class Api:
),
TriplesStreamEndpoint(
pulsar_host=self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
auth = self.auth,
),
GraphEmbeddingsStreamEndpoint(
pulsar_host=self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
auth = self.auth,
),
TriplesLoadEndpoint(
pulsar_host=self.pulsar_host,
auth = self.auth,
pulsar_api_key=self.pulsar_api_key,
),
GraphEmbeddingsLoadEndpoint(
pulsar_host=self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
auth = self.auth,
),
MuxEndpoint(
pulsar_host=self.pulsar_host,
auth = self.auth,
services = self.services,
pulsar_api_key=self.pulsar_api_key,
),
]
@ -225,6 +232,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(
'--port',

14
trustgraph-flow/trustgraph/gateway/subscriber.py
View file

@ -6,9 +6,10 @@ import time
class Subscriber:
def __init__(self, pulsar_host, topic, subscription, consumer_name,
def __init__(self, pulsar_host, topic, subscription, consumer_name, pulsar_api_key=None,
schema=None, max_size=100):
self.pulsar_host = pulsar_host
self.pulsar_api_key = pulsar_api_key
self.topic = topic
self.subscription = subscription
self.consumer_name = consumer_name
@ -28,9 +29,16 @@ class Subscriber:
try:
client = pulsar.Client(
if self.pulsar_api_key:
auth = pulsar.AuthenticationToken(self.pulsar_api_key)
client = pulsar.Client(
self.pulsar_host,
)
authentication=auth,
)
else:
client = pulsar.Client(
self.pulsar_host,
)
consumer = client.subscribe(
topic=self.topic,

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

@ -7,10 +7,11 @@ from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
class TextCompletionRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth):
def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
super(TextCompletionRequestor, self).__init__(
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
request_queue=text_completion_request_queue,
response_queue=text_completion_response_queue,
request_schema=TextCompletionRequest,

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

@ -8,10 +8,11 @@ from . sender import ServiceSender
from . serialize import to_subgraph
class TextLoadSender(ServiceSender):
def __init__(self, pulsar_host):
def __init__(self, pulsar_host, pulsar_api_key=None):
super(TextLoadSender, self).__init__(
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
request_queue=text_ingest_queue,
request_schema=TextDocument,
)

6
trustgraph-flow/trustgraph/gateway/triples_load.py
View file

@ -14,17 +14,19 @@ from . serialize import to_subgraph
class TriplesLoadEndpoint(SocketEndpoint):
def __init__(self, pulsar_host, auth, path="/api/v1/load/triples"):
def __init__(self, pulsar_host, auth, path="/api/v1/load/triples", pulsar_api_key=None):
super(TriplesLoadEndpoint, self).__init__(
endpoint_path=path, auth=auth,
)
self.pulsar_host=pulsar_host
self.pulsar_api_key=pulsar_api_key
self.publisher = Publisher(
self.pulsar_host, triples_store_queue,
schema=JsonSchema(Triples)
schema=JsonSchema(Triples),
pulsar_api_key=self.pulsar_api_key
)
async def start(self):

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

@ -8,10 +8,11 @@ from . requestor import ServiceRequestor
from . serialize import to_value, serialize_subgraph
class TriplesQueryRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth):
def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
super(TriplesQueryRequestor, self).__init__(
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
request_queue=triples_request_queue,
response_queue=triples_response_queue,
request_schema=TriplesQueryRequest,

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

@ -13,17 +13,19 @@ from . serialize import serialize_triples
class TriplesStreamEndpoint(SocketEndpoint):
def __init__(self, pulsar_host, auth, path="/api/v1/stream/triples"):
def __init__(self, pulsar_host, auth, path="/api/v1/stream/triples", pulsar_api_key=None):
super(TriplesStreamEndpoint, self).__init__(
endpoint_path=path, auth=auth,
)
self.pulsar_host=pulsar_host
self.pulsar_api_key=pulsar_api_key
self.subscriber = Subscriber(
self.pulsar_host, 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

10
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):
@ -125,12 +126,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',

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,