server.py

from fastapi import FastAPI, HTTPException, Request, Body
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")

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)