testDocker / api /main.py
thuonguyenvan's picture
Update api/main.py
fedb0f8 verified
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import google.generativeai as genai
app = FastAPI()
genai.configure(api_key='AIzaSyBvhkco5aRtkqEyPFFaTKaaDmJkgU12MJU')
llm_model = genai.GenerativeModel('gemini-1.5-pro')
class InputText(BaseModel):
text: str
@app.post("/generate")
async def generate_text(input_text: InputText):
try:
response = llm_model.generate_content(input_text.text)
return {"generated_text": response.text}
except Exception as e:
raise HTTPException(status_code=500, detail=f"Error generating content: {e}")