EVAL/api/main.py

129 lines
3.4 KiB
Python
Raw Normal View History

2023-03-17 15:55:15 +00:00
import re
from multiprocessing import Process
from tempfile import NamedTemporaryFile
from typing import List, TypedDict
2023-03-17 15:55:15 +00:00
2023-04-03 07:43:34 +00:00
import uvicorn
from fastapi import FastAPI, Request, UploadFile
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from api.container import agent_manager, file_handler, reload_dirs, templates, uploader
from api.worker import get_task_result, start_worker, task_execute
2023-04-03 07:43:34 +00:00
from env import settings
2023-03-17 15:55:15 +00:00
app = FastAPI()
2023-03-18 12:42:14 +00:00
app.mount("/static", StaticFiles(directory=uploader.path), name="static")
2023-03-17 15:55:15 +00:00
class ExecuteRequest(BaseModel):
session: str
prompt: str
2023-03-18 06:05:02 +00:00
files: List[str]
2023-03-17 15:55:15 +00:00
class ExecuteResponse(TypedDict):
answer: str
2023-03-18 06:05:02 +00:00
files: List[str]
2023-03-17 15:55:15 +00:00
@app.get("/", response_class=HTMLResponse)
async def index(request: Request):
return templates.TemplateResponse("index.html", {"request": request})
@app.get("/dashboard", response_class=HTMLResponse)
async def dashboard(request: Request):
return templates.TemplateResponse("dashboard.html", {"request": request})
2023-03-17 15:55:15 +00:00
@app.post("/upload")
async def create_upload_file(files: List[UploadFile]):
urls = []
for file in files:
extension = "." + file.filename.split(".")[-1]
with NamedTemporaryFile(suffix=extension) as tmp_file:
tmp_file.write(file.file.read())
tmp_file.flush()
urls.append(uploader.upload(tmp_file.name))
return {"urls": urls}
@app.post("/api/execute")
async def execute(request: ExecuteRequest) -> ExecuteResponse:
query = request.prompt
2023-03-17 15:55:15 +00:00
files = request.files
session = request.session
2023-03-17 15:55:15 +00:00
executor = agent_manager.create_executor(session)
promptedQuery = "\n".join([file_handler.handle(file) for file in files])
2023-03-18 06:05:02 +00:00
promptedQuery += query
2023-03-17 15:55:15 +00:00
try:
res = executor({"input": promptedQuery})
2023-03-17 15:55:15 +00:00
except Exception as e:
return {"answer": str(e), "files": []}
files = re.findall(r"\[file://\S*\]", res["output"])
files = [file[1:-1] for file in files]
2023-03-17 15:55:15 +00:00
return {
"answer": res["output"],
"files": [uploader.upload(file) for file in files],
2023-03-17 15:55:15 +00:00
}
@app.post("/api/execute/async")
async def execute_async(request: ExecuteRequest):
query = request.prompt
files = request.files
session = request.session
promptedQuery = "\n".join([file_handler.handle(file) for file in files])
promptedQuery += query
execution = task_execute.delay(session, promptedQuery)
return {"id": execution.id}
@app.get("/api/execute/async/{execution_id}")
async def execute_async(execution_id: str):
execution = get_task_result(execution_id)
result = {}
if execution.status == "SUCCESS" and execution.result:
output = execution.result.get("output", "")
files = re.findall(r"\[file://\S*\]", output)
files = [file[1:-1] for file in files]
result = {
"answer": output,
"files": [uploader.upload(file) for file in files],
}
return {
"status": execution.status,
"info": execution.info,
"result": result,
}
def serve():
p = Process(target=start_worker, args=[])
p.start()
uvicorn.run("api.main:app", host="0.0.0.0", port=settings["EVAL_PORT"])
def dev():
p = Process(target=start_worker, args=[])
p.start()
uvicorn.run(
"api.main:app",
host="0.0.0.0",
port=settings["EVAL_PORT"],
reload=True,
reload_dirs=reload_dirs,
)