File size: 2,468 Bytes
e23fa9f a029b6f 97bde60 a029b6f e23fa9f cc24ffb e23fa9f 0248d33 aacab3f a029b6f e23fa9f aacab3f e23fa9f aacab3f a029b6f aacab3f f4997a0 aacab3f a029b6f 8ccf9dc a029b6f cc24ffb a029b6f cc24ffb a029b6f aacab3f a029b6f e23fa9f a029b6f e23fa9f 0248d33 a029b6f aacab3f a029b6f |
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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
from typing import Any, Dict, List, Optional
import os
import google.generativeai as genai
# === Load environment variables ===
API_KEY = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
MODEL_NAME = os.getenv("GEMINI_MODEL", "gemini-2.5-pro")
TEMPERATURE = float(os.getenv("TEMPERATURE", "0.7"))
TOP_P = float(os.getenv("TOP_P", "0.95"))
MAX_TOKENS = int(os.getenv("MAX_OUTPUT_TOKENS", "1024"))
SYSTEM_PROMPT = os.getenv("SYSTEM_PROMPT", "You are a helpful assistant.")
genai.configure(api_key=API_KEY)
def _extract_text(resp: Any) -> str:
if getattr(resp, "text", None):
return resp.text
try:
for c in getattr(resp, "candidates", []):
for p in getattr(c, "content", {}).get("parts", []):
if p.get("text"):
return p["text"]
except Exception:
pass
return ""
class EndpointHandler:
def __init__(self, path: str = ""):
print("[handler:init] Loading Gemini model...", flush=True)
self.model = genai.GenerativeModel(
MODEL_NAME,
system_instruction=SYSTEM_PROMPT
)
print("[handler:init] Model ready ✅", flush=True)
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
print("[handler:call] Incoming data:", data, flush=True)
# 1️⃣ Try flexible input parsing
text = data.get("message") or data.get("inputs") or ""
if isinstance(text, dict):
text = text.get("content") or ""
text = str(text).strip()
if not text:
return {"text": "(empty input)"}
params = data.get("parameters", {})
gen_cfg = {
"temperature": float(params.get("temperature", TEMPERATURE)),
"top_p": float(params.get("top_p", TOP_P)),
"max_output_tokens": int(params.get("max_output_tokens", MAX_TOKENS))
}
history = data.get("history") or []
try:
if history:
chat = self.model.start_chat(history=history)
resp = chat.send_message(text, generation_config=gen_cfg)
else:
resp = self.model.generate_content(text, generation_config=gen_cfg)
reply = _extract_text(resp)
if not reply:
reply = "(no response)"
return {"text": reply}
except Exception as e:
print("[handler:error]", e, flush=True)
return {"text": f"Error: {e}"}
|