harsimran726/FineTune_Data_Generation_Agent

1import fastapi2from fastapi import FastAPI, HTTPException, Request3from fastapi.responses import FileResponse, JSONResponse4from fastapi.staticfiles import StaticFiles5import json6from fastapi.responses import HTMLResponse7from fastapi.templating import Jinja2Templates8import uvicorn9from pydantic import BaseModel10from typing import List, Optional11from Data_Geneartion_Agent import generate_data_agent12import os13from dotenv import load_dotenv14 15app = FastAPI()16from fastapi.middleware.cors import CORSMiddleware17 18app.add_middleware(19    CORSMiddleware,20    allow_origins=["*"],  # or your Netlify frontend URL21    allow_credentials=True,22    allow_methods=["*"],23    allow_headers=["*"],24)25# Mount static files26templates = Jinja2Templates(directory="templates")27# app.mount("/static", StaticFiles(directory="static"), name="static")28 29class Query(BaseModel):30    # api_key: str31    query: str32 33class Result(BaseModel):34    status: str35    message: str36    csv_file: Optional[str] = None37 38def save_api_key(api_key: str):39    """Save API key to .env file"""40    # with open(".env", "w") as f:41    #     f.write(f"GOOGLE_API_KEY={api_key}\n")42    # with open("runtime_env.json", "w") as f:43    #     json.dump({"API_KEY": api_key}, f)44    # Reload environment variables45    load_dotenv()46 47@app.get("/",response_class=HTMLResponse)48async def root(request: Request):49    return templates.TemplateResponse("index.html",{"request":request})50 51# class ApiKey(BaseModel):52#     api_key: str53 54# @app.post("/api/save_api_key")55# async def save_api_key(api_key: ApiKey):56#     with open(".env", "w") as f:57#         f.write(f"GOOGLE_API_KEY={api_key.api_key}")58#     return JSONResponse(content={"message": "API key saved successfully"}, status_code=200)59 60@app.post("/generate")61async def generate(query: Query):62    try:63        # Save the API key to .env file64        # print(f"Here is api_key in main.py {query.api_key}")65 66        # save_api_key(query.api_key)67        68        # Generate data using the agent with the new API key69        result = generate_data_agent(query.query)70        print(f"Here is the final result {result}")71        return result72    except Exception as e:73        raise HTTPException(status_code=500, detail=str(e))74 75@app.get("/download/{filename}")76async def download_file(filename: str):77    try:78        print(f"Here is file filename in main.py  {filename}")79        return FileResponse(80            path=filename,81            filename=filename,82            media_type="text/csv"83        )84    except Exception as e:85        raise HTTPException(status_code=404, detail="File not found")86 87if __name__ == "__main__":88    uvicorn.run(app, host="0.0.0.0", port=8000)89