73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
from typing import Optional
|
|
|
|
from aiokit import AioThing
|
|
from grpc.experimental.aio import insecure_channel
|
|
from idm.api.proto.chat_manager_service_pb2 import Chat as ChatPb
|
|
from nexus.hub.proto.delivery_service_pb2 import \
|
|
StartDeliveryRequest as StartDeliveryRequestPb
|
|
from nexus.hub.proto.delivery_service_pb2 import \
|
|
StartDeliveryResponse as StartDeliveryResponsePb
|
|
from nexus.hub.proto.delivery_service_pb2_grpc import DeliveryStub
|
|
from nexus.hub.proto.submitter_service_pb2 import \
|
|
SubmitRequest as SubmitRequestPb
|
|
from nexus.hub.proto.submitter_service_pb2 import \
|
|
SubmitResponse as SubmitResponsePb
|
|
from nexus.hub.proto.submitter_service_pb2_grpc import SubmitterStub
|
|
from nexus.models.proto.typed_document_pb2 import \
|
|
TypedDocument as TypedDocumentPb
|
|
|
|
|
|
class HubGrpcClient(AioThing):
|
|
def __init__(
|
|
self,
|
|
base_url: str,
|
|
):
|
|
super().__init__()
|
|
self.channel = insecure_channel(base_url, [
|
|
('grpc.dns_min_time_between_resolutions_ms', 1000),
|
|
('grpc.initial_reconnect_backoff_ms', 1000),
|
|
('grpc.lb_policy_name', 'round_robin'),
|
|
('grpc.min_reconnect_backoff_ms', 1000),
|
|
('grpc.max_reconnect_backoff_ms', 2000),
|
|
])
|
|
self.delivery_stub = DeliveryStub(self.channel)
|
|
self.submitter_stub = SubmitterStub(self.channel)
|
|
|
|
async def start(self):
|
|
await self.channel.channel_ready()
|
|
|
|
async def stop(self):
|
|
await self.channel.close()
|
|
|
|
async def start_delivery(
|
|
self,
|
|
typed_document_pb: TypedDocumentPb,
|
|
chat: ChatPb,
|
|
request_id: Optional[str],
|
|
session_id: Optional[str],
|
|
) -> StartDeliveryResponsePb:
|
|
return await self.delivery_stub.start_delivery(
|
|
StartDeliveryRequestPb(
|
|
typed_document=typed_document_pb,
|
|
chat=chat,
|
|
),
|
|
metadata=(('request-id', request_id), ('session-id', session_id))
|
|
)
|
|
|
|
async def submit(
|
|
self,
|
|
telegram_document: bytes,
|
|
telegram_file_id: str,
|
|
chat: ChatPb,
|
|
request_id: Optional[str] = None,
|
|
session_id: Optional[str] = None,
|
|
) -> SubmitResponsePb:
|
|
return await self.submitter_stub.submit(
|
|
SubmitRequestPb(
|
|
telegram_document=telegram_document,
|
|
telegram_file_id=telegram_file_id,
|
|
chat=chat,
|
|
),
|
|
metadata=(('request-id', request_id), ('session-id', session_id))
|
|
)
|