main.py

import os

import dotenv
from openai import OpenAI, BadRequestError

dotenv.load_dotenv()

client = OpenAI(
    api_key=os.getenv("OPENAI_KEY"),
    base_url=os.getenv("OPENAI_BASE_URL")
)

MODEL = "gpt-4o-mini"

completion_tools = [
    {
        "type": "function",
        "function": {
            "name": "exit_conversation",
            "description": "Signal for the application to exit the conversation",
            "parameters": {
                "type": "object",
                "properties": {
                    "call_id": {
                        "type": "string",
                        "description": "The Call Id of the termination function"
                    }
                },
                "required": ["call_id"],
                "additionalProperties": False
            },
            "strict": True
        }
    }
]


def exit_conversation(call_id):
    print(call_id)


def calculate_cost(input_tokens, output_tokens):
    if MODEL != "gpt-4o-mini":
        raise ValueError(f"Pricing data is currently unavailable for model: {MODEL}")

    input_token_cost = 0.15
    output_token_cost = 0.60

    input_cost = (input_tokens / 1_000_000) * input_token_cost
    output_cost = (output_tokens / 1_000_000) * output_token_cost

    return input_cost + output_cost


def main():
    should_continue = True
    messages = [
        {
            "role": "system",
            "content": "You are a helpful assistant for a simple CLI chat. Only respond with text messages. "
                       "Get creative with the answers!"
        },
    ]

    while should_continue:
        query = input("Enter a message: ")
        messages.append({"role": "user", "content": query})

        try:
            completion = client.chat.completions.create(
                model=MODEL,
                messages=messages,
                tools=completion_tools
            )
        except BadRequestError as e:
            print(f"Request failed: {e.message}")
            continue

        message = completion.choices[0].message
        messages.append(message)
        usage = completion.usage

        for tool_call in (message.tool_calls or []):
            if tool_call.type == "function":
                try:
                    if tool_call.function.name == "exit_conversation":
                        exit_conversation(tool_call.id)
                        should_continue = False
                    else:
                        print(f"Unknown tool: {tool_call.function.name}")
                except Exception as e:
                    print(f"An error occurred: {str(e)}")

        print(f"You: {query}")
        print(f"Assistant: {message.content}")
        try:
            total_cost = calculate_cost(usage.prompt_tokens, usage.completion_tokens)
            print(f"Cost: ${total_cost:.6f}\n")
        except ValueError as e:
            print(f"Unable to calculate cost: {e}\n")


if __name__ == "__main__":
    main()