Update utils/mistral.py

utils/mistral.py

@@ -44,20 +44,34 @@ def call_llm(messages, max_tokens=2048, temperature=0.3):
 
 
 # Function to clean model output
+# def Data_Cleaner(text):
+#     pattern = r".*?format:"
+#     result = re.split(pattern, text, maxsplit=1)
+#     if len(result) > 1:
+#         text_after_format = result[1].strip().strip('`').strip('json')
+#     else:
+#         text_after_format = text.strip().strip('`').strip('json')
+
+#     try:
+#         json.loads(text_after_format)
+#         return text_after_format
+#     except json.JSONDecodeError:
+#         logging.error("Data cleaning led to invalid JSON")
+#         return text
 def Data_Cleaner(text):
-    pattern = r".*?format:"
-    result = re.split(pattern, text, maxsplit=1)
-    if len(result) > 1:
-        text_after_format = result[1].strip().strip('`').strip('json')
-    else:
-        text_after_format = text.strip().strip('`').strip('json')
-
     try:
-        json.loads(text_after_format)
-        return text_after_format
-    except json.JSONDecodeError:
-        logging.error("Data cleaning led to invalid JSON")
-        return text
+        # Extract JSON block using regex
+        json_match = re.search(r'\{.*\}', text, re.DOTALL)
+        if json_match:
+            json_str = json_match.group(0)
+            return json_str
+        else:
+            raise ValueError("No JSON found in response")
+
+    except Exception as e:
+        logging.error(f"JSON extraction failed: {e}")
+        return None
+
 
 
 # Function to call LLM and process output
@@ -95,6 +109,10 @@ def Model_ProfessionalDetails_Output(resume, client=None):
     try:
         response = call_llm([system_role, user_prompt], max_tokens=3000, temperature=0.35)
         clean_response = Data_Cleaner(response)
+
+        if not clean_response:
+            raise ValueError("Empty or invalid LLM response")
+
         parsed_response = json.loads(clean_response)
     except Exception as e:
         logging.error(f"LLM Error: {e}")
@@ -129,6 +147,10 @@ def Model_PersonalDetails_Output(resume, client=None):
     try:
         response = call_llm([system_role, user_prompt], max_tokens=2000, temperature=0.35)
         clean_response = Data_Cleaner(response)
+
+        if not clean_response:
+            raise ValueError("Empty or invalid LLM response")
+
         parsed_response = json.loads(clean_response)
     except Exception as e:
         print("JSON Decode Error:", e)