test_agent.py

import unittest
import os
import sys
from core.text_processor import TextProcessor
from core.character_extractor import CharacterExtractor
from core.character_analyzer import CharacterAnalyzer
from core.character_agent import CharacterAgent
from utils.text_utils import TextUtils
from utils.cache_manager import CacheManager

class TestLargeScaleAgent(unittest.TestCase):
    """大规模文本角色Agent测试"""
    
    @classmethod
    def setUpClass(cls):
        """准备测试数据"""
        cls.sample_text = """
Chapter 1: The Boy Who Lived

Mr. and Mrs. Dursley, of number four, Privet Drive, were proud to say that they were perfectly normal, thank you very much. They were the last people you'd expect to be involved in anything strange or mysterious, because they just didn't hold with such nonsense.

Harry Potter rolled over inside his blankets without waking up. One small hand closed on the letter beside him and he slept on, not knowing he was special, not knowing he was famous.

"Harry! Harry Potter!" said Hermione Granger excitedly. She was a clever girl with bushy brown hair.

Ron Weasley grinned at Harry. "Blimey, Harry, everyone knows who you are!" Ron was Harry's best friend, a tall boy with red hair.

Professor Dumbledore smiled wisely. "Harry, you have your mother's eyes," he said gently. Dumbledore was the headmaster of Hogwarts, known for his wisdom and kindness.

Harry looked confused. "I don't understand, Professor," he said quietly. Harry was a brave boy, though he often felt uncertain.

Hermione rolled her eyes. "Honestly, Harry, you need to read more!" she said in her know-it-all voice.

Ron laughed. "Don't worry, mate. We'll figure it out together," he said loyally.

""" * 50  # 重复50次模拟长文本
        
        cls.processor = TextProcessor()
        cls.chunks = cls.processor.chunk_text(cls.sample_text)
    
    def test_01_text_utils(self):
        """测试文本工具"""
        print("\n" + "="*70)
        print("测试 1: 文本工具函数")
        print("="*70)
        
        utils = TextUtils()
        
        # 测试 token 计数
        tokens = utils.count_tokens(self.sample_text[:1000])
        print(f"Token 数量: {tokens}")
        self.assertGreater(tokens, 0)
        
        # 测试语言检测
        lang = utils.detect_language(self.sample_text)
        print(f"检测语言: {lang}")
        self.assertEqual(lang, "en")
        
        # 测试句子分割
        sentences = utils.split_into_sentences(self.sample_text[:500])
        print(f"句子数量: {len(sentences)}")
        self.assertGreater(len(sentences), 0)
        
        # 测试关键词提取
        keywords = utils.extract_keywords(self.sample_text[:1000], top_n=5)
        print(f"关键词: {keywords}")
        self.assertGreater(len(keywords), 0)
    
    def test_02_cache_manager(self):
        """测试缓存管理器"""
        print("\n" + "="*70)
        print("测试 2: 缓存管理器")
        print("="*70)
        
        cache = CacheManager()
        
        # 测试缓存设置和获取
        test_key = "test_key"
        test_data = {"name": "Harry", "age": 11}
        
        cache.set(test_key, test_data)
        cached_data = cache.get(test_key)
        
        print(f"缓存数据: {cached_data}")
        self.assertEqual(cached_data, test_data)
        
        # 测试缓存存在性
        self.assertTrue(cache.exists(test_key))
        
        # 测试 JSON 缓存
        cache.save_json("test_json", test_data)
        json_data = cache.load_json("test_json")
        print(f"JSON 数据: {json_data}")
        self.assertEqual(json_data, test_data)
        
        # 测试缓存信息
        info = cache.get_cache_info()
        print(f"缓存信息: {info}")
        self.assertGreater(info['count'], 0)
        
        # 清理测试缓存
        cache.delete(test_key)
    
    def test_03_text_processing(self):
        """测试文本处理"""
        print("\n" + "="*70)
        print("测试 3: 文本处理和分块")
        print("="*70)
        
        stats = self.processor.get_statistics(self.sample_text)
        
        print(f"总字符数: {stats['total_length']:,}")
        print(f"Token数: {stats['total_tokens']:,}")
        print(f"段落数: {stats['paragraphs']}")
        print(f"句子数: {stats['sentences']}")
        print(f"语言: {stats['language']}")
        print(f"分块数: {len(self.chunks)}")
        
        self.assertGreater(len(self.chunks), 0)
        self.assertGreater(stats['total_length'], 1000)
        self.assertEqual(stats['language'], 'en')
    
    def test_04_character_extraction(self):
        """测试角色提取"""
        print("\n" + "="*70)
        print("测试 4: 角色提取")
        print("="*70)
        
        extractor = CharacterExtractor()
        characters = extractor.extract_main_characters(
            self.chunks,
            text_sample=self.sample_text[:1000],
            language="en"
        )
        
        print(f"\n提取到 {len(characters)} 个主要角色:")
        for i, char in enumerate(characters[:5], 1):
            print(f"  {i}. {char['name']}: 出现 {char['info']['count']} 次")
        
        self.assertGreater(len(characters), 0)
        
        # 检查是否包含预期角色
        char_names = [c['name'] for c in characters]
        expected_chars = ['Harry', 'Hermione', 'Ron']
        
        for expected in expected_chars:
            found = any(expected in name for name in char_names)
            if found:
                print(f"✓ 找到角色: {expected}")
            self.assertTrue(found, f"应该找到角色: {expected}")
    
    def test_05_character_analysis(self):
        """测试角色分析"""
        print("\n" + "="*70)
        print("测试 5: 角色性格分析")
        print("="*70)
        
        extractor = CharacterExtractor()
        characters = extractor.extract_main_characters(
            self.chunks,
            language="en"
        )
        
        if not characters:
            self.skipTest("没有提取到角色")
        
        # 分析第一个角色
        character = characters[0]
        print(f"\n分析角色: {character['name']}")
        
        analyzer = CharacterAnalyzer()
        representative_chunks = analyzer.select_representative_chunks(
            self.chunks,
            character['info']['chunks'][:10]
        )
        
        print(f"选择了 {len(representative_chunks)} 个代表性片段")
        
        profile = analyzer.analyze_character_batch(
            character['name'],
            representative_chunks
        )
        
        print(f"\n角色档案:")
        print(f"  名字: {profile['name']}")
        print(f"  核心特质: {profile.get('core_traits', [])}")
        print(f"  说话风格: {profile.get('speaking_style', 'N/A')}")
        print(f"  性格总结: {profile.get('personality_summary', 'N/A')[:100]}...")
        
        self.assertEqual(profile['name'], character['name'])
        self.assertIn('core_traits', profile)
        self.assertIn('speaking_style', profile)
    
    def test_06_agent_creation(self):
        """测试Agent创建"""
        print("\n" + "="*70)
        print("测试 6: 创建对话Agent")
        print("="*70)
        
        test_profile = {
            'name': 'Harry Potter',
            'core_traits': ['勇敢', '善良', '忠诚', '谦逊'],
            'speaking_style': '谦逊、真诚，有时会感到困惑',
            'behavior_patterns': '面对困难时表现出勇气',
            'values': '友谊、正义、牺牲精神',
            'emotional_style': '内敛但情感丰富',
            'relationship_style': '忠诚于朋友',
            'background': '年轻的巫师，霍格沃茨学生',
            'personality_summary': 'Harry是一个勇敢善良的年轻巫师',
            'key_quotes': ['I don\'t understand', 'What do you mean?']
        }
        
        agent = CharacterAgent(test_profile)
        
        print("\n" + agent.get_character_info())
        
        self.assertIsNotNone(agent)
        self.assertEqual(agent.character_profile['name'], 'Harry Potter')
        self.assertIsNotNone(agent.system_prompt)
    
    def test_07_conversation(self):
        """测试对话功能"""
        print("\n" + "="*70)
        print("测试 7: 对话测试")
        print("="*70)
        
        test_profile = {
            'name': 'Hermione Granger',
            'core_traits': ['聪明', '好学', '正直', '有时完美主义'],
            'speaking_style': '自信、条理清晰，有时显得严肃',
            'behavior_patterns': '总是提前完成作业，遇到问题查阅书籍',
            'values': '知识就是力量，规则应该被遵守',
            'emotional_style': '理性为主，但也有感性的一面',
            'relationship_style': '对朋友忠诚，但不容忍懒惰',
            'background': '麻瓜出身的女巫，格兰芬多学生',
            'personality_summary': 'Hermione是个极其聪明、勤奋的学生',
            'key_quotes': ['Honestly!', 'You need to read more!', 'It\'s in Hogwarts: A History']
        }
        
        agent = CharacterAgent(test_profile)
        
        test_messages = [
            "Hi Hermione, how are you?",
            "Can you help me with my homework?",
            "Do you ever just relax?"
        ]
        
        for msg in test_messages:
            print(f"\n🧑 你: {msg}")
            response = agent.chat(msg, use_memory=False)
            print(f"🎭 Hermione: {response}")
            
            self.assertIsInstance(response, str)
            self.assertGreater(len(response), 0)
        
        print(f"\n对话历史: {len(agent.conversation_history)} 条")
        self.assertGreater(len(agent.conversation_history), 0)
    
    def test_08_conversation_management(self):
        """测试对话管理功能"""
        print("\n" + "="*70)
        print("测试 8: 对话管理")
        print("="*70)
        
        test_profile = {
            'name': 'Ron Weasley',
            'core_traits': ['忠诚', '幽默', '勇敢'],
            'speaking_style': '随意、幽默',
            'personality_summary': 'Ron是Harry的忠实朋友',
            'key_quotes': ['Blimey!', 'Brilliant!']
        }
        
        agent = CharacterAgent(test_profile)
        
        # 进行对话
        agent.chat("Hello Ron!")
        agent.chat("How's it going?")
        
        initial_count = len(agent.conversation_history)
        print(f"对话数量: {initial_count}")
        self.assertGreater(initial_count, 0)
        
        # 测试重置
        agent.reset_conversation()
        print(f"重置后对话数量: {len(agent.conversation_history)}")
        self.assertEqual(len(agent.conversation_history), 0)
        
        # 测试保存和加载
        agent.chat("Test message")
        test_file = "test_conversation.json"
        agent.save_conversation(test_file)
        
        self.assertTrue(os.path.exists(test_file))
        
        # 清理
        if os.path.exists(test_file):
            os.remove(test_file)

