MedResearcher/experiment_runner.py

#!/usr/bin/env python3
"""
数字提取实验运行器

基于Agno框架的SubAgent系统，从文本中提取数字并进行解释
展示SubAgent系统的核心功能：配置读取、模型创建、动态prompt构建、JSON解析
"""

import sys
import os
from typing import List

# 添加src路径到Python路径
sys.path.append(os.path.join(os.path.dirname(__file__), 'src'))

from src.agent_system import SubAgent, create_json_agent
from number_extraction_models import NumberExtractionResult


def create_number_extraction_agent() -> SubAgent:
    """创建数字提取专用SubAgent"""
    
    instructions = [
        "你是一个专业的数字提取专家",
        "你的任务是从给定文本中识别和提取所有的数字",
        "重点关注：整数、小数、百分比、比率、统计数据等",
        "对每个数字提供准确的上下文和清晰的解释"
    ]
    
    prompt_template = """
请仔细分析以下文本，提取其中的所有数字：

【输入文本】
{input_text}

【提取要求】
1. 识别文本中的所有数字（整数、小数、百分比、比率等）
2. 对每个数字进行分类和解释
3. 记录数字出现的完整上下文
4. 识别数字的单位（如果有）
5. 提供对整个文本中数字的总体摘要
"""

    try:
        agent = SubAgent(
            provider="aliyun",
            model_name="qwen-max",
            name="number_extractor",
            description="专业的文本数字提取和解释系统",
            instructions=instructions,
            prompt_template=prompt_template,
            response_model=NumberExtractionResult
        )
        
        print("✅ 数字提取Agent初始化成功")
        return agent
        
    except Exception as e:
        print(f"❌ Agent初始化失败: {e}")
        raise


def display_results(result: NumberExtractionResult):
    """展示提取结果"""
    print("\n" + "="*50)
    print("🔍 数字提取结果")
    print("="*50)
    
    print(f"摘要: {result.summary}")
    print(f"总数: {result.total_count}")
    
    if result.extractions:
        for i, item in enumerate(result.extractions, 1):
            print(f"\n数字{i}: {item.number}")
            if item.unit:
                print(f"  单位: {item.unit}")
            print(f"  上下文: {item.context}")
            print(f"  解释: {item.explanation}")
    
    print("="*50)


def get_user_input() -> str:
    """获取用户输入"""
    print("\n请输入要分析的文本:")
    print("(输入'quit'退出，输入'END'结束输入)")
    
    lines = []
    while True:
        try:
            line = input(">>> ").strip()
            if line.lower() == 'quit':
                return 'quit'
            elif line.upper() == 'END':
                break
            else:
                lines.append(line)
                break
        except (KeyboardInterrupt, EOFError):
            return 'quit'
    
    return '\n'.join(lines).strip()


def main():
    """主函数"""
    print("🎉 数字提取工具启动!")
    
    try:
        agent = create_number_extraction_agent()
    except Exception as e:
        print(f"初始化失败: {e}")
        return
    
    print("✅ 系统就绪")
    
    while True:
        user_input = get_user_input()
        
        if user_input == 'quit':
            print("再见!")
            break
        
        if not user_input:
            continue
        
        try:
            print(f"分析中... (文本长度: {len(user_input)}字符)")
            result = agent.run(template_vars={'input_text': user_input})
            display_results(result)
            
        except Exception as e:
            print(f"提取失败: {e}")
            print("请尝试重新输入")


def test_basic():
    """基础测试"""
    try:
        agent = create_number_extraction_agent()
        
        # 测试prompt构建
        prompt = agent.build_prompt({'input_text': '测试数字95.2%'})
        print(f"✅ 测试通过，prompt长度: {len(prompt)}")
        return True
        
    except Exception as e:
        print(f"❌ 测试失败: {e}")
        return False


if __name__ == "__main__":
    main()