kbsearch-extension/scripts/llm.js

// scripts/llm.js
const LLM_TOKEN_API = "http://llm.wisdompulse.cn/api/request-token";
const LLM_WORKSPACES_API = "http://llm.wisdompulse.cn/api/workspaces";
const LLM_STREAM_API_BASE = "http://llm.wisdompulse.cn/api/workspace";
const LLM_TOKEN_EXPIRE_MS = 23 * 60 * 60 * 1000;

async function getLLMPassword() {
    return new Promise((resolve) => {
        chrome.storage.local.get(["llmPassword"], (res) => {
            resolve(res.llmPassword || "");
        });
    });
}

async function getLLMToken(forceRefresh = false) {
    return new Promise((resolve, reject) => {
        chrome.storage.local.get(["llmToken", "llmTokenAt"], async (res) => {
            const now = Date.now();
            if (!forceRefresh && res.llmToken && res.llmTokenAt && now - res.llmTokenAt < LLM_TOKEN_EXPIRE_MS) {
                log("LLM token 命中缓存");
                resolve(res.llmToken);
            } else {
                try {
                    const pwd = await getLLMPassword();
                    const response = await fetch(LLM_TOKEN_API, {
                        method: "POST",
                        headers: { "Content-Type": "application/json" },
                        body: JSON.stringify({ password: pwd }),
                    });
                    const data = await response.json();
                    if (data.token) {
                        chrome.storage.local.set({ llmToken: data.token, llmTokenAt: now });
                        resolve(data.token);
                    } else {
                        reject("LLM token获取失败");
                    }
                } catch (e) {
                    reject("LLM token接口异常");
                }
            }
        });
    });
}

async function getLLMWorkspaceSlug(token, forceRefresh = false) {
    return new Promise((resolve, reject) => {
        chrome.storage.local.get(["llmWorkspaceSlug", "llmWorkspaceAt"], async (res) => {
            const now = Date.now();
            if (!forceRefresh && res.llmWorkspaceSlug && res.llmWorkspaceAt && now - res.llmWorkspaceAt < LLM_TOKEN_EXPIRE_MS) {
                resolve(res.llmWorkspaceSlug);
            } else {
                try {
                    const response = await fetch(LLM_WORKSPACES_API, {
                        method: "GET",
                        headers: { Authorization: `Bearer ${token}` },
                    });
                    const data = await response.json();
                    if (data.workspaces && data.workspaces.length > 0 && data.workspaces[0].slug) {
                        chrome.storage.local.set({ llmWorkspaceSlug: data.workspaces[0].slug, llmWorkspaceAt: now });
                        resolve(data.workspaces[0].slug);
                    } else {
                        reject("未获取到LLM工作区slug");
                    }
                } catch (e) {
                    reject("LLM工作区接口异常");
                }
            }
        });
    });
}

async function fetchLLMAnswer(question, container, contentDiv) {
    if (!question || !question.trim()) return;

    const userItem = document.createElement("div");
    userItem.className = "llm-msg-item llm-user";
    userItem.innerHTML = `<span style="color:#2a63c8;font-weight:bold;">我：</span> ${question}`;
    contentDiv.appendChild(userItem);
    container.classList.remove("hidden");

    const msgBox = document.createElement("div");
    msgBox.className = "llm-msg-item llm-assistant";
    contentDiv.appendChild(msgBox);
    contentDiv.scrollTop = contentDiv.scrollHeight;

    try {
        const token = await getLLMToken();
        const workspace = await getLLMWorkspaceSlug(token);
        const url = `${LLM_STREAM_API_BASE}/${workspace}/stream-chat`;

        const res = await fetch(url, {
            method: "POST",
            headers: {
                "Authorization": `Bearer ${token}`,
                "Content-Type": "application/json"
            },
            body: JSON.stringify({ message: question, attachments: [] })
        });

        if (!res.ok || !res.body) {
            msgBox.innerHTML = "<span style='color:red;'>LLM接口请求失败</span>";
            return;
        }

        const reader = res.body.getReader();
        const decoder = new TextDecoder("utf-8");
        let buffer = "", done = false;
        let thinkStr = "", answerStr = "", inThink = false;

        while (!done) {
            const { value, done: readerDone } = await reader.read();
            done = readerDone;
            if (value) {
                buffer += decoder.decode(value, { stream: true });
                const lines = buffer.split("\n");
                buffer = lines.pop();
                for (let line of lines) {
                    if (line.startsWith("data: ")) {
                        const data = line.slice(6);
                        const chunk = JSON.parse(data);
                        let content = chunk.textResponse || "";
                        if (content.includes("<think>")) {
                            inThink = true;
                            content = content.replace("<think>", "");
                        }
                        if (content.includes("</think>")) {
                            inThink = false;
                            content = content.replace("</think>", "");
                        }
                        if (inThink) {
                            thinkStr += content.replace(/\n/g, "<br>");
                        } else {
                            answerStr += content.replace(/\n/g, "<br>");
                        }
                        msgBox.innerHTML =
                            (thinkStr ? `<div class="llm-think">${thinkStr}</div>` : "") +
                            `<span class="assistant-label">智能助手：</span>${answerStr}`;
                        contentDiv.scrollTop = contentDiv.scrollHeight;
                    }
                }
            }
        }
    } catch (e) {
        msgBox.innerHTML += `<span style='color:red;'>LLM异常：${e}</span>`;
    }
}