#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/3/16 09:46
# @Author  : old-tom
# @File    : llm_agent
# @Project : llmFunctionCallDemo
# @Desc    : llm代理，负责初始化模型

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

# 默认系统提示词
DEFAULT_SYS_PROMPT = ''

parser = StrOutputParser()

# 模型初始化，注意替换成自己的模型和api_key
llm = ChatOpenAI(
    model='deepseek-r1-250120', api_key='4eefc827-187f-4756-9637-7e0153c93d81',
    base_url='https://ark.cn-beijing.volces.com/api/v3/', max_tokens=4096, temperature=0.5, streaming=True
)

# 工具函数定义,deepseek 官方定义的格式（最多支持 128 个 function），参考 https://api-docs.deepseek.com/zh-cn/guides/function_calling
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "get_current_weather",
            "description": "Get weather of an location, the user shoud supply a location first",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "The city and state, e.g. San Francisco, CA",
                    }
                },
                "required": ["location"]
            },
        }
    },
    {
        "type": "function",
        "function": {
            "name": "play_video",
            "description": "播放、查看、打开实时视频",
            "parameters": {
                "type": "object",
                "properties": {
                    "camera_name": {
                        "type": "string",
                        "description": "相机名称,例如：南卡口1号相机",
                    }
                },
                "required": ["camera_name"]
            },
        }
    },
    {
        "type": "function",
        "function": {
            "name": "split_screen",
            "description": "切换分屏",
            "parameters": {
                "type": "object",
                "properties": {
                    "split_n": {
                        "type": "int",
                        "description": "要切换的分屏数量,整数并且大于0，例如：1分屏、2分屏",
                    }
                },
                "required": ["split_n"]
            },
        }
    },
    {
        "type": "function",
        "function": {
            "name": "play_video_record",
            "description": "播放、打开录像",
            "parameters": {
                "type": "object",
                "properties": {
                    "camera_name": {
                        "type": "string",
                        "description": "相机名称，例如：南卡口1号相机",
                    },
                    "start_time": {
                        "type": "string",
                        "description": "录像开始时间,格式为yyyy-MM-dd hh:mm:ss，例 2025-03-16 01:00:00",
                    },
                    "end_time": {
                        "type": "string",
                        "description": "录像结束时间,格式为yyyy-MM-dd hh:mm:ss，例 2025-03-16 02:09:31",
                    }
                },
                "required": ["camera_name", "start_time", "end_time"]
            },
        }
    },
    {
        "type": "function",
        "function": {
            "name": "switch_page",
            "description": "打开、跳转页面",
            "parameters": {
                "type": "object",
                "properties": {
                    "page_name": {
                        "type": "string",
                        "description": "页面中文名称或者缩写，例：人员核查、系统日志、设备管理、首页",
                    }
                },
                "required": ["page_name"]
            },
        }
    },
    {
        "type": "function",
        "function": {
            "name": "zoom_in",
            "description": "放大电子地图",
            "parameters": {
                "type": "object",
                "properties": {
                    "level_n": {
                        "type": "int",
                        "description": "放大等级,整数并且大于0小于5，例如：放大1级、放大2级",
                    }
                },
                "required": ["level_n"]
            },
        }
    },
    {
        "type": "function",
        "function": {
            "name": "view_flight_details",
            "description": "查询指定机场指定航班及时间的出入境人员明细",
            "parameters": {
                "type": "object",
                "properties": {
                    "airport_name": {
                        "type": "string",
                        "description": "机场名称，简体中文，可以是缩写，例如：成都天府机场、天府机场、长水机场、上海浦东机场",
                    },
                    "flight_code": {
                        "type": "string",
                        "description": "航班编号,由字母+数字组成的完整编号,若编号包含多余字符（如标点符号），需过滤后保留有效部分",
                    },
                    "flight_date": {
                        "type": "string",
                        "description": "提取完整日期（年月日），自动补零至标准格式 yyyy-MM-dd, 例：2025-03-16",
                    },
                    "ie_type": {
                        "type": "string",
                        "description": "出入境类型，仅识别'入境'或'出境'两种类型",
                    }
                },
                "required": ["airport_name", "flight_code", "flight_date"]
            },
        }
    }
]


class DeepSeekR1Agent(object):
    def __init__(self, system_prompt: str = DEFAULT_SYS_PROMPT):
        """
        :param system_prompt: 系统提示词
        """
        self.prompt = ChatPromptTemplate(
            [
                ("system", system_prompt),
                ("human", "{user_input}")
            ]
        )

    def invoke(self, user_input: str):
        """
        请求模型并一次性返回
        :param user_input: 用户输入
        :return:
        """
        chain = self.prompt | llm
        return chain.invoke({
            'user_input': user_input
        }).content

    def invoke_by_stream(self, user_input: str):
        """
        请求模型并流式返回（同步流）
        :param user_input: 用户输入
        :return:
        """
        chain = self.prompt | llm | parser
        response = chain.stream({'user_input': user_input})
        for chunk in response:
            print(chunk, flush=True, end='')

    @staticmethod
    def invoke_with_tool(user_input: str):
        """
        工具链调用,function calling时system prompt不会生效
        :param user_input:
        :return: 这里返回的是LLM推理出的tool信息，格式如下：
            [{'name': 'get_current_weather', 'args': {'location': 'Beijing, China'}, 'id': 'call_xeeq4q52fw9x61lkrqwy9cr6', 'type': 'tool_call'}]
        """
        llm_with_tools = llm.bind_tools(TOOLS)
        return llm_with_tools.invoke(user_input).tool_calls