app.py

import os
import hashlib
import base64
import struct
import threading
import httpx
import logging
import gradio as gr
from fastapi import FastAPI, Request
from fastapi.responses import PlainTextResponse
from openai import OpenAI
from Crypto.Cipher import AES
import xml.etree.ElementTree as ET
import uvicorn

# ── 日志 ──────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger("hermes")

# ── 环境变量 ──────────────────────────────────────────────────
NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY", "")
WECOM_TOKEN    = os.getenv("WECOM_TOKEN", "").strip()
WECOM_AES_KEY  = os.getenv("WECOM_AES_KEY", "").strip()
WECOM_CORPID   = os.getenv("CORPID", "").strip()
WECOM_SECRET   = os.getenv("WECOM_SECRET", "").strip()
WECOM_AGENTID  = os.getenv("WECOM_AGENTID", "1000003").strip() or "1000003"

# ── LLM ───────────────────────────────────────────────────────
llm = OpenAI(base_url="https://integrate.api.nvidia.com/v1", api_key=NVIDIA_API_KEY)

SYSTEM_PROMPT = (
    "You are Hermes, a helpful AI assistant. "
    "Answer concisely and accurately. Reply in the same language as the user."
)

def chat_with_hermes(message: str, history: list) -> str:
    log.info("[WEB] user: %s", message)
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for h in history:
        if isinstance(h, (list, tuple)) and len(h) == 2:
            if h[0]:
                messages.append({"role": "user", "content": str(h[0])})
            if h[1]:
                messages.append({"role": "assistant", "content": str(h[1])})
    messages.append({"role": "user", "content": message})
    resp = llm.chat.completions.create(
       # model="meta/llama-3.1-8b-instruct",
        model="minimaxai/minimax-m2.7",
        messages=messages,
        max_tokens=1024,
        temperature=0.7,
    )
    reply = resp.choices[0].message.content
    log.info("[WEB] assistant: %s", reply)
    return reply

# ── 企业微信 ──────────────────────────────────────────────────
_wx_cache = {"token": "", "expires": 0}
_lock = threading.Lock()

def get_wecom_token() -> str:
    import time
    with _lock:
        if time.time() < _wx_cache["expires"] - 60:
            return _wx_cache["token"]
        r = httpx.get(
            "https://qyapi.weixin.qq.com/cgi-bin/gettoken",
            params={"corpid": WECOM_CORPID, "corpsecret": WECOM_SECRET},
            timeout=10,
        ).json()
        log.info("[WECOM] get_token: %s", r)
        _wx_cache["token"] = r.get("access_token", "")
        _wx_cache["expires"] = time.time() + r.get("expires_in", 7200)
        return _wx_cache["token"]

def send_wecom_message(to_user: str, content: str):
    token = get_wecom_token()
    r = httpx.post(
        f"https://qyapi.weixin.qq.com/cgi-bin/message/send?access_token={token}",
        json={"touser": to_user, "msgtype": "text",
              "agentid": int(WECOM_AGENTID), "text": {"content": content}},
        timeout=10,
    ).json()
    log.info("[WECOM] send to=%s result=%s", to_user, r)

def decrypt_msg(encrypted: str) -> str:
    # 企业微信 AESKey = Base64(43位) 解码后恰好 32 字节，不需要补 =
    aes_key = base64.b64decode(WECOM_AES_KEY + "=")  # 补一个 = 凑成44位合法base64
    cipher = AES.new(aes_key, AES.MODE_CBC, aes_key[:16])
    raw = cipher.decrypt(base64.b64decode(encrypted))
    # 手动去除 PKCS7 padding（企业微信标准）
    pad_len = raw[-1]
    raw = raw[:-pad_len]
    # 格式: 16字节随机 + 4字节消息长度(big-endian) + 消息内容 + corpid
    msg_len = struct.unpack(">I", raw[16:20])[0]
    return raw[20: 20 + msg_len].decode("utf-8")

def verify_sig(timestamp: str, nonce: str, encrypt: str = "") -> str:
    parts = [WECOM_TOKEN, timestamp, nonce] + ([encrypt] if encrypt else [])
    return hashlib.sha1("".join(sorted(parts)).encode()).hexdigest()

# ── FastAPI ───────────────────────────────────────────────────
fastapi_app = FastAPI()

@fastapi_app.get("/gateway/wecom")
async def wecom_verify(request: Request):
    p = request.query_params
    log.info("[WECOM] GET params: %s", dict(p))
    try:
        sig = verify_sig(p.get("timestamp", ""), p.get("nonce", ""), p.get("echostr", ""))
        expected = p.get("msg_signature", "")
        log.info("[WECOM] sig=%s expected=%s match=%s", sig, expected, sig == expected)
        if sig != expected:
            return PlainTextResponse("")
        plain = decrypt_msg(p.get("echostr", ""))
        log.info("[WECOM] verify ok plain=%s", plain)
        return PlainTextResponse(plain)
    except Exception as e:
        log.exception("[WECOM] verify error: %s", e)
        return PlainTextResponse("")

@fastapi_app.post("/gateway/wecom")
async def wecom_message(request: Request):
    p = request.query_params
    log.info("[WECOM] POST params: %s", dict(p))
    try:
        body = await request.body()
        log.info("[WECOM] body: %s", body.decode())
        encrypt = ET.fromstring(body.decode()).findtext("Encrypt", "")
        sig = verify_sig(p.get("timestamp", ""), p.get("nonce", ""), encrypt)
        expected = p.get("msg_signature", "")
        log.info("[WECOM] sig=%s expected=%s match=%s", sig, expected, sig == expected)
        if sig != expected:
            return PlainTextResponse("success")
        root = ET.fromstring(decrypt_msg(encrypt))
        msg_type = root.findtext("MsgType")
        from_user = root.findtext("FromUserName", "")
        log.info("[WECOM] msg_type=%s from=%s", msg_type, from_user)
        if msg_type == "text":
            content = root.findtext("Content", "").strip()
            log.info("[WECOM] content=%s", content)
            if content:
                def reply():
                    try:
                        answer = chat_with_hermes(content, [])
                        log.info("[WECOM] reply=%s", answer)
                        send_wecom_message(from_user, answer)
                    except Exception as e:
                        log.exception("[WECOM] reply error: %s", e)
                threading.Thread(target=reply, daemon=True).start()
    except Exception as e:
        log.exception("[WECOM] message error: %s", e)
    return PlainTextResponse("success")

# ── Gradio 挂载到 FastAPI ─────────────────────────────────────
demo = gr.ChatInterface(
    fn=chat_with_hermes,
    title="🤖 Hermes Agent",
    description="Powered by NVIDIA NIM · 企业微信同步接入",
)
app = gr.mount_gradio_app(fastapi_app, demo, path="/")

# ── 启动 ──────────────────────────────────────────────────────
if __name__ == "__main__":
    log.info("WECOM_TOKEN=%s AES_KEY=%s CORPID=%s SECRET=%s NVIDIA=%s AGENTID=%s",
             bool(WECOM_TOKEN), bool(WECOM_AES_KEY),
             bool(WECOM_CORPID), bool(WECOM_SECRET), bool(NVIDIA_API_KEY), WECOM_AGENTID)
    uvicorn.run(app, host="0.0.0.0", port=7860)