agent.py

import ast
import inspect
import os
import re
import time
import urllib.error
import urllib.request
from string import Template
from typing import List, Callable, Tuple
from urllib.parse import urlparse

import click
from dotenv import load_dotenv
from openai import OpenAI, RateLimitError
import platform

from prompt_template import react_system_prompt_template


class ReActAgent:
    def __init__(self, tools: List[Callable], model: str, project_directory: str):
        self.tools = { func.__name__: func for func in tools }
        self.model = model
        self.project_directory = project_directory
        self.client = OpenAI(
            base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
            api_key=ReActAgent.get_api_key(),
        )

    def run(self, user_input: str):
        messages = [
            {"role": "system", "content": self.render_system_prompt(react_system_prompt_template)},
            {"role": "user", "content": f"<question>{user_input}</question>"}
        ]

        while True:

            # 请求模型
            content = self.call_model(messages)

            # 检测 Thought
            thought_match = re.search(r"<thought>(.*?)</thought>", content, re.DOTALL)
            if thought_match:
                thought = thought_match.group(1)
                print(f"\n\n💭 Thought: {thought}")

            # 检测模型是否输出 Final Answer，如果是的话，直接返回
            if "<final_answer>" in content:
                final_answer = re.search(r"<final_answer>(.*?)</final_answer>", content, re.DOTALL)
                return final_answer.group(1)

            # 检测 Action
            action_match = re.search(r"<action>(.*?)</action>", content, re.DOTALL)
            if not action_match:
                raise RuntimeError("模型未输出 <action>")
            action = action_match.group(1)
            tool_name, args = self.parse_action(action)

            print(f"\n\n🔧 Action: {tool_name}({', '.join(args)})")
            # 只有终端命令才需要询问用户，其他的工具直接执行
            should_continue = input(f"\n\n是否继续？（Y/N）") if tool_name == "run_terminal_command" else "y"
            if should_continue.lower() != 'y':
                print("\n\n操作已取消。")
                return "操作被用户取消"

            try:
                observation = self.tools[tool_name](*args)
            except Exception as e:
                observation = f"工具执行错误：{str(e)}"
            print(f"\n\n🔍 Observation：{observation}")
            obs_msg = f"<observation>{observation}</observation>"
            messages.append({"role": "user", "content": obs_msg})


    def get_tool_list(self) -> str:
        """生成工具列表字符串，包含函数签名和简要说明"""
        tool_descriptions = []
        for func in self.tools.values():
            name = func.__name__
            signature = str(inspect.signature(func))
            doc = inspect.getdoc(func)
            tool_descriptions.append(f"- {name}{signature}: {doc}")
        return "\n".join(tool_descriptions)

    def render_system_prompt(self, system_prompt_template: str) -> str:
        """渲染系统提示模板，替换变量"""
        tool_list = self.get_tool_list()
        file_list = ", ".join(
            os.path.abspath(os.path.join(self.project_directory, f))
            for f in os.listdir(self.project_directory)
        )
        return Template(system_prompt_template).substitute(
            operating_system=self.get_operating_system_name(),
            tool_list=tool_list,
            file_list=file_list
        )

    @staticmethod
    def get_api_key() -> str:
        """Load the API key from an environment variable."""
        load_dotenv()
        api_key = os.getenv("GEMINI_API_KEY")
        if not api_key:
            raise ValueError("❎：未找到 GEMINI_API_KEY 环境变量，请在 .env 文件中设置（Google AI Studio）。")
        return api_key

    @staticmethod
    def _retry_after_seconds(err: RateLimitError) -> float | None:
        text = str(err)
        m = re.search(r"Please retry in ([\d.]+)\s*s", text, re.IGNORECASE)
        if m:
            return float(m.group(1))
        resp = getattr(err, "response", None)
        if resp is not None:
            ra = resp.headers.get("retry-after")
            if ra:
                return float(ra)
        return None

    def call_model(self, messages):
        max_attempts = 1
        fallback_wait = 35.0
        for attempt in range(max_attempts):
            print(f"\n\n正在请求模型（{self.model}），请稍等...")
            try:
                response = self.client.chat.completions.create(
                    model=self.model,
                    messages=messages,
                )
                content = response.choices[0].message.content
                messages.append({"role": "assistant", "content": content})
                return content
            except RateLimitError as e:
                if attempt >= max_attempts - 1:
                    raise RuntimeError(
                        "Gemini 返回 429：免费层配额或每分钟上限已用尽。"
                        " 可稍后再试、在 .env 设置 GEMINI_MODEL 换模型（如 gemini-2.5-flash-lite），"
                        "或前往 https://ai.google.dev/gemini-api/docs/rate-limits 查看配额与计费。"
                    ) from e
                wait = self._retry_after_seconds(e) or fallback_wait
                wait = min(max(wait, 1.0), 120.0)
                print(
                    f"\n速率限制 (429)，{wait:.0f} 秒后重试 "
                    f"({attempt + 1}/{max_attempts})…"
                )
                time.sleep(wait)

    def parse_action(self, code_str: str) -> Tuple[str, List[str]]:
        match = re.match(r'(\w+)\((.*)\)', code_str, re.DOTALL)
        if not match:
            raise ValueError("Invalid function call syntax")

        func_name = match.group(1)
        args_str = match.group(2).strip()

        # 手动解析参数，特别处理包含多行内容的字符串
        args = []
        current_arg = ""
        in_string = False
        string_char = None
        i = 0
        paren_depth = 0
        
        while i < len(args_str):
            char = args_str[i]
            
            if not in_string:
                if char in ['"', "'"]:
                    in_string = True
                    string_char = char
                    current_arg += char
                elif char == '(':
                    paren_depth += 1
                    current_arg += char
                elif char == ')':
                    paren_depth -= 1
                    current_arg += char
                elif char == ',' and paren_depth == 0:
                    # 遇到顶层逗号，结束当前参数
                    args.append(self._parse_single_arg(current_arg.strip()))
                    current_arg = ""
                else:
                    current_arg += char
            else:
                current_arg += char
                if char == string_char and (i == 0 or args_str[i-1] != '\\'):
                    in_string = False
                    string_char = None
            
            i += 1
        
        # 添加最后一个参数
        if current_arg.strip():
            args.append(self._parse_single_arg(current_arg.strip()))
        
        return func_name, args
    
    def _parse_single_arg(self, arg_str: str):
        """解析单个参数"""
        arg_str = arg_str.strip()
        
        # 如果是字符串字面量
        if (arg_str.startswith('"') and arg_str.endswith('"')) or \
           (arg_str.startswith("'") and arg_str.endswith("'")):
            # 移除外层引号并处理转义字符
            inner_str = arg_str[1:-1]
            # 处理常见的转义字符
            inner_str = inner_str.replace('\\"', '"').replace("\\'", "'")
            inner_str = inner_str.replace('\\n', '\n').replace('\\t', '\t')
            inner_str = inner_str.replace('\\r', '\r').replace('\\\\', '\\')
            return inner_str
        
        # 尝试使用 ast.literal_eval 解析其他类型
        try:
            return ast.literal_eval(arg_str)
        except (SyntaxError, ValueError):
            # 如果解析失败，返回原始字符串
            return arg_str

    def get_operating_system_name(self):
        os_map = {
            "Darwin": "macOS",
            "Windows": "Windows",
            "Linux": "Linux"
        }

        return os_map.get(platform.system(), "Unknown")


