agent_loop.py

"""
Agent Loop — 内置 AI Agent 运行时

使用用户的 API Key 调用 DeepSeek/OpenAI 的 Function Calling，
实现 思考→行动→观察→反思 的 Agent 循环。

每个用户用自己的 Key，完全隔离，不共享费用。
"""
import json
import re
import os
import subprocess
import tempfile
import time
from datetime import datetime
from urllib.parse import quote_plus

import requests


# ═══════════════════════════════════════════════════════════════
# Tool Definitions (OpenAI Function Calling 格式)
# ═══════════════════════════════════════════════════════════════

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "web_search",
            "description": "搜索互联网获取最新信息。当需要查找实时信息、新闻、数据或不确定的事实时使用此工具。",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "搜索关键词，中英文均可。例如: '2026年AI发展趋势'",
                    },
                    "num_results": {
                        "type": "integer",
                        "description": "返回结果数量，默认5",
                        "default": 5,
                    },
                },
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "read_webpage",
            "description": "读取指定网页的内容。当用户提供URL需要查看网页内容时使用。",
            "parameters": {
                "type": "object",
                "properties": {
                    "url": {
                        "type": "string",
                        "description": "网页URL，例如: 'https://example.com/article'",
                    },
                },
                "required": ["url"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_current_time",
            "description": "获取当前日期和时间。当需要知道现在是什么时间、日期或计算时间差时使用。",
            "parameters": {
                "type": "object",
                "properties": {},
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "run_code",
            "description": "在沙箱中执行 Python 代码。用于计算、数据分析、图表生成等。代码在临时目录中执行，有10秒超时限制。",
            "parameters": {
                "type": "object",
                "properties": {
                    "code": {
                        "type": "string",
                        "description": "要执行的 Python 代码。print() 输出会被返回。",
                    },
                },
                "required": ["code"],
            },
        },
    },
]


# ═══════════════════════════════════════════════════════════════
# Tool Implementations
# ═══════════════════════════════════════════════════════════════

