inference.py

import os
import sys
import pandas as pd
import httpx
from openai import OpenAI
from env.environment import DataCleaningEnv

# ------------------ ENV VARIABLES ------------------
API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")

# ------------------ OPENAI CLIENT ------------------
client = None

try:
    if not HF_TOKEN:
        raise RuntimeError("API_KEY or HF_TOKEN environment variable is not set")

    client = OpenAI(
        base_url=API_BASE_URL,
        api_key=HF_TOKEN,
        http_client=httpx.Client(timeout=30.0, trust_env=False)
    )
except Exception as e:
    print(f"[WARN] OpenAI client init failed: {e}", file=sys.stderr)
    client = None

MAX_STEPS = 6

# ------------------ LOGGING ------------------
def clean_line(value):
    return str(value).replace("\r", " ").replace("\n", " ")

def log_start(task, env, model):
    print(f"[START] task={clean_line(task)} env={clean_line(env)} model={clean_line(model)}")

def log_step(step, action, reward, done, error):
    error_val = clean_line(error) if error else "null"
    print(f"[STEP] step={step} action={clean_line(action)} reward={reward:.2f} done={str(done).lower()} error={error_val}")

def log_end(success, steps, rewards, score=None):
    rewards_str = ",".join(f"{r:.2f}" for r in rewards)

    if score is None:
        print(f"[END] success={str(success).lower()} steps={steps} rewards={rewards_str}")
    else:
        print(f"[END] success={str(success).lower()} steps={steps} score={score:.4f} rewards={rewards_str}")

# ------------------ LLM DECISION ------------------
def fallback_action(history):
    actions = [
        {"type": "fill_nulls", "column": "city"},
        {"type": "deduplicate", "column": "customer_id"},
        {"type": "convert_types", "column": "age"},
        {"type": "trim_whitespace", "column": "city"},
        {"type": "normalize", "column": "income"},
        {"type": "remove_nulls", "column": "customer_id"},
    ]

    for action in actions:
        if str(action) not in history:
            return action

    return {"type": "deduplicate", "column": "customer_id"}

def normalize_action(action, df, history):
    valid_actions = {
        "fill_nulls",
        "remove_nulls",
        "deduplicate",
        "convert_types",
        "trim_whitespace",
        "normalize",
    }

    if not isinstance(action, dict):
        return fallback_action(history)

    action_type = action.get("type")
    column = action.get("column")

    if action_type not in valid_actions:
        return fallback_action(history)

    if action_type != "deduplicate" and column not in df.columns:
        return fallback_action(history)

    return action

def get_action_from_llm(dataset, history):
    if client is None:
        return fallback_action(history)

    prompt = f"""
You are an intelligent data cleaning agent.

Actions:
fill_nulls, remove_nulls, deduplicate, convert_types, trim_whitespace, normalize

Previous actions:
{history}

Dataset:
{dataset}

Return ONLY:
action_type,column_name
"""

    try:
        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.3,
            max_tokens=50
        )

        output = response.choices[0].message.content.strip()

        first_line = output.splitlines()[0].strip()
        parts = first_line.split(",", 1)

        if len(parts) < 2:
            return fallback_action(history)

        action_type, column = parts[0], parts[1]
        return {"type": action_type.strip(), "column": column.strip()}

    except Exception as e:
        print(f"[WARN] LLM call failed: {e}", file=sys.stderr)
        return fallback_action(history)

# ------------------ MAIN ------------------
def main():
    env = None
    rewards = []
    steps_taken = 0
    history = []
    success = False
    score = None

    log_start("task1", "data_cleaning", MODEL_NAME)

    try:
        for task_id in [1, 2, 3]:
           env = DataCleaningEnv(task=task_id)
        obs = env.reset()

        for step in range(1, MAX_STEPS + 1):

            df = pd.DataFrame(obs["dataset"])
            col_info = {}

            for col in df.columns:
                col_info[col] = {
                    "nulls": float(df[col].isnull().mean()),
                    "dtype": str(df[col].dtype),
                    "unique": int(df[col].nunique())
                }

            summary = f"""
Columns: {list(df.columns)}
Column Info: {col_info}
Duplicates: {df.duplicated().sum()}
Sample: {df.head(3).to_dict()}
"""

            action = get_action_from_llm(summary, history)
            action = normalize_action(action, df, history)

            if str(action) in history:
                action = fallback_action(history)
                action = normalize_action(action, df, history)

            history.append(str(action))

            error = None

            try:
                obs, reward, done, _ = env.step(action)
            except Exception as e:
                error = str(e)
                reward = 0.0
                done = True

            rewards.append(reward)
            steps_taken = step

            log_step(step, str(action), reward, done, error)

            if done:
                break

        final = env.submit_cleaned_data(env.dirty_df)
        score = final["final_score"]
        success = score > 0.3

    except Exception as e:
        print(f"[WARN] inference failed: {e}", file=sys.stderr)
        success = False
        score = None

    finally:
        if env is not None and hasattr(env, "close"):
            try:
                env.close()
            except Exception as e:
                print(f"[WARN] env.close failed: {e}", file=sys.stderr)

        log_end(success, steps_taken, rewards, score)

# ------------------ RUN ------------------
if __name__ == "__main__":
    main()