app.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
TreeChat - 树状 AI 对话
-----------------------
一个零依赖（仅用 Python 标准库）的本地服务器，作为浏览器与各 AI API 之间的代理：
  - 浏览器只跟 localhost 通信，避免 CORS 问题；
  - API key 由前端从 localStorage 恢复并 POST 过来，服务端只在内存中持有，不读环境变量；
  - 支持流式（SSE）输出。

用法：
  1) python app.py
  2) 浏览器自动打开 http://127.0.0.1:8000，首次使用在设置里填写 API Key
"""

import argparse
import gzip
import json
import os
import socket
import subprocess
import sys
import threading
import time
import urllib.error
import urllib.request
import webbrowser
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Optional, Dict, Union

# 只读资源（index.html/static/VERSION）所在目录
HERE: str = os.path.dirname(os.path.abspath(__file__))

# 持久化数据（日志等）不能用 HERE：PyInstaller onefile 模式下 __file__ 指向运行时
# 临时解压目录（sys._MEIPASS），进程退出后通常被清理；只有可执行文件自身所在的目录
# 才是用户能找到、跨次运行持久保存的位置
EXE_DIR: str = (
    os.path.dirname(sys.executable) if getattr(sys, "frozen", False) else HERE
)

DO_NOT_MASK_KEY: bool = False

_VERSION_CACHE: Optional[Dict[str, str]] = None


def _get_version() -> Dict[str, str]:
    """返回 {"version": str, "type": "tag"|"commit"|"unknown"}"""

    global _VERSION_CACHE
    if _VERSION_CACHE is not None:
        return _VERSION_CACHE

    def _is_commit_hash(s: str) -> bool:
        return 7 <= len(s) <= 40 and all(c in "0123456789abcdef" for c in s.lower())

    # 打包二进制：读捆绑的 VERSION 文件，按内容判断类型
    if getattr(sys, "frozen", False):
        vfile = os.path.join(HERE, "VERSION")
        if os.path.isfile(vfile):
            try:
                v = open(vfile, encoding="utf-8").read().strip()
                if v:
                    t = "commit" if _is_commit_hash(v) else "tag"
                    _VERSION_CACHE = {"version": v, "type": t}
                    return _VERSION_CACHE
            except Exception:
                pass
        _VERSION_CACHE = {"version": "未知", "type": "unknown"}
        return _VERSION_CACHE

    # 源码运行：先查 git tag，再查 commit hash
    try:
        tag = (
            subprocess.check_output(
                ["git", "describe", "--tags", "--exact-match", "HEAD"],
                stderr=subprocess.DEVNULL,
                cwd=HERE,
            )
            .decode()
            .strip()
        )
        if tag:
            _VERSION_CACHE = {"version": tag, "type": "tag"}
            return _VERSION_CACHE
    except Exception:
        pass
    try:
        commit = (
            subprocess.check_output(
                ["git", "rev-parse", "--short", "HEAD"],
                stderr=subprocess.DEVNULL,
                cwd=HERE,
            )
            .decode()
            .strip()
        )
        if commit:
            _VERSION_CACHE = {"version": commit, "type": "commit"}
            return _VERSION_CACHE
    except Exception:
        pass
    _VERSION_CACHE = {"version": "未知", "type": "unknown"}
    return _VERSION_CACHE


# ---------- 日志 ----------
# 单个 latest.log 超过这个大小就轮转，不再只在进程启动时轮转一次——server 模式一次运行可能
# 持续很久，聊天记录（尤其带图片附件的多模态对话，见 _redact_for_log）不加限制地写下去
# 会让日志文件无限增长
LOG_MAX_BYTES = 10 * 1024 * 1024


class _Logger:
    def __init__(self, log_dir: str, log_to_file: bool) -> None:
        self._log_dir = log_dir
        self._f = None
        # ThreadingHTTPServer 下每个请求各自的线程都会调 write()，不加锁的话文件写入可能
        # 交错，产生损坏/错位的日志行（screen 那半没这个问题，但同一把锁一起保护更简单）
        self._lock = threading.Lock()
        if log_to_file:
            os.makedirs(log_dir, exist_ok=True)
            self._rotate_if_exists()
            self._f = open(os.path.join(log_dir, "latest.log"), "a", encoding="utf-8")

    def _rotate_if_exists(self) -> None:
        latest = os.path.join(self._log_dir, "latest.log")
        if not os.path.exists(latest):
            return
        today = time.strftime("%Y-%m-%d")
        n = 1
        while True:
            arc = os.path.join(self._log_dir, "%s-%d.log.gz" % (today, n))
            if not os.path.exists(arc):
                break
            n += 1
        with open(latest, "rb") as fin:
            with gzip.open(arc, "wb") as fout:
                fout.writelines(fin)
        os.remove(latest)

    def write(self, msg: str, screen: Optional[str] = None) -> None:
        ts = time.strftime("%Y-%m-%d %H:%M:%S")
        scr = "[%s] %s" % (ts, screen if screen is not None else msg)
        with self._lock:
            sys.stderr.write(scr + "\n")
            if self._f:
                self._f.write("[%s] %s\n" % (ts, msg))
                self._f.flush()
                if self._f.tell() > LOG_MAX_BYTES:
                    self._f.close()
                    self._rotate_if_exists()
                    self._f = open(
                        os.path.join(self._log_dir, "latest.log"), "a", encoding="utf-8"
                    )

    def close(self) -> None:
        if self._f:
            self._f.close()


def _mask_key(k: str) -> str:
    if DO_NOT_MASK_KEY:
        return k

    if not k or len(k) < 8:
        return "***"
    return k[:4] + "***" + k[-4:]


def _redact_for_log(obj):
    """递归拷贝一份 messages 结构用于写日志：多模态附件（图片/PDF）在 content 里是
    data: 开头的完整 base64，动辄几十 KB 到几 MB，原样写进日志既让文件体积暴涨，也是
    明文落盘的隐私风险（完整图片内容）——这里只截断超长的 data: URI，不影响真正发给
    上游 API 的 payload（这个函数的返回值只用于日志，调用方传进来的原始 msgs 不会被改）"""
    if isinstance(obj, str):
        if obj.startswith("data:") and len(obj) > 200:
            return "data:<%d bytes redacted>" % len(obj)
        return obj
    if isinstance(obj, list):
        return [_redact_for_log(v) for v in obj]
    if isinstance(obj, dict):
        return {k: _redact_for_log(v) for k, v in obj.items()}
    return obj


PROVIDERS: Dict[str, str] = {
    "deepseek": "https://api.deepseek.com/chat/completions",
    "openai": "https://api.openai.com/v1/chat/completions",
    "groq": "https://api.groq.com/openai/v1/chat/completions",
    "xai": "https://api.x.ai/v1/chat/completions",
    "mistral": "https://api.mistral.ai/v1/chat/completions",
    "kimi": "https://api.moonshot.cn/v1/chat/completions",
    "zhipu": "https://open.bigmodel.cn/api/paas/v4/chat/completions",
    "qwen": "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions",
    # Claude/Gemini 官方都提供了 OpenAI 兼容层（同样的 /chat/completions 请求/响应格式、
    # Authorization: Bearer 鉴权），走这套兼容端点就能直接复用这里现成的透传代理，不用为
    # 它们各自的原生 API（/v1/messages、GenerateContent）单独写请求/响应格式转换逻辑
    "claude": "https://api.anthropic.com/v1/chat/completions",
    "gemini": "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
}

# 联网搜索：Tavily 是专为 LLM/Agent 场景设计的搜索 API，单次 POST 请求、Authorization: Bearer
# 鉴权，跟上面这些 LLM provider 的调用方式一致——不需要为它单独写一套请求逻辑
TAVILY_SEARCH_URL: str = "https://api.tavily.com/search"


# index.html 磁盘原文里这一行恒为 false；只有这个函数真正把文件发给浏览器时才替换成 true。
# 前端 init() 靠这个变量判断"是不是真的连着这个后端"，而不是看请求协议——直接双击
# index.html、拖进 IDE 自带的预览、随便起一个静态文件服务器指到这个目录，走的都是
# http(s)/file 协议，但都不会经过这个函数，读到的永远是磁盘上原始的 false
_BACKEND_FLAG_FALSE = b"window.__TREECHAT_BACKEND__ = false;"
_BACKEND_FLAG_TRUE = b"window.__TREECHAT_BACKEND__ = true;"


def read_index() -> bytes:
    path = os.path.join(HERE, "index.html")
    with open(path, "rb") as f:
        data = f.read()
    if _BACKEND_FLAG_FALSE not in data:
        # index.html 被改动导致标记对不上，只打个日志到 stderr，不影响正常访问——
        # 这只是个 UX 提示用的标记，不是安全边界，没必要为它把整个页面搞挂
        print(
            "警告: index.html 里没找到 __TREECHAT_BACKEND__ 标记，"
            "「未连接后端」提示会失效",
            file=sys.stderr,
        )
        return data
    return data.replace(_BACKEND_FLAG_FALSE, _BACKEND_FLAG_TRUE, 1)


class _QuietHTTPServer(ThreadingHTTPServer):
    """客户端在服务端读完请求之前就断开连接（关标签页、切后台、代理/网关超时、
    弱网、甚至端口扫描器探测），会在 socketserver 内部读 socket 时抛出
    ConnectionResetError/BrokenPipeError 等——这发生在我们自己的 do_GET/do_POST
    执行之前（stdlib 的 handle_one_request 里），没法用 try/except 包在业务代码里挡住。
    ThreadingHTTPServer 默认的 handle_error() 会把完整 traceback 摔到 stderr，看着像
    程序炸了，其实是正常的网络状况，在暴露局域网/公网的 --server 模式下尤其常见。
    这里覆盖 handle_error()，把这类"预期内的连接异常"降级成一行安静日志；其它类型的
    异常（真正的程序 bug）仍然走默认行为、打印完整 traceback，不能被这次改动误伤"""

    _QUIET_EXCS = (
        ConnectionResetError,
        BrokenPipeError,
        ConnectionAbortedError,
        TimeoutError,
    )

    def __init__(self, server_address, handler_cls, logger: _Logger) -> None:
        self._logger = logger
        super().__init__(server_address, handler_cls)

    def handle_error(self, request, client_address) -> None:
        exc = sys.exc_info()[1]
        if isinstance(exc, self._QUIET_EXCS):
            self._logger.write(
                "[%s:%d] 连接中断（%s），已忽略"
                % (client_address[0], client_address[1], type(exc).__name__)
            )
            return
        super().handle_error(request, client_address)


class Handler(BaseHTTPRequestHandler):
    def __init__(self, logger: _Logger, *args, **kwargs) -> None:
        self._logger = logger
        super().__init__(*args, **kwargs)

    def log_message(self, fmt: str, *args) -> None:
        pass  # 由 _log_request 统一处理

    def _log(self, msg: str, screen: Optional[str] = None) -> None:
        self._logger.write(
            "[%s] %s" % (self.address_string(), msg),
            "[%s] %s" % (self.address_string(), screen) if screen is not None else None,
        )

    def _send(
        self,
        code: int,
        body: Union[dict, bytes],
        ctype: str = "application/json; charset=utf-8",
    ) -> None:
        body_bytes: bytes = (
            json.dumps(body, ensure_ascii=False).encode("utf-8")
            if isinstance(body, dict)
            else body
        )
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(body_bytes)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        try:
            self.wfile.write(body_bytes)
        except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
            pass

    def _read_json(self) -> dict:
        length = int(self.headers.get("Content-Length", 0))
        raw = self.rfile.read(length) if length else b""
        if not raw:
            return {}
        return json.loads(raw.decode("utf-8"))

    # ---------- GET ----------
    def do_GET(self) -> None:
        if self.path in ("/", "/index.html"):
            try:
                self._send(200, read_index(), "text/html; charset=utf-8")
            except FileNotFoundError:
                self._send(500, {"error": "index.html not found next to app.py"})
        elif self.path == "/api/status":
            self._send(200, {"ok": True})
        elif self.path == "/api/version":
            self._send(200, _get_version())
        elif self.path.startswith("/static/"):
            self._serve_static(self.path)
        else:
            self._send(404, {"error": "not found"})

    # 托管本地静态资源（第三方库 / PWA manifest 与图标等），仅限 static/ 目录，防目录穿越
    def _serve_static(self, path: str) -> None:
        rel = path.split("?", 1)[0].lstrip("/")
        full = os.path.normpath(os.path.join(HERE, rel))
        staticdir = os.path.normpath(os.path.join(HERE, "static"))
        if not full.startswith(staticdir + os.sep) or not os.path.isfile(full):
            return self._send(404, {"error": "not found"})
        ext = os.path.splitext(full)[1].lower()
        ctype = {
            ".js": "application/javascript; charset=utf-8",
            ".css": "text/css; charset=utf-8",
            ".woff2": "font/woff2",
            ".woff": "font/woff",
            ".ttf": "font/ttf",
            ".json": "application/json; charset=utf-8",
            ".png": "image/png",
            ".ico": "image/x-icon",
            ".svg": "image/svg+xml",
        }.get(ext, "application/octet-stream")
        with open(full, "rb") as f:
            data = f.read()
        self.send_response(200)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(data)))
        self.send_header("Cache-Control", "max-age=86400")
        self.end_headers()
        try:
            self.wfile.write(data)
        except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
            pass

    # ---------- POST ----------
    def do_POST(self) -> None:
        if self.path == "/api/test-key":
            return self._test_key()

        if self.path == "/api/chat":
            return self._proxy_chat()

        if self.path == "/api/search":
            return self._web_search()

        self._send(404, {"error": "not found"})

    # ---------- Key 验证（发一条最小非流式请求，只看状态码） ----------
    def _test_key(self) -> None:
        try:
            data = self._read_json()
        except Exception as e:
            return self._send(400, {"error": str(e)})
        provider = (data.get("provider", "")).strip()
        api_key = (data.get("key", "")).strip()
        model = (data.get("model", "")).strip()

        # 日志
        self._log(
            "POST /api/test-key  provider=%s  model=%s  key=%s"
            % (provider, model, _mask_key(api_key))
        )

        if provider not in PROVIDERS:
            return self._send(400, {"error": "缺少或未知的提供商"})
        if not api_key:
            return self._send(400, {"error": "Key 未设置"})
        payload = json.dumps(
            {
                "model": model,
                "messages": [{"role": "user", "content": "hi"}],
                "stream": False,
                "max_tokens": 1,
            }
        ).encode("utf-8")
        req = urllib.request.Request(
            PROVIDERS[provider],
            data=payload,
            headers={
                "Content-Type": "application/json",
                "Authorization": "Bearer " + api_key,
            },
            method="POST",
        )
        try:
            resp = urllib.request.urlopen(req, timeout=15)
            resp.read()
            return self._send(200, {"ok": True})
        except urllib.error.HTTPError as e:
            try:
                body = _safe_json(e.read().decode("utf-8", "replace"))
            except Exception:
                body = str(e)
            return self._send(200, {"ok": False, "status": e.code, "body": body})
        except Exception as e:
            return self._send(200, {"ok": False, "body": str(e)})

    # ---------- 联网搜索（Tavily，非流式，单次请求即可） ----------
    def _web_search(self) -> None:
        try:
            data = self._read_json()
        except Exception as e:
            return self._send(400, {"error": "bad json: %s" % e})

        query = (data.get("query", "")).strip()
        api_key = (data.get("key", "")).strip()

        self._log("POST /api/search  key=%s  query=%s" % (_mask_key(api_key), query))

        if not query:
            return self._send(400, {"error": "缺少 query 参数"})
        if not api_key:
            return self._send(401, {"error": "未设置 Tavily API key，请在设置中填写。"})

        payload = json.dumps({"query": query, "max_results": 5}).encode("utf-8")
        req = urllib.request.Request(
            TAVILY_SEARCH_URL,
            data=payload,
            headers={
                "Content-Type": "application/json",
                "Authorization": "Bearer " + api_key,
            },
            method="POST",
        )
        try:
            resp = urllib.request.urlopen(req, timeout=20)
            body = json.loads(resp.read().decode("utf-8"))
            return self._send(
                200,
                {
                    "ok": True,
                    "results": body.get("results", []),
                    "answer": body.get("answer", ""),
                },
            )
        except urllib.error.HTTPError as e:
            try:
                body = _safe_json(e.read().decode("utf-8", "replace"))
            except Exception:
                body = str(e)
            return self._send(200, {"ok": False, "status": e.code, "body": body})
        except Exception as e:
            return self._send(502, {"ok": False, "body": str(e)})

    # ---------- 流式代理 ----------
    def _proxy_chat(self) -> None:
        try:
            data = self._read_json()
        except Exception as e:
            return self._send(400, {"error": "bad json: %s" % e})

        provider = (data.get("provider", "")).strip()
        if provider not in PROVIDERS:
            return self._send(400, {"error": "缺少或未知的提供商"})
        api_key = (data.get("key", "")).strip()
        if not api_key:
            return self._send(
                401, {"error": "未设置 %s 的 API key，请在设置中填写。" % provider}
            )

        # 日志：屏幕截断对话预览，文件记录完整对话
        msgs = data.get("messages", [])
        last_user = ""
        for m in reversed(msgs):
            if isinstance(m, dict) and m.get("role") == "user":
                c = m.get("content", "")
                last_user = c if isinstance(c, str) else str(c)
                break
        preview = (
            last_user[:80] + ("..." if len(last_user) > 80 else "")
            if last_user
            else "(无对话)"
        )
        self._log(
            "POST /api/chat  provider=%s  model=%s  key=%s  messages=%s"
            % (
                provider,
                data.get("model", "?"),
                _mask_key(api_key),
                json.dumps(_redact_for_log(msgs), ensure_ascii=False),
            ),
            "POST /api/chat  provider=%s  model=%s  key=%s  msgs=%d  preview=%s"
            % (
                provider,
                data.get("model", "?"),
                _mask_key(api_key),
                len(msgs),
                preview,
            ),
        )

        # 透传前端 payload，去掉仅供路由用的字段
        payload = {k: v for k, v in data.items() if k not in ("provider", "key")}
        payload["stream"] = True
        if not payload.get("model"):
            return self._send(400, {"error": "缺少 model 参数"})
        payload.setdefault("messages", [])

        req = urllib.request.Request(
            PROVIDERS[provider],
            data=json.dumps(payload).encode("utf-8"),
            headers={
                "Content-Type": "application/json",
                "Authorization": "Bearer " + api_key,
                "Accept": "text/event-stream",
            },
            method="POST",
        )

        try:
            upstream = urllib.request.urlopen(req, timeout=300)
        except urllib.error.HTTPError as e:
            # 把上游错误体透传给前端（仍以 SSE 形式，方便前端统一处理）
            try:
                err_body = e.read().decode("utf-8", "replace")
            except Exception:
                err_body = str(e)
            self.send_response(200)
            self.send_header("Content-Type", "text/event-stream; charset=utf-8")
            self.send_header("Cache-Control", "no-store")
            self.end_headers()
            msg = {"error": {"status": e.code, "body": _safe_json(err_body)}}
            self._sse_write("data: " + json.dumps(msg, ensure_ascii=False) + "\n\n")
            return
        except Exception as e:
            return self._send(502, {"error": "无法连接 %s: %s" % (provider, e)})

        # 成功：把上游 SSE 字节流原样转发给浏览器
        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream; charset=utf-8")
        self.send_header("Cache-Control", "no-cache")
        self.send_header("X-Accel-Buffering", "no")  # 禁止 nginx 缓冲 SSE
        self.send_header("Connection", "close")  # 流结束后立即关闭连接
        self.end_headers()
        try:
            while True:
                chunk = upstream.read(1024)
                if not chunk:
                    break
                self.wfile.write(chunk)
                self.wfile.flush()
            # 主动发送 [DONE]，确保在连接关闭前到达浏览器
            self._sse_write("data: [DONE]\n\n")
        except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
            pass  # 浏览器断开（例如用户停止），正常结束
        finally:
            try:
                upstream.close()
            except Exception:
                pass

    def _sse_write(self, text: str) -> None:
        try:
            self.wfile.write(text.encode("utf-8"))
            self.wfile.flush()
        except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
            pass


def _safe_json(s: str) -> Union[dict, str]:
    try:
        return json.loads(s)
    except Exception:
        return s


def main() -> None:
    parser = argparse.ArgumentParser(description="TreeChat local server")
    parser.add_argument(
        "--host", default="127.0.0.1", help="bind host (default: 127.0.0.1)"
    )
    parser.add_argument(
        "--port", type=int, default=8000, help="bind port (default: 8000)"
    )
    parser.add_argument(
        "--server",
        action="store_true",
        help="服务器模式：监听 0.0.0.0（局域网/公网可访问），而不是仅本机",
    )
    parser.add_argument(
        "--do-not-mask-key",
        action="store_true",
        help="不隐藏 API Key（仅用于调试）",
    )
    parser.add_argument(
        "--do-not-log-file",
        action="store_true",
        help="不写日志文件（logs/ 目录），仍会打印到屏幕",
    )
    args = parser.parse_args()

    host = "0.0.0.0" if args.server else args.host
    port = args.port

    # 尝试绑定端口，失败说明端口已被占用（可能是本程序已在运行）
    probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        probe.bind((host, port))
    except OSError:
        url = "http://127.0.0.1:%d" % port
        print("TreeChat already running at %s — opening browser." % url)
        webbrowser.open(url)
        return
    finally:
        probe.close()

    # 日志文件默认开启（与是否 --server 无关），--do-not-log-file 显式关闭
    _logger = _Logger(
        os.path.join(EXE_DIR, "logs"), log_to_file=not args.do_not_log_file
    )

    # 是否隐藏 API Key
    global DO_NOT_MASK_KEY
    DO_NOT_MASK_KEY = args.do_not_mask_key

    server = _QuietHTTPServer(
        (host, port),
        lambda *_args, **_kwargs: Handler(_logger, *_args, **_kwargs),
        _logger,
    )
    url = "http://127.0.0.1:%d" % port

    if host == "0.0.0.0":
        # 服务器模式：找一个局域网/公网 IP 提示用户
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.connect(("8.8.8.8", 80))
            lan_ip = s.getsockname()[0]
            s.close()
        except Exception:
            lan_ip = "<your-ip>"
        print("TreeChat running in SERVER mode")
        print("  Local:   http://127.0.0.1:%d" % port)
        print("  Network: http://%s:%d" % (lan_ip, port))
    else:
        print("TreeChat running at %s" % url)
        threading.Timer(0.6, lambda: webbrowser.open(url)).start()

    print("  Press Ctrl+C to stop.")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nbye")
        _logger.close()
        server.shutdown()


if __name__ == "__main__":
    main()