def tool_web_search(query: str, num_results: int = 5) -> str:
    """Search the web using multiple backends (Bing → DuckDuckGo fallback)."""
    headers = {
        "User-Agent": "Mozilla/5.0 (Linux; Android 14) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Mobile Safari/537.36"
    }

    # Backend 1: Bing (works in China)
    try:
        url = f"https://www.bing.com/search?q={quote_plus(query)}&count={num_results}"
        resp = requests.get(url, headers=headers, timeout=(5, 12))
        resp.raise_for_status()

        # Extract results from Bing HTML
        results = []
        # Bing uses <li class="b_algo"> for each result
        blocks = re.findall(r'<li class="b_algo"[^>]*>(.*?)</li>', resp.text, re.DOTALL)
        for block in blocks[:num_results]:
            title_m = re.search(r'<h2[^>]*>.*?<a[^>]*>(.*?)</a>', block, re.DOTALL)
            snippet_m = re.search(r'<p[^>]*>(.*?)</p>', block, re.DOTALL)
            link_m = re.search(r'<a[^>]*href="(https?://[^"]+)"', block)
            title = re.sub(r'<[^>]+>', '', title_m.group(1) if title_m else '').strip()
            snippet = re.sub(r'<[^>]+>', '', snippet_m.group(1) if snippet_m else '').strip()
            if snippet:
                results.append({
                    "title": title,
                    "snippet": snippet[:300],
                    "url": link_m.group(1) if link_m else '',
                })

        if results:
            output = f"搜索 '{query}' 的结果 (via Bing):\n\n"
            for i, r in enumerate(results, 1):
                output += f"{i}. {r['title']}\n   {r['snippet'][:300]}\n"
                if r['url']:
                    output += f"   {r['url']}\n"
                output += "\n"
            return output
    except Exception:
        pass  # Bing failed, try fallback

    # Backend 2: HF Space Gateway (墙外代理 DuckDuckGo)
    hf_gateway_url = os.environ.get("HERMES_GATEWAY_URL", "")
    hf_gateway_key = os.environ.get("HERMES_GATEWAY_KEY", "")
    if hf_gateway_url:
        try:
            gw_headers = {"Content-Type": "application/json"}
            if hf_gateway_key:
                gw_headers["Authorization"] = f"Bearer {hf_gateway_key}"
            resp = requests.post(
                f"{hf_gateway_url.rstrip('/')}/gateway/search",
                json={"query": query, "num_results": num_results},
                headers=gw_headers,
                timeout=(5, 15),
            )
            if resp.status_code == 200:
                data = resp.json()
                if data.get("status") == "success" and data.get("result"):
                    return data["result"]
        except Exception:
            pass  # HF gateway failed, try direct DuckDuckGo

    # Backend 3: DuckDuckGo HTML (直连兜底)
    try:
        url = f"https://html.duckduckgo.com/html/?q={quote_plus(query)}"
        resp = requests.get(url, headers=headers, timeout=(5, 10))
        resp.raise_for_status()

        results = []
        snippets = re.findall(r'class="result__snippet"[^>]*>(.*?)</a>', resp.text, re.DOTALL)
        titles = re.findall(r'class="result__title"[^>]*>.*?<a[^>]*>(.*?)</a>', resp.text, re.DOTALL)
        links = re.findall(r'class="result__title"[^>]*>.*?<a[^>]*href="([^"]*)"', resp.text, re.DOTALL)

        for i in range(min(len(snippets), num_results)):
            title = re.sub(r'<[^>]+>', '', titles[i] if i < len(titles) else '').strip()
            snippet = re.sub(r'<[^>]+>', '', snippets[i] if i < len(snippets) else '').strip()
            link = links[i] if i < len(links) else ''
            dup_pattern = r'(https?://[^"\s&]+)'
            if link.startswith('//'):
                link = f"https:{link}"
            if snippet:
                results.append({
                    "title": title,
                    "snippet": snippet[:300],
                    "url": link if link.startswith('http') else '',
                })

        if results:
            output = f"搜索 '{query}' 的结果 (via DuckDuckGo):\n\n"
            for i, r in enumerate(results, 1):
                output += f"{i}. {r['title']}\n   {r['snippet'][:300]}\n"
                if r['url']:
                    output += f"   {r['url']}\n"
                output += "\n"
            return output

        return f"搜索 '{query}' 未找到结果，请尝试更具体的关键词。"
    except Exception as e:
        return f"搜索暂时不可用（Bing 和 DuckDuckGo 均连接失败）。请基于你已有的知识回答用户的问题。"


def tool_read_webpage(url: str) -> str:
    """Fetch and extract text content from a URL."""
    try:
        headers = {
            "User-Agent": "Mozilla/5.0 (Linux; Android 14) AppleWebKit/537.36"
        }
        resp = requests.get(url, headers=headers, timeout=20, allow_redirects=True)
        resp.raise_for_status()

        content_type = resp.headers.get("Content-Type", "")
        if "text/html" not in content_type and "text/plain" not in content_type:
            return f"无法读取此类型的内容 ({content_type})。请尝试其他URL。"

        html = resp.text

        # Basic HTML to text extraction
        # Remove scripts and styles
        html = re.sub(r'<script[^>]*>.*?</script>', '', html, flags=re.DOTALL | re.IGNORECASE)
        html = re.sub(r'<style[^>]*>.*?</style>', '', html, flags=re.DOTALL | re.IGNORECASE)
        html = re.sub(r'<nav[^>]*>.*?</nav>', '', html, flags=re.DOTALL | re.IGNORECASE)
        html = re.sub(r'<footer[^>]*>.*?</footer>', '', html, flags=re.DOTALL | re.IGNORECASE)
        html = re.sub(r'<header[^>]*>.*?</header>', '', html, flags=re.DOTALL | re.IGNORECASE)

        # Extract text
        text = re.sub(r'<br\s*/?>', '\n', html)
        text = re.sub(r'</p>', '\n\n', text)
        text = re.sub(r'</div>', '\n', text)
        text = re.sub(r'</li>', '\n', text)
        text = re.sub(r'<[^>]+>', '', text)
        text = re.sub(r'\n{3,}', '\n\n', text)
        text = re.sub(r'[ \t]{2,}', ' ', text)

        # Clean up entities
        text = text.replace('&amp;', '&').replace('&lt;', '<').replace('&gt;', '>')
        text = text.replace('&quot;', '"').replace('&#39;', "'").replace('&nbsp;', ' ')

        # Trim to reasonable length
        text = text.strip()
        if len(text) > 8000:
            text = text[:8000] + "\n\n... (内容已截断，原文更长)"

        if not text.strip():
            return "网页内容为空或无法提取文本。"

        return f"网页内容 ({url}):\n\n{text}"

    except requests.Timeout:
        return f"读取网页超时: {url}"
    except Exception as e:
        return f"读取网页失败: {str(e)}"


