client = openai.AsyncClient()
ENVIRONMENT = "propositional_logic"
model = "gpt-4.1-mini"
semaphore = asyncio.Semaphore(50)
num_datapoints = 200
system_prompt = (
    system_prompt
    + """You will also be provided the real answer. Your thinking should eventually result in producing the real answer."""
)

dataloader = create_dataset(name=ENVIRONMENT, size=num_datapoints)

@backoff.on_exception(backoff.expo, openai.RateLimitError)
async def generate_response(item):
    async with semaphore:
        messages = [
            {"role": "system", "content": system_prompt},
            {
                "role": "user",
                "content": f"""
    Question: {item['question']}
    Metadata: {item['metadata']}
    Answer: {item['answer']}
                    """,
            },
        ]
        response = await client.chat.completions.create(messages=messages, model=model)
        return {
            "question": item["question"],
            "metadata": item["metadata"],
            "answer": item["answer"],
            "response": response.choices[0].message.content,
        }

async def main():
    responses = await asyncio.gather(*[generate_response(item) for item in dataloader])
    fname = f"responses_{ENVIRONMENT}_{model}.json"
    json.dump(responses, open(fname, "w"), indent=4)
    print(f"Saved responses to {fname}")

if __name__ == "__main__":
    asyncio.run(main())