mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-17 09:11:03 +02:00
40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
|
|
from .. schema import TextCompletionRequest, TextCompletionResponse
|
|
from .. schema import text_completion_request_queue
|
|
from .. schema import text_completion_response_queue
|
|
|
|
from . endpoint import ServiceEndpoint
|
|
from . requestor import ServiceRequestor
|
|
|
|
class TextCompletionRequestor(ServiceRequestor):
|
|
def __init__(self, pulsar_host, timeout, auth):
|
|
|
|
super(TextCompletionRequestor, self).__init__(
|
|
pulsar_host=pulsar_host,
|
|
request_queue=text_completion_request_queue,
|
|
response_queue=text_completion_response_queue,
|
|
request_schema=TextCompletionRequest,
|
|
response_schema=TextCompletionResponse,
|
|
timeout=timeout,
|
|
)
|
|
|
|
def to_request(self, body):
|
|
return TextCompletionRequest(
|
|
system=body["system"],
|
|
prompt=body["prompt"]
|
|
)
|
|
|
|
def from_response(self, message):
|
|
return { "response": message.response }
|
|
|
|
class TextCompletionEndpoint(ServiceEndpoint):
|
|
def __init__(self, pulsar_host, timeout, auth):
|
|
|
|
super(TextCompletionEndpoint, self).__init__(
|
|
endpoint_path="/api/v1/text-completion",
|
|
auth=auth,
|
|
requestor = TextCompletionRequestor(
|
|
pulsar_host=pulsar_host, timeout=timeout, auth=auth
|
|
)
|
|
)
|
|
|