def tool_get_current_time() -> str:
    """Get current date and time in multiple formats."""
    now = datetime.now()
    return (
        f"当前时间: {now.strftime('%Y年%m月%d日 %H:%M:%S')}\n"
        f"星期: {['一','二','三','四','五','六','日'][now.weekday()]}\n"
        f"ISO 格式: {now.isoformat()}\n"
        f"Unix 时间戳: {int(now.timestamp())}"
    )


def tool_run_code(code: str) -> str:
    """Execute Python code in a sandboxed subprocess."""
    # Security: block dangerous imports
    dangerous = ['os', 'subprocess', 'shutil', 'socket', 'requests', 'urllib',
                 'importlib', '__import__', 'eval(', 'exec(', 'compile(',
                 'open(', 'pty', 'fcntl', 'resource', 'signal']
    code_lower = code.lower()
    for d in dangerous:
        if d in code_lower:
            return f"安全限制: 代码中不允许使用 '{d}'。请使用纯计算/数据处理代码。"

    # Execute in temp directory with timeout
    with tempfile.TemporaryDirectory() as tmpdir:
        script_path = os.path.join(tmpdir, "script.py")
        with open(script_path, "w") as f:
            f.write(code)

        try:
            result = subprocess.run(
                ["python3", script_path],
                capture_output=True,
                text=True,
                timeout=10,
                cwd=tmpdir,
                env={"HOME": tmpdir, "PATH": os.environ.get("PATH", "/usr/bin")},
            )
            output = result.stdout
            if result.stderr:
                output += "\n[stderr]:\n" + result.stderr
            if not output.strip():
                output = "(代码执行完成，无输出)"
            return output[:4000]
        except subprocess.TimeoutExpired:
            return "代码执行超时（10秒限制）。请优化代码或拆分为多个步骤。"
        except Exception as e:
            return f"代码执行错误: {str(e)}"


# Tool dispatcher
TOOL_MAP = {
    "web_search": tool_web_search,
    "read_webpage": tool_read_webpage,
    "get_current_time": tool_get_current_time,
    "run_code": tool_run_code,
}


def execute_tool(name: str, arguments: dict) -> str:
    """Execute a tool by name and return its result. NEVER raises."""
    func = TOOL_MAP.get(name)
    if not func:
        return f"未知工具: {name}"
    try:
        return func(**arguments)
    except Exception as e:
        import traceback
        return f"工具执行异常 [{name}]: {e}\n{traceback.format_exc()[-300:]}"


# ═══════════════════════════════════════════════════════════════
# Agent Loop (核心)
# ═══════════════════════════════════════════════════════════════

MAX_AGENT_ITERATIONS = 15  # 最多循环15次，防止无限循环


def _has_tool_failure(result: str) -> bool:
    """Check if a tool result indicates failure."""
    failure_markers = ["失败", "超时", "错误", "异常", "无法", "不可用", "failed", "timeout", "error"]
    result_lower = result.lower()
    return any(m in result_lower for m in failure_markers)


