vedkdev/FlakyTestSleuthOpenEnvRL

1import os2import httpx3 4API_KEY = os.environ.get("GEMINI_API_KEY", "your-api-key-here")5MODEL_NAME = os.environ.get("MODEL_NAME", "gemini-2.0-flash")6BASE_URL = "https://generativelanguage.googleapis.com/v1beta"7 8def chat(messages: list[dict], **kwargs) -> str:9    headers = {10        "Content-Type": "application/json"11    }12    13    # Convert messages to Gemini format14    contents = []15    for msg in messages:16        role = "user" if msg["role"] == "user" else "model"17        contents.append({18            "role": role,19            "parts": [{"text": msg["content"]}]20        })21    22    payload = {23        "contents": contents,24        "generationConfig": {25            "temperature": kwargs.get("temperature", 1.0),26            "maxOutputTokens": kwargs.get("max_tokens", 2048),27            "topP": kwargs.get("top_p", 0.95),28            "topK": kwargs.get("top_k", 40)29        }30    }31    32    url = f"{BASE_URL}/models/{MODEL_NAME}:generateContent?key={API_KEY}"33    34    response = httpx.post(35        url,36        json=payload,37        headers=headers,38        timeout=60.039    )40    41    if response.status_code != 200:42        raise Exception(f"Error: {response.status_code} - {response.text}")43    44    result = response.json()45    return result["candidates"][0]["content"]["parts"][0]["text"]46 47if __name__ == "__main__":48    response = chat([49        {"role": "user", "content": "Hello! What model are you?"}50    ])51    print(response)