ritual/projects/gpt4/container/src/app.py
2024-06-06 13:18:48 -04:00

111 lines
3.5 KiB
Python

import logging
import os
from typing import Any, cast
from eth_abi import decode, encode # type: ignore
from infernet_ml.utils.css_mux import (
ConvoMessage,
CSSCompletionParams,
CSSRequest,
Provider,
)
from infernet_ml.utils.service_models import InfernetInput
from infernet_ml.utils.service_models import JobLocation
from infernet_ml.workflows.inference.css_inference_workflow import CSSInferenceWorkflow
from quart import Quart, request
log = logging.getLogger(__name__)
def create_app() -> Quart:
app = Quart(__name__)
workflow = CSSInferenceWorkflow(
api_keys={Provider.OPENAI: os.environ["OPENAI_API_KEY"]}
)
workflow.setup()
@app.route("/")
def index() -> str:
"""
Utility endpoint to check if the service is running.
"""
return "GPT4 Example Program"
@app.route("/service_output", methods=["POST"])
async def inference() -> Any:
req_data = await request.get_json()
"""
InfernetInput has the format:
source: (0 on-chain, 1 off-chain)
data: dict[str, Any]
"""
infernet_input: InfernetInput = InfernetInput(**req_data)
match infernet_input:
case InfernetInput(source=JobLocation.OFFCHAIN):
prompt = cast(dict[str, Any], infernet_input.data).get("prompt")
case InfernetInput(source=JobLocation.ONCHAIN):
# On-chain requests are sent as a generalized hex-string which we will
# decode to the appropriate format.
(prompt,) = decode(
["string"], bytes.fromhex(cast(str, infernet_input.data))
)
case _:
raise ValueError("Invalid source")
result = workflow.inference(
CSSRequest(
provider=Provider.OPENAI,
endpoint="completions",
model="gpt-4-0613",
params=CSSCompletionParams(
messages=[
ConvoMessage(
role="system", content="you are a helpful " "assistant."
),
ConvoMessage(role="user", content=cast(str, prompt)),
]
),
)
)
match infernet_input:
case InfernetInput(destination=JobLocation.OFFCHAIN):
"""
In case of an off-chain request, the result is returned as is.
"""
return {"message": result}
case InfernetInput(destination=JobLocation.ONCHAIN):
"""
In case of an on-chain request, the result is returned in the format:
{
"raw_input": str,
"processed_input": str,
"raw_output": str,
"processed_output": str,
"proof": str,
}
refer to: https://docs.ritual.net/infernet/node/containers for more
info.
"""
return {
"raw_input": "",
"processed_input": "",
"raw_output": encode(["string"], [result]).hex(),
"processed_output": "",
"proof": "",
}
case _:
raise ValueError("Invalid destination")
return app
if __name__ == "__main__":
"""
Utility to run the app locally. For development purposes only.
"""
create_app().run(port=3000)