|
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)
|
|
|
