- adding robust parsing strategy for LLM output

- including token counter for all models for price analysis

Messaging_system/LLM.py

@@ -114,6 +114,11 @@ class LLM:
                     'total_tokens': response.usage.total_tokens
                 }
 
+                # validating the JSON
+                self.Core.total_tokens['prompt_tokens'] += tokens['prompt_tokens']
+                self.Core.total_tokens['completion_tokens'] += tokens['completion_tokens']
+                self.Core.temp_token_counter += tokens['total_tokens']
+
                 try:
                     content = response.choices[0].message.content
 
@@ -132,10 +137,6 @@ class LLM:
                                 f"'header' or 'message' is more than specified characters in response on attempt {attempt + 1}. Retrying...")
                             continue
 
-                    # validating the JSON
-                    self.Core.total_tokens['prompt_tokens'] += tokens['prompt_tokens']
-                    self.Core.total_tokens['completion_tokens'] += tokens['completion_tokens']
-                    self.Core.temp_token_counter += tokens['total_tokens']
                     return output
 
                 except json.JSONDecodeError:
@@ -172,6 +173,17 @@ class LLM:
                     ))
 
                 # output = json.loads(str(response.text))
+                tokens = {
+                    'prompt_tokens': response.usage_metadata.prompt_token_count,
+                    'completion_tokens': response.usage_metadata.candidates_token_count,
+                    'total_tokens': response.usage_metadata.total_token_count
+                }
+
+                # validating the JSON
+                self.Core.total_tokens['prompt_tokens'] += tokens['prompt_tokens']
+                self.Core.total_tokens['completion_tokens'] += tokens['completion_tokens']
+                self.Core.temp_token_counter += tokens['total_tokens']
+
                 output = self.preprocess_and_parse_json(response.text)
 
                 if 'message' not in output or 'header' not in output:
@@ -214,7 +226,7 @@ class LLM:
                             {"role": "system", "content": instructions},
                             {"role": "user", "content": prompt}
                         ],
-                        reasoning_effort="medium",
+                        reasoning_effort="minimal",
                         n=1,
                     )
 
@@ -236,6 +248,11 @@ class LLM:
                     'total_tokens': response.usage.total_tokens
                 }
 
+                # validating the JSON
+                self.Core.total_tokens['prompt_tokens'] += tokens['prompt_tokens']
+                self.Core.total_tokens['completion_tokens'] += tokens['completion_tokens']
+                self.Core.temp_token_counter += tokens['total_tokens']
+
                 try:
                     content = response.choices[0].message.content
 
@@ -254,10 +271,6 @@ class LLM:
                                 f"'header' or 'message' is more than specified characters in response on attempt {attempt + 1}. Retrying...")
                             continue
 
-                    # validating the JSON
-                    self.Core.total_tokens['prompt_tokens'] += tokens['prompt_tokens']
-                    self.Core.total_tokens['completion_tokens'] += tokens['completion_tokens']
-                    self.Core.temp_token_counter += tokens['total_tokens']
                     return output
 
                 except json.JSONDecodeError:
@@ -359,11 +372,15 @@ class LLM:
                 response = message.content[0].text
 
                 tokens = {
-                    'prompt_tokens': 0,
-                    'completion_tokens': 0,
-                    'total_tokens': 0
+                    'prompt_tokens': message.usage.input_tokens,
+                    'completion_tokens': message.usage.output_tokens,
+                    'total_tokens': message.usage.output_tokens + message.usage.input_tokens
                 }
 
+                self.Core.total_tokens['prompt_tokens'] += tokens['prompt_tokens']
+                self.Core.total_tokens['completion_tokens'] += tokens['completion_tokens']
+                self.Core.temp_token_counter += tokens['total_tokens']
+
                 try:
                     output = self.preprocess_and_parse_json_claude(response)
                     if output is None:
@@ -394,30 +411,56 @@ class LLM:
 
     def preprocess_and_parse_json(self, response: str):
         """
-        Cleans an LLM response by removing <think> tags and extracting JSON
-        from ```json ... ``` fences (or bare text if no fence is found),
-        then returns the parsed object or None on failure.
+        Remove <think> blocks, extract JSON (from ```json fences or first {...} block),
+        and parse. Includes a repair pass to handle common LLM issues like trailing commas.
         """
-        # 1) Remove all <think>...</think> blocks
-        cleaned = re.sub(r'<think>.*?</think>', '', response, flags=re.DOTALL)
-
-        # 2) Look for a ```json ... ``` fenced block
-        fence_pattern = re.compile(r'```json(.*?)```', flags=re.DOTALL)
-        fence_match = fence_pattern.search(cleaned)
-        if fence_match:
-            json_text = fence_match.group(1).strip()
-        else:
-            # No fence; assume whole cleaned text is JSON
-            json_text = cleaned.strip()
 
-        # 3) Attempt to parse
+        def extract_json(text: str) -> str:
+            # Remove <think>...</think>
+            text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL).strip()
+
+            # Prefer fenced code if present
+            fence = re.search(r'```(?:json)?(.*?)```', text, flags=re.DOTALL | re.IGNORECASE)
+            if fence:
+                return fence.group(1).strip()
+
+            # Otherwise, grab the first {...} block
+            brace = re.search(r'\{.*\}', text, flags=re.DOTALL)
+            return brace.group(0).strip() if brace else text.strip()
+
+        def normalize_quotes(text: str) -> str:
+            return (text
+                    .replace('\ufeff', '')  # strip BOM if present
+                    .replace('“', '"').replace('”', '"')
+                    .replace('‘', "'").replace('’', "'"))
+
+        def strip_comments(text: str) -> str:
+            # Remove // line comments and /* block comments */
+            text = re.sub(r'//.*?$', '', text, flags=re.MULTILINE)
+            text = re.sub(r'/\*.*?\*/', '', text, flags=re.DOTALL)
+            return text
+
+        def remove_trailing_commas(text: str) -> str:
+            # Remove commas before } or ]
+            return re.sub(r',(\s*[}\]])', r'\1', text)
+
+        raw = extract_json(response)
+        raw = normalize_quotes(raw)
+
         try:
-            return json.loads(json_text)
-        except json.JSONDecodeError as e:
-            print(f"Failed to parse JSON: {e}")
-            # Optionally, log the offending text for debugging:
-            # print("Offending text:", json_text)
-            return None
+            return json.loads(raw)
+        except json.JSONDecodeError:
+            # Repair pass
+            repaired = strip_comments(raw)
+            repaired = remove_trailing_commas(repaired)
+            repaired = repaired.strip()
+
+            try:
+                return json.loads(repaired)
+            except json.JSONDecodeError as e:
+                print(f"Failed to parse JSON: {e}")
+                # print('Offending text:', repaired)
+                return None
 
     # ===============================================================
     # def preprocess_and_parse_json_claude(self, response: str):