File size: 1,188 Bytes
77057cf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
from fastapi import FastAPI, HTTPException
from fastapi.responses import PlainTextResponse
from pydantic import BaseModel
import os
from google import genai
from google.genai import types
app = FastAPI()
class GenerateRequest(BaseModel):
system_prompt: str
user_input: str
@app.post("/generate", response_class=PlainTextResponse)
def generate(req: GenerateRequest):
client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
model = "gemini-flash-latest"
contents = [
types.Content(
role="user",
parts=[types.Part.from_text(text=req.user_input)]
)
]
config = types.GenerateContentConfig(
thinking_config=types.ThinkingConfig(thinking_budget=0),
response_mime_type="application/json",
system_instruction=[types.Part.from_text(text=req.system_prompt)]
)
output = ""
try:
for chunk in client.models.generate_content_stream(
model=model,
contents=contents,
config=config
):
output += chunk.text
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
return output
|