task_worker_v2_portable/platforms/wenxin.py

import base64
import hashlib
import json
import re
import time
from urllib.parse import quote


class WenxinClient:
    BASE_URL = "https://chat.baidu.com"

    def __init__(self, cookies=None):
        self._cookie_str = ""
        if cookies:
            self._cookie_str = "; ".join(f'{c["name"]}={c["value"]}' for c in cookies)

    async def chat(self, question, thinking=False):
        from curl_cffi.requests import AsyncSession

        async with AsyncSession(impersonate="chrome136", verify=False) as session:
            if self._cookie_str:
                session.headers["cookie"] = self._cookie_str

            resp = await session.get(self.BASE_URL + "/", timeout=20)
            match = re.search(r'\{"token"\s*:\s*"([^"]+)"\s*,\s*"lid"\s*:\s*"(\d+)"', resp.text)
            if not match:
                return "", ""
            token, lid = match.group(1), match.group(2)

            ts = int(time.time() * 1000)
            query_md5 = hashlib.md5(question.encode("utf-8")).hexdigest()
            raw_token = f"{token}|{query_md5}|{ts}|{lid}"
            chat_token = f"{base64.b64encode(raw_token.encode()).decode()}-{lid}-3"

            payload = {
                "message": {
                    "inputMethod": "chat_search", "isRebuild": False,
                    "content": {
                        "query": "", "qtype": 0,
                        "agentInfo": {"agent_id": [""], "params": json.dumps({"agt_rk": 1, "agt_sess_cnt": 1})},
                        "agentInfoList": [],
                    },
                    "searchInfo": {
                        "srcid": "", "order": "", "tplname": "", "dqaKey": "",
                        "re_rank": "1", "ori_lid": "", "sa": "bkb", "enter_type": "chat_url",
                        "chatParams": {"setype": "csaitab", "chat_samples": "WISE_NEW_CSAITAB",
                                       "chat_token": chat_token, "scene": ""},
                        "isPrivateChat": False,
                        "usedModel": {
                            "modelName": "smartMode",
                            "modelFunction": {"deepSearch": "1" if thinking else "0", "thinkMode": "0"},
                            "showModelName": "smartMode",
                        },
                        "landingPageSwitch": "", "landingPage": "aitab", "isInnovate": 2,
                        "showMindMap": False,
                    },
                    "from": "", "source": "pc_csaitab",
                    "query": [{"type": "TEXT", "data": {"text": {"query": question, "extData": "{}", "text_type": ""}}}],
                    "anti_ext": {"inputT": None},
                },
                "sa": "bkb", "setype": "csaitab", "rank": 1,
            }

            headers = {
                "Accept": "text/event-stream", "Content-Type": "application/json",
                "Origin": self.BASE_URL, "Referer": f"{self.BASE_URL}/",
                "source": "pc_csaitab", "isDeepseek": "1",
                "X-Chat-Message": f"query:{quote(question)},anti_ext:%7B%22inputT%22%3Anull%7D,enter_type:chat_url,re_rank:1,modelName:smartMode,sa:bkb",
            }

            resp = await session.post(
                self.BASE_URL + "/aichat/api/conversation",
                headers=headers, json=payload, stream=True, timeout=180,
            )

            raw_text = ""
            req_body = json.dumps(payload, ensure_ascii=False)
            async for chunk in resp.aiter_content():
                raw_text += chunk.decode("utf-8", errors="replace")

            return raw_text, req_body


async def extract_wenxin_cookies():
    from playwright.async_api import async_playwright
    from .base import CDP_URL
    try:
        async with async_playwright() as pw:
            browser = await pw.chromium.connect_over_cdp(CDP_URL)
            ctx = browser.contexts[0]
            cookies = await ctx.cookies(["https://chat.baidu.com", "https://www.baidu.com", "https://baidu.com"])
            if any(c["name"] in ("BDUSS", "BDUSS_BFESS") for c in cookies):
                return cookies
    except Exception:
        pass
    return None


async def get_wenxin_share_url(lid, wx_page=None):
    from playwright.async_api import async_playwright
    from .base import CDP_URL
    page = wx_page
    pw_ctx = None
    try:
        if not page:
            pw_ctx = await async_playwright().start()
            browser = await pw_ctx.chromium.connect_over_cdp(CDP_URL)
            page = next((p for p in browser.contexts[0].pages if "baidu.com" in p.url), None)
            if not page:
                if pw_ctx:
                    await pw_ctx.stop()
                return ""

        await page.goto(f"https://chat.baidu.com/search/{lid}", wait_until="domcontentloaded", timeout=20000)
        await asyncio.sleep(3)
        await page.evaluate("""() => {
            window.__sc = [];
            const of = window.fetch;
            window.fetch = async function(...a) {
                const r = await of.apply(this, a);
                try { const c=r.clone(); const t=await c.text();
                    if (t.includes('short_url')) window.__sc.push(t);
                } catch {}
                return r;
            };
        }""")
        await page.click(".cos-icon-share1", timeout=8000)
        await asyncio.sleep(1)
        await page.click('button:has-text("复制链接")', timeout=5000)
        await asyncio.sleep(2)
        captures = await page.evaluate("() => window.__sc")
        import json as _json
        for c in captures:
            d = _json.loads(c)
            url = d.get("data", {}).get("short_url", "")
            if url:
                if pw_ctx:
                    await pw_ctx.stop()
                return url
        if pw_ctx:
            await pw_ctx.stop()
    except Exception as e:
        print(f"    ⚠ 文心分享链接获取失败 (lid={lid}): {e}")
        if pw_ctx:
            try:
                await pw_ctx.stop()
            except Exception:
                pass
    return ""