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

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

21
trustgraph-base/trustgraph/base/base_processor.py
View file

@ -11,6 +11,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):
@ -28,14 +29,22 @@ class BaseProcessor:
})
pulsar_host = params.get("pulsar_host", self.default_pulsar_host)
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,
logger=pulsar.ConsoleLogger(log_level.to_pulsar())
)
)
def __del__(self):
@ -51,6 +60,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(
'-l', '--log-level',

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(

15
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,
logger=pulsar.ConsoleLogger(log_level.to_pulsar())
)
else:
self.client = pulsar.Client(
pulsar_host,
logger=pulsar.ConsoleLogger(log_level.to_pulsar())
)
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,