Update app.py

Update the logic to separate the <think> block and the JSON block in the generated 'response'.

app.py

@@ -133,7 +133,39 @@ class BasicModel:
         response = outputs[0]["generated_text"]
         assert isinstance(response, str)
 
-        # 3. Optionally map SmolAgents’ stop_sequences → HF pipeline’s 'stop'
+        # 3. The response consists of a <think></think> block and a JSON block. we want to separate the two blocks.
+        ## Save the <think></think> block to print in the log
+        try:
+            # This regex finds the content between <think> and </think>
+            thought_pattern = r"<think>(.*?)</think>"
+            # re.DOTALL allows the '.' to match newline characters
+            match = regex.search(thought_pattern, response, flags=regex.DOTALL)
+            if match:
+                # .group(1) gets the content captured by the (.*?) part of the pattern
+                thought_content = match.group(1).strip()
+                print(f"--- CAPTURED THOUGHT ---\n{thought_content}\n------------------------")
+        except Exception as e:
+            # This is just for logging, so we don't want it to crash the main process.
+            print(f"Could not extract <think> block: {e}")
+            
+        ## Pass the JSON block as the 'response'
+        try:
+            # Find the first '{' and the last '}'
+            start_index = response.find('{')
+            end_index = response.rfind('}')
+            # If both braces are found, extract the substring
+            if start_index != -1 and end_index != -1 and end_index > start_index:
+                json_substring = response[start_index : end_index + 1]
+                # Try to parse the substring to ensure it's valid JSON
+                json.loads(json_substring)
+                # If it's valid, we replace the original response with just the clean JSON part.
+                response = json_substring  
+        # If no valid JSON is found, we proceed with the original response string.
+        except (json.JSONDecodeError, TypeError):
+            # Pass silently and let the framework handle the potentially malformed string.
+            pass
+
+        # 4. Optionally map SmolAgents’ stop_sequences → HF pipeline’s 'stop'
         if stop_sequences:
             # find the earliest occurrence of any stop token
             cuts = [response.find(s) for s in stop_sequences if response.find(s) != -1]
@@ -142,7 +174,7 @@ class BasicModel:
 
         print(f"\n--- RAW MODEL OUTPUT ---\n{response}\n------------------------")
 
-        # 4. NEW: Parse, Fix, and Re-serialize the agent's code output
+        # 5. NEW: Parse, Fix, and Re-serialize the agent's code output
         try:
             # The agent's response is expected to be a JSON string, according to the system prompt
             # Parse it into a Python dictionary
@@ -178,7 +210,7 @@ class BasicModel:
 
         print(f"\n--- FIXED AND FINAL OUTPUT ---\n{response}\n----------------------------")
         
-        # 5. Wrap back into a chat message dict
+        # 6. Wrap back into a chat message dict
         return ChatMessage(role="assistant", content=response)
         # return {
         #     "role": 'assistant',