Feature/pulsar api key support (#308)

* 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

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(

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,

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

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

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

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

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

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

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

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,