Kye
2 years ago
parent
14d48289fe
commit
f381962d91
@ -1,130 +1,61 @@
|
||||
import os
|
||||
import re
|
||||
from multiprocessing import Process
|
||||
from tempfile import NamedTemporaryFile
|
||||
from pathlib import Path
|
||||
from typing import Dict, List
|
||||
|
||||
from typing import List, TypedDict
|
||||
import uvicorn
|
||||
from fastapi import FastAPI, Request, UploadFile
|
||||
from fastapi.responses import HTMLResponse
|
||||
from fastapi.templating import Jinja2Templates
|
||||
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from pydantic import BaseModel
|
||||
from swarms import Swarms
|
||||
from swarms.utils.utils import BaseHandler, FileHandler, FileType, StaticUploader, CsvToDataframe
|
||||
|
||||
from api.container import agent_manager, file_handler, reload_dirs, templates, uploader
|
||||
from api.worker import get_task_result, start_worker, task_execute
|
||||
from swarms.tools.main import BaseToolSet, ExitConversation, RequestsGet, CodeEditor, Terminal
|
||||
|
||||
|
||||
app = FastAPI()
|
||||
BASE_DIR = Path(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
os.chdir(BASE_DIR / os.getenv("PLAYGROUND_DIR"))
|
||||
|
||||
app.mount("/static", StaticFiles(directory=uploader.path), name="static")
|
||||
api_key = os.getenv("OPENAI_API_KEY")
|
||||
|
||||
toolsets: List[BaseToolSet] = [
|
||||
Terminal(),
|
||||
CodeEditor(),
|
||||
RequestsGet(),
|
||||
ExitConversation(),
|
||||
]
|
||||
handlers: Dict[FileType, BaseHandler] = {FileType.DATAFRAME: CsvToDataframe()}
|
||||
|
||||
class ExecuteRequest(BaseModel):
|
||||
session: str
|
||||
prompt: str
|
||||
files: List[str]
|
||||
if os.getenv("USE_GPU") == "True":
|
||||
import torch
|
||||
|
||||
from swarms.tools.main import ImageCaptioning
|
||||
from swarms.tools.main import (
|
||||
ImageEditing,
|
||||
InstructPix2Pix,
|
||||
Text2Image,
|
||||
VisualQuestionAnswering,
|
||||
)
|
||||
|
||||
class ExecuteResponse(TypedDict):
|
||||
answer: str
|
||||
files: List[str]
|
||||
if torch.cuda.is_available():
|
||||
toolsets.extend(
|
||||
[
|
||||
Text2Image("cuda"),
|
||||
ImageEditing("cuda"),
|
||||
InstructPix2Pix("cuda"),
|
||||
VisualQuestionAnswering("cuda"),
|
||||
]
|
||||
)
|
||||
handlers[FileType.IMAGE] = ImageCaptioning("cuda")
|
||||
|
||||
swarms = Swarms(api_key)
|
||||
|
||||
@app.get("/", response_class=HTMLResponse)
|
||||
async def index(request: Request):
|
||||
return templates.TemplateResponse("index.html", {"request": request})
|
||||
file_handler = FileHandler(handlers=handlers, path=BASE_DIR)
|
||||
|
||||
templates = Jinja2Templates(directory=BASE_DIR / "api" / "templates")
|
||||
|
||||
@app.get("/dashboard", response_class=HTMLResponse)
|
||||
async def dashboard(request: Request):
|
||||
return templates.TemplateResponse("dashboard.html", {"request": request})
|
||||
uploader = StaticUploader(
|
||||
static_dir=BASE_DIR / "static",
|
||||
endpoint="static",
|
||||
public_url=os.getenv("PUBLIC_URL")
|
||||
)
|
||||
|
||||
|
||||
@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
|
||||
files = request.files
|
||||
session = request.session
|
||||
|
||||
executor = agent_manager.create_executor(session)
|
||||
|
||||
promptedQuery = "\n".join([file_handler.handle(file) for file in files])
|
||||
promptedQuery += query
|
||||
|
||||
try:
|
||||
res = executor({"input": promptedQuery})
|
||||
except Exception as e:
|
||||
return {"answer": str(e), "files": []}
|
||||
|
||||
files = re.findall(r"\[file://\S*\]", res["output"])
|
||||
files = [file[1:-1].split("file://")[1] for file in files]
|
||||
|
||||
return {
|
||||
"answer": res["output"],
|
||||
"files": [uploader.upload(file) for file in files],
|
||||
}
|
||||
|
||||
|
||||
@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].split("file://")[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=os.getenv["EVAL_PORT"])
|
||||
|
||||
|
||||
def dev():
|
||||
p = Process(target=start_worker, args=[])
|
||||
p.start()
|
||||
uvicorn.run(
|
||||
"api.main:app",
|
||||
host="0.0.0.0",
|
||||
port=os.getenv["EVAL_PORT"],
|
||||
reload=True,
|
||||
reload_dirs=reload_dirs,
|
||||
)
|
||||
reload_dirs = [BASE_DIR / "swarms", BASE_DIR / "api"]
|
||||
|
||||
Loading…
Reference in new issue