Upload alpha_factory/infra/llm_client.py

alpha_factory/infra/llm_client.py

@@ -1,7 +1,8 @@
 """
-LLM Client v3.1 — unified interface with token budget enforcement,
+LLM Client v3.2 — unified interface with token budget enforcement,
 retry logic for transient failures, and proper error classification.
 No more error-amplification via 3-fallback strategy.
+Fixed thinking-tag regex that was mangling JSON responses.
 """
 import asyncio
 import json
@@ -256,24 +257,20 @@ class LLMClient:
         )
         return self._parse_json_response(result["content"], schema)
 
-    def _check_token_count(self, budget: int) -> bool:
-        """Check if accumulated token count has exceeded budget. Sets flag if so."""
-        if self._token_count >= budget:
-            self._budget_exceeded = True
-            return True
-        return False
-
     def _parse_json_response(self, content: str, schema: type[T]) -> T:
         """
         Parse JSON from LLM response, handling common issues:
         - Markdown code blocks
         - Leading/trailing text
         - Thinking tags (<think>...</think>, <thinking>...</thinking>)
+        
+        CAUTION: Thinking-tag regex must be SPECIFIC to known tag formats.
+        Previous broad regex `\s*.*?\s*` was stripping all whitespace!
         """
         text = content.strip()
 
-        # Remove thinking tags (Qwen/DeepSeek R1 style)
-        # Match <think>...</think> or <thinking>...</thinking>
+        # Remove thinking tags (DeepSeek/Qwen R1 style)
+        # ONLY match known XML-style thinking tags, not arbitrary whitespace
         text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL).strip()
         text = re.sub(r'<thinking>.*?</thinking>', '', text, flags=re.DOTALL).strip()