apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-04-26 00:46:22 +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,