Update app.py

app.py

@@ -8,14 +8,34 @@ import re
 API_URL = "https://deployment.datasaur.ai/api/deployment/8/2723/chat/completions"
 API_TOKEN = os.environ["DATASAUR_API_KEY"]
 
-def extract_suggestion_with_regex(xml_string):
-    """Extract content of <suggestion> tag using regex."""
-    pattern = r"<suggestion>(.*?)</suggestion>"
-    match = re.search(pattern, xml_string, re.DOTALL)
+import re
+import json
+
+def extract_suggestion_value(text):
+    """
+    Extracts the value of the 'suggestion' key from a given string that may contain JSON-like content.
+
+    Parameters:
+        text (str): Input text containing a JSON-like 'suggestion' field.
+
+    Returns:
+        dict: Dictionary with the 'suggestion' key and its extracted value, or None if not found.
+    """
+    # First, try to directly parse as JSON
+    try:
+        parsed = json.loads(text)
+        if "suggestion" in parsed:
+            return {"suggestion": parsed["suggestion"]}
+    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)
     if match:
-        return match.group(1).strip()
-    else:
-        raise ValueError("No <suggestion> tag found")
+        return {"suggestion": match.group(1)}
+
+    return None
+
         
 def magic_function(input_text):
     """
@@ -26,7 +46,7 @@ def magic_function(input_text):
         "Authorization": f"Bearer {API_TOKEN}",
     }
     data = {
-        "messages": [{"role": "user", "content": f"`{input_text}`"}]
+        "messages": [{"role": "user", "content": f"Input: `{input_text}`"}]
     }
     
     try:
@@ -39,7 +59,7 @@ 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_with_regex(content)
+        content = extract_suggestion_value(content)
 
         return content.strip()