test_thinking.py

#!/usr/bin/env python3
"""
Test script to verify thinking extraction works correctly
"""

import re
from typing import Tuple

def extract_thinking_and_response(text: str) -> Tuple[str, str]:
    """Extract thinking process and clean response from AI output"""
    thinking = ""
    response = text
    
    # Extract thinking content
    thinking_match = re.search(r'<think>(.*?)</think>', text, re.DOTALL)
    if thinking_match:
        thinking = thinking_match.group(1).strip()
        response = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
    
    # Clean up the response
    response = re.sub(r'^(Assistant:|AI:|Response:|Answer:)\s*', '', response.strip())
    response = re.sub(r'\[INST\].*?\[\/INST\]', '', response)
    response = re.sub(r'<\|.*?\|>', '', response)
    
    return thinking.strip(), response.strip()

# Test cases
test_cases = [
    # Basic thinking extraction
    "<think>I need to analyze this query and generate search terms.</think>Here are the search queries:\n1. AI developments 2024\n2. artificial intelligence news",
    
    # No thinking
    "Here are the search queries:\n1. AI developments 2024\n2. artificial intelligence news",
    
    # Complex thinking
    "<think>The user is asking about quantum computing breakthroughs. I should focus on recent developments, key research areas, and practical applications. Let me think about the best search terms...</think>Based on your question, here are optimized search queries for quantum computing breakthroughs.",
    
    # Thinking with newlines
    "<think>\nThis is a complex question about climate change.\nI need to consider multiple aspects:\n1. Ocean currents\n2. Temperature changes\n3. Recent research\n</think>Climate change affects ocean currents in several ways..."
]

print("Testing thinking extraction...")
for i, test in enumerate(test_cases, 1):
    thinking, response = extract_thinking_and_response(test)
    print(f"\nTest {i}:")
    print(f"Original: {test[:100]}...")
    print(f"Thinking: {thinking[:100]}..." if thinking else "Thinking: None")
    print(f"Response: {response[:100]}...")
    print("-" * 50)

print("✅ Thinking extraction test completed!")