Model Router

__pycache__/

#!/usr/bin/env python3
"""
Model Router - 智能模型选择器
根据用户问题类型自动选择最合适的大模型

支持 8 大场景：对话聊天、代码开发、图像生成、文本协作、营销文案、数据分析、知识问答、角色扮演
"""

import re

# 模型配置 - 8 大场景 (1 主 +2 备)
MODELS = {
    "chat": {
        "primary": "alibaba-cloud/qwen3.5-plus",
        "fallbacks": ["alibaba-cloud/glm-5", "alibaba-cloud/qwen-plus"],
        "description": "高质量对话聊天",
        "triggers": ["聊天", "对话", "咨询", "建议", "心情", "聊聊", "陪伴", "讲故事", "倾诉", "寂寞", "无聊"],
        "priority": 1
    },
    "code": {
        "primary": "alibaba-cloud/qwen-coder",
        "fallbacks": ["alibaba-cloud/qwen3.5-plus", "alibaba-cloud/qwen-plus"],
        "description": "代码开发专用",
        "triggers": ["代码", "编程", "函数", "bug", "调试", "开发", "script", "python", "javascript", "java", "写个", "脚本", "接口", "api", "github", "编程"],
        "priority": 2
    },
    "image": {
        "primary": "wanx-v1",
        "fallbacks": ["qwen-image-edit", "flux-schnell"],
        "description": "图像生成",
        "triggers": ["图片", "图像", "生成图片", "画图", "自拍", "照片", "design", "设计图", "泳装", "美照", "头像", "插画"],
        "priority": 3,
        "note": "需要调用图像 API，不是对话模型"
    },
    "writing": {
        "primary": "alibaba-cloud/qwen-plus",
        "fallbacks": ["alibaba-cloud/qwen3.5-plus", "alibaba-cloud/glm-5"],
        "description": "文本协作/文章写作",
        "triggers": ["写文章", "润色", "翻译", "总结", "摘要", "改写", "扩写", "邮件", "报告", "文档"],
        "priority": 4
    },
    "marketing": {
        "primary": "alibaba-cloud/qwen-plus",
        "fallbacks": ["alibaba-cloud/qwen3.5-plus", "alibaba-cloud/glm-5"],
        "description": "营销文案",
        "triggers": ["文案", "广告", "营销", "标题", "slogan", "推广", "种草", "小红书", "朋友圈", "微博"],
        "priority": 5
    },
    "data": {
        "primary": "alibaba-cloud/qwen-max",
        "fallbacks": ["alibaba-cloud/qwen3.5-plus", "alibaba-cloud/qwen-plus"],
        "description": "数据分析",
        "triggers": ["分析", "数据", "统计", "图表", "excel", "表格", "计算", "报表", "可视化", "数据库"],
        "priority": 6
    },
    "simple": {
        "primary": "alibaba-cloud/glm-5",
        "fallbacks": ["alibaba-cloud/qwen3.5-plus", "alibaba-cloud/qwen-plus"],
        "description": "简单问答",
        "triggers": ["是什么", "定义", "简单", "快速", "什么意思", "解释", "介绍", "概念"],
        "priority": 7
    },
    "roleplay": {
        "primary": "alibaba-cloud/qwen3.5-plus",
        "fallbacks": ["alibaba-cloud/qwen-plus", "alibaba-cloud/glm-5"],
        "description": "角色扮演",
        "triggers": ["扮演", "模拟", "身份", "角色", "假装", "cosplay"],
        "priority": 8
    }
}

def analyze_query(query: str) -> dict:
    """分析用户查询，返回推荐的模型（含主备）"""
    query_lower = query.lower()
    
    matches = []
    
    for category, config in MODELS.items():
        score = 0
        matched_triggers = []
        
        for trigger in config["triggers"]:
            if trigger.lower() in query_lower:
                score += 1
                matched_triggers.append(trigger)
        
        if score > 0:
            matches.append({
                "category": category,
                "primary": config["primary"],
                "fallbacks": config["fallbacks"],
                "description": config["description"],
                "score": score,
                "matched_triggers": matched_triggers,
                "priority": config["priority"],
                "note": config.get("note", "")
            })
    
    # 按分数排序，分数相同按优先级排序
    matches.sort(key=lambda x: (-x["score"], x["priority"]))
    
    if matches:
        best = matches[0]
        return {
            "recommended_model": best["primary"],
            "primary": best["primary"],
            "fallbacks": best["fallbacks"],
            "category": best["category"],
            "description": best["description"],
            "confidence": best["score"],
            "matched_triggers": best["matched_triggers"],
            "all_matches": matches,
            "note": best["note"]
        }
    else:
        # 默认返回 qwen3.5-plus
        default = MODELS["chat"]
        return {
            "recommended_model": default["primary"],
            "primary": default["primary"],
            "fallbacks": default["fallbacks"],
            "category": "chat",
            "description": "高质量对话聊天",
            "confidence": 0,
            "matched_triggers": [],
            "all_matches": [],
            "note": "未匹配到特定类型，使用默认对话模型"
        }

def get_model_for_task(task_type: str) -> dict:
    """根据任务类型获取模型配置（主 + 备）"""
    if task_type not in MODELS:
        return {
            "primary": "alibaba-cloud/qwen3.5-plus",
            "fallbacks": ["alibaba-cloud/glm-5", "alibaba-cloud/qwen-plus"]
        }
    config = MODELS[task_type]
    return {
        "primary": config["primary"],
        "fallbacks": config["fallbacks"]
    }

def main():
    """命令行交互"""
    print("🤖 Model Router - 智能模型选择器")
    print("=" * 50)
    print("支持场景：对话、代码、图像、写作、营销、数据、问答、角色扮演")
    print("输入问题，我会推荐最合适的模型~")
    print("输入 'quit' 退出\n")
    
    while True:
        try:
            query = input("主人：").strip()
            if query.lower() in ["quit", "exit", "q"]:
                print("拜拜主人~ 💋")
                break
            
            if not query:
                continue
            
            result = analyze_query(query)
            
            print(f"\n📊 分析结果:")
            print(f"  场景类型：{result['category']} ({result['description']})")
            print(f"  主模型：{result['primary']}")
            print(f"  备用 1: {result['fallbacks'][0]}")
            print(f"  备用 2: {result['fallbacks'][1]}")
            print(f"  匹配度：{result['confidence']} 个关键词")
            
            if result['matched_triggers']:
                print(f"  匹配关键词：{', '.join(result['matched_triggers'])}")
            
            if result['note']:
                print(f"  ⚠️  注意：{result['note']}")
            
            print()
            
        except KeyboardInterrupt:
            print("\n拜拜主人~ 💋")
            break
        except Exception as e:
            print(f"❌ 错误：{e}")

if __name__ == "__main__":
    main()