package llm
|
|
import (
|
"bytes"
|
"encoding/json"
|
"fmt"
|
"io"
|
"net/http"
|
"os"
|
"path/filepath"
|
"regexp"
|
"strings"
|
"time"
|
|
"gopkg.in/yaml.v3"
|
)
|
|
// LLMConfig 单个 LLM 配置
|
type LLMConfig struct {
|
APIBase string `yaml:"api_base"`
|
APIKey string `yaml:"api_key"`
|
Model string `yaml:"model"`
|
Temperature float64 `yaml:"temperature"`
|
MaxTokens int `yaml:"max_tokens"`
|
DisableThinking bool `yaml:"disable_thinking"`
|
}
|
|
// Config 配置文件结构
|
type Config struct {
|
LLM struct {
|
Primary LLMConfig `yaml:"primary"`
|
Fallback *LLMConfig `yaml:"fallback"`
|
} `yaml:"llm"`
|
KnowledgeBase struct {
|
VaultPath string `yaml:"vault_path"`
|
DBPath string `yaml:"db_path"`
|
} `yaml:"knowledge_base"`
|
Draft struct {
|
ReviewDir string `yaml:"review_dir"`
|
DefaultType string `yaml:"default_type"`
|
} `yaml:"draft"`
|
}
|
|
// Client LLM 客户端
|
type Client struct {
|
config Config
|
usePrimary bool
|
}
|
|
// ExtractResult 提取结果
|
type ExtractResult struct {
|
Tags []string `json:"tags"`
|
RelatedDocs []string `json:"related_docs"`
|
}
|
|
// MergeHint 合并指示
|
type MergeHint struct {
|
Recommendation struct {
|
Action string `json:"action"`
|
Target string `json:"target"`
|
Reason string `json:"reason"`
|
Confidence string `json:"confidence"`
|
} `json:"recommendation"`
|
Analysis []struct {
|
Path string `json:"path"`
|
Relevance string `json:"relevance"`
|
Reason string `json:"reason"`
|
} `json:"analysis"`
|
}
|
|
// SearchCandidate 搜索候选
|
type SearchCandidate struct {
|
Title string `json:"title"`
|
Path string `json:"path"`
|
Score float64 `json:"score"`
|
}
|
|
// NewClient 创建 LLM 客户端
|
func NewClient() *Client {
|
config := loadConfig()
|
return &Client{config: config, usePrimary: true}
|
}
|
|
// loadConfig 加载配置文件
|
func loadConfig() Config {
|
var config Config
|
|
// 配置文件路径
|
homeDir, _ := os.UserHomeDir()
|
configPath := filepath.Join(homeDir, ".kb-cli", "config.yaml")
|
|
// 读取配置文件
|
data, err := os.ReadFile(configPath)
|
if err != nil {
|
// 如果配置文件不存在，使用默认值
|
fmt.Fprintf(os.Stderr, "警告: 无法读取配置文件 %s，使用默认值\n", configPath)
|
return getDefaultConfig()
|
}
|
|
// 解析 YAML
|
if err := yaml.Unmarshal(data, &config); err != nil {
|
fmt.Fprintf(os.Stderr, "警告: 解析配置文件失败: %v，使用默认值\n", err)
|
return getDefaultConfig()
|
}
|
|
// 展开 ~ 路径
|
config.LLM.Primary.APIBase = expandPath(config.LLM.Primary.APIBase)
|
config.KnowledgeBase.VaultPath = expandPath(config.KnowledgeBase.VaultPath)
|
config.KnowledgeBase.DBPath = expandPath(config.KnowledgeBase.DBPath)
|
if config.LLM.Fallback != nil {
|
config.LLM.Fallback.APIBase = expandPath(config.LLM.Fallback.APIBase)
|
}
|
|
return config
|
}
|
|
// getDefaultConfig 获取默认配置
|
func getDefaultConfig() Config {
|
var config Config
|
config.LLM.Primary.APIBase = "http://192.168.3.246:1127/v1"
|
config.LLM.Primary.APIKey = "sk-local"
|
config.LLM.Primary.Model = "qwen3.6-35b-a3b"
|
config.LLM.Primary.Temperature = 0.3
|
config.LLM.Primary.MaxTokens = 2000
|
config.LLM.Primary.DisableThinking = true
|
config.KnowledgeBase.VaultPath = "~/aisim/note/001/笔记001"
|
config.KnowledgeBase.DBPath = "~/.cache/kb-cli/kb.db"
|
config.Draft.ReviewDir = "待审阅"
|
config.Draft.DefaultType = "售后"
|
return config
|
}
|
|
// expandPath 展开路径中的 ~
|
func expandPath(path string) string {
|
if strings.HasPrefix(path, "~/") {
|
homeDir, _ := os.UserHomeDir()
|
return filepath.Join(homeDir, path[2:])
|
}
|
return path
|
}
|
|
// getCurrentConfig 获取当前使用的 LLM 配置
|
func (c *Client) getCurrentConfig() LLMConfig {
|
if c.usePrimary {
|
return c.config.LLM.Primary
|
}
|
if c.config.LLM.Fallback != nil {
|
return *c.config.LLM.Fallback
|
}
|
return c.config.LLM.Primary
|
}
|
|
// switchToNext 切换到下一个可用的 LLM
|
func (c *Client) switchToNext() bool {
|
if c.usePrimary && c.config.LLM.Fallback != nil {
|
fmt.Fprintf(os.Stderr, "主 LLM 不可用，切换到备用 LLM\n")
|
c.usePrimary = false
|
return true
|
}
|
return false
|
}
|
|
// ExtractTags 提取 tags 和相关文档
|
func (c *Client) ExtractTags(content string) (*ExtractResult, error) {
|
prompt := fmt.Sprintf(`分析以下知识库草稿，提取：
|
1. tags（5-10个）：
|
- 核心问题标签（如"充不进气"、"档案下载失败"）
|
- 扩展词（不同人可能的描述，如"充气慢"、"进气不足"）
|
- 平台/设备标签（如"电子秤平台"、"智能枪"）
|
2. related_docs（0-5个）：相关文档标题（用于创建链接）
|
|
输出 JSON：
|
{
|
"tags": ["充不进气", "充气慢", "进气不足", "智能枪", "电子秤平台"],
|
"related_docs": ["智能枪通气杆卡住漏气", "角阀充装功率不足"]
|
}
|
|
草稿内容：
|
%s`, content)
|
|
response, err := c.callLLMWithRetry(prompt)
|
if err != nil {
|
return nil, err
|
}
|
|
// 解析 JSON
|
var result ExtractResult
|
if err := parseJSON(response, &result); err != nil {
|
return nil, fmt.Errorf("解析 LLM 响应失败: %w", err)
|
}
|
|
return &result, nil
|
}
|
|
// GenerateMergeHint 生成合并指示
|
func (c *Client) GenerateMergeHint(draftContent string, candidates []SearchCandidate) (*MergeHint, error) {
|
prompt := fmt.Sprintf(`你是一个知识库管理助手。请判断以下草稿与哪些知识库文档相关。
|
|
## 草稿内容
|
%s
|
|
## 候选文档（Top 3）
|
`, draftContent)
|
|
for i, cand := range candidates {
|
prompt += fmt.Sprintf("%d. 标题: %s, 路径: %s, 相关度: %.2f\n", i+1, cand.Title, cand.Path, cand.Score)
|
}
|
|
prompt += `
|
## 任务
|
请分析草稿与每个候选文档的相关性，输出 JSON 格式：
|
|
{
|
"recommendation": {
|
"action": "merge|new|split",
|
"target": "目标路径（如果 action=merge）",
|
"reason": "判断理由",
|
"confidence": "high|medium|low"
|
},
|
"analysis": [
|
{
|
"path": "文档路径",
|
"relevance": "high|medium|low",
|
"reason": "相关性说明"
|
}
|
]
|
}
|
|
判断标准：
|
- action=merge: 草稿内容与某个候选文档高度相关，应该合并
|
- action=new: 草稿内容是全新的，应该新建文档
|
- action=split: 草稿内容包含多个独立主题，应该拆分
|
- confidence=high: 判断很确定
|
- confidence=medium: 判断比较确定
|
- confidence=low: 判断不太确定
|
|
只输出 JSON，不要其他内容。`
|
|
response, err := c.callLLMWithRetry(prompt)
|
if err != nil {
|
return nil, err
|
}
|
|
var result MergeHint
|
if err := parseJSON(response, &result); err != nil {
|
return nil, fmt.Errorf("解析 LLM 响应失败: %w", err)
|
}
|
|
return &result, nil
|
}
|
|
// ClassifyDocuments 分类文档
|
func (c *Client) ClassifyDocuments(prompt string) (string, error) {
|
response, err := c.callLLMWithRetry(prompt)
|
if err != nil {
|
return "", err
|
}
|
|
// 解析 JSON
|
var result []struct {
|
Path string `json:"path"`
|
Platform string `json:"platform"`
|
Device string `json:"device"`
|
ContentType string `json:"content_type"`
|
Confidence float64 `json:"confidence"`
|
}
|
|
if err := parseJSON(response, &result); err != nil {
|
return "", fmt.Errorf("解析 LLM 响应失败: %w", err)
|
}
|
|
// 转换回 JSON 字符串
|
jsonBytes, err := json.Marshal(result)
|
if err != nil {
|
return "", err
|
}
|
|
return string(jsonBytes), nil
|
}
|
|
// ClassifyDocuments 分类文档
|
|
// callLLMWithRetry 带重试的 LLM 调用
|
func (c *Client) callLLMWithRetry(prompt string) (string, error) {
|
// 第一次尝试
|
response, err := c.callLLM(prompt)
|
if err == nil {
|
return response, nil
|
}
|
|
// 判断是否是需要切换的错误
|
if shouldSwitchLLM(err) {
|
// 尝试切换到备用 LLM
|
if c.switchToNext() {
|
fmt.Fprintf(os.Stderr, "重试备用 LLM...\n")
|
response, err = c.callLLM(prompt)
|
if err == nil {
|
return response, nil
|
}
|
}
|
}
|
|
return "", err
|
}
|
|
// shouldSwitchLLM 判断是否应该切换到备用 LLM
|
func shouldSwitchLLM(err error) bool {
|
errMsg := err.Error()
|
// 网络错误、超时、认证失败等应该切换
|
return strings.Contains(errMsg, "connection") ||
|
strings.Contains(errMsg, "timeout") ||
|
strings.Contains(errMsg, "401") ||
|
strings.Contains(errMsg, "403") ||
|
strings.Contains(errMsg, "500") ||
|
strings.Contains(errMsg, "502") ||
|
strings.Contains(errMsg, "503") ||
|
strings.Contains(errMsg, "504")
|
}
|
|
// callLLM 调用 LLM API
|
func (c *Client) callLLM(prompt string) (string, error) {
|
llmConfig := c.getCurrentConfig()
|
url := llmConfig.APIBase + "/chat/completions"
|
|
payload := map[string]interface{}{
|
"model": llmConfig.Model,
|
"messages": []map[string]string{
|
{"role": "user", "content": prompt},
|
},
|
"temperature": llmConfig.Temperature,
|
"max_tokens": llmConfig.MaxTokens,
|
}
|
|
// 如果需要禁用思考模式
|
if llmConfig.DisableThinking {
|
payload["chat_template_kwargs"] = map[string]bool{
|
"enable_thinking": false,
|
}
|
}
|
|
jsonData, err := json.Marshal(payload)
|
if err != nil {
|
return "", fmt.Errorf("序列化请求失败: %w", err)
|
}
|
|
req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData))
|
if err != nil {
|
return "", fmt.Errorf("创建请求失败: %w", err)
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
if llmConfig.APIKey != "" {
|
req.Header.Set("Authorization", "Bearer "+llmConfig.APIKey)
|
}
|
|
client := &http.Client{Timeout: 60 * time.Second}
|
resp, err := client.Do(req)
|
if err != nil {
|
return "", fmt.Errorf("请求失败: %w", err)
|
}
|
defer resp.Body.Close()
|
|
body, err := io.ReadAll(resp.Body)
|
if err != nil {
|
return "", fmt.Errorf("读取响应失败: %w", err)
|
}
|
|
if resp.StatusCode != http.StatusOK {
|
return "", fmt.Errorf("API 返回错误状态: %d, body: %s", resp.StatusCode, string(body))
|
}
|
|
// 解析响应
|
var result struct {
|
Choices []struct {
|
Message struct {
|
Content string `json:"content"`
|
} `json:"message"`
|
} `json:"choices"`
|
}
|
|
if err := json.Unmarshal(body, &result); err != nil {
|
return "", fmt.Errorf("解析响应失败: %w", err)
|
}
|
|
if len(result.Choices) == 0 {
|
return "", fmt.Errorf("API 返回空结果")
|
}
|
|
return result.Choices[0].Message.Content, nil
|
}
|
|
// parseJSON 解析 JSON（支持 markdown 代码块）
|
func parseJSON(content string, v interface{}) error {
|
// 尝试提取 ```json 代码块
|
re := regexp.MustCompile("(?s)```json\\s*(.*?)\\s*```")
|
matches := re.FindStringSubmatch(content)
|
if len(matches) > 1 {
|
content = matches[1]
|
}
|
|
if err := json.Unmarshal([]byte(content), v); err != nil {
|
return err
|
}
|
|
return nil
|
}
|
