import os

# 可以指定一个绝对路径，统一存放所有的Embedding和LLM模型。
# 每个模型可以是一个单独的目录，也可以是某个目录下的二级子目录。
# 如果模型目录名称和 MODEL_PATH 中的 key 或 value 相同，程序会自动检测加载，无需修改 MODEL_PATH 中的路径。
MODEL_ROOT_PATH = "/home/gc/gangyan/models"

# 选用的 Embedding 名称
EMBEDDING_MODEL = "bge-m3-api"

# 知识库表里或旧环境可能仍为本地键名；迁移内网 embedding API 后统一映射到 ONLINE_LLM_MODEL 的键
EMBED_MODEL_ALIASES = {
    "bge_m3": "bge-m3-api",
}

# LLM 名称别名映射：用于旧前端/历史数据仍请求无权限模型时的兼容与兜底
# 例如内网网关令牌不允许 Qwen2-72B-Instruct，则统一转到 deepseek-v3
LLM_MODEL_ALIASES = {
    "Qwen2-72B-Instruct": "deepseek-v3",
}

# Embedding 模型运行设备。设为 "auto" 会自动检测(会有警告)，也可手动设定为 "cuda","mps","cpu","xpu" 其中之一。
EMBEDDING_DEVICE = "cpu"

# mivlus混合检索条件
EXPR = ""

# 选用的reranker模型
RERANKER_MODEL = "bge-reranker"
# 是否启用reranker模型
USE_RERANKER = False
RERANKER_MAX_LENGTH = 1024

# 如果需要在 EMBEDDING_MODEL 中增加自定义的关键字时配置
EMBEDDING_KEYWORD_FILE = "keywords.txt"
EMBEDDING_MODEL_OUTPUT_PATH = "output"

#模型数据库的定义
POLICY_KNOWLEDGE_BASE = "t_policy_total_bge_new_v2" 
# REPORT_KNOWLEDGE_BASE = "t_strategy_report_bge_v2"
REPORT_KNOWLEDGE_BASE = "gydemo_report_v2"
JOURNAL_KNOWLEDGE_BASE = "t_journal_article_bge_v1"

# 润色改写等功能相关模型设置
STRATEGY_MODEL_DICT = {
     'DEFAULT_EXTRACT_KEYWORDS_MODEL_NAME': "deepseek-v3",
     'DEFAULT_QUERY_REWRITE_MODEL_NAME': "deepseek-v3",
     'DEFAULT_SUMMARY_MODEL_NAME': "deepseek-v3",
     'DEFAULT_RRQJ_MODEL_NAME': "deepseek-v3",
     'CONTINUE_WRITE_MODEL_NAME': "deepseek-v3",
     'REWRITE_MODEL_NAME': "deepseek-v3",
     'EXPAND_WRITE_MODEL_NAME': "deepseek-v3",
     'ABB_REWRITE_MODEL_NAME': "deepseek-v3",
     'EMBELLISH_MODEL_NAME': "deepseek-v3",
     'CHI_TO_ENS_MODEL_NAME': "deepseek-v3",
     'ENS_TO_CHI_MODEL_NAME': "deepseek-v3",
     'FORMAL_STYLE_MODEL_NAME': "deepseek-v3",
     'PARTY_STYLE_MODEL_NAME': "deepseek-v3",
     'COLLOQUIAL_STYLE_MODEL_NAME': "deepseek-v3",
}

# 要运行的 LLM 名称，可以包括本地模型和在线模型。列表中本地模型将在启动项目时全部加载。
# 列表中第一个模型将作为 API 和 WEBUI 的默认模型。
# 在这里，使用目前主流的两个离线模型，其中，chatglm3-6b 为默认加载模型。
# 如果你的显存不足，可使用 Qwen-1_8B-Chat, 该模型 FP16 仅需 3.8G显存。
LLM_MODELS = ["deepseek-v3", "deepseek-r1", "deepseek-chat", "qwen-max", "Qwen2-72B-Instruct"]
Agent_MODEL = None

# LLM 模型运行设备。设为"auto"会自动检测(会有警告)，也可手动设定为 "cuda","mps","cpu","xpu" 其中之一。
LLM_DEVICE = "cuda"

HISTORY_LEN = 20

MAX_TOKENS = None
MAX_CUT_TOKENS = 30 * 1024

