fix: make K2 Think JSON extraction and repair extremely robust

app/services/k2_think_engine.py

@@ -181,135 +181,173 @@ DO NOT USE <think> TAGS. DO NOT CONVERSE.
             elif "<think>" in processed_content:
                 processed_content = re.sub(r'<think>.*', '', processed_content, flags=re.DOTALL)
 
-            candidates = []
-            
-            # Stratégie 0 : Merge ALL Markdown blocks (Highest Priority)
-            # If the LLM splits its JSON into multiple blocks, we want all of them.
-            json_blocks = re.findall(r'```(?:json)?\s*(\{.*?\})\s*```', processed_content, re.DOTALL | re.IGNORECASE)
-            if json_blocks:
-                # We will try to parse and merge them later
-                pass
-            
-            # Stratégie Prioritaire : Le plus grand bloc d'accolades (souvent le plus fiable si </think> est bien géré)
-            start_idx = processed_content.find('{')
-            end_idx = processed_content.rfind('}')
-            if start_idx != -1 and end_idx != -1 and end_idx > start_idx:
-                candidates.append(processed_content[start_idx:end_idx + 1])
-            
-            # Stratégie 1 : Chercher le vrai début du JSON (via comparative_analysis) - FIRST occurrence
-            first_idx = processed_content.find('"comparative_analysis"')
-            if first_idx != -1:
-                start_idx = processed_content.rfind('{', 0, first_idx)
-                if start_idx != -1:
-                    cand = processed_content[start_idx:]
-                    end_idx = cand.rfind('}')
-                    if end_idx != -1:
-                        candidates.append(cand[:end_idx + 1])
-                    else:
-                        candidates.append(cand) # Truncated
-            
-            # Stratégie 2 : Blocs Markdown existants
-            if json_blocks:
-                candidates.append(json_blocks[-1].strip())
-                if len(json_blocks) > 1:
-                    candidates.append(json_blocks[0].strip())
-                    
-            # Stratégie 3 : Balise [RESULT]
-            result_match = re.search(r'\[RESULT\]\s*(.*)', processed_content, re.DOTALL | re.IGNORECASE)
-            if result_match:
-                cand = result_match.group(1).strip()
-                json_inner = re.search(r'(\{.*\})', cand, re.DOTALL)
-                if json_inner:
-                    candidates.append(json_inner.group(1).strip())
-
             # Fonction utilitaire pour réparer le JSON
             def repair_json(text):
                 repaired = text.strip()
                 
-                # Heuristic: if odd number of quotes, we are inside an unclosed string. Remove it.
-                if repaired.count('"') % 2 != 0:
-                    last_quote = repaired.rfind('"')
-                    if last_quote != -1:
-                        repaired = repaired[:last_quote].strip()
+                # If it doesn't start with '{', find the first '{'
+                first_brace = repaired.find('{')
+                if first_brace != -1:
+                    repaired = repaired[first_brace:]
                 
-                # Strip trailing commas or colons before closing
-                repaired = repaired.rstrip(',:').strip()
+                if not repaired:
+                    return repaired
+                    
+                # Fix unclosed quotes
+                in_string = False
+                escape = False
+                last_good_pos = len(repaired)
                 
+                for i, char in enumerate(repaired):
+                    if escape:
+                        escape = False
+                        continue
+                    if char == '\\':
+                        escape = True
+                        continue
+                    if char == '"':
+                        in_string = not in_string
+                        if not in_string:
+                            last_good_pos = i + 1
+                            
+                if in_string:
+                    repaired = repaired[:last_good_pos].strip()
+                
+                # Remove trailing incomplete keys
+                # e.g., ',"key"' or ',"key":'
+                for suffix in ['', '}', ']']:
+                    escaped_suffix = re.escape(suffix)
+                    repaired = re.sub(r',\s*"[^"]*"\s*:\s*' + escaped_suffix + r'$', suffix, repaired)
+                    repaired = re.sub(r',\s*"[^"]*"\s*' + escaped_suffix + r'$', suffix, repaired)
+                    repaired = re.sub(r'\{\s*"[^"]*"\s*:\s*' + escaped_suffix + r'$', '{' + suffix, repaired)
+                    repaired = re.sub(r'\{\s*"[^"]*"\s*' + escaped_suffix + r'$', '{' + suffix, repaired)
+                
+                # Strip trailing punctuation/spaces
+                repaired = re.sub(r'[\s,:+]+$', '', repaired)
+
+                # Balance braces and brackets
                 open_braces = repaired.count('{')
                 close_braces = repaired.count('}')
                 open_brackets = repaired.count('[')
                 close_brackets = repaired.count(']')
+                
                 while open_brackets > close_brackets:
                     repaired += ']'
                     close_brackets += 1
                 while open_braces > close_braces:
                     repaired += '}'
                     close_braces += 1
+                    
                 return repaired
 
