Files
u-desktop/knowledge.go

311 lines
7.8 KiB
Go
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
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()
}