TEMPERATURE = 0.7
DEEPSEEK_MODELS = ["deepseek-reasoner", "deepseek-chat"]
CAST_MODELS = ["kexie_0.5b"]
ONLINE_LLM_MODEL = {
    # 本地部署的大模型 API (10.102.24.75:3000)
    "bge-m3-api": {
        "model_name": "bge-m3",
        "api_base_url": "http://10.102.24.75:3000/v1",
        "api_key": "sk-nDr7vDHOxJQOGFxbmSE6g2wCK0WELtyZst4kD3eo4383P6j5",
        "provider": "OpenAIWorker",
    },
    "bge-reranker": {
        "model_name": "bge-reranker",
        "api_base_url": "http://10.102.24.75:3000/v1",
        "api_key": "sk-nDr7vDHOxJQOGFxbmSE6g2wCK0WELtyZst4kD3eo4383P6j5",
        "provider": "OpenAIWorker",
    },
    "deepseek-v3": {
        "model_name": "deepseek-v3",
        "api_base_url": "http://10.102.24.75:3000/v1",
        "api_key": "sk-BlQIGRrotbVDWE5mXCPBFjVWIvJ83hldzz67xInNwzVo7pPb",
    },
    "deepseek-r1": {
        "model_name": "deepseek-r1",
        "api_base_url": "http://10.102.24.75:3000/v1",
        "api_key": "sk-BlQIGRrotbVDWE5mXCPBFjVWIvJ83hldzz67xInNwzVo7pPb",
    },
    "Qwen2-72B-Instruct": {
        "model_name": "Qwen2-72B-Instruct",
        "api_base_url": "http://10.102.24.75:3000/v1",
        "api_key": "sk-BlQIGRrotbVDWE5mXCPBFjVWIvJ83hldzz67xInNwzVo7pPb",
    },
    # 阿里云通义千问
    # 文档参考 https://help.aliyun.com/zh/model-studio/qwen-api-reference/
    "qwen-max":{
        "model_name":"qwen-max",
        "api_base_url":"https://dashscope.aliyuncs.com/compatible-mode/v1",
        "api_key":"sk-672f9d1fc4404674bf1a713dfd130a14",
    },
    # deepseek
    "deepseek-chat":{
        # DeepSeek-V3.2
        "model_name":"deepseek-chat",
        "api_base_url":"https://api.deepseek.com/v1",
        "api_key":"sk-26858b50690a49828766fcfcf3290de9",
    },
    "deepseek-reasoner":{
        # DeepSeek-V3.2
        "model_name":"deepseek-reasoner",
        "api_base_url":"https://api.deepseek.com/v1",
        "api_key":"sk-26858b50690a49828766fcfcf3290de9",
    },
    # 智谱AI API,具体注册及api key获取请前往 http://open.bigmodel.cn
    "zhipu-api": {
        "api_key": "http://open.bigmodel.cn",
        "version": "glm-4",
        "provider": "ChatGLMWorker",
    },
}

# 在以下字典中修改属性值，以指定本地embedding模型存储位置。支持3种设置方法：
# 1、将对应的值修改为模型绝对路径
# 2、不修改此处的值（以 text2vec 为例）：
#       2.1 如果{MODEL_ROOT_PATH}下存在如下任一子目录：
#           - text2vec
#           - GanymedeNil/text2vec-large-chinese
#           - text2vec-large-chinese
#       2.2 如果以上本地路径不存在，则使用huggingface模型

MODEL_PATH = {
    "embed_model": {
        "bge_m3": "bge-m3",
    },
    "llm_model": {
        "deepseek-v3": "deepseek-v3",
        "deepseek-r1": "deepseek-r1",
        "qwen-max": "qwen-max",
        "deepseek-chat": "deepseek-chat",
        "deepseek-reasoner": "deepseek-reasoner",
        "Qwen2-72B-Instruct": "Qwen2-72B-Instruct"
    },
    "reranker": {
        "bge-reranker": "bge-reranker",
        "bge-reranker-large": "bge-reranker-large",
    }
}

# 通常情况下不需要更改以下内容
# nltk 模型存储路径
NLTK_DATA_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "nltk_data")

# 使用VLLM可能导致模型推理能力下降，无法完成Agent任务
VLLM_MODEL_DICT = {
    "chatglm3-6b": "chatglm3-6b",
}

SUPPORT_AGENT_MODEL = [
    "Qwen",  # 所有Qwen系列本地模型
    "chatglm3-6b"
]