Added missing api_key references

This commit is contained in:
Tyler O 2025-02-11 15:30:53 +00:00
parent 5bdb9c1919
commit 3165bfee8d
5 changed files with 13 additions and 6 deletions

View file

@ -15,7 +15,7 @@ from . serialize import to_subgraph
class DocumentEmbeddingsLoadEndpoint(SocketEndpoint):
def __init__(
self, pulsar_host, auth, path="/api/v1/load/document-embeddings",
self, pulsar_host, auth, path="/api/v1/load/document-embeddings", pulsar_api_key=None
):
super(DocumentEmbeddingsLoadEndpoint, self).__init__(
@ -23,10 +23,12 @@ class DocumentEmbeddingsLoadEndpoint(SocketEndpoint):
)
self.pulsar_host=pulsar_host
self.pulsar_api_key=pulsar_api_key
self.publisher = Publisher(
self.pulsar_host, document_embeddings_store_queue,
schema=JsonSchema(DocumentEmbeddings)
schema=JsonSchema(DocumentEmbeddings),
pulsar_api_key=self.pulsar_api_key
)
async def start(self):

View file

@ -14,7 +14,7 @@ from . serialize import serialize_document_embeddings
class DocumentEmbeddingsStreamEndpoint(SocketEndpoint):
def __init__(
self, pulsar_host, auth, path="/api/v1/stream/document-embeddings"
self, pulsar_host, auth, path="/api/v1/stream/document-embeddings", pulsar_api_key=None
):
super(DocumentEmbeddingsStreamEndpoint, self).__init__(
@ -22,11 +22,13 @@ class DocumentEmbeddingsStreamEndpoint(SocketEndpoint):
)
self.pulsar_host=pulsar_host
self.pulsar_api_key=pulsar_api_key
self.subscriber = Subscriber(
self.pulsar_host, 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):

View file

@ -7,10 +7,11 @@ from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor
class DocumentRagRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth):
def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
super(DocumentRagRequestor, self).__init__(
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
request_queue=document_rag_request_queue,
response_queue=document_rag_response_queue,
request_schema=DocumentRagQuery,

View file

@ -27,7 +27,7 @@ class GraphEmbeddingsLoadEndpoint(SocketEndpoint):
self.publisher = Publisher(
self.pulsar_host, graph_embeddings_store_queue,
self.pulsar_api_key,
pulsar_api_key=self.pulsar_api_key,
schema=JsonSchema(GraphEmbeddings)
)

View file

@ -69,6 +69,7 @@ class Processing:
params = {
"pulsar_host": self.pulsar_host,
"pulsar_api_key": self.pulsar_api_key,
"log_level": str(self.log_level),
}
@ -161,6 +162,7 @@ def run():
try:
p = Processing(
pulsar_host=args.pulsar_host,
pulsar_api_key=args.pulsar_api_key,
file=args.file,
log_level=args.log_level,
)