Generate customer bank transaction data with the following columns:
- customer_name: Full names in Western format
- transaction_date: Dates within the last 30 days
- transaction_amount: Dollar amounts between $1-$10,000
- transaction_type: Either 'debit' or 'credit'
- transaction_category: Common banking categories like 'dining', 'retail', 'utilities'
- account_balance: Running balance after each transaction

CREATE TABLE transactions (
    customer_name VARCHAR(100),
    customer_id CHAR(8),
    transaction_date DATE,
    transaction_amount DECIMAL(10,2),
    transaction_type VARCHAR(6),
    transaction_category VARCHAR(50),
    account_balance DECIMAL(10,2)
);

# !pip install -Uqq openai==1.52.0
import json
import os

from openai import OpenAI

client = OpenAI(
    base_url="https://<url>/v1/inference/oai/v1",
    api_key="XXX",
)

model = "gretelai/auto"
message = """Generate a mock dataset for users from the Foo company based in France.
Each user should have the following columns:
* first_name: traditional French first names.
* last_name: traditional French surnames. 
* email: formatted as the first letter of their first name followed by their last name @foo.io (e.g., jdupont@foo.io)
* gender: Male/Female
* city: a city in France
* country: always 'France'.
"""

chat_completion_stream = client.chat.completions.create(
    messages=[
        {
            "role": "user",
            "content": message,
        }
    ],
    stream=True,
    n=10,
    model=model,
)
for chunk in chat_completion_stream:
    if content := chunk.choices[0].delta.content:
        # each chunk streams back as a json row
        print(json.loads(content)["table_data"])

    if usage := chunk.usage:
        print(usage.model_dump())

# !pip install gretel-client
import pandas as pd

from gretel_client import Gretel

PROMPT = """Generate a mock dataset for users from the Foo company based in France.
Each user should have the following columns:
* first_name: traditional French first names.
* last_name: traditional French surnames. 
* email: formatted as the first letter of their first name followed by their last name @foo.io (e.g., jdupont@foo.io)
* gender: Male/Female
* city: a city in France
* country: always 'France'.
"""

EDIT_PROMPT = """Edit the table and add the following columns:
* occupation: a random occupation
* education level: make it relevant to the occupation
"""

table_headers = ["first_name", "last_name", "email", "gender", "city", "country"]
table_data = [
    {
        "first_name": "Lea",
        "last_name": "Martin",
        "email": "lmartin@foo.io",
        "gender": "Female",
        "city": "Lyon",
        "country": "France",
    }
]

SAMPLE_DATA = pd.DataFrame(table_data, columns=table_headers)

STREAM = True

from openai import OpenAI

# client = AzureOpenAI(azure_endpoint="http://localhost:8000/v1/inference/oai/v1", api_key="abc", api_version="foo")

# This is used for local testing since the AzureOpenAI class requires an actual deployed Azure OpenAI endpoint
client = OpenAI(
    base_url="https://<endpoint>/v1/inference/oai/v1",
    api_key="abc"
)

MODEL = "gretelai-azure/gpt-3.5-turbo"

azure_open_ai = Gretel.create_navigator_azure_oai_adapter(client)

usage, generated_df = azure_open_ai.generate(
    MODEL,
    PROMPT,
    num_records=10,
    sample_data=SAMPLE_DATA,
    stream=STREAM,
)

print(generated_df)
print("*****")
print(usage)

# Let's edit this now

EDIT_PROMPT = """Edit the table and add the following columns:
* Occupation: a random occupation
* Education level: make it relevant to the occupation
"""

usage, edited_df = azure_open_ai.edit(
    MODEL,
    EDIT_PROMPT,
    seed_data=generated_df,
    stream=STREAM
)
print(edited_df)
print("*****")
print(usage)