311 lines
7.8 KiB
Go
311 lines
7.8 KiB
Go
package main
|
||
|
||
import (
|
||
"bytes"
|
||
"database/sql"
|
||
"encoding/json"
|
||
"fmt"
|
||
"log"
|
||
"math/rand"
|
||
"net/http"
|
||
"path/filepath"
|
||
"strings"
|
||
"time"
|
||
"unicode/utf8"
|
||
|
||
_ "modernc.org/sqlite"
|
||
)
|
||
|
||
const cpaURL = "https://cpa.1216.top/v1/chat/completions"
|
||
const cpaModel = "glm-4.5-air"
|
||
const minKnowledgeRunes = 80
|
||
|
||
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 ""
|
||
}
|
||
rows, err := knowledgeDB.Query(
|
||
"SELECT content FROM knowledge_cards WHERE keyword = ? ORDER BY RANDOM() LIMIT 20",
|
||
keyword,
|
||
)
|
||
if err != nil {
|
||
return ""
|
||
}
|
||
defer rows.Close()
|
||
for rows.Next() {
|
||
var content string
|
||
if rows.Scan(&content) == nil {
|
||
content = normalizeKnowledgeContent(content)
|
||
if isQualityKnowledgeCard(content) {
|
||
return content
|
||
}
|
||
}
|
||
}
|
||
return ""
|
||
}
|
||
|
||
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 := buildKnowledgePrompt(keyword, cfg.KnowledgePrompt)
|
||
messages := []map[string]string{
|
||
{
|
||
"role": "system",
|
||
"content": "你是严谨的中文知识卡片作者,输出必须具体、准确、有信息密度。不要写空泛鸡汤,不要只给一句定义。",
|
||
},
|
||
{"role": "user", "content": basePrompt},
|
||
}
|
||
|
||
body := map[string]interface{}{
|
||
"model": cpaModel,
|
||
"max_tokens": 512,
|
||
"temperature": 0.55,
|
||
"messages": messages,
|
||
}
|
||
content := requestKnowledgeCompletion(body)
|
||
content = normalizeKnowledgeContent(content)
|
||
if isQualityKnowledgeCard(content) {
|
||
return content
|
||
}
|
||
|
||
log.Printf("知识卡片质量不足,重试: %q", content)
|
||
messages = append(messages, map[string]string{
|
||
"role": "user",
|
||
"content": fmt.Sprintf(
|
||
"上一条太短或信息密度不足。请重写一条「%s」知识卡片:120-180个中文字符,必须包含一个明确机制/原理、一个具体例子或应用场景、一个结论。只输出正文。",
|
||
keyword,
|
||
),
|
||
})
|
||
body["messages"] = messages
|
||
content = requestKnowledgeCompletion(body)
|
||
content = normalizeKnowledgeContent(content)
|
||
if isQualityKnowledgeCard(content) {
|
||
return content
|
||
}
|
||
|
||
return ""
|
||
}
|
||
|
||
func buildKnowledgePrompt(keyword, customPrompt string) string {
|
||
basePrompt := fmt.Sprintf(`围绕关键词「%s」生成一条桌面知识卡片。
|
||
硬性要求:
|
||
1. 120-180个中文字符,分成2-3句;
|
||
2. 必须讲清一个具体机制、原理、权衡或实践经验;
|
||
3. 必须包含一个具体例子、场景或判断标准;
|
||
4. 避免“很重要、非常有用、提升效率”这类空泛表述;
|
||
5. 不要标题、序号、Markdown、表情,直接输出正文。`, keyword)
|
||
if customPrompt != "" {
|
||
basePrompt += "\n附加风格要求(不能覆盖上面的字数和质量要求):" + customPrompt
|
||
}
|
||
return basePrompt
|
||
}
|
||
|
||
func requestKnowledgeCompletion(body map[string]interface{}) string {
|
||
jsonData, _ := json.Marshal(body)
|
||
|
||
req, err := http.NewRequest("POST", cpaURL, bytes.NewReader(jsonData))
|
||
if err != nil {
|
||
log.Println("知识API请求创建失败:", err)
|
||
return ""
|
||
}
|
||
key := loadConfig().cpaKey()
|
||
if key == "" {
|
||
log.Println("未配置知识卡片 API Key")
|
||
return ""
|
||
}
|
||
req.Header.Set("Authorization", "Bearer "+key)
|
||
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 ""
|
||
}
|
||
|
||
// normalizeKnowledgeContent 清洗 LLM 输出的格式残留(Markdown标记、标题前缀、多余空白)
|
||
func normalizeKnowledgeContent(content string) string {
|
||
content = strings.TrimSpace(content)
|
||
content = strings.Trim(content, "` \t\r\n")
|
||
content = strings.TrimPrefix(content, "知识卡片:")
|
||
content = strings.TrimPrefix(content, "知识卡片:")
|
||
content = strings.TrimPrefix(content, "正文:")
|
||
content = strings.TrimPrefix(content, "正文:")
|
||
content = strings.ReplaceAll(content, "\r\n", "\n")
|
||
lines := strings.FieldsFunc(content, func(r rune) bool {
|
||
return r == '\n' || r == '\r' || r == '\t'
|
||
})
|
||
content = strings.Join(lines, " ")
|
||
content = strings.Join(strings.Fields(content), " ")
|
||
return strings.TrimSpace(content)
|
||
}
|
||
|
||
// isQualityKnowledgeCard 检查字数 ≥80 且空泛表述 <3 条
|
||
func isQualityKnowledgeCard(content string) bool {
|
||
if utf8.RuneCountInString(content) < minKnowledgeRunes {
|
||
return false
|
||
}
|
||
weakPhrases := []string{"很重要", "非常重要", "很有用", "提升效率", "值得关注", "可以帮助", "广泛应用"}
|
||
weakHits := 0
|
||
for _, phrase := range weakPhrases {
|
||
if strings.Contains(content, phrase) {
|
||
weakHits++
|
||
}
|
||
}
|
||
if weakHits >= 3 {
|
||
return false
|
||
}
|
||
return strings.ContainsAny(content, "。;;::,,")
|
||
}
|
||
|
||
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()
|
||
}
|