projects/repobot/v4/tools.py

"""RepoBot v3 的工具：查文档（检索、按行读）和翻源码（搜索、按行读）。

工具的写法和 05 模块第 2 课一样：普通函数 + 装饰器生成的说明书。
"""
import inspect
import os
import re
import subprocess
from pathlib import Path

HERE = Path(__file__).parent
DOCS_DIR = Path(os.environ.get("REPOBOT_DOCS", HERE / "../../../data/httpx-docs")).resolve()
SOURCE_DIR = Path(os.environ.get("REPOBOT_SOURCE", HERE / ".cache/httpx")).resolve()
SOURCE_REPO = "https://github.com/encode/httpx"

TOOLS = {}
retriever = None  # repobot.py 启动时设置成 retrieval.Retriever 的实例


def tool(description, **params):
    def register(fn):
        sig = inspect.signature(fn)
        properties = {name: {"type": "integer" if p.annotation is int else "string", "description": params[name]}
                      for name, p in sig.parameters.items()}
        required = [name for name, p in sig.parameters.items() if p.default is inspect.Parameter.empty]
        TOOLS[fn.__name__] = {"fn": fn, "schema": {"type": "function", "function": {
            "name": fn.__name__, "description": description,
            "parameters": {"type": "object", "properties": properties, "required": required}}}}
        return fn
    return register


def ensure_source():
    """第一次运行时把 httpx 的源码克隆下来。"""
    if not (SOURCE_DIR / "httpx").is_dir():
        print(f"第一次运行，正在从 {SOURCE_REPO} 下载 httpx 源码……", flush=True)
        SOURCE_DIR.parent.mkdir(parents=True, exist_ok=True)
        subprocess.run(["git", "clone", "--depth", "1", "--quiet", SOURCE_REPO, str(SOURCE_DIR)], check=True)


def inside(base, path):
    """把相对路径转成绝对路径，并确认它没有跑出 base 目录。跑出去了就返回 None。"""
    target = (base / path).resolve()
    return target if base in target.parents and target.is_file() else None


def numbered(target, start, end, limit=80):
    lines = target.read_text(errors="replace").splitlines()
    start = max(1, start)
    end = min(end, start + limit - 1, len(lines))
    return "\n".join(f"{n}: {lines[n - 1]}" for n in range(start, end + 1))


@tool("在 httpx 官方文档里检索和问题最相关的 5 段（关键词 + 向量混合检索），每段带文件名。"
      "用户问 httpx 怎么用、某个功能是什么时，先用它。",
      query="英文检索词，写出问题的英文表述和可能的参数名、类名，例如 follow redirects default follow_redirects")
def search_docs(query):
    parts = []
    for i, (file, text) in enumerate(retriever.search(query, k=5), 1):
        parts.append(f"[{i}] 文档 {file}\n{text[:800]}")
    return "\n\n".join(parts)


@tool("读取 httpx 文档里某个文件的指定行，返回带行号的内容，一次最多 80 行。search_docs 的片段不够完整时使用。",
      path="文档文件路径，例如 advanced/timeouts.md", start="起始行号，从 1 开始", end="结束行号")
def read_doc(path, start: int = 1, end: int = 80):
    target = inside(DOCS_DIR, path)
    if target is None:
        return f"错误：文档里没有 {path}。文档路径来自 search_docs 的结果，例如 advanced/timeouts.md"
    return numbered(target, start, end)


@tool("在 httpx 的 Python 源码里搜索一段文本或正则表达式，返回文件路径、行号和那一行，最多 30 条。"
      "文档里找不到答案时使用，比如某个参数的默认值、某个异常在什么情况下抛出、某个函数的内部逻辑。",
      pattern="要搜索的文本或正则表达式，例如 DEFAULT_MAX_REDIRECTS 或 def raise_for_status")
def grep_source(pattern):
    try:
        regex = re.compile(pattern)
    except re.error:
        regex = re.compile(re.escape(pattern))
    hits = []
    for p in sorted((SOURCE_DIR / "httpx").rglob("*.py")):
        for n, line in enumerate(p.read_text(errors="replace").splitlines(), 1):
            if regex.search(line):
                hits.append(f"{p.relative_to(SOURCE_DIR)}:{n}: {line.strip()[:120]}")
    if not hits:
        return f"源码里没有找到 {pattern}，换个写法再试，比如只搜函数名或常量名"
    more = f"\n……还有 {len(hits) - 30} 条没有显示，请用更具体的搜索词" if len(hits) > 30 else ""
    return "\n".join(hits[:30]) + more


@tool("读取 httpx 源码文件的指定行，返回带行号的内容，一次最多 80 行。通常在 grep_source 找到行号之后使用。",
      path="源码文件路径，来自 grep_source 的结果，例如 httpx/_config.py", start="起始行号", end="结束行号")
def read_source(path, start: int = 1, end: int = 80):
    target = inside(SOURCE_DIR, path)
    if target is None:
        return f"错误：源码里没有 {path}。路径应该来自 grep_source 的结果，例如 httpx/_client.py"
    return numbered(target, start, end)