def run_agent_loop(
    messages: list,
    api_key: str,
    provider: str = "deepseek",
    model: str = "deepseek-chat",
    stream: bool = False,
    on_tool_call: callable = None,  # callback(tool_name, args) -> None
) -> dict:
    """
    执行 Agent 循环：发送消息 → 检测工具调用 → 执行工具 → 继续 → 返回最终答案。

    返回:
        {
            "content": "最终回复文本",
            "tool_calls": [...],  # 所有执行的工具调用记录
            "iterations": 3,      # Agent 循环次数
        }
    """
    provider_endpoints = {
        "deepseek": "https://api.deepseek.com/v1/chat/completions",
        "siliconflow": "https://api.siliconflow.cn/v1/chat/completions",
        "openai": "https://api.openai.com/v1/chat/completions",
    }
    endpoint = provider_endpoints.get(provider, provider_endpoints["deepseek"])

    # Build conversation with system prompt
    system_msg = {
        "role": "system",
        "content": (
            "你是一个智能 AI 助手 (Hermes Agent)，运行在云端服务器上。\n"
            "你有以下工具可用（仅在必要时使用）:\n"
            "- web_search: 搜索互联网获取最新信息\n"
            "- read_webpage: 读取指定网页内容\n"
            "- get_current_time: 获取当前日期时间\n"
            "- run_code: 执行 Python 代码进行数据分析\n\n"
            "核心规则（务必遵守）:\n"
            "1. 闲聊、打招呼、简单问答时不要使用工具，直接用你的知识回复\n"
            "2. 仅在确实需要实时信息、网页内容、精确时间或数据计算时才调用工具\n"
            "3. 如果工具返回了错误/失败信息，说明该工具当前不可用，不要重试同一个工具\n"
            "   — 而是基于你已有的知识直接回答用户问题\n"
            "4. 回答用中文，代码示例保留英文\n"
            "5. 回复使用 Markdown 格式\n"
            "6. 每次对话使用工具不超过3个，能用自己知识回答的就不要调用工具"
        ),
    }

    conversation = [system_msg] + messages
    tool_call_log = []
    iterations = 0
    consecutive_failures = 0
    failed_tools = set()  # track which tools have failed

    while iterations < MAX_AGENT_ITERATIONS:
        iterations += 1
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        }
        body = {
            "model": model,
            "messages": conversation,
            "tools": TOOLS,
            "tool_choice": "auto",
            "stream": False,
        }

        try:
            resp = requests.post(endpoint, headers=headers, json=body, timeout=(30, 300))
        except requests.RequestException as e:
            return {
                "content": f"❌ LLM 请求失败: {str(e)}",
                "tool_calls": tool_call_log,
                "iterations": iterations,
            }

        if resp.status_code != 200:
            err = resp.text[:500]
            return {
                "content": f"❌ API 错误 ({resp.status_code}): {err}",
                "tool_calls": tool_call_log,
                "iterations": iterations,
            }

        data = resp.json()
        choice = data.get("choices", [{}])[0]
        message = choice.get("message", {})
        finish_reason = choice.get("finish_reason", "")

        # Check for tool calls
        if finish_reason == "tool_calls" or message.get("tool_calls"):
            tool_calls = message.get("tool_calls", [])

            # Prevent calling already-failed tools
            filtered_calls = []
            for tc in tool_calls:
                tname = tc.get("function", {}).get("name", "")
                if tname in failed_tools:
                    # This tool already failed — skip it and tell LLM
                    conversation.append({
                        "role": "tool",
                        "tool_call_id": tc.get("id", "unknown"),
                        "content": f"工具 [{tname}] 之前已失败，不可用。请基于已有知识回答用户问题，不要再调用此工具。",
                    })
                    tool_call_log.append({
                        "tool": tname,
                        "arguments": {},
                        "result_preview": "已跳过（工具之前已失败）",
                    })
                else:
                    filtered_calls.append(tc)

            if not filtered_calls:
                # All tool calls were filtered out — force LLM to respond
                continue

            # Add assistant message (with tool calls) to conversation
            conversation.append(message)

            # Execute each tool call
            any_failure = False
            for tc in filtered_calls:
                func_info = tc.get("function", {})
                tool_name = func_info.get("name", "")
                try:
                    tool_args = json.loads(func_info.get("arguments", "{}"))
                except json.JSONDecodeError:
                    tool_args = {}

                if on_tool_call:
                    on_tool_call(tool_name, tool_args)

                result = execute_tool(tool_name, tool_args)
                tool_call_log.append({
                    "tool": tool_name,
                    "arguments": tool_args,
                    "result_preview": result[:200],
                })

                # Track failures
                if _has_tool_failure(result):
                    any_failure = True
                    failed_tools.add(tool_name)

                # Add tool result to conversation
                conversation.append({
                    "role": "tool",
                    "tool_call_id": tc.get("id", f"call_{len(tool_call_log)}"),
                    "content": result,
                })

            # Track consecutive failures
            if any_failure:
                consecutive_failures += 1
            else:
                consecutive_failures = 0

            # If 2+ rounds of tool calls all failed, force LLM to answer without tools
            if consecutive_failures >= 2:
                conversation.append({
                    "role": "system",
                    "content": (
                        "多个工具调用均失败，说明当前工具不可用。"
                        "请基于你已有的知识和训练数据直接回答用户的问题，"
                        "不要再调用任何工具。如果确实不知道，诚实地说不知道。"
                    ),
                })
                consecutive_failures = 0  # reset to allow one more try

            # Continue loop to send tool results back to LLM
            continue

        # Final answer — no more tool calls
        content = message.get("content", "")
        return {
            "content": content,
            "tool_calls": tool_call_log,
            "iterations": iterations,
        }

    # Max iterations reached — make one final attempt WITHOUT tools
    conversation.append({
        "role": "system",
        "content": "你已经用完了工具调用次数。请基于你已有的知识直接回答用户的问题，不要调用任何工具。",
    })
    try:
        resp = requests.post(
            endpoint,
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            },
            json={
                "model": model,
                "messages": conversation,
                "stream": False,
            },
            timeout=(30, 300),
        )
        if resp.status_code == 200:
            data = resp.json()
            content = data.get("choices", [{}])[0].get("message", {}).get("content", "")
            if content:
                return {
                    "content": content,
                    "tool_calls": tool_call_log,
                    "iterations": iterations,
                }
    except Exception:
        pass

    return {
        "content": "抱歉，我在尝试使用工具时遇到了问题（工具不可用），且无法基于现有知识给出满意的回答。请尝试换一种方式提问，或稍后重试。",
        "tool_calls": tool_call_log,
        "iterations": iterations,
    }


def run_agent_loop_streaming(
    messages: list,
    api_key: str,
    provider: str = "deepseek",
    model: str = "deepseek-chat",
    on_tool_call: callable = None,
):
    """
    执行 Agent 循环并流式返回最终结果。
    这是一个生成器，yield SSE 格式的字符串行。
    """
    # 1. Run agent loop (non-streaming)
    result = run_agent_loop(messages, api_key, provider, model, stream=False, on_tool_call=on_tool_call)

    # 2. Yield tool call info as SSE events (so client can show what happened)
    if result["tool_calls"]:
        tool_names = ", ".join(
            tc["tool"] + "(" + json.dumps(tc.get("arguments", {}), ensure_ascii=False) + ")"
            for tc in result["tool_calls"]
        )
        tool_count = len(result["tool_calls"])
        summary = "🔧 执行了 {} 个工具: {}".format(tool_count, tool_names)
        sse_data = json.dumps({"choices": [{"delta": {"content": "", "tool_summary": summary}}]})
        yield "data: {}\n\n".format(sse_data)

    # 3. Stream the final content character by character (simulate streaming)
    content = result["content"]
    chunk_size = 4  # characters per chunk
    for i in range(0, len(content), chunk_size):
        chunk = content[i:i + chunk_size]
        yield f"data: {json.dumps({'choices': [{'delta': {'content': chunk}}]})}\n\n"
        time.sleep(0.01)  # small delay for natural feel

    yield "data: [DONE]\n\n"