+            candidates = []
+            
+            # Stratégie 0 : Extraction par blocs markdown
+            for block in re.findall(r'```(?:json)?\s*(.*?)\s*```', processed_content, re.DOTALL | re.IGNORECASE):
+                b_start = block.find('{')
+                b_end = block.rfind('}')
+                if b_start != -1 and b_end != -1 and b_end > b_start:
+                    candidates.append(block[b_start:b_end + 1])
+                else:
+                    candidates.append(block)
+
+            # Stratégie 1 : Recherche par comptage de parenthèses (très robuste)
+            positions = [m.start() for m in re.finditer(r'\{', processed_content)]
+            for start_idx in positions:
+                brace_count = 0
+                in_string = False
+                escape = False
+                end_idx = -1
+                for i in range(start_idx, len(processed_content)):
+                    char = processed_content[i]
+                    if escape:
+                        escape = False
+                        continue
+                    if char == '\\':
+                        escape = True
+                        continue
+                    if char == '"':
+                        in_string = not in_string
+                        continue
+                    if not in_string:
+                        if char == '{':
+                            brace_count += 1
+                        elif char == '}':
+                            brace_count -= 1
+                            if brace_count == 0:
+                                end_idx = i
+                                break
+                if end_idx != -1:
+                    candidates.append(processed_content[start_idx:end_idx + 1])
+                else:
+                    candidates.append(processed_content[start_idx:])  # Truncated fallback
+
+            # Stratégie 2 : Blocs d'accolades globaux
+            start_idx = processed_content.find('{')
+            end_idx = processed_content.rfind('}')
+            if start_idx != -1 and end_idx != -1 and end_idx > start_idx:
+                candidates.append(processed_content[start_idx:end_idx + 1])
+
+            # Stratégie 3 : Recherche de "comparative_analysis"
+            first_idx = processed_content.find('"comparative_analysis"')
+            if first_idx != -1:
+                start_idx = processed_content.rfind('{', 0, first_idx)
+                if start_idx != -1:
+                    candidates.append(processed_content[start_idx:])
+
+            # Prioritize candidates containing key schema terms
+            schema_keywords = ["comparative_analysis", "proposed_protocol", "reasoning_summary", "research_gaps"]
+            
+            # Sort candidates by:
+            # 1. Matches at least one keyword (Boolean)
+            # 2. Length (longer is better for completeness)
+            def candidate_key(c):
+                has_keyword = any(kw in c for kw in schema_keywords)
+                return (1 if has_keyword else 0, len(c))
+                
+            candidates.sort(key=candidate_key, reverse=True)
+
             # 5. Parsing
             import pathlib
             debug_dir = pathlib.Path(__file__).parent.parent.parent
             with open(debug_dir / "k2_debug_raw.txt", "w", encoding="utf-8") as f:
                 f.write(raw_content)
+
+            # Try parsing each candidate
+            for i, cand in enumerate(candidates):
+                if not cand.strip():
+                    continue
+                with open(debug_dir / f"k2_debug_cand_{i}.txt", "w", encoding="utf-8") as f:
+                    f.write(cand)
                 
-            # FIRST: Try merging all markdown blocks if multiple exist
-            if json_blocks:
-                merged_dict = {}
-                for block in json_blocks:
-                    try:
-                        parsed = json.loads(block.strip())
-                        if isinstance(parsed, dict):
-                            merged_dict.update(parsed)
-                    except:
+                # 5.1 Direct parse
+                try:
+                    k2_analysis = json.loads(cand)
+                    clean_json = cand
+                    break
+                except json.JSONDecodeError as e:
+                    if "Extra data" in str(e) and hasattr(e, "pos"):
                         try:
-                            parsed = json.loads(repair_json(block))
-                            if isinstance(parsed, dict):
-                                merged_dict.update(parsed)
-                        except:
+                            k2_analysis = json.loads(cand[:e.pos].strip())
+                            clean_json = cand[:e.pos].strip()
+                            break
+                        except Exception:
                             pass
-                # Check if merge yielded a decent result
-                if merged_dict and ("comparative_analysis" in merged_dict or "research_gaps" in merged_dict):
-                    k2_analysis = merged_dict
-                    clean_json = json.dumps(merged_dict)
-                    logger.info("JSON successfully merged from multiple Markdown blocks.")
-
-            if not k2_analysis:
-                for i, cand in enumerate(candidates):
-                    if not cand.strip():
-                        continue
-                    with open(debug_dir / f"k2_debug_cand_{i}.txt", "w", encoding="utf-8") as f:
-                        f.write(cand)
-                    
-                    # Tentative directe
-                    try:
-                        k2_analysis = json.loads(cand)
-                        clean_json = cand
-                        break
-                    except json.JSONDecodeError as e:
-                        # Troncature si "Extra data"
-                        if "Extra data" in str(e) and hasattr(e, "pos"):
-                            try:
-                                k2_analysis = json.loads(cand[:e.pos].strip())
-                                clean_json = cand[:e.pos].strip()
-                                break
-                            except Exception:
-                                pass
-                    
-                    # Tentative avec réparation
-                    repaired_cand = repair_json(cand)
+                
+                # 5.2 Parse with repair
+                repaired_cand = repair_json(cand)
+                try:
+                    k2_analysis = json.loads(repaired_cand)
+                    clean_json = repaired_cand
+                    break
+                except Exception:
+                    # 5.3 Parse python literal dict fallback
                     try:
-                        k2_analysis = json.loads(repaired_cand)
+                        k2_analysis = ast.literal_eval(repaired_cand)
                         clean_json = repaired_cand
                         break
                     except Exception:
-                        # Tentative avec ast.literal_eval pour Python dicts
-                        try:
-                            k2_analysis = ast.literal_eval(repaired_cand)
-                            clean_json = repaired_cand
-                            break
-                        except Exception:
-                            pass
+                        pass
             
             if k2_analysis:
                 logger.info("JSON successfully extracted and parsed.")