Skip to content

Commit

Permalink
server: refactor server for modularity and readability
Browse files Browse the repository at this point in the history
  • Loading branch information
Arieg419 committed Sep 27, 2024
1 parent 7547d94 commit be5a816
Show file tree
Hide file tree
Showing 7 changed files with 251 additions and 229 deletions.
227 changes: 0 additions & 227 deletions backend/server.py

This file was deleted.

Empty file added backend/server/__init__.py
Empty file.
129 changes: 129 additions & 0 deletions backend/server/server.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,129 @@
import json
import os
from typing import Dict, List

from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect, File, UploadFile, Header
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel

from backend.server.server_utils import generate_report_files
from backend.server.websocket_manager import WebSocketManager
from multi_agents.main import run_research_task
from gpt_researcher.document.document import DocumentLoader
from gpt_researcher.master.actions import stream_output
from backend.server.server_utils import (
sanitize_filename, handle_start_command, handle_human_feedback,
generate_report_files, send_file_paths, get_config_dict,
update_environment_variables, handle_file_upload, handle_file_deletion,
execute_multi_agents, handle_websocket_communication, extract_command_data
)

# Models
class ResearchRequest(BaseModel):
task: str
report_type: str
agent: str

class ConfigRequest(BaseModel):
ANTHROPIC_API_KEY: str
TAVILY_API_KEY: str
LANGCHAIN_TRACING_V2: str
LANGCHAIN_API_KEY: str
OPENAI_API_KEY: str
DOC_PATH: str
RETRIEVER: str
GOOGLE_API_KEY: str = ''
GOOGLE_CX_KEY: str = ''
BING_API_KEY: str = ''
SEARCHAPI_API_KEY: str = ''
SERPAPI_API_KEY: str = ''
SERPER_API_KEY: str = ''
SEARX_URL: str = ''

# App initialization
app = FastAPI()

# Static files and templates
app.mount("/site", StaticFiles(directory="./frontend"), name="site")
app.mount("/static", StaticFiles(directory="./frontend/static"), name="static")
templates = Jinja2Templates(directory="./frontend")

# WebSocket manager
manager = WebSocketManager()

# Middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["http://localhost:3000"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)

# Constants
DOC_PATH = os.getenv("DOC_PATH", "./my-docs")

# Startup event
@app.on_event("startup")
def startup_event():
os.makedirs("outputs", exist_ok=True)
app.mount("/outputs", StaticFiles(directory="outputs"), name="outputs")
os.makedirs(DOC_PATH, exist_ok=True)

# Routes
@app.get("/")
async def read_root(request: Request):
return templates.TemplateResponse("index.html", {"request": request, "report": None})

@app.get("/getConfig")
async def get_config(
langchain_api_key: str = Header(None),
openai_api_key: str = Header(None),
tavily_api_key: str = Header(None),
google_api_key: str = Header(None),
google_cx_key: str = Header(None),
bing_api_key: str = Header(None),
searchapi_api_key: str = Header(None),
serpapi_api_key: str = Header(None),
serper_api_key: str = Header(None),
searx_url: str = Header(None)
):
return get_config_dict(
langchain_api_key, openai_api_key, tavily_api_key,
google_api_key, google_cx_key, bing_api_key,
searchapi_api_key, serpapi_api_key, serper_api_key, searx_url
)

@app.get("/files/")
async def list_files():
files = os.listdir(DOC_PATH)
print(f"Files in {DOC_PATH}: {files}")
return {"files": files}

@app.post("/api/multi_agents")
async def run_multi_agents():
return await execute_multi_agents(manager)

@app.post("/setConfig")
async def set_config(config: ConfigRequest):
update_environment_variables(config.dict())
return {"message": "Config updated successfully"}

@app.post("/upload/")
async def upload_file(file: UploadFile = File(...)):
return await handle_file_upload(file, DOC_PATH)

@app.delete("/files/{filename}")
async def delete_file(filename: str):
return await handle_file_deletion(filename, DOC_PATH)

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await manager.connect(websocket)
try:
await handle_websocket_communication(websocket, manager)
except WebSocketDisconnect:
await manager.disconnect(websocket)
Loading

0 comments on commit be5a816

Please sign in to comment.