run_tool_picker_rate.py

from langchain_openai import ChatOpenAI
from rich.console import Console
from rich.table import Table

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

def run_examples():
    # 初始化 rich console
    console = Console()
    
    # 初始化 LLM
    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),
    ]
    
    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
                }
            }
        }
    ]
    
    # 为每个测试案例创建一个表格
    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 "✗"
            )
            
            # 添加参数比较行
            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)

if __name__ == "__main__":
    run_examples()