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