80 lines
2.5 KiB
Python
Executable file
80 lines
2.5 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
import google.generativeai as genai
|
|
import sys
|
|
import json
|
|
import base64
|
|
|
|
def log(msg):
|
|
"""Print to stderr for logging"""
|
|
print(msg, file=sys.stderr, flush=True)
|
|
|
|
def query_gemini(api_key, prompt, image_path=None):
|
|
log("\n=== Gemini API Request ===")
|
|
log(f"API Key: {api_key[:4]}...{api_key[-4:]}") # Show only first/last 4 chars
|
|
log(f"Prompt:\n{prompt}")
|
|
|
|
# Configure the Gemini API
|
|
genai.configure(api_key=api_key)
|
|
model = genai.GenerativeModel('gemini-2.0-flash-exp')
|
|
|
|
try:
|
|
if image_path:
|
|
log(f"\nProcessing image: {image_path}")
|
|
with open(image_path, 'rb') as f:
|
|
image_data = f.read()
|
|
log(f"Image size: {len(image_data):,} bytes")
|
|
|
|
image_parts = [
|
|
{
|
|
"mime_type": "image/jpeg",
|
|
"data": base64.b64encode(image_data).decode('utf-8')
|
|
}
|
|
]
|
|
response = model.generate_content([prompt, image_parts[0]])
|
|
else:
|
|
log("\nProcessing text-only request")
|
|
response = model.generate_content(prompt)
|
|
|
|
log("\n=== Gemini API Response ===")
|
|
log(f"Response:\n{response.text}")
|
|
|
|
# Create the JSON response
|
|
result = {
|
|
"success": True,
|
|
"text": response.text
|
|
}
|
|
|
|
# Log the result to stderr
|
|
log("\n=== JSON Response ===")
|
|
log(json.dumps(result, indent=2))
|
|
|
|
# Return the compact JSON to stdout without any extra newlines
|
|
return json.dumps(result, separators=(',', ':'))
|
|
except Exception as e:
|
|
log("\n=== Gemini API Error ===")
|
|
log(f"Error: {str(e)}")
|
|
|
|
error_result = {
|
|
"success": False,
|
|
"error": str(e)
|
|
}
|
|
|
|
# Log the error to stderr
|
|
log("\n=== JSON Error Response ===")
|
|
log(json.dumps(error_result, indent=2))
|
|
|
|
# Return the compact JSON to stdout without any extra newlines
|
|
return json.dumps(error_result, separators=(',', ':'))
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) < 3:
|
|
log("Usage: query_gemini_api.py <api_key> <prompt> [image_path]")
|
|
sys.exit(1)
|
|
|
|
api_key = sys.argv[1]
|
|
prompt = sys.argv[2]
|
|
image_path = sys.argv[3] if len(sys.argv) > 3 else None
|
|
|
|
# Get the result and print it to stdout without any extra newlines
|
|
result = query_gemini(api_key, prompt, image_path)
|
|
print(result, end='', flush=True)
|