| 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() | |