u-desktop/knowledge.go

package main

import (
	"bytes"
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"math/rand"
	"net/http"
	"path/filepath"
	"time"

	_ "modernc.org/sqlite"
)

const cpaURL = "https://cpa.1216.top/v1/chat/completions"
const cpaKey = "alink-shared-key-1"
const cpaModel = "glm-4.5-air"

type knowledgeData struct {
	Content string `json:"content"`
	Keyword string `json:"keyword"`
}

var knowledgeDB *sql.DB

func initKnowledgeDB() {
	dbPath := filepath.Join(configDir(), "knowledge.db")
	var err error
	knowledgeDB, err = sql.Open("sqlite", dbPath)
	if err != nil {
		log.Println("知识库打开失败:", err)
		return
	}
	knowledgeDB.SetMaxOpenConns(1)
	_, err = knowledgeDB.Exec(`CREATE TABLE IF NOT EXISTS knowledge_cards (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		keyword TEXT NOT NULL,
		content TEXT NOT NULL,
		created_at DATETIME DEFAULT CURRENT_TIMESTAMP
	)`)
	if err != nil {
		log.Println("知识库建表失败:", err)
	}
}

func saveKnowledgeCard(keyword, content string) {
	if knowledgeDB == nil {
		return
	}
	_, err := knowledgeDB.Exec("INSERT INTO knowledge_cards (keyword, content) VALUES (?, ?)", keyword, content)
	if err != nil {
		log.Println("知识保存失败:", err)
	}
}

func getRandomKnowledgeCard(keyword string) string {
	if knowledgeDB == nil {
		return ""
	}
	var content string
	err := knowledgeDB.QueryRow(
		"SELECT content FROM knowledge_cards WHERE keyword = ? ORDER BY RANDOM() LIMIT 1",
		keyword,
	).Scan(&content)
	if err != nil {
		return ""
	}
	return content
}

func getKnowledgeCardCount(keyword string) int {
	if knowledgeDB == nil {
		return 0
	}
	var count int
	err := knowledgeDB.QueryRow(
		"SELECT COUNT(*) FROM knowledge_cards WHERE keyword = ?",
		keyword,
	).Scan(&count)
	if err != nil {
		return 0
	}
	return count
}

func fetchKnowledgeFromLLM(keyword string, cfg *Config) string {
	basePrompt := fmt.Sprintf(
		"根据关键词「%s」，生成一条有趣的知识小卡片。要求：控制在80字以内，简洁有趣，有知识性。直接输出内容，不要加标题、序号或其他格式。",
		keyword,
	)
	if cfg.KnowledgePrompt != "" {
		basePrompt += "\n附加要求：" + cfg.KnowledgePrompt
	}

	body := map[string]interface{}{
		"model":      cpaModel,
		"max_tokens": 256,
		"messages": []map[string]string{
			{"role": "user", "content": basePrompt},
		},
	}
	jsonData, _ := json.Marshal(body)

	req, err := http.NewRequest("POST", cpaURL, bytes.NewReader(jsonData))
	if err != nil {
		log.Println("知识API请求创建失败:", err)
		return ""
	}
	req.Header.Set("Authorization", "Bearer "+cpaKey)
	req.Header.Set("Content-Type", "application/json")

	resp, err := httpClient.Do(req)
	if err != nil {
		log.Println("知识API请求失败:", err)
		return ""
	}
	defer resp.Body.Close()

	var result struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
	}
	if json.NewDecoder(resp.Body).Decode(&result) != nil {
		log.Println("知识API响应解析失败")
		return ""
	}
	if len(result.Choices) > 0 {
		return result.Choices[0].Message.Content
	}
	return ""
}

func pushKnowledgeJSON(content, keyword string) {
	data, _ := json.Marshal(knowledgeData{Content: content, Keyword: keyword})
	evalJS(fmt.Sprintf(`if(window.updateKnowledgeFromGo) window.updateKnowledgeFromGo(%s)`, string(data)))
}

func fetchAndPushKnowledge() {
	cfg := loadConfig()
	keyword := cfg.KnowledgeKeyword
	if keyword == "" {
		return
	}

	var content string

	count := getKnowledgeCardCount(keyword)
	if count > 0 && rand.Intn(10) < 3 {
		content = getRandomKnowledgeCard(keyword)
	}

	if content == "" {
		content = fetchKnowledgeFromLLM(keyword, cfg)
		if content != "" {
			saveKnowledgeCard(keyword, content)
		}
	}

	if content == "" && count > 0 {
		content = getRandomKnowledgeCard(keyword)
	}

	if content == "" {
		return
	}

	pushKnowledgeJSON(content, keyword)
	preview := content
	if len(preview) > 30 {
		preview = preview[:30] + "..."
	}
	log.Println("知识卡片已推送:", preview)
}

func pushKnowledgeLoading(keyword string) {
	pushKnowledgeJSON("加载中...", keyword)
}

func pushKnowledgePlaceholder() {
	pushKnowledgeJSON("请设置知识关键字", "")
}

func knowledgeLoop() {
	initKnowledgeDB()

	cfg := loadConfig()
	if cfg.KnowledgeKeyword != "" && !cfg.HideKnowledge {
		if cached := getRandomKnowledgeCard(cfg.KnowledgeKeyword); cached != "" {
			pushKnowledgeJSON(cached, cfg.KnowledgeKeyword)
		} else {
			pushKnowledgeLoading(cfg.KnowledgeKeyword)
		}
	} else if cfg.KnowledgeKeyword == "" {
		pushKnowledgePlaceholder()
	}

	time.Sleep(3 * time.Second)

	cfg = loadConfig()
	if cfg.KnowledgeKeyword != "" && !cfg.HideKnowledge {
		fetchAndPushKnowledge()
	}

	ticker := time.NewTicker(30 * time.Minute)
	for range ticker.C {
		cfg := loadConfig()
		if cfg.KnowledgeKeyword != "" && !cfg.HideKnowledge {
			fetchAndPushKnowledge()
		}
	}
}

func triggerKnowledgeRefresh() {
	go fetchAndPushKnowledge()
}