def run_quick_demo():
    """快速演示"""
    print("\n" + "="*70)
    print("🎭 快速演示：与 Hermione Granger 对话")
    print("="*70)
    
    profile = {
        'name': 'Hermione Granger',
        'core_traits': ['聪明绝顶', '勤奋好学', '正义感强', '有时完美主义'],
        'speaking_style': '自信、条理清晰，喜欢引用书本知识，有时显得严肃',
        'behavior_patterns': '总是提前完成作业，遇到问题首先想到查阅书籍',
        'values': '知识就是力量，规则应该被遵守',
        'emotional_style': '理性为主，但面对朋友会展现温暖的一面',
        'relationship_style': '对朋友忠诚，但不容忍懒惰和不负责任',
        'background': '麻瓜出身的年轻女巫，格兰芬多学院学生',
        'key_quotes': [
            'Honestly, don\'t you two read?',
            'It\'s LeviOsa, not LeviosA!',
            'When in doubt, go to the library.',
            'Fear of a name only increases fear of the thing itself.'
        ],
        'personality_summary': 'Hermione是一个极其聪明、勤奋的学生，对知识充满渴望。她有强烈的正义感，虽然有时显得有点自以为是，但内心深处对朋友极其忠诚。'
    }
    
    agent = CharacterAgent(profile)
    print(agent.get_character_info())
    
    demo_conversations = [
        "Hi Hermione! How's your day going?",
        "I'm really struggling with my Potions homework...",
        "You always seem so confident. Don't you ever feel stressed?",
        "What's your secret to being so good at everything?"
    ]
    
    print("\n💬 开始对话演示\n")
    print("提示: 这是自动演示模式，每个问题后会自动继续")
    print("-"*70 + "\n")
    
    for i, msg in enumerate(demo_conversations, 1):
        print(f"[{i}/{len(demo_conversations)}]")
        print(f"🧑 你: {msg}")
        print("⏳ Hermione 正在思考...", end='', flush=True)
        
        response = agent.chat(msg, use_memory=False)
        
        print(f"\r🎭 Hermione: {response}\n")
        
        if i < len(demo_conversations):
            try:
                input("[按 Enter 继续下一个问题...]")
            except:
                break
    
    print("\n" + "="*70)
    print("演示完成！")
    print("="*70)

