import openai

import os
from openai import OpenAI

# Function to call the AI/ML API with your key and input prompt
def call_aiml_api(prompt, max_tokens=2000):
    # Define the AI/ML API key and base URL
    aiml_api_key = os.getenv('AIML_API_KEY')  # Make sure you set this in your environment
    base_url = "https://api.aimlapi.com/"

    # Initialize the OpenAI client with the custom base URL and API key
    client = OpenAI(
        api_key=aiml_api_key,
        base_url=base_url
    )

    # Create a chat completion request using the provided model and prompt
    chat_completion = client.chat.completions.create(
        model="o1-mini",
        messages=[
            {"role": "user", "content": prompt},
        ],
        max_tokens=max_tokens,
    )

    # Extract the response from the API response object
    response = chat_completion.choices[0].message.content

    return response