mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-05-19 12:25:13 +02:00
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:
parent
f7df2df266
commit
617eb7efd5
40 changed files with 173 additions and 21 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue