run_tool_picker_rate.py

from langchain_openai import ChatOpenAI
from rich.console import Console
from rich.table import Table
import time
import sys,os

sys.path.append("../")


from src.server.tool_picker import ToolPicker
from src.agent.tool_rate import RegionRateTool, RankingRateTool
from src.agent.tool_monitor import MonitorPointTool
from src.agent.tool_warn import WarningTool

def create_tool_picker():
    llm = ChatOpenAI(
        openai_api_key="xxxxxx",
        openai_api_base="http://192.168.10.14:8000/v1",
        model_name="Qwen2-7B",
        verbose=True
    )

    base_url = "http://172.30.0.37:30007"
    # 初始化工具
    tools = [
        RegionRateTool(base_url=base_url),
        RankingRateTool(base_url=base_url),
        MonitorPointTool(base_url=base_url),
        WarningTool(base_url=base_url),
    ]
    return ToolPicker(llm, tools)


def run_examples():
    # 初始化 rich console
    console = Console()
    
    # 初始化 LLM
    
    base_url = "http://172.30.0.37:30007"
    # 初始化工具
    tools = [
        RegionRateTool(base_url=base_url),
        RankingRateTool(base_url=base_url),
        MonitorPointTool(base_url=base_url),
        WarningTool(base_url=base_url),
    ]
    
    tool_dict = {tool.name: tool for tool in tools}
    # 初始化 ToolPicker
    picker = ToolPicker(llm, tools)
    
    # 测试案例和预期结果
    test_cases = [
        {
            "query": "请分析下今天全国各地区在线率情况",
            "expected": {
                "tool": "region_online_rate",
                "params": {
                    "start_time": "2024-11-19",
                    "end_time": "2024-11-19",
                    "region_name": "",
                    "month_statistics": False
                }
            }
        },
        {
            "query": "请分析下今天甘肃省设备在线率情况",
            "expected": {
                "tool": "region_online_rate",
                "params": {
                    "start_time": "2024-11-19",
                    "end_time": "2024-11-19",
                    "region_name": "甘肃省",
                    "month_statistics": False
                }
            }
        },
        {
            "query": "查询2024年11月13日各地区排名情况",
            "expected": {
                "tool": "online_rate_ranking",
                "params": {
                    "rate_type": 1
                }
            }
        },
        {
            "query": "查询各厂商在线率排名情况",
            "expected": {
                "tool": "online_rate_ranking",
                "params": {
                    "rate_type": 2
                }
            }
        },
        {
            "query": "2023年甘肃省每月的设备在线率分别是多少？",
            "expected": {
                "tool": "region_online_rate",
                "params": {
                    "start_time": "2023-01-01",
                    "end_time": "2023-12-31",
                    "region_name": "甘肃省",
                    "month_statistics": True
                }
            }
        },
        {
            "query": "查询2024年甘肃省各月在线率",
            "expected": {
                "tool": "region_online_rate",
                "params": {
                    "start_time": "2024-01-01",
                    "end_time": "2024-12-31",
                    "region_name": "甘肃省",
                    "month_statistics": True
                }
            }
        },
        {
            "query": "2024年10月15日，兰州市的设备在线率是多少？",
            "expected": {
                "tool": "region_online_rate",
                "params": {
                    "start_time": "2024-10-15",
                    "end_time": "2024-10-15",
                    "region_name": "兰州市",
                    "month_statistics": False
                }
            }
        },
        {
            "query": "2024年，兰州市榆中县的设备在线率是多少？",
            "expected": {
                "tool": "region_online_rate",
                "params": {
                    "start_time": "2024-11-26",  "region_name": "兰州市榆中县", "month_statistics": False
                }
            }
        },
        {
            "query": "2023年甘肃省每月的设备在线率分别是多少？",
            "expected": {
                "tool": "region_online_rate",
                "params": {
                    "start_time": "2023-01-01", "end_time": "2023-12-31", "region_name": "甘肃省", "month_statistics": True
                }
            }
        },
        {
            "query": "2023年甘肃省按月统计设备在线率？",
            "expected": {
                "tool": "region_online_rate",
                "params": {
                    "start_time": "2023-01-01", "region_name": "甘肃省", "month_statistics": True
                }
            }
        },
        {
            "query": "2023年全国每个月的设备在线率",
            "expected": {
                "tool": "region_online_rate",
                "params": {
                    "start_time": "2023-01-01", "region_name": "", "month_statistics": True
                }
            }
        },
        {
            "query": "2023年1月-2023年12月期间青海实验点在线率是多少？",
            "expected": {
                "tool": "region_online_rate",
                "params": {
                    "start_time": "2023-01-01", "end_time": "2023-12-31", "region_name": "青海", "month_statistics": False
                }
            }
        },
        {
            "query": "2023年1月-2023年12月期间青海实验点各月在线率是多少？",
            "expected": {
                "tool": "region_online_rate",
                "params": {
                    "start_time": "2023-01-01", "end_time": "2023-12-31", "region_name": "青海", "month_statistics": True
                }
            }
        },
        {
            "query": "2022年各个月设备在线率统计；",
            "expected": {
                "tool": "region_online_rate",
                "params": {
                    "start_time": "2022-01-01", "end_time": "2022-12-31", "region_name": "", "month_statistics": True
                }
            }
        },
        {
            "query": "2023年1月-2023年12月期间西藏实验点每个月在线率是多少？",
            "expected": {
                "tool": "region_online_rate",
                "params": {
                    "start_time": "2023-01-01", "end_time": "2023-12-31", "region_name": "西藏", "month_statistics": True
                }
            }
        },
        {
            "query": "2023年1月-2023年12月期间西藏实验点每个月在线率是多少？",
            "expected": {
                "tool": "region_online_rate",
                "params": {
                    "start_time": "2023-01-01", "end_time": "2023-12-31", "region_name": "西藏", "month_statistics": True
                }
            }
        },
        {
            "query": "贵阳市各县区设备在线率是多少",
            "expected": {
                "tool": "region_online_rate",   
                "params": {
                    "region_name": "贵阳市"
                }
            }
        }
    ]
    

    tool_selected_success = 0
    # 为每个测试案例创建一个表格
    for i, case in enumerate(test_cases, 1):
        console.print(f"\n[bold cyan]=== 测试案例 {i} ===[/bold cyan]")
        
        table = Table(title=f"查询: {case['query']}")
        table.add_column("项目", style="cyan")
        table.add_column("预期结果", style="green")
        table.add_column("实际结果", style="yellow")
        table.add_column("是否匹配", style="magenta")
        
        try:
            result = picker.pick(case["query"])
            
            # 添加工具比较行
            expected_tool = case["expected"]["tool"]
            actual_tool = result["tool"]
            table.add_row(
                "选择的工具",
                expected_tool,
                actual_tool,
                "✓" if expected_tool == actual_tool else "✗"
            )
            
            if expected_tool == actual_tool:
                tool_selected_success += 1
            
            # 添加参数比较行
            for param_key in case["expected"]["params"]:
                expected_value = str(case["expected"]["params"][param_key])
                actual_value = str(result["params"].get(param_key, "未提供"))
                table.add_row(
                    f"参数: {param_key}",
                    expected_value,
                    actual_value,
                    "✓" if expected_value == actual_value else "✗"
                )

            tool = tool_dict[result["tool"]]
            params = result["params"]

            # result =   tool.invoke(params)
            # print(result)
            
        except Exception as e:
            table.add_row("错误", "", str(e), "✗")
        
        console.print(table)
        console.print("=" * 80)

    print(f"工具选择成功率: {tool_selected_success} / {len(test_cases)} = {tool_selected_success / len(test_cases) * 100:.2f}%")

def run_case_17():
    console = Console()
    picker = create_tool_picker()

    

    success_count = 0
    total_time = 0
    test_cases_count = 50
    for i in range(test_cases_count):
        start_time = time.time()
        query = "贵阳市各县区设备在线率是多少"
        expected = {
            "tool": "region_online_rate",
            "params": {
                "region_name": "贵阳市"
            }
        }

        result = picker.pick(query)
        exec_time = time.time() - start_time
        total_time += exec_time
        table = Table(title=f"第 {i+1:02d} 次测试结果 【{query}】")
        table.add_column("查询", style="cyan")
        table.add_column("预期工具", style="green")
        table.add_column("实际工具", style="yellow")
        table.add_column("是否匹配", style="magenta")
        table.add_column("耗时", style="magenta")
        table.add_row(query, expected["tool"], result["tool"], "✓" if result["tool"] == expected["tool"] else "✗", f"{exec_time:.2f}秒")

        if result["tool"] == expected["tool"]:
            success_count += 1
        console.print(table)
    
    print(f"工具选择成功率: {success_count} / {test_cases_count} = {success_count / test_cases_count * 100:.2f}%")
    print(f"平均耗时: {total_time / test_cases_count:.2f}秒 总共耗时: {total_time:.2f}秒")
if __name__ == "__main__":
    # run_examples()
    run_case_17()