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


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


app = FastAPI(
    title='CommandRLLMAPI'
)


@app.post("/completion/")
async def get_answer(question: PromptType = None):
    answer = await get_answer_from_llm(question=question.prompt)
    return answer


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