import asyncio
import gc
import os
from typing import List

from memory_profiler import profile
from openai import AsyncOpenAI
from openai.lib._parsing import type_to_response_format_param
from pydantic import Field, create_model


StepModel = create_model(
    "Step",
    explanation=(str, Field()),
    output=(str, Field()),
)


def create_new_model():
    """This sounds useless as it is. In our business case, I'm generating a model that slightly different at each call, hence the use of create_model. This illustrates of a model that seems to always be the same keeps on adding up in the memory."""
    return create_model(
        "MathResponse",
        steps=(List[StepModel], Field()),
        final_answer=(str, Field()),
    )


@profile()
async def leaking_call(client, new_model):
    await client.beta.chat.completions.parse(
        model="gpt-4o-2024-08-06",
        messages=[
            {"role": "system", "content": "You are a helpful math tutor."},
            {"role": "user", "content": "solve 8x + 31 = 2"},
        ],
        response_format=new_model,
    )


async def non_leaking_call(client, new_model):
    await client.chat.completions.create(
        model="gpt-4o-2024-08-06",
        messages=[
            {"role": "system", "content": "You are a helpful math tutor."},
            {"role": "user", "content": "solve 8x + 31 = 2"},
        ],
        response_format=type_to_response_format_param(new_model),
    )


async def main():
    client = AsyncOpenAI()

    for _ in range(200):
        # You can switch to `non_leaking_call` and see that the memory is correctly emptied
        await leaking_call(client, create_new_model())

        # We wanted to thoroughly check the memory usage, hence memory profiler + gc
        gc.collect()
        print(len(gc.get_objects()))


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