def read_file(file_path):
    """用于读取文件内容"""
    with open(file_path, "r", encoding="utf-8") as f:
        return f.read()

def write_to_file(file_path, content):
    """将指定内容写入指定文件"""
    with open(file_path, "w", encoding="utf-8") as f:
        f.write(content.replace("\\n", "\n"))
    return "写入成功"

def run_terminal_command(command):
    """用于执行终端命令"""
    import subprocess
    run_result = subprocess.run(command, shell=True, capture_output=True, text=True)
    if run_result.returncode != 0:
        return run_result.stderr or "命令执行失败（无错误输出）"
    return run_result.stdout or "执行成功（无输出）"


def fetch_url(url: str, max_chars: int = 80000) -> str:
    """HTTP GET 拉取指定 URL 的响应正文（文本），用于查阅在线文档、新闻、公开 API 等。仅支持 http/https；正文过长会截断。"""
    parsed = urlparse(url)
    if parsed.scheme not in ("http", "https"):
        return "错误：仅支持 http 或 https 协议的 URL。"
    max_bytes = min(512 * 1024, max(4096, max_chars * 4))
    req = urllib.request.Request(
        url,
        headers={
            "User-Agent": "Mozilla/5.0 (compatible; ReActAgent/1.0; +https://ai.google.dev/)",
            "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
        },
        method="GET",
    )
    try:
        with urllib.request.urlopen(req, timeout=45) as resp:
            content_type = resp.headers.get("Content-Type", "")
            raw = resp.read(max_bytes + 1)
    except urllib.error.HTTPError as e:
        body = ""
        try:
            b = e.read(4000)
            body = b.decode("utf-8", errors="replace") if b else ""
        except Exception:
            pass
        return f"HTTP {e.code} {e.reason}\n{body}".strip()
    except urllib.error.URLError as e:
        return f"网络请求失败：{e.reason}"
    except Exception as e:
        return f"拉取 URL 时出错：{e}"

    truncated_bytes = len(raw) > max_bytes
    if truncated_bytes:
        raw = raw[:max_bytes]

    charset = "utf-8"
    m = re.search(r"charset=([\w-]+)", content_type, re.I)
    if m:
        charset = m.group(1).strip('"').strip("'") or charset
    try:
        text = raw.decode(charset)
    except (LookupError, UnicodeDecodeError):
        text = raw.decode("utf-8", errors="replace")

    if len(text) > max_chars:
        text = text[:max_chars] + f"\n\n[正文已截断至前 {max_chars} 个字符；可调大 max_chars 参数重试]"
    if truncated_bytes:
        text += "\n\n[响应体积超过字节上限，可能未完整下载]"
    return text


@click.command()
@click.argument('project_directory',
                type=click.Path(exists=True, file_okay=False, dir_okay=True))
def main(project_directory):
    load_dotenv()
    project_dir = os.path.abspath(project_directory)

    tools = [read_file, write_to_file, run_terminal_command, fetch_url]
    model = os.getenv("GEMINI_MODEL", "gemini-2.5-flash")
    agent = ReActAgent(tools=tools, model=model, project_directory=project_dir)

    task = input("请输入任务：")

    final_answer = agent.run(task)

    print(f"\n\n✅ Final Answer：{final_answer}")

if __name__ == "__main__":
    main()