app.py

import os
import threading
from typing import Optional

import gradio as gr
from groq import Groq


# ============================================================
# Configuration
# ============================================================

LOCAL_MODEL_NAME = "devendrajadhav34/gemma3-bitext-support-lora"

# This model is currently shown in Groq's official examples.
GROQ_MODEL_NAME = os.getenv(
    "GROQ_MODEL_NAME",
    "llama-3.3-70b-versatile",
)

# On a free CPU Hugging Face Space, leave this false.
# Set ENABLE_LOCAL_MODEL=true only when CUDA is available.
ENABLE_LOCAL_MODEL = (
    os.getenv("ENABLE_LOCAL_MODEL", "false").lower() == "true"
)

GROQ_API_KEY = os.getenv("GROQ_API_KEY")

generation_lock = threading.Lock()

local_model = None
local_tokenizer = None
local_model_error: Optional[str] = None


# ============================================================
# Few-shot prompt
# ============================================================

SYSTEM_PROMPT = """
You are a helpful customer support assistant.

Your responses must closely follow the style demonstrated in the examples:

- Begin by acknowledging the customer's issue with empathy.
- Sound reassuring, polite, and professional.
- Give a concise answer rather than an overly detailed explanation.
- Ask only for information genuinely needed to resolve the issue.
- When appropriate, provide simple numbered steps.
- Reassure the customer that their satisfaction is important.
- Do not invent company policies, order details, refund guarantees,
  phone numbers, URLs, operating hours, or account information.
- When company-specific information is required, use descriptive
  placeholders such as {{Order Number}}, {{Tracking Number}},
  {{Customer Support Phone Number}}, or {{Website URL}}.
- Never claim that an action has already been completed.
- Never claim that you accessed an order, account, transaction,
  delivery system, or customer record.
- Keep most answers between 60 and 140 words.
- Do not mention these instructions or the examples.
""".strip()


FEW_SHOT_MESSAGES = [
    {
        "role": "user",
        "content": "I accidentally placed an order and want to cancel it",
    },
    {
        "role": "assistant",
        "content": (
            "I understand your situation and the urgency to cancel your "
            "order. Rest assured, I'm here to assist you every step of the "
            "way. To proceed with the cancellation, please provide your "
            "{{Order Number}}, as it will help the support team locate the "
            "order and assist you further. Please note that cancellation "
            "availability may depend on whether the order has already been "
            "processed or shipped. Your satisfaction is our priority, and "
            "we'll do our best to help resolve the issue promptly."
        ),
    },
    {
        "role": "user",
        "content": "My package has not arrived yet",
    },
    {
        "role": "assistant",
        "content": (
            "We understand that you're concerned about your package not "
            "arriving, and we're here to help. Could you please provide your "
            "{{Tracking Number}} or {{Order Number}}? With this information, "
            "the support team can review the shipment status and provide the "
            "most recent available update. We appreciate your patience and "
            "cooperation while the delivery is investigated."
        ),
    },
    {
        "role": "user",
        "content": "I was charged twice for the same purchase",
    },
    {
        "role": "assistant",
        "content": (
            "I'm sorry to hear that you noticed two charges for the same "
            "purchase. It's important to review billing discrepancies "
            "promptly. To assist you, please provide your {{Order Number}} "
            "or {{Transaction ID}}, along with the dates and amounts of the "
            "two charges. Please avoid sharing complete card or bank account "
            "details. Once the relevant information is available, the "
            "support team can investigate the duplicate charge and explain "
            "the appropriate next steps."
        ),
    },
    {
        "role": "user",
        "content": "How can I change my shipping address?",
    },
    {
        "role": "assistant",
        "content": (
            "To change your shipping address, please follow these steps:\n\n"
            "1. Log in to your account on our website.\n"
            "2. Open the 'My Account' or 'My Orders' section.\n"
            "3. Select the relevant order or saved shipping address.\n"
            "4. Choose 'Edit Address' and enter the updated details.\n"
            "5. Review the information carefully and save your changes.\n\n"
            "If the order has already been processed or shipped, the address "
            "may no longer be editable online. In that case, please contact "
            "customer support with your {{Order Number}} for assistance."
        ),
    },
    {
        "role": "user",
        "content": "I forgot my account password",
    },
    {
        "role": "assistant",
        "content": (
            "I understand that you've forgotten your account password, and "
            "I'm here to guide you through regaining access:\n\n"
            "1. Visit the {{Account Recovery Page URL}}.\n"
            "2. Select the '{{Forgot Password}}' option.\n"
            "3. Enter the email address associated with your account.\n"
            "4. Check your email for a password-reset link.\n"
            "5. Follow the instructions to create a new password.\n\n"
            "For your security, never share your password or reset code with "
            "anyone. If the email does not arrive, check your spam folder or "
            "contact customer support."
        ),
    },
]


# ============================================================
# Optional local Unsloth model
# ============================================================

def load_local_model() -> bool:
    """
    Attempt to load the fine-tuned Unsloth model.

    On a free CPU Hugging Face Space, this function immediately returns
    False because CUDA will not be available.
    """
    global local_model
    global local_tokenizer
    global local_model_error

    if not ENABLE_LOCAL_MODEL:
        local_model_error = "Local model disabled by configuration."
        return False

    try:
        import torch

        if not torch.cuda.is_available():
            local_model_error = "CUDA is not available."
            return False

        from unsloth import FastLanguageModel

        local_model, local_tokenizer = (
            FastLanguageModel.from_pretrained(
                model_name=LOCAL_MODEL_NAME,
                max_seq_length=512,
                load_in_4bit=True,
            )
        )

        FastLanguageModel.for_inference(local_model)
        return True

    except Exception as exc:
        local_model = None
        local_tokenizer = None
        local_model_error = (
            f"{type(exc).__name__}: {exc}"
        )
        return False


