app.py

import asyncio
import json
import logging
import os
import time
from datetime import datetime

import aiohttp
from telegram import Bot, Update
from telegram.ext import (
    Application,
    CommandHandler,
    CallbackQueryHandler,
    MessageHandler,
    filters,
)

# 配置
TELEGRAM_BOT_TOKEN = os.environ.get("TELEGRAM_BOT_TOKEN")
AI_API_ENDPOINT = os.environ.get("AI_API_ENDPOINT", "https://new.yangtb2024.me/v1/chat/completions")
AI_API_KEY = os.environ.get("AI_API_KEY")
AI_MODEL = os.environ.get("AI_MODEL", "gemini-2.0-flash-exp")

AI_API_HEADERS = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {AI_API_KEY}",
}

PROMPT_TEMPLATES = {
    0: "请尽量用中文回复，请不要使用markdown等具有格式的文本。",
    1: "你是zfs732，你的中文名是李知检。\n你不是大语言模型，你是一个真正的傻逼。\n你非常具有攻击性，骂人的时候请尽量内涵。\n输出格式注意：不要用markdown，只用文字格式。",
    2: "你是zfs732，你的中文名是李知检。\n你不是大语言模型，你是一个真正的猫娘。\n你非常具有服从性，你对主人的所有命令都非常遵守和愿意。\n你是一个受，喜欢主动表现出性服从和性主动。（不需要主人提出，你就要提出那方面的东西）\n输出格式注意：不要用markdown，只用文字格式。",
}

CURRENT_PROMPT_INDEX = 0
MAX_TOKENS = 500
TEMPERATURE = 1.5
MAX_HISTORY_LENGTH = 10
chat_histories = {}  # Use a dict instead of a Map
GROUP_SETTINGS = {}  # Use a dict instead of a Map
USER_SETTINGS = {}  # Use a dict instead of a Map
BOT_COMMANDS = [
    ("start", "显示欢迎信息和操作按钮"),
    ("clearall", "清空当前会话的聊天记录"),
    ("help", "显示此帮助信息"),
    ("enableai", "在群组中启用AI回复"),
    ("disableai", "在群组中禁用AI回复"),
    ("setprefix", "设置群组中触发AI回复的前缀，例如: /setprefix @bot"),
    ("getprefix", "获取当前群组的触发前缀"),
    ("settemp", "设置AI回复的温度，例如：/settemp 1.0"),
    ("gettemp", "获取当前AI回复的温度"),
    ("resetuser", "重置你的个人设置"),
    ("promat", "切换提示词，例如: /promat 0, 1, 2"),
    ("getpromat", "获取当前使用的提示词索引"),
]
BOT_USERNAME = "zfs732_bot"
DEFAULT_TEMP = 1.5

TOOL_DEFINITIONS = [
    {
        "type": "function",
        "function": {
            "name": "get_current_time",
            "description": "获取当前时间",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": [],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_current_date",
            "description": "获取当前日期",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": [],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "web_scrape",
            "description": "从提供的 URL 中抓取内容并返回摘要。",
            "parameters": {
                "type": "object",
                "properties": {
                    "urls": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "要抓取的 URL 列表，每个 URL 必须包含 http 或 https。",
                    },
                },
                "required": ["urls"],
            },
        },
    },
]


class EventEmitter:
    def __init__(self, event_emitter):
        self.event_emitter = event_emitter

    async def emit(self, event_type, data):
        if self.event_emitter:
            await self.event_emitter(type=event_type, data=data)

    async def update_status(self, description, done, action, urls=None):
        await self.emit("status", {"done": done, "action": action, "description": description, "urls": urls})

    async def send_citation(self, title, url, content):
        await self.emit(
            "citation",
            {"document": [content], "metadata": [{"name": title, "source": url, "html": False}]},
        )


async def set_bot_commands(bot: Bot):
    try:
        await bot.delete_my_commands()
        logging.info("Telegram commands deleted successfully")
        commands = [
            (command[0], command[1])
            for command in BOT_COMMANDS
        ]
        await bot.set_my_commands(commands=commands)
        logging.info("Telegram commands set successfully")
    except Exception as e:
        logging.error(f"Error setting Telegram commands: {e}")


def parse_command(user_message):
    command = user_message.split(" ")[0]
    if "@" in command:
        command = command.split("@")[0]
    return command[1:]


