server.py

from fastapi import FastAPI, HTTPException, Request, Body
from fastapi.middleware.cors import CORSMiddleware
from langserve import add_routes
from run import chain
import os
import uvicorn
from pydantic import BaseModel
from titleGenerator import generate_title

app = FastAPI(title="Code Vulnerability AI")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows all origins for testing purposes
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

add_routes(app, chain)

class MessageInput(BaseModel):
    message: str

@app.post("/generate_title")
async def generate_title_endpoint(request: Request, body: MessageInput = Body(...)):
    message = body.message

    if not message:
        raise HTTPException(status_code=400, detail="Message is required")

    generated_title = generate_title(message)
    
    return {"title": generated_title}

if __name__ == "__main__":
    host = os.environ.get('STAGING', '0.0.0.0')
    uvicorn.run(app, host=host, port=4000)