Update app.py

app.py

@@ -5,48 +5,55 @@ import re
 
 # 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 = "https://deployment.datasaur.ai/api/deployment/8/2723/chat/completions"
+API_URL = "https://deployment.datasaur.ai/api/deployment/8/2738/chat/completions"
 API_TOKEN = os.environ["DATASAUR_API_KEY"]
 
 import re
 import json
 
-def extract_suggestion_value(text):
+def extract_json_from_text(text):
     """
-    Extracts the value of the 'suggestion' key from a given string that may contain JSON-like content.
+    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 containing a JSON-like 'suggestion' field.
+        text (str): Input text that may contain JSON content.
 
     Returns:
-        dict: Dictionary with the 'suggestion' key and its extracted value, or None if not found.
+        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 "suggestion" in parsed:
-            return parsed["suggestion"]
+        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 suggestion value
-    match = re.search(r'"suggestion"\s*:\s*"([^"]+)"', text)
+    # Fallback: use regex to extract JSON content between { and }
+    match = re.search(r'\{[^{}]*\}', text)
     if match:
-        return match.group(1)
+        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 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"Input: `{input_text}`"}]
+        "messages": [{"role": "user", "content": f"Time Entry Input: `{input_text}`"}]
     }
     
     try:
@@ -59,38 +66,59 @@ def magic_function(input_text):
         # 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 = extract_suggestion_value(content)
-
-        return content.strip()
+        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 f"API Request Error: {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 f"Error processing API response: {response.text}"
+        return {"category_fix": "ERROR", "suggestion": f"Error processing API response: {response.text}"}
 
-def handle_magic_click(current_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 the improved text has no "[" or "]" characters, make text area light green and hide previous text.
+    If category_fix == "COMPLIANT", make text area light green, hide previous text, lock text area, and change button to Reset.
     """
-    improved_text = magic_function(current_text)
+    if current_button_text == "Reset":
+        # Reset to initial state
+        return (
+            gr.update(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
+            "Your Text"  # text_area label - reset
+        )
+    
+    result = magic_function(current_text)
+    category_fix = result.get("category_fix", "")
+    suggestion = result.get("suggestion", "")
     
-    # Check if improved text contains "[" or "]"
-    if "[" not in improved_text and "]" not in improved_text:
-        # No brackets found - make text area light green and hide previous text
+    if category_fix == "COMPLIANT":
+        # COMPLIANT case - no changes to input, green background, lock, add checklist, change button to Reset
         return (
-            gr.update(value=improved_text, elem_classes="success-text"),
-            current_text,
-            gr.update(visible=False)
+            gr.update(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
+            "Your Text ✓"  # text_area label - add checklist symbol
         )
     else:
-        # Brackets found - reset text area to normal appearance and show previous text
+        # Non-compliant case - show suggestion, show previous text, show category
         return (
-            gr.update(value=improved_text, elem_classes=""),
-            current_text,
-            gr.update(visible=True)
+            gr.update(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
+            "Your Text"  # text_area label - normal
         )
 
 
@@ -102,11 +130,12 @@ with gr.Blocks(css=".success-text { background-color: #d4edda !important; }") as
             previous_text_area = gr.Textbox(label="Previous Text", lines=5, visible=False)
         with gr.Column(scale=1):
             magic_button = gr.Button("Magic Button")
+            category_label = gr.Markdown("", visible=False)
 
     magic_button.click(
         fn=handle_magic_click,
-        inputs=text_area,
-        outputs=[text_area, previous_text_area, previous_text_area]
+        inputs=[text_area, magic_button],
+        outputs=[text_area, previous_text_area, previous_text_area, category_label, magic_button, text_area]
     )
 
 if __name__ == "__main__":