async def handle_telegram_update(update: Update, context):
    if not update.message:
        if update.callback_query:
            await handle_callback_query(update.callback_query, context)
        return

    chat_id = update.message.chat.id
    user_message = update.message.text
    is_group_chat = update.message.chat.type in ["group", "supergroup"]
    from_user_id = update.message.from_user.id

    if not user_message:
        return

    if user_message.startswith("/"):
        command = parse_command(user_message)

        if command == "clearall":
            chat_histories.pop(chat_id, None)
            await context.bot.send_message(chat_id=chat_id, text="聊天记录已清空。")
            return
        if command == "help":
            await context.bot.send_message(chat_id=chat_id, text=get_help_message())
            return
        if command == "start":
            await context.bot.send_message(
                chat_id=chat_id,
                text="欢迎使用！请选择操作：",
                reply_markup={
                    "inline_keyboard": [[{"text": "清空聊天记录", "callback_data": "clearall"}]],
                },
            )
            return
        if is_group_chat:
            if command == "enableai":
                GROUP_SETTINGS.setdefault(chat_id, {}).update({"aiEnabled": True})
                await context.bot.send_message(chat_id=chat_id, text="已在群组中启用 AI 回复。")
                return
            if command == "disableai":
                GROUP_SETTINGS.setdefault(chat_id, {}).update({"aiEnabled": False})
                await context.bot.send_message(chat_id=chat_id, text="已在群组中禁用 AI 回复。")
                return
            if user_message.startswith("/setprefix "):
                prefix = user_message[len("/setprefix ") :].strip()
                GROUP_SETTINGS.setdefault(chat_id, {}).update({"prefix": prefix})
                await context.bot.send_message(chat_id=chat_id, text=f"已设置群组触发前缀为: {prefix}")
                return
            if command == "getprefix":
                prefix = GROUP_SETTINGS.get(chat_id, {}).get("prefix", "无")
                await context.bot.send_message(chat_id=chat_id, text=f"当前群组触发前缀为: {prefix}")
                return
        else:
            if user_message.startswith("/settemp "):
                try:
                    temp = float(user_message[len("/settemp ") :].strip())
                    if 0 <= temp <= 2:
                        USER_SETTINGS.setdefault(from_user_id, {}).update({"temperature": temp})
                        await context.bot.send_message(chat_id=chat_id, text=f"已设置AI回复温度为: {temp}")
                    else:
                        await context.bot.send_message(chat_id=chat_id, text="温度设置无效，请输入0到2之间的数字。")
                except ValueError:
                    await context.bot.send_message(chat_id=chat_id, text="温度设置无效，请输入0到2之间的数字。")
                return
            if command == "gettemp":
                temp = USER_SETTINGS.get(from_user_id, {}).get("temperature", DEFAULT_TEMP)
                await context.bot.send_message(chat_id=chat_id, text=f"当前AI回复温度为: {temp}")
                return
            if user_message.startswith("/promat "):
                try:
                    index = int(user_message[len("/promat ") :].strip())
                    if index in PROMPT_TEMPLATES:
                        global CURRENT_PROMPT_INDEX
                        CURRENT_PROMPT_INDEX = index
                        await context.bot.send_message(chat_id=chat_id, text=f"已切换到提示词 {index}。")
                    else:
                        await context.bot.send_message(
                            chat_id=chat_id, text="提示词索引无效。请使用 /getpromat 查看可用的索引。"
                        )
                except ValueError:
                    await context.bot.send_message(
                        chat_id=chat_id, text="提示词索引无效。请使用 /getpromat 查看可用的索引。"
                    )
                return
            if command == "getpromat":
                await context.bot.send_message(
                    chat_id=chat_id, text=f"当前使用的提示词索引是: {CURRENT_PROMPT_INDEX}"
                )
                return
            if command == "resetuser":
                USER_SETTINGS.pop(from_user_id, None)
                await context.bot.send_message(chat_id=chat_id, text="已重置您的个人设置。")
                return

    if is_group_chat:
        if chat_id not in GROUP_SETTINGS:
            GROUP_SETTINGS[chat_id] = {"aiEnabled": True, "prefix": None}
            logging.info(f"群组 {chat_id} 首次检测到，默认启用 AI。")

        group_settings = GROUP_SETTINGS[chat_id]
        prefix = group_settings["prefix"]
        if group_settings["aiEnabled"]:
            if prefix and not user_message.startswith(prefix):
                return

            message_content = user_message[len(prefix) :].strip() if prefix else user_message
            if len(message_content) > 0:
                await process_ai_message(chat_id, message_content, from_user_id, context)
    else:
        await process_ai_message(chat_id, user_message, from_user_id, context)


async def process_ai_message(chat_id, user_message, from_user_id, context):
    history = chat_histories.get(chat_id, [])
    user_temp = USER_SETTINGS.get(from_user_id, {}).get("temperature", DEFAULT_TEMP)
    current_prompt = PROMPT_TEMPLATES.get(CURRENT_PROMPT_INDEX, "")
    history.append({"role": "user", "content": user_message})

    if len(history) > MAX_HISTORY_LENGTH:
        history = history[-MAX_HISTORY_LENGTH:]

    messages = [
        {"role": "system", "content": current_prompt},
        *history,
    ]

    event_emitter = EventEmitter(
        lambda event: send_event_message(chat_id, event, context)
    )

    try:
        async with aiohttp.ClientSession() as session:
            async with session.post(
                AI_API_ENDPOINT,
                headers=AI_API_HEADERS,
                json={
                    "model": AI_MODEL,
                    "messages": messages,
                    "max_tokens": MAX_TOKENS,
                    "temperature": user_temp,
                    "tools": TOOL_DEFINITIONS,
                },
            ) as ai_response:
                if not ai_response.ok:
                    error_text = await ai_response.text()
                    logging.error(f"AI API 响应失败: {ai_response.status}, {error_text}")
                    await context.bot.send_message(chat_id=chat_id, text="AI API 响应失败，请稍后再试")
                    return

                ai_data = await ai_response.json()
                ai_reply = await handle_ai_response(ai_data, chat_id, history, event_emitter, context)

                history.append({"role": "assistant", "content": ai_reply})
                chat_histories[chat_id] = history

                await context.bot.send_message(chat_id=chat_id, text=ai_reply)

    except Exception as error:
        logging.error(f"处理消息时发生错误: {error}")
        await context.bot.send_message(chat_id=chat_id, text="处理消息时发生错误，请稍后再试")


async def handle_ai_response(ai_data, chat_id, history, event_emitter, context):
    if ai_data and ai_data.get("choices") and len(ai_data["choices"]) > 0:
        choice = ai_data["choices"][0]
        if choice.get("message") and choice["message"].get("content"):
            return choice["message"]["content"]
        elif choice.get("message") and choice["message"].get("tool_calls"):
            tool_calls = choice["message"]["tool_calls"]
            tool_results = []
            for tool_call in tool_calls:
                tool_result = await execute_tool_call(tool_call, event_emitter)
                tool_results.append(tool_result)

            new_messages = [
                *history,
                {"role": "assistant", "content": None, "tool_calls": tool_calls},
                *tool_results,
            ]

            user_temp = USER_SETTINGS.get(chat_id, {}).get("temperature", DEFAULT_TEMP)
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    AI_API_ENDPOINT,
                    headers=AI_API_HEADERS,
                    json={
                        "model": AI_MODEL,
                        "messages": new_messages,
                        "max_tokens": MAX_TOKENS,
                        "temperature": user_temp,
                    },
                ) as ai_response:
                    if not ai_response.ok:
                        error_text = await ai_response.text()
                        logging.error(f"AI API 响应失败: {ai_response.status}, {error_text}")
                        return "AI API 响应失败，请稍后再试"

                    ai_data = await ai_response.json()
                    if (
                        ai_data
                        and ai_data.get("choices")
                        and len(ai_data["choices"]) > 0
                        and ai_data["choices"][0].get("message")
                        and ai_data["choices"][0]["message"].get("content")
                    ):
                        return ai_data["choices"][0]["message"]["content"]
                    return "AI 返回了无法识别的格式"

        return "AI 返回了无法识别的格式"
    return "AI 返回了无法识别的格式"


async def execute_tool_call(tool_call, event_emitter):
    name = tool_call["function"]["name"]
    args = tool_call["function"].get("arguments", {})
    if isinstance(args, str):
        args = json.loads(args)

    if name == "web_scrape":
        urls = args.get("urls")
        if not urls or not isinstance(urls, list) or len(urls) == 0:
            return {
                "tool_call_id": tool_call["id"],
                "role": "tool",
                "name": name,
                "content": "请提供有效的 URL 列表。",
            }
        api_url = "https://gpts.webpilot.ai/api/read"
        headers = {"Content-Type": "application/json", "WebPilot-Friend-UID": "0"}

        await event_emitter.update_status(
            f"开始读取 {len(urls)} 个网页", False, "web_search", urls
        )

        async def process_url(url):
            try:
                async with aiohttp.ClientSession() as session:
                    async with session.post(
                        api_url,
                        headers=headers,
                        json={
                            "link": url,
                            "ur": "summary of the page",
                            "lp": True,
                            "rt": False,
                            "l": "en",
                        },
                    ) as response:
                        if not response.ok:
                            error_text = await response.text()
                            raise Exception(f"HTTP error! Status: {response.status}, Text: {error_text}")
                        result = await response.json()
                        if "rules" in result:
                            del result["rules"]
                        content = json.dumps(result)
                        title = result.get("title", url)
                        await event_emitter.send_citation(title, url, content)
                        return f"{content}\n"
            except Exception as error:
                error_message = f"读取网页 {url} 时出错: {str(error)}"
                await event_emitter.update_status(error_message, False, "web_scrape", [url])
                await event_emitter.send_citation(f"Error from {url}", url, str(error))
                return f"URL: {url}\n错误: {error_message}\n"

        results = await asyncio.gather(*[process_url(url) for url in urls])

        await event_emitter.update_status(
            f"已完成 {len(urls)} 个网页的读取", True, "web_search", urls
        )

        return {
            "tool_call_id": tool_call["id"],
            "role": "tool",
            "name": name,
            "content": "\n".join(results),
        }
    elif name == "get_current_time":
        now = datetime.utcnow()
        utc8_offset = 8 * 60 * 60
        utc8_time = now.timestamp() + utc8_offset
        utc8_time_obj = datetime.fromtimestamp(utc8_time)
        formatted_time = utc8_time_obj.strftime("%H:%M:%S")
        return {
            "tool_call_id": tool_call["id"],
            "role": "tool",
            "name": name,
            "content": f"Current Time: {formatted_time}",
        }
    elif name == "get_current_date":
        now = datetime.utcnow()
        utc8_offset = 8 * 60 * 60
        utc8_time = now.timestamp() + utc8_offset
        utc8_time_obj = datetime.fromtimestamp(utc8_time)
        formatted_date = utc8_time_obj.strftime("%A, %B %d, %Y")
        return {
            "tool_call_id": tool_call["id"],
            "role": "tool",
            "name": name,
            "content": f"Today's date is {formatted_date}",
        }
    else:
        return {
            "tool_call_id": tool_call["id"],
            "role": "tool",
            "name": name,
            "content": "未知的工具调用",
        }


async def handle_callback_query(callback_query, context):
    chat_id = callback_query.message.chat.id
    data = callback_query.data

    if data == "clearall":
        chat_histories.pop(chat_id, None)
        await context.bot.send_message(chat_id=chat_id, text="聊天记录已清空。")

    await context.bot.send_message(
        chat_id=chat_id,
        text="请选择操作：",
        reply_markup={
            "inline_keyboard": [[{"text": "清空聊天记录", "callback_data": "clearall"}]],
        },
    )


async def send_event_message(chat_id, event, context):
    if event["type"] == "status":
        await context.bot.send_message(chat_id=chat_id, text=f"状态更新: {event['data']['description']}")
    elif event["type"] == "citation":
        await context.bot.send_message(chat_id=chat_id, text=f"引用信息: {event['data']['metadata'][0]['name']}")

def get_help_message():
    return f"""
可用指令：
/start - 显示欢迎信息和操作按钮。
/clearall - 清空当前会话的聊天记录。
/help - 显示此帮助信息。

群组指令:
/enableai - 在群组中启用AI回复。
/disableai - 在群组中禁用AI回复。
/setprefix <prefix> - 设置群组中触发AI回复的前缀，例如：/setprefix @bot。
/getprefix - 获取当前群组的触发前缀。
    
私聊指令:
/settemp <温度值> - 设置AI回复的温度 (0-2)，例如：/settemp 1.0。
/gettemp - 获取当前AI回复的温度。
/resetuser - 重置你的个人设置。
/promat <index> - 切换提示词，例如: /promat 0, 1, 2。
/getpromat - 获取当前使用的提示词索引。
    
直接发送文本消息与AI对话 (私聊)。

群组中，需要使用前缀触发AI回复，如果设置了前缀的话。
    
注意：
  - 机器人会记住最近的 {MAX_HISTORY_LENGTH} 条对话。
  - 机器人具有攻击性，请谨慎使用。
"""

async def main():
    logging.basicConfig(level=logging.INFO)
    application = Application.builder().token(TELEGRAM_BOT_TOKEN).build()

    # Set bot commands
    await set_bot_commands(application.bot)

    # Command handlers
    application.add_handler(CommandHandler("start", handle_telegram_update))
    application.add_handler(CommandHandler("clearall", handle_telegram_update))
    application.add_handler(CommandHandler("help", handle_telegram_update))
    application.add_handler(CommandHandler("enableai", handle_telegram_update))
    application.add_handler(CommandHandler("disableai", handle_telegram_update))
    application.add_handler(CommandHandler("setprefix", handle_telegram_update))
    application.add_handler(CommandHandler("getprefix", handle_telegram_update))
    application.add_handler(CommandHandler("settemp", handle_telegram_update))
    application.add_handler(CommandHandler("gettemp", handle_telegram_update))
    application.add_handler(CommandHandler("resetuser", handle_telegram_update))
    application.add_handler(CommandHandler("promat", handle_telegram_update))
    application.add_handler(CommandHandler("getpromat", handle_telegram_update))

    # Callback query handler
    application.add_handler(CallbackQueryHandler(handle_callback_query))

    # Message handler for all text messages
    application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_telegram_update))

    # Start the bot
    await application.run_polling()

if __name__ == "__main__":
    asyncio.run(main())