def run_interactive_demo():
    """交互式演示"""
    print("\n" + "="*70)
    print("🎭 交互式演示：与 Harry Potter 对话")
    print("="*70)
    
    profile = {
        'name': 'Harry Potter',
        'core_traits': ['勇敢', '善良', '忠诚', '谦逊', '有正义感'],
        'speaking_style': '真诚、谦逊，说话时常带有不确定感',
        'behavior_patterns': '面对危险时挺身而出，但对赞美感到不自在',
        'values': '友谊高于一切，愿意为朋友牺牲',
        'emotional_style': '内敛但情感深沉，容易被朋友的遭遇触动',
        'relationship_style': '对朋友极度忠诚，珍视每一段友谊',
        'background': '失去父母的年轻巫师，被预言为击败伏地魔的人',
        'key_quotes': [
            'I don\'t know what you\'re talking about.',
            'We can figure this out together.',
            'I had a lot of help.',
            'It\'s not about being brave, it\'s about doing what\'s right.'
        ],
        'personality_summary': 'Harry是一个勇敢善良但又谦逊的年轻巫师。虽然承担着巨大的使命，但他从不认为自己特殊。他最看重的是友谊，愿意为朋友付出一切。'
    }
    
    agent = CharacterAgent(profile)
    print(agent.get_character_info())
    
    print("\n💬 开始交互式对话")
    print("-"*70)
    print("提示:")
    print("  • 输入你的问题开始对话")
    print("  • 输入 'quit' 退出")
    print("  • 输入 'reset' 重置对话")
    print("="*70 + "\n")
    
    while True:
        try:
            user_input = input("🧑 你: ").strip()
            
            if not user_input:
                continue
            
            if user_input.lower() in ['quit', 'exit', 'q']:
                print(f"\n🎭 Harry: Goodbye! It was nice talking to you.\n")
                break
            
            if user_input.lower() == 'reset':
                agent.reset_conversation()
                continue
            
            print("⏳ Harry 正在思考...", end='', flush=True)
            response = agent.chat(user_input, use_memory=False)
            print(f"\r🎭 Harry: {response}\n")
            
        except KeyboardInterrupt:
            print(f"\n\n🎭 Harry: Goodbye!\n")
            break
        except Exception as e:
            print(f"\n❌ 错误: {e}\n")

if __name__ == "__main__":
    import sys
    
    if len(sys.argv) > 1:
        if sys.argv[1] == "demo":
            run_quick_demo()
        elif sys.argv[1] == "interactive":
            run_interactive_demo()
        else:
            print("用法:")
            print("  python test_agent.py          - 运行所有测试")
            print("  python test_agent.py demo     - 快速演示")
            print("  python test_agent.py interactive - 交互式演示")
    else:
        # 运行单元测试
        unittest.main(verbosity=2)