test_length.py

#!/usr/bin/env python3
"""
Test script to verify length instructions are working
"""

import sys
import os

# Add current directory to Python path
script_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, script_dir)

def test_length_instructions():
    """Test that length instructions are properly formatted"""
    print("🧪 Testing Length Instructions...")
    
    try:
        from src.autocomplete import SmartAutoComplete
        
        # Create mock settings
        class MockSettings:
            def __init__(self):
                self.OPENAI_API_KEY = "test-key"
                self.ANTHROPIC_API_KEY = ""
                self.DEFAULT_PROVIDER = "openai"
                self.CACHE_TTL = 3600
                self.CACHE_MAX_SIZE = 100
        
        # Create mock API client that captures the messages
        class MockAPIClient:
            def __init__(self, settings=None):
                self.last_messages = None
                
            def get_completion(self, messages, temperature=0.7, max_tokens=150, provider=None):
                self.last_messages = messages
                print(f"\n📝 API called with max_tokens: {max_tokens}")
                print(f"📝 System prompt: {messages[0]['content'][:200]}...")
                print(f"📝 User message: {messages[1]['content'][:200]}...")
                return f"Mock completion response ({max_tokens} tokens requested)"
        
        # Create mock cache
        class MockCacheManager:
            def __init__(self, settings=None):
                pass
            def get(self, key):
                return None
            def set(self, key, value):
                pass
        
        # Test setup
        settings = MockSettings()
        autocomplete = SmartAutoComplete(settings)
        autocomplete.api_client = MockAPIClient(settings)
        autocomplete.cache_manager = MockCacheManager(settings)
        
        # Test different token lengths
        test_cases = [
            (50, "short"),
            (150, "medium"), 
            (300, "long"),
            (500, "very long")
        ]
        
        for max_tokens, description in test_cases:
            print(f"\n🔍 Testing {description} output ({max_tokens} tokens):")
            
            suggestions = autocomplete.get_suggestions(
                text="Dear Mr. Johnson,",
                context="email",
                max_tokens=max_tokens,
                user_context="Meeting about quarterly budget"
            )
            
            # Check if the messages contain the token count
            messages = autocomplete.api_client.last_messages
            system_prompt = messages[0]['content']
            user_message = messages[1]['content']
            
            # Verify token count is mentioned
            token_in_system = str(max_tokens) in system_prompt
            token_in_user = str(max_tokens) in user_message
            
            print(f"   ✅ Token count in system prompt: {token_in_system}")
            print(f"   ✅ Token count in user message: {token_in_user}")
            print(f"   ✅ Length instruction present: {'IMPORTANT' in user_message}")
            
            # Check for appropriate length guidance
            if max_tokens <= 100:
                expected_guidance = "concise and brief"
            elif max_tokens <= 200:
                expected_guidance = "moderate length"
            elif max_tokens <= 300:
                expected_guidance = "detailed response"
            else:
                expected_guidance = "comprehensive and detailed"
            
            guidance_present = expected_guidance in user_message
            print(f"   ✅ Appropriate guidance ({expected_guidance}): {guidance_present}")
            
            if not (token_in_system or token_in_user):
                print(f"   ❌ Token count not found in prompts!")
                return False
        
        print("\n✅ All length instruction tests passed!")
        return True
        
    except Exception as e:
        print(f"❌ Length instruction test failed: {str(e)}")
        import traceback
        traceback.print_exc()
        return False

def main():
    """Main test function"""
    print("🚀 Smart Auto-Complete Length Test")
    print("=" * 50)
    
    if test_length_instructions():
        print("\n🎉 Length instructions are working correctly!")
        print("\n💡 The AI should now generate responses of the requested length.")
        print("   - 50 tokens: ~1-2 sentences")
        print("   - 150 tokens: ~3-4 sentences") 
        print("   - 300 tokens: ~1-2 paragraphs")
        print("   - 500 tokens: ~2-3 paragraphs")
    else:
        print("\n❌ Length instruction tests failed.")
        return 1
    
    return 0

if __name__ == "__main__":
    sys.exit(main())