File size: 1,295 Bytes
67e93c9 | 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 | import os
import sys
import json
import google.generativeai as genai
from dotenv import load_dotenv
load_dotenv()
genai.configure(api_key=os.environ.get("GOOGLE_API_KEY"))
model = genai.GenerativeModel("gemini-3.1-flash-lite-preview")
def main():
try:
# Promptfoo passes the prompt string as JSON to stdin
input_data = sys.stdin.read()
if not input_data:
print(json.dumps({"error": "No input provided on stdin"}))
sys.exit(1)
# The prompt is a simple JSON string or object containing '{ "prompt": "..." }'
parsed = json.loads(input_data)
# Depending on how promptfoo calls this, it may be a direct string or a dict
if isinstance(parsed, dict) and 'prompt' in parsed:
prompt = parsed['prompt']
elif isinstance(parsed, str):
prompt = parsed
else:
prompt = str(parsed)
response = model.generate_content(prompt)
# Promptfoo expects the response inside a JSON object: { "output": "..." }
print(json.dumps({"output": response.text}))
sys.exit(0)
except Exception as e:
print(json.dumps({"error": str(e)}))
sys.exit(1)
if __name__ == "__main__":
main()
|