app.py

import gradio as gr
import requests
import os
import re

from run_model import run

# WARNING: It is not recommended to hardcode sensitive data like API tokens in code.
# Consider using environment variables or other secure methods for production applications.
API_URL = os.environ["DATASAUR_API_URL"]
API_TOKEN = os.environ["DATASAUR_API_KEY"]

import re
import json

def extract_json_from_text(text):
    """
    Extracts JSON content from text that starts with { and ends with }.
    Returns the parsed JSON object with category_fix and suggestion keys.

    Parameters:
        text (str): Input text that may contain JSON content.

    Returns:
        dict: Dictionary with 'category_fix' and 'suggestion' keys, or None if not found.
    """
    # First, try to directly parse as JSON
    try:
        parsed = json.loads(text)
        if "category_fix" in parsed and "suggestion" in parsed:
            return parsed
    except json.JSONDecodeError:
        pass  # Fall back to regex if not valid JSON

    # Fallback: use regex to extract JSON content between { and }
    match = re.search(r'\{[^{}]*\}', text)
    if match:
        try:
            json_str = match.group(0)
            parsed = json.loads(json_str)
            if "category_fix" in parsed and "suggestion" in parsed:
                return parsed
        except json.JSONDecodeError:
            pass

    return None

        
def magic_function(input_text):
    """
    Sends text to the Datasaur deployment API and returns the processed JSON with category_fix and suggestion.
    """
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_TOKEN}",
    }
    data = {
        "messages": [{"role": "user", "content": f"Time Entry Input: `{input_text}`"}]
    }
    
    try:
        """
        response = requests.post(API_URL, headers=headers, json=data)
        response.raise_for_status()  # Raise an exception for bad status codes (4xx or 5xx)
        
        response_json = response.json()
        
        # Extract content from a standard chat completion response structure.
        # This may need adjustment if the API has a different format.
        content = response_json.get("choices", [{}])[0].get("message", {}).get("content", "Error: Could not parse response.")
        """
        
        content = run(input_text)

        print(content)

        result = extract_json_from_text(content)
        
        if result:
            return result
        else:
            return {"category_fix": "ERROR", "suggestion": "Could not parse response"}
        
    except requests.exceptions.RequestException as e:
        return {"category_fix": "ERROR", "suggestion": f"API Request Error: {e}"}
    except (ValueError, KeyError, IndexError):
        # Handle cases where response is not valid JSON or structure is unexpected
        return {"category_fix": "ERROR", "suggestion": f"Error processing API response: {response.text}"}

def handle_magic_click(current_text, current_button_text):
    """
    When the magic button is clicked, this function gets the improved text,
    and returns the new and previous text to update the UI.
    If category_fix == "COMPLIANT", make text area light green, hide previous text, lock text area, and change button to Reset.
    """
    if current_button_text == "Reset":
        # Reset to initial state
        return (
            gr.update(label="Your Text", value="", elem_classes="", interactive=True),  # text_area - reset and unlock
            "",  # previous_text_area - clear
            gr.update(visible=False),  # previous_text_area visibility - hide
            gr.update(value="", visible=False),  # category_label - hide
            gr.update(value="Magic Button")  # button text - reset
        )
    
    result = magic_function(current_text)
    category_fix = result.get("category_fix", "")
    suggestion = result.get("suggestion", "")

    print(current_text)
    print(category_fix)
    print(suggestion)
    print()
    
    if category_fix == "COMPLIANT":
        # COMPLIANT case - no changes to input, green background, lock, add checklist, change button to Reset
        return (
            gr.update(label="Your Text ✓", value=current_text, elem_classes="success-text", interactive=False),  # text_area - keep original, green, lock
            current_text,  # previous_text_area - store current
            gr.update(visible=False),  # previous_text_area visibility - hide
            gr.update(value=f"**{category_fix}**", visible=True),  # category_label - show category in bold
            gr.update(value="Reset")  # button text - change to Reset
        )
    else:
        # Non-compliant case - show suggestion, show previous text, show category
        return (
            gr.update(label="Your Text", value=suggestion, elem_classes="", interactive=True),  # text_area - show suggestion, normal style, unlocked
            current_text,  # previous_text_area - store current
            gr.update(visible=True),  # previous_text_area visibility - show
            gr.update(value=f"**{category_fix}**", visible=True),  # category_label - show category in bold
            gr.update(value="Magic Button")  # button text - keep as Magic Button
        )


with gr.Blocks(
        theme=gr.themes.Default(primary_hue="blue", secondary_hue="sky"), 
        css="""
        .success-text { background-color: #d4edda !important; } 
        """
    ) as demo:
    gr.Markdown("# Time Entry Improvement")
    with gr.Row():
        with gr.Column(scale=4):
            text_area = gr.Textbox(label="Your Text", lines=5)
            previous_text_area = gr.Textbox(label="Previous Text", lines=5, visible=False)
            category_label = gr.Markdown("", visible=False)
        with gr.Column(scale=1):
            magic_button = gr.Button("Magic Button")
            
    magic_button.click(
        fn=handle_magic_click,
        inputs=[text_area, magic_button],
        outputs=[text_area, previous_text_area, previous_text_area, category_label, magic_button]
    )

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