app.py

##########====================================================================################
##########====================PRODUCTION VERSION -- vLLM, GRADIO=====================###########
##########====================================================================################
import os, time
import requests
from typing import List, Dict, Tuple
from datetime import datetime
from anthropic import Anthropic
from openai import OpenAI
import gradio as gr
from tqdm import tqdm

ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
assert ANTHROPIC_API_KEY, "Set ANTHROPIC_API_KEY in Space settings"

VLLM_API = "http://localhost:8000/v1"

QWEN_MODEL = "Qwen/Qwen1.5-4B-Chat-AWQ"
CLAUDE_MODEL = "claude-3-5-haiku-latest"

open_source_client = OpenAI(api_key="EMPTY", base_url=VLLM_API)
claude_client = Anthropic(api_key=ANTHROPIC_API_KEY)


def wait_for_vllm_ready(timeout=900):
    start = time.time() 
    while time.time() - start < timeout:
        try:
            r = requests.get("http://localhost:8000/health", timeout=3)
            if r.status_code == 200:
                return True
        except Exception:
            pass
        time.sleep(2)
    raise RuntimeError("vLLM did not start within timeout")

def invoke_messages(
        rows_num: int,
        business_category: str,
        columns: str,
        instruction: str,
) -> List[Dict[str, str]]:
    system_message = """
        You are a helpful assistant generating synthetic mockup dataset as per
        user's request across all types of businesses and sorts.
        User's specific request for the data niche, data column types, and all
        other details and your job is to create wonderful mockup data for them
        to use for their demo apps or develop in a testing environment.
    """.strip()

    user_prompt = f"""
        Generate a synthetic mockup data that fits the following instruction:
        - Number of rows: {rows_num}
        - Business area: {business_category}
        - Columns: {columns}
        - Other instruction: {instruction}
        ㅡ Make sure to deliver only the markdown content without any additional comments
    """.strip()

    system_message = system_message + """
        In the case of sql file selection as an output, make sure to
        contain the full sql file format, including CREATE TABLE command.
    """.strip()

    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": user_prompt}
    ]

    return messages


def pass_claude_msg(file_format: str, content: str) -> Tuple[str, str]:
    claude_sys_msg = """
        You are a helpful assistant, converting generated outputs (done by other model)
        into the format of chosen type:
        example: csv, sql, or json format.
        NOTE: generate the result output that only includes the markdown content
        without any addtional comments!
    """.strip()
    claude_user_msg = f"""
        Convert the output into the {file_format} format for the following content:
        ----------------------------------------------------------------------
        {content}
    """.strip()

    return claude_sys_msg, claude_user_msg


def generate_output(messages):
    
    resp = open_source_client.chat.completions.create(
        model=QWEN_MODEL,
        messages=messages,
        max_tokens=400,
        temperature=0.2,
        stream=False
    )

    return resp.choices[0].message.content


def launch_claude_api(sys_msg, user_msg):

    if not claude_client:
        return None
    
    response = claude_client.messages.create(
        model=CLAUDE_MODEL,
        system=sys_msg,
        max_tokens=400,
        temperature=0.1,
        messages=[
            {"role": "user", "content": user_msg}
        ]
    )
    return response.content[0].text


###============= Gradio Function =============###

def generate_mockup_data(category, num_data_rows, columns, a_instruction,
                         progress=gr.Progress()):
    progress(0.2, desc="Generating...")
    msg = invoke_messages(
        rows_num=int(num_data_rows or 10),
        business_category=category,
        columns=columns,
        instruction=a_instruction
    )

    result = generate_output(msg)
    progress(1.0, desc="Done")

    return result


def show_hidden_row():
    return gr.update(visible=True)


def make_file(btn_sort: str, category: str, content: str):
    '''
    btn_sort: one of the 3 download file tpes from the buttons -- download csv, sql, json 
    category: Business category or area that the data is associated with.
    content: LLM generated text output to write in a file
    '''

    if not content or not content.strip():
        raise gr.Error("The result content is empty. Cannot create a file.")

    if not claude_client:
        raise gr.Error("File formatting requires ANTHROPIC_API_KEY.")

    try:
        sys_msg, user_msg = pass_claude_msg(btn_sort, content)
        claude_output = launch_claude_api(sys_msg, user_msg)

        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
        filepath = f"/tmp/{category}_mockup_{ts}.{btn_sort}"

        with open(filepath, "w") as f:
            f.write(claude_output)

        return filepath
    except Exception as e:
        raise gr.Error("Failed to format or create the file.")


###============= Gradio UI =============###

def render_interface():

    with gr.Blocks(title="Mockup Data Generator", css="footer {visibility:hidden}") as demo:
        category = gr.Textbox(
            label="Business Area/Category",
            placeholder="e.g. HR, Sales, Hospitality, Senior Care, E-commerce, Finance",
        )
        num_data_rows = gr.Number(
            label="Number of Rows",
            placeholder="Type number...",
            minimum=10,
            maximum=50,
            step=10,
            precision=0
        )
        columns = gr.Textbox(
            label="Insert Columns",
            placeholder="Comma, separated..."
        )
        a_instruction = gr.Textbox(
            label="Additional Instruction",
            placeholder="Any additional instruction. Leave blank if none.",
            lines=5
        )
        btn = gr.Button(
            value="Generate"
        )
        out = gr.Textbox(label="Result shown here.")

        buttons_row = gr.Row(visible=False)

        with buttons_row:
            btn_csv = gr.DownloadButton(label="Download csv", size="md", elem_classes=["download-btn"])
            btn_sql = gr.DownloadButton(label="Download sql", size="md", elem_classes=["download-btn"])
            btn_json = gr.DownloadButton(label="Download json", size="md", elem_classes=["download-btn"])

        chain = btn.click(
            fn=generate_mockup_data,
            inputs=[category, num_data_rows, columns, a_instruction],
            outputs=out,
            queue=True
        )

        chain = chain.then(
            fn=show_hidden_row,
            inputs=None,
            outputs=buttons_row,
        )

        btn_csv.click(
            lambda category, data: make_file("csv", category, data),
            inputs=[category, out],
            outputs=btn_csv
        )

        btn_sql.click(
            lambda category, data: make_file("sql", category, data),
            inputs=[category, out],
            outputs=btn_sql
        )

        btn_json.click(
            lambda category, data: make_file("json", category, data),
            inputs=[category, out],
            outputs=btn_json
        )
    return demo


if __name__ == "__main__":
    wait_for_vllm_ready(900)
    app = render_interface()
    app.queue(default_concurrency_limit=1)
    app.launch(server_name="0.0.0.0", server_port=7860)