| from typing import Any, Dict, List, Optional |
| import os |
| import google.generativeai as genai |
|
|
| |
| 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) |
|
|
| |
| 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}"} |
|
|
