- Pass single Pulsar client to everything in gateway & librarian

- Pulsar listener-name support in gateway
- PDF and text load working in librarian
This commit is contained in:
Cyber MacGeddon 2025-02-12 21:24:19 +00:00
parent b5f51731bd
commit d8abbf3d0b
28 changed files with 111 additions and 108 deletions

View file

@ -6,14 +6,13 @@ import threading
class Publisher: class Publisher:
def __init__(self, pulsar_host, topic, schema=None, max_size=10, def __init__(self, pulsar_client, topic, schema=None, max_size=10,
chunking_enabled=True, listener=None): chunking_enabled=True):
self.pulsar_host = pulsar_host self.client = pulsar_client
self.topic = topic self.topic = topic
self.schema = schema self.schema = schema
self.q = queue.Queue(maxsize=max_size) self.q = queue.Queue(maxsize=max_size)
self.chunking_enabled = chunking_enabled self.chunking_enabled = chunking_enabled
self.listener_name = listener
self.running = True self.running = True
def start(self): def start(self):
@ -33,11 +32,8 @@ class Publisher:
try: try:
client = pulsar.Client( print(self.chunking_enabled)
self.pulsar_host, listener_name=self.listener_name producer = self.client.create_producer(
)
producer = client.create_producer(
topic=self.topic, topic=self.topic,
schema=self.schema, schema=self.schema,
chunking_enabled=self.chunking_enabled, chunking_enabled=self.chunking_enabled,

View file

@ -6,9 +6,9 @@ import time
class Subscriber: class Subscriber:
def __init__(self, pulsar_host, topic, subscription, consumer_name, def __init__(self, pulsar_client, topic, subscription, consumer_name,
schema=None, max_size=100, listener=None): schema=None, max_size=100):
self.pulsar_host = pulsar_host self.client = pulsar_client
self.topic = topic self.topic = topic
self.subscription = subscription self.subscription = subscription
self.consumer_name = consumer_name self.consumer_name = consumer_name
@ -17,7 +17,6 @@ class Subscriber:
self.full = {} self.full = {}
self.max_size = max_size self.max_size = max_size
self.lock = threading.Lock() self.lock = threading.Lock()
self.listener_name = listener
self.running = True self.running = True
def start(self): def start(self):
@ -36,12 +35,7 @@ class Subscriber:
try: try:
client = pulsar.Client( consumer = self.client.subscribe(
self.pulsar_host,
listener_name=self.listener_name,
)
consumer = client.subscribe(
topic=self.topic, topic=self.topic,
subscription_name=self.subscription, subscription_name=self.subscription,
consumer_name=self.consumer_name, consumer_name=self.consumer_name,

View file

@ -7,10 +7,10 @@ from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor from . requestor import ServiceRequestor
class AgentRequestor(ServiceRequestor): class AgentRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth): def __init__(self, pulsar_client, timeout, auth):
super(AgentRequestor, self).__init__( super(AgentRequestor, self).__init__(
pulsar_host=pulsar_host, pulsar_client=pulsar_client,
request_queue=agent_request_queue, request_queue=agent_request_queue,
response_queue=agent_response_queue, response_queue=agent_response_queue,
request_schema=AgentRequest, request_schema=AgentRequest,

View file

@ -7,10 +7,10 @@ from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor from . requestor import ServiceRequestor
class DbpediaRequestor(ServiceRequestor): class DbpediaRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth): def __init__(self, pulsar_client, timeout, auth):
super(DbpediaRequestor, self).__init__( super(DbpediaRequestor, self).__init__(
pulsar_host=pulsar_host, pulsar_client=pulsar_client,
request_queue=dbpedia_lookup_request_queue, request_queue=dbpedia_lookup_request_queue,
response_queue=dbpedia_lookup_response_queue, response_queue=dbpedia_lookup_response_queue,
request_schema=LookupRequest, request_schema=LookupRequest,

View file

@ -15,17 +15,17 @@ 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_client, auth, path="/api/v1/load/document-embeddings",
): ):
super(DocumentEmbeddingsLoadEndpoint, self).__init__( super(DocumentEmbeddingsLoadEndpoint, self).__init__(
endpoint_path=path, auth=auth, endpoint_path=path, auth=auth,
) )
self.pulsar_host=pulsar_host self.pulsar_client=pulsar_client
self.publisher = Publisher( self.publisher = Publisher(
self.pulsar_host, document_embeddings_store_queue, self.pulsar_client, document_embeddings_store_queue,
schema=JsonSchema(DocumentEmbeddings) schema=JsonSchema(DocumentEmbeddings)
) )

View file

@ -14,17 +14,18 @@ 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_client, auth,
path="/api/v1/stream/document-embeddings"
): ):
super(DocumentEmbeddingsStreamEndpoint, self).__init__( super(DocumentEmbeddingsStreamEndpoint, self).__init__(
endpoint_path=path, auth=auth, endpoint_path=path, auth=auth,
) )
self.pulsar_host=pulsar_host self.pulsar_client=pulsar_client
self.subscriber = Subscriber( self.subscriber = Subscriber(
self.pulsar_host, document_embeddings_store_queue, self.pulsar_client, document_embeddings_store_queue,
"api-gateway", "api-gateway", "api-gateway", "api-gateway",
schema=JsonSchema(DocumentEmbeddings) schema=JsonSchema(DocumentEmbeddings)
) )

View file

@ -8,10 +8,10 @@ from . sender import ServiceSender
from . serialize import to_subgraph from . serialize import to_subgraph
class DocumentLoadSender(ServiceSender): class DocumentLoadSender(ServiceSender):
def __init__(self, pulsar_host): def __init__(self, pulsar_client):
super(DocumentLoadSender, self).__init__( super(DocumentLoadSender, self).__init__(
pulsar_host=pulsar_host, pulsar_client=pulsar_client,
request_queue=document_ingest_queue, request_queue=document_ingest_queue,
request_schema=Document, request_schema=Document,
) )

View file

@ -7,10 +7,10 @@ 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_client, timeout, auth):
super(DocumentRagRequestor, self).__init__( super(DocumentRagRequestor, self).__init__(
pulsar_host=pulsar_host, pulsar_client=pulsar_client,
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

@ -7,10 +7,10 @@ from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor from . requestor import ServiceRequestor
class EmbeddingsRequestor(ServiceRequestor): class EmbeddingsRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth): def __init__(self, pulsar_client, timeout, auth):
super(EmbeddingsRequestor, self).__init__( super(EmbeddingsRequestor, self).__init__(
pulsar_host=pulsar_host, pulsar_client=pulsar_client,
request_queue=embeddings_request_queue, request_queue=embeddings_request_queue,
response_queue=embeddings_response_queue, response_queue=embeddings_response_queue,
request_schema=EmbeddingsRequest, request_schema=EmbeddingsRequest,

View file

@ -7,10 +7,10 @@ from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor from . requestor import ServiceRequestor
class EncyclopediaRequestor(ServiceRequestor): class EncyclopediaRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth): def __init__(self, pulsar_client, timeout, auth):
super(EncyclopediaRequestor, self).__init__( super(EncyclopediaRequestor, self).__init__(
pulsar_host=pulsar_host, pulsar_client=pulsar_client,
request_queue=encyclopedia_lookup_request_queue, request_queue=encyclopedia_lookup_request_queue,
response_queue=encyclopedia_lookup_response_queue, response_queue=encyclopedia_lookup_response_queue,
request_schema=LookupRequest, request_schema=LookupRequest,

View file

@ -15,17 +15,17 @@ from . serialize import to_subgraph, to_value
class GraphEmbeddingsLoadEndpoint(SocketEndpoint): class GraphEmbeddingsLoadEndpoint(SocketEndpoint):
def __init__( def __init__(
self, pulsar_host, auth, path="/api/v1/load/graph-embeddings", self, pulsar_client, auth, path="/api/v1/load/graph-embeddings",
): ):
super(GraphEmbeddingsLoadEndpoint, self).__init__( super(GraphEmbeddingsLoadEndpoint, self).__init__(
endpoint_path=path, auth=auth, endpoint_path=path, auth=auth,
) )
self.pulsar_host=pulsar_host self.pulsar_client=pulsar_client
self.publisher = Publisher( self.publisher = Publisher(
self.pulsar_host, graph_embeddings_store_queue, self.pulsar_client, graph_embeddings_store_queue,
schema=JsonSchema(GraphEmbeddings) schema=JsonSchema(GraphEmbeddings)
) )

View file

@ -8,10 +8,10 @@ from . requestor import ServiceRequestor
from . serialize import serialize_value from . serialize import serialize_value
class GraphEmbeddingsQueryRequestor(ServiceRequestor): class GraphEmbeddingsQueryRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth): def __init__(self, pulsar_client, timeout, auth):
super(GraphEmbeddingsQueryRequestor, self).__init__( super(GraphEmbeddingsQueryRequestor, self).__init__(
pulsar_host=pulsar_host, pulsar_client=pulsar_client,
request_queue=graph_embeddings_request_queue, request_queue=graph_embeddings_request_queue,
response_queue=graph_embeddings_response_queue, response_queue=graph_embeddings_response_queue,
request_schema=GraphEmbeddingsRequest, request_schema=GraphEmbeddingsRequest,

View file

@ -14,17 +14,17 @@ from . serialize import serialize_graph_embeddings
class GraphEmbeddingsStreamEndpoint(SocketEndpoint): class GraphEmbeddingsStreamEndpoint(SocketEndpoint):
def __init__( def __init__(
self, pulsar_host, auth, path="/api/v1/stream/graph-embeddings" self, pulsar_client, auth, path="/api/v1/stream/graph-embeddings"
): ):
super(GraphEmbeddingsStreamEndpoint, self).__init__( super(GraphEmbeddingsStreamEndpoint, self).__init__(
endpoint_path=path, auth=auth, endpoint_path=path, auth=auth,
) )
self.pulsar_host=pulsar_host self.pulsar_client=pulsar_client
self.subscriber = Subscriber( self.subscriber = Subscriber(
self.pulsar_host, graph_embeddings_store_queue, self.pulsar_client, graph_embeddings_store_queue,
"api-gateway", "api-gateway", "api-gateway", "api-gateway",
schema=JsonSchema(GraphEmbeddings) schema=JsonSchema(GraphEmbeddings)
) )

View file

@ -7,10 +7,10 @@ from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor from . requestor import ServiceRequestor
class GraphRagRequestor(ServiceRequestor): class GraphRagRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth): def __init__(self, pulsar_client, timeout, auth):
super(GraphRagRequestor, self).__init__( super(GraphRagRequestor, self).__init__(
pulsar_host=pulsar_host, pulsar_client=pulsar_client,
request_queue=graph_rag_request_queue, request_queue=graph_rag_request_queue,
response_queue=graph_rag_response_queue, response_queue=graph_rag_response_queue,
request_schema=GraphRagQuery, request_schema=GraphRagQuery,

View file

@ -7,10 +7,10 @@ from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor from . requestor import ServiceRequestor
class InternetSearchRequestor(ServiceRequestor): class InternetSearchRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth): def __init__(self, pulsar_client, timeout, auth):
super(InternetSearchRequestor, self).__init__( super(InternetSearchRequestor, self).__init__(
pulsar_host=pulsar_host, pulsar_client=pulsar_client,
request_queue=internet_search_request_queue, request_queue=internet_search_request_queue,
response_queue=internet_search_response_queue, response_queue=internet_search_response_queue,
request_schema=LookupRequest, request_schema=LookupRequest,

View file

@ -9,10 +9,10 @@ from . serialize import serialize_document_package, serialize_document_info
from . serialize import to_document_package, to_document_info, to_criteria from . serialize import to_document_package, to_document_info, to_criteria
class LibrarianRequestor(ServiceRequestor): class LibrarianRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth): def __init__(self, pulsar_client, timeout, auth):
super(LibrarianRequestor, self).__init__( super(LibrarianRequestor, self).__init__(
pulsar_host=pulsar_host, pulsar_client=pulsar_client,
request_queue=librarian_request_queue, request_queue=librarian_request_queue,
response_queue=librarian_response_queue, response_queue=librarian_response_queue,
request_schema=LibrarianRequest, request_schema=LibrarianRequest,
@ -22,17 +22,19 @@ class LibrarianRequestor(ServiceRequestor):
def to_request(self, body): def to_request(self, body):
print("TRR")
if "document" in body: if "document" in body:
dp = to_document_package(body["document"]) dp = to_document_package(body["document"])
else: else:
dp = None dp = None
print("GOT")
if "criteria" in body: if "criteria" in body:
criteria = to_criteria(body["criteria"]) criteria = to_criteria(body["criteria"])
else: else:
criteria = None criteria = None
limit = int(body.get("limit", 10000)) print("ASLDKJ")
return LibrarianRequest( return LibrarianRequest(
operation = body.get("operation", None), operation = body.get("operation", None),

View file

@ -18,7 +18,7 @@ MAX_QUEUE_SIZE = 10
class MuxEndpoint(SocketEndpoint): class MuxEndpoint(SocketEndpoint):
def __init__( def __init__(
self, pulsar_host, auth, self, pulsar_client, auth,
services, services,
path="/api/v1/socket", path="/api/v1/socket",
): ):

View file

@ -9,10 +9,10 @@ from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor from . requestor import ServiceRequestor
class PromptRequestor(ServiceRequestor): class PromptRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth): def __init__(self, pulsar_client, timeout, auth):
super(PromptRequestor, self).__init__( super(PromptRequestor, self).__init__(
pulsar_host=pulsar_host, pulsar_client=pulsar_client,
request_queue=prompt_request_queue, request_queue=prompt_request_queue,
response_queue=prompt_response_queue, response_queue=prompt_response_queue,
request_schema=PromptRequest, request_schema=PromptRequest,

View file

@ -14,7 +14,7 @@ class ServiceRequestor:
def __init__( def __init__(
self, self,
pulsar_host, pulsar_client,
request_queue, request_schema, request_queue, request_schema,
response_queue, response_schema, response_queue, response_schema,
subscription="api-gateway", consumer_name="api-gateway", subscription="api-gateway", consumer_name="api-gateway",
@ -22,12 +22,12 @@ class ServiceRequestor:
): ):
self.pub = Publisher( self.pub = Publisher(
pulsar_host, request_queue, pulsar_client, request_queue,
schema=JsonSchema(request_schema) schema=JsonSchema(request_schema),
) )
self.sub = Subscriber( self.sub = Subscriber(
pulsar_host, response_queue, pulsar_client, response_queue,
subscription, consumer_name, subscription, consumer_name,
JsonSchema(response_schema) JsonSchema(response_schema)
) )
@ -53,9 +53,11 @@ class ServiceRequestor:
q = self.sub.subscribe(id) q = self.sub.subscribe(id)
print("BOUT TO SEDN")
await asyncio.to_thread( await asyncio.to_thread(
self.pub.send, id, self.to_request(request) self.pub.send, id, self.to_request(request)
) )
print("SENT")
while True: while True:

View file

@ -15,13 +15,13 @@ class ServiceSender:
def __init__( def __init__(
self, self,
pulsar_host, pulsar_client,
request_queue, request_schema, request_queue, request_schema,
): ):
self.pub = Publisher( self.pub = Publisher(
pulsar_host, request_queue, pulsar_client, request_queue,
schema=JsonSchema(request_schema) schema=JsonSchema(request_schema),
) )
async def start(self): async def start(self):
@ -53,4 +53,3 @@ class ServiceSender:
return err return err

View file

@ -126,7 +126,7 @@ def to_document_package(x):
return DocumentPackage( return DocumentPackage(
metadata = to_subgraph(x["metadata"]), metadata = to_subgraph(x["metadata"]),
document = base64.b64decode(x["document"].encode("utf-8")), document = x.get("document", None),
kind = x.get("kind", None), kind = x.get("kind", None),
user = x.get("user", None), user = x.get("user", None),
collection = x.get("collection", None), collection = x.get("collection", None),

View file

@ -73,6 +73,11 @@ class Api:
self.port = int(config.get("port", default_port)) self.port = int(config.get("port", default_port))
self.timeout = int(config.get("timeout", default_timeout)) self.timeout = int(config.get("timeout", default_timeout))
self.pulsar_host = config.get("pulsar_host", default_pulsar_host) self.pulsar_host = config.get("pulsar_host", default_pulsar_host)
self.pulsar_listener = config.get("pulsar_listener", None)
self.pulsar_client = pulsar.Client(
self.pulsar_host, listener_name=self.pulsar_listener
)
self.prometheus_url = config.get( self.prometheus_url = config.get(
"prometheus_url", default_prometheus_url, "prometheus_url", default_prometheus_url,
@ -91,58 +96,58 @@ class Api:
self.services = { self.services = {
"text-completion": TextCompletionRequestor( "text-completion": TextCompletionRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout, pulsar_client=self.pulsar_client, timeout=self.timeout,
auth = self.auth, auth = self.auth,
), ),
"prompt": PromptRequestor( "prompt": PromptRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout, pulsar_client=self.pulsar_client, timeout=self.timeout,
auth = self.auth, auth = self.auth,
), ),
"graph-rag": GraphRagRequestor( "graph-rag": GraphRagRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout, pulsar_client=self.pulsar_client, timeout=self.timeout,
auth = self.auth, auth = self.auth,
), ),
"document-rag": DocumentRagRequestor( "document-rag": DocumentRagRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout, pulsar_client=self.pulsar_client, timeout=self.timeout,
auth = self.auth, auth = self.auth,
), ),
"triples-query": TriplesQueryRequestor( "triples-query": TriplesQueryRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout, pulsar_client=self.pulsar_client, timeout=self.timeout,
auth = self.auth, auth = self.auth,
), ),
"graph-embeddings-query": GraphEmbeddingsQueryRequestor( "graph-embeddings-query": GraphEmbeddingsQueryRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout, pulsar_client=self.pulsar_client, timeout=self.timeout,
auth = self.auth, auth = self.auth,
), ),
"embeddings": EmbeddingsRequestor( "embeddings": EmbeddingsRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout, pulsar_client=self.pulsar_client, timeout=self.timeout,
auth = self.auth, auth = self.auth,
), ),
"agent": AgentRequestor( "agent": AgentRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout, pulsar_client=self.pulsar_client, timeout=self.timeout,
auth = self.auth, auth = self.auth,
), ),
"librarian": LibrarianRequestor( "librarian": LibrarianRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout, pulsar_client=self.pulsar_client, timeout=self.timeout,
auth = self.auth, auth = self.auth,
), ),
"encyclopedia": EncyclopediaRequestor( "encyclopedia": EncyclopediaRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout, pulsar_client=self.pulsar_client, timeout=self.timeout,
auth = self.auth, auth = self.auth,
), ),
"dbpedia": DbpediaRequestor( "dbpedia": DbpediaRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout, pulsar_client=self.pulsar_client, timeout=self.timeout,
auth = self.auth, auth = self.auth,
), ),
"internet-search": InternetSearchRequestor( "internet-search": InternetSearchRequestor(
pulsar_host=self.pulsar_host, timeout=self.timeout, pulsar_client=self.pulsar_client, timeout=self.timeout,
auth = self.auth, auth = self.auth,
), ),
"document-load": DocumentLoadSender( "document-load": DocumentLoadSender(
pulsar_host=self.pulsar_host, pulsar_client=self.pulsar_client,
), ),
"text-load": TextLoadSender( "text-load": TextLoadSender(
pulsar_host=self.pulsar_host, pulsar_client=self.pulsar_client,
), ),
} }
@ -205,31 +210,31 @@ class Api:
requestor = self.services["text-load"], requestor = self.services["text-load"],
), ),
TriplesStreamEndpoint( TriplesStreamEndpoint(
pulsar_host=self.pulsar_host, pulsar_client=self.pulsar_client,
auth = self.auth, auth = self.auth,
), ),
GraphEmbeddingsStreamEndpoint( GraphEmbeddingsStreamEndpoint(
pulsar_host=self.pulsar_host, pulsar_client=self.pulsar_client,
auth = self.auth, auth = self.auth,
), ),
DocumentEmbeddingsStreamEndpoint( DocumentEmbeddingsStreamEndpoint(
pulsar_host=self.pulsar_host, pulsar_client=self.pulsar_client,
auth = self.auth, auth = self.auth,
), ),
TriplesLoadEndpoint( TriplesLoadEndpoint(
pulsar_host=self.pulsar_host, pulsar_client=self.pulsar_client,
auth = self.auth, auth = self.auth,
), ),
GraphEmbeddingsLoadEndpoint( GraphEmbeddingsLoadEndpoint(
pulsar_host=self.pulsar_host, pulsar_client=self.pulsar_client,
auth = self.auth, auth = self.auth,
), ),
DocumentEmbeddingsLoadEndpoint( DocumentEmbeddingsLoadEndpoint(
pulsar_host=self.pulsar_host, pulsar_client=self.pulsar_client,
auth = self.auth, auth = self.auth,
), ),
MuxEndpoint( MuxEndpoint(
pulsar_host=self.pulsar_host, pulsar_client=self.pulsar_client,
auth = self.auth, auth = self.auth,
services = self.services, services = self.services,
), ),
@ -266,6 +271,11 @@ def run():
help=f'Pulsar host (default: {default_pulsar_host})', help=f'Pulsar host (default: {default_pulsar_host})',
) )
parser.add_argument(
'--pulsar-listener',
help=f'Pulsar listener (default: none)',
)
parser.add_argument( parser.add_argument(
'-m', '--prometheus-url', '-m', '--prometheus-url',
default=default_prometheus_url, default=default_prometheus_url,

View file

@ -7,10 +7,10 @@ from . endpoint import ServiceEndpoint
from . requestor import ServiceRequestor from . requestor import ServiceRequestor
class TextCompletionRequestor(ServiceRequestor): class TextCompletionRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth): def __init__(self, pulsar_client, timeout, auth):
super(TextCompletionRequestor, self).__init__( super(TextCompletionRequestor, self).__init__(
pulsar_host=pulsar_host, pulsar_client=pulsar_client,
request_queue=text_completion_request_queue, request_queue=text_completion_request_queue,
response_queue=text_completion_response_queue, response_queue=text_completion_response_queue,
request_schema=TextCompletionRequest, request_schema=TextCompletionRequest,

View file

@ -8,10 +8,10 @@ from . sender import ServiceSender
from . serialize import to_subgraph from . serialize import to_subgraph
class TextLoadSender(ServiceSender): class TextLoadSender(ServiceSender):
def __init__(self, pulsar_host): def __init__(self, pulsar_client):
super(TextLoadSender, self).__init__( super(TextLoadSender, self).__init__(
pulsar_host=pulsar_host, pulsar_client=pulsar_client,
request_queue=text_ingest_queue, request_queue=text_ingest_queue,
request_schema=TextDocument, request_schema=TextDocument,
) )

View file

@ -14,16 +14,16 @@ from . serialize import to_subgraph
class TriplesLoadEndpoint(SocketEndpoint): class TriplesLoadEndpoint(SocketEndpoint):
def __init__(self, pulsar_host, auth, path="/api/v1/load/triples"): def __init__(self, pulsar_client, auth, path="/api/v1/load/triples"):
super(TriplesLoadEndpoint, self).__init__( super(TriplesLoadEndpoint, self).__init__(
endpoint_path=path, auth=auth, endpoint_path=path, auth=auth,
) )
self.pulsar_host=pulsar_host self.pulsar_client=pulsar_client
self.publisher = Publisher( self.publisher = Publisher(
self.pulsar_host, triples_store_queue, self.pulsar_client, triples_store_queue,
schema=JsonSchema(Triples) schema=JsonSchema(Triples)
) )

View file

@ -8,10 +8,10 @@ from . requestor import ServiceRequestor
from . serialize import to_value, serialize_subgraph from . serialize import to_value, serialize_subgraph
class TriplesQueryRequestor(ServiceRequestor): class TriplesQueryRequestor(ServiceRequestor):
def __init__(self, pulsar_host, timeout, auth): def __init__(self, pulsar_client, timeout, auth):
super(TriplesQueryRequestor, self).__init__( super(TriplesQueryRequestor, self).__init__(
pulsar_host=pulsar_host, pulsar_client=pulsar_client,
request_queue=triples_request_queue, request_queue=triples_request_queue,
response_queue=triples_response_queue, response_queue=triples_response_queue,
request_schema=TriplesQueryRequest, request_schema=TriplesQueryRequest,

View file

@ -13,16 +13,16 @@ from . serialize import serialize_triples
class TriplesStreamEndpoint(SocketEndpoint): class TriplesStreamEndpoint(SocketEndpoint):
def __init__(self, pulsar_host, auth, path="/api/v1/stream/triples"): def __init__(self, pulsar_client, auth, path="/api/v1/stream/triples"):
super(TriplesStreamEndpoint, self).__init__( super(TriplesStreamEndpoint, self).__init__(
endpoint_path=path, auth=auth, endpoint_path=path, auth=auth,
) )
self.pulsar_host=pulsar_host self.pulsar_client=pulsar_client
self.subscriber = Subscriber( self.subscriber = Subscriber(
self.pulsar_host, triples_store_queue, self.pulsar_client, triples_store_queue,
"api-gateway", "api-gateway", "api-gateway", "api-gateway",
schema=JsonSchema(Triples) schema=JsonSchema(Triples)
) )

View file

@ -7,6 +7,7 @@ from functools import partial
import asyncio import asyncio
import threading import threading
import queue import queue
import base64
from pulsar.schema import JsonSchema from pulsar.schema import JsonSchema
@ -94,32 +95,27 @@ class Processor(ConsumerProducer):
) )
self.document_load = Publisher( self.document_load = Publisher(
self.pulsar_host, document_load_queue, JsonSchema(Document), self.client, document_load_queue, JsonSchema(Document),
listener=self.pulsar_listener,
) )
self.text_load = Publisher( self.text_load = Publisher(
self.pulsar_host, text_load_queue, JsonSchema(TextDocument), self.client, text_load_queue, JsonSchema(TextDocument),
listener=self.pulsar_listener,
) )
self.triples_brk = Subscriber( self.triples_brk = Subscriber(
self.pulsar_host, triples_store_queue, self.client, triples_store_queue,
"librarian", "librarian", "librarian", "librarian",
schema=JsonSchema(Triples), schema=JsonSchema(Triples),
listener=self.pulsar_listener,
) )
self.graph_embeddings_brk = Subscriber( self.graph_embeddings_brk = Subscriber(
self.pulsar_host, graph_embeddings_store_queue, self.client, graph_embeddings_store_queue,
"librarian", "librarian", "librarian", "librarian",
schema=JsonSchema(GraphEmbeddings), schema=JsonSchema(GraphEmbeddings),
listener=self.pulsar_listener,
) )
self.document_embeddings_brk = Subscriber( self.document_embeddings_brk = Subscriber(
self.pulsar_host, document_embeddings_store_queue, self.client, document_embeddings_store_queue,
"librarian", "librarian", "librarian", "librarian",
schema=JsonSchema(DocumentEmbeddings), schema=JsonSchema(DocumentEmbeddings),
listener=self.pulsar_listener,
) )
self.triples_reader = threading.Thread( self.triples_reader = threading.Thread(
@ -243,6 +239,9 @@ class Processor(ConsumerProducer):
async def load_text(self, id, document): async def load_text(self, id, document):
text = base64.b64decode(document.document)
text = text.decode("utf-8")
doc = TextDocument( doc = TextDocument(
metadata = Metadata( metadata = Metadata(
id = id, id = id,
@ -250,7 +249,7 @@ class Processor(ConsumerProducer):
user = document.user, user = document.user,
collection = document.collection collection = document.collection
), ),
text = document.document text = text,
) )
self.text_load.send(None, doc) self.text_load.send(None, doc)