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): class DocumentEmbeddingsLoadEndpoint(SocketEndpoint):
def __init__( 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__( super(DocumentEmbeddingsLoadEndpoint, self).__init__(
@ -23,10 +23,12 @@ class DocumentEmbeddingsLoadEndpoint(SocketEndpoint):
) )
self.pulsar_host=pulsar_host self.pulsar_host=pulsar_host
self.pulsar_api_key=pulsar_api_key
self.publisher = Publisher( self.publisher = Publisher(
self.pulsar_host, document_embeddings_store_queue, self.pulsar_host, document_embeddings_store_queue,
schema=JsonSchema(DocumentEmbeddings) schema=JsonSchema(DocumentEmbeddings),
pulsar_api_key=self.pulsar_api_key
) )
async def start(self): async def start(self):

View file

@ -14,7 +14,7 @@ from . serialize import serialize_document_embeddings
class DocumentEmbeddingsStreamEndpoint(SocketEndpoint): class DocumentEmbeddingsStreamEndpoint(SocketEndpoint):
def __init__( 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__( super(DocumentEmbeddingsStreamEndpoint, self).__init__(
@ -22,11 +22,13 @@ class DocumentEmbeddingsStreamEndpoint(SocketEndpoint):
) )
self.pulsar_host=pulsar_host self.pulsar_host=pulsar_host
self.pulsar_api_key=pulsar_api_key
self.subscriber = Subscriber( self.subscriber = Subscriber(
self.pulsar_host, document_embeddings_store_queue, self.pulsar_host, document_embeddings_store_queue,
"api-gateway", "api-gateway", "api-gateway", "api-gateway",
schema=JsonSchema(DocumentEmbeddings) schema=JsonSchema(DocumentEmbeddings),
pulsar_api_key=self.pulsar_api_key
) )
async def listener(self, ws, running): async def listener(self, ws, running):

View file

@ -7,10 +7,11 @@ from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor from . requestor import ServiceRequestor
class DocumentRagRequestor(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__( super(DocumentRagRequestor, self).__init__(
pulsar_host=pulsar_host, pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
request_queue=document_rag_request_queue, request_queue=document_rag_request_queue,
response_queue=document_rag_response_queue, response_queue=document_rag_response_queue,
request_schema=DocumentRagQuery, request_schema=DocumentRagQuery,

View file

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

View file

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