import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel
from model import get_answer_from_llm


class Prompt(BaseModel):
    prompt: str = ''


app = FastAPI(
    title='CommandRLLMAPI'
)


@app.post("/completion/")
def get_answer(question: Prompt.prompt):
    answer = get_answer_from_llm(question)
    return answer


if __name__ == '__main__':
    uvicorn.run(
        app,
        host='0.0.0.0',
        port=8081
    )