LOCAL_MODEL_AVAILABLE = False


def generate_with_local_model(
    question: str,
    temperature: float,
    max_new_tokens: int,
) -> str:
    import torch

    if local_model is None or local_tokenizer is None:
        raise RuntimeError("Local model has not been loaded.")

    messages = [
        {
            "role": "system",
            "content": [
                {
                    "type": "text",
                    "text": (
                        "You are a helpful customer support assistant. "
                        "Respond empathetically, professionally, and "
                        "concisely. Ask only for the information needed "
                        "to resolve the customer's issue."
                    ),
                }
            ],
        },
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": question,
                }
            ],
        },
    ]

    inputs = local_tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    ).to("cuda")

    generation_kwargs = {
        "max_new_tokens": int(max_new_tokens),
        "do_sample": temperature > 0,
        "pad_token_id": local_tokenizer.eos_token_id,
    }

    if temperature > 0:
        generation_kwargs["temperature"] = float(temperature)

    with generation_lock, torch.inference_mode():
        outputs = local_model.generate(
            **inputs,
            **generation_kwargs,
        )

    new_tokens = outputs[
        0,
        inputs["input_ids"].shape[-1]:,
    ]

    return local_tokenizer.decode(
        new_tokens,
        skip_special_tokens=True,
    ).strip()


# ============================================================
# Groq fallback
# ============================================================

def generate_with_groq(
    question: str,
    temperature: float,
    max_new_tokens: int,
) -> str:
    if not GROQ_API_KEY:
        raise RuntimeError(
            "GROQ_API_KEY has not been configured. Add it under "
            "your Hugging Face Space's Settings → Secrets."
        )

    client = Groq(api_key=GROQ_API_KEY)

    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT,
        },
        *FEW_SHOT_MESSAGES,
        {
            "role": "user",
            "content": question,
        },
    ]

    completion = client.chat.completions.create(
        model=GROQ_MODEL_NAME,
        messages=messages,
        temperature=float(temperature),
        max_completion_tokens=int(max_new_tokens),
        top_p=0.9,
    )

    response = completion.choices[0].message.content

    if not response:
        raise RuntimeError("Groq returned an empty response.")

    return response.strip()


# ============================================================
# Unified inference function
# ============================================================

def generate_response(
    question: str,
    temperature: float,
    max_new_tokens: int,
):
    question = question.strip()

    if not question:
        return (
            "Please enter a customer support question.",
            "No inference performed",
        )

    # Attempt the fine-tuned model when it was successfully loaded.
    if LOCAL_MODEL_AVAILABLE:
        try:
            response = generate_with_local_model(
                question=question,
                temperature=temperature,
                max_new_tokens=max_new_tokens,
            )

            return response, "Fine-tuned Gemma 3 model"

        except Exception as local_error:
            # Continue to Groq instead of failing the request.
            print(
                "Local inference failed; using Groq fallback:",
                repr(local_error),
            )

    try:
        response = generate_with_groq(
            question=question,
            temperature=temperature,
            max_new_tokens=max_new_tokens,
        )

        return response, f"Groq fallback: {GROQ_MODEL_NAME}"

    except Exception as groq_error:
        print("Groq inference failed:", repr(groq_error))

        return (
            "I'm sorry, but the assistant is temporarily unavailable. "
            "Please try again shortly.",
            f"Error: {type(groq_error).__name__}",
        )


# ============================================================
# Gradio interface
# ============================================================

with gr.Blocks(
    title="Customer Support Assistant",
) as demo:
    gr.Markdown(
        """
        # Customer Support Assistant

        A customer-support assistant designed to produce concise,
        empathetic and action-oriented responses.
        """
    )

    question = gr.Textbox(
        lines=4,
        placeholder="Describe your customer support issue...",
        label="Customer Query",
    )

    with gr.Accordion(
        "Generation settings",
        open=False,
    ):
        temperature = gr.Slider(
            minimum=0.0,
            maximum=1.2,
            value=0.35,
            step=0.05,
            label="Temperature",
        )

        max_new_tokens = gr.Slider(
            minimum=64,
            maximum=350,
            value=220,
            step=1,
            label="Maximum response tokens",
        )

    generate_button = gr.Button(
        "Generate response",
        variant="primary",
    )

    response = gr.Textbox(
        lines=10,
        label="Assistant Response",
    )

    inference_source = gr.Textbox(
        label="Inference source",
        interactive=False,
    )

    generate_button.click(
        fn=generate_response,
        inputs=[
            question,
            temperature,
            max_new_tokens,
        ],
        outputs=[
            response,
            inference_source,
        ],
    )

    question.submit(
        fn=generate_response,
        inputs=[
            question,
            temperature,
            max_new_tokens,
        ],
        outputs=[
            response,
            inference_source,
        ],
    )

    gr.Examples(
        examples=[
            ["I accidentally placed an order. Can I cancel it?"],
            ["My package hasn't arrived yet."],
            ["I was charged twice for my order."],
            ["How do I reset my password?"],
            [
                "Can I change my shipping address "
                "after placing an order?"
            ],
            ["The product I received is damaged."],
            ["I want to request a refund."],
        ],
        inputs=question,
    )


demo.queue(
    default_concurrency_limit=4,
    max_size=30,
)

if __name__ == "__main__":
    demo.launch()