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

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

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)
)