84 lines
3.0 KiB
Python
84 lines
3.0 KiB
Python
import os
|
|
import json
|
|
import io
|
|
import logging
|
|
from PIL import Image
|
|
from google import genai
|
|
from google.genai import types
|
|
from ..config_loader import get_config
|
|
|
|
log = logging.getLogger("ainventory")
|
|
|
|
def get_best_models():
|
|
# Priority list based on aliases and future-gen models found in user logs
|
|
return [
|
|
"gemini-flash-latest", # Points to the best stable Flash version
|
|
"gemini-3.1-flash-lite-preview", # Cutting edge 3.1 Lite
|
|
"gemini-flash-lite-latest", # Best Lite alias
|
|
"gemini-pro-latest" # Pro fallback alias
|
|
]
|
|
|
|
def extract(image_bytes: bytes, prompt: str):
|
|
config = get_config()
|
|
api_key = config.get("ai", {}).get("gemini_api_key")
|
|
if not api_key:
|
|
log.error("CRITICAL: gemini_api_key is MISSING in configuration!")
|
|
return None
|
|
|
|
# Log partial key for safety debug
|
|
key_hint = f"{api_key[:4]}...{api_key[-4:]}" if len(api_key) > 8 else "too short"
|
|
log.info(f"🔑 Using API Key: {key_hint}")
|
|
|
|
try:
|
|
# Switching to v1beta which is required for many experimental/new models
|
|
client = genai.Client(api_key=api_key, http_options={'api_version': 'v1beta'})
|
|
|
|
# DEBUG: List allowed models for this key
|
|
log.debug("🔍 Checking available models for your key...")
|
|
try:
|
|
for m in client.models.list():
|
|
log.debug(f" - Found: {m.name}")
|
|
except Exception as list_e:
|
|
log.warning(f" ⚠️ Could not list models: {list_e}")
|
|
|
|
models_to_try = get_best_models()
|
|
|
|
# Try models in order
|
|
for model_name in models_to_try:
|
|
try:
|
|
log.info(f"🚀 AI Launching (v2 SDK): {model_name}...")
|
|
|
|
# In the new SDK, we pass a list of parts (text string and image bytes)
|
|
response = client.models.generate_content(
|
|
model=model_name,
|
|
contents=[
|
|
prompt,
|
|
types.Part.from_bytes(data=image_bytes, mime_type="image/jpeg")
|
|
]
|
|
)
|
|
|
|
if not response or not response.text:
|
|
log.warning(f"⚠️ Gemini {model_name} returned NO text.")
|
|
continue
|
|
|
|
text = response.text.strip()
|
|
log.info(f"✅ AI Response Received ({len(text)} bytes)")
|
|
log.debug(f"FULL RESPONSE:\n{text}")
|
|
|
|
# Extract JSON block
|
|
if "```json" in text:
|
|
text = text.split("```json")[1].split("```")[0].strip()
|
|
elif "```" in text:
|
|
text = text.split("```")[1].strip()
|
|
|
|
return json.loads(text)
|
|
|
|
except Exception as e:
|
|
log.error(f"❌ Gemini {model_name} failed: {e}")
|
|
continue
|
|
|
|
except Exception as outer_e:
|
|
log.error(f"❌ Gemini Client Init failed: {outer_e}")
|
|
|
|
return None
|