151 lines
4.0 KiB
Go
151 lines
4.0 KiB
Go
package tool
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"log"
|
||
"net/http"
|
||
"time"
|
||
|
||
"tcm-agent/internal/config"
|
||
)
|
||
|
||
// ========================================================================
|
||
// MaxKB 知识库客户端
|
||
// ========================================================================
|
||
// MaxKB 是一个开源的企业级 RAG 知识库平台：
|
||
// - 文档管理（PDF/Word/Markdown 等）
|
||
// - 智能分段 + 向量化（pgvector）
|
||
// - 语义检索 + LLM 增强生成
|
||
// - 兼容 OpenAI 协议
|
||
//
|
||
// 官方文档：https://maxkb.cn/docs/
|
||
// Docker 一键部署：docker run -d --name maxkb -p 8080:8080 1panel/maxkb
|
||
// ========================================================================
|
||
|
||
// MaxKBClient MaxKB 知识库客户端
|
||
type MaxKBClient struct {
|
||
BaseURL string `json:"base_url"`
|
||
APIKey string `json:"api_key"`
|
||
AppID string `json:"app_id"`
|
||
Client *http.Client `json:"-"`
|
||
}
|
||
|
||
// NewMaxKBClient 创建 MaxKB 客户端
|
||
//
|
||
// 参数：
|
||
// cfg - MaxKB 配置（BaseURL/APIKey/AppID）
|
||
//
|
||
// 返回：
|
||
// 初始化完成的客户端
|
||
func NewMaxKBClient(cfg config.MaxKBConfig) *MaxKBClient {
|
||
return &MaxKBClient{
|
||
BaseURL: cfg.BaseURL,
|
||
APIKey: cfg.APIKey,
|
||
AppID: cfg.AppID,
|
||
Client: &http.Client{Timeout: 60 * time.Second},
|
||
}
|
||
}
|
||
|
||
// Chat 调用 MaxKB 的对话接口（兼容 OpenAI 格式）
|
||
//
|
||
// 这是最核心的方法：把问题发给 MaxKB，
|
||
// 它会自动做 RAG 检索 + LLM 生成，返回融合知识库的回答。
|
||
//
|
||
// 参数：
|
||
// ctx - 上下文（支持超时取消）
|
||
// query - 用户问题或检索关键词
|
||
//
|
||
// 返回：
|
||
// MaxKB 生成的回答（已融合知识库检索结果）
|
||
func (c *MaxKBClient) Chat(ctx context.Context, query string) (string, error) {
|
||
body := map[string]any{
|
||
"model": "maxkb-model",
|
||
"messages": []map[string]string{
|
||
{"role": "user", "content": query},
|
||
},
|
||
"stream": false,
|
||
}
|
||
|
||
buf, _ := json.Marshal(body)
|
||
url := fmt.Sprintf("%s/api/application/%s/chat/completions", c.BaseURL, c.AppID)
|
||
|
||
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(buf))
|
||
if err != nil {
|
||
return "", fmt.Errorf("[MaxKB] 创建请求失败: %w", err)
|
||
}
|
||
req.Header.Set("Authorization", "Bearer "+c.APIKey)
|
||
req.Header.Set("Content-Type", "application/json")
|
||
|
||
resp, err := c.Client.Do(req)
|
||
if err != nil {
|
||
return "", fmt.Errorf("[MaxKB] 调用失败: %w", err)
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
if resp.StatusCode != 200 {
|
||
data, _ := io.ReadAll(resp.Body)
|
||
return "", fmt.Errorf("[MaxKB] 返回错误 %d: %s", resp.StatusCode, string(data))
|
||
}
|
||
|
||
var result struct {
|
||
Choices []struct {
|
||
Message struct {
|
||
Content string `json:"content"`
|
||
} `json:"message"`
|
||
} `json:"choices"`
|
||
}
|
||
json.NewDecoder(resp.Body).Decode(&result)
|
||
|
||
if len(result.Choices) == 0 {
|
||
return "", fmt.Errorf("[MaxKB] 返回空结果")
|
||
}
|
||
|
||
log.Printf("[MaxKB] ✅ 检索成功 | 查询: %.50s...", query)
|
||
return result.Choices[0].Message.Content, nil
|
||
}
|
||
|
||
// Search 仅做知识检索（不生成，返回原始片段）
|
||
//
|
||
// 适合需要"引用来源"的场景。
|
||
// 返回 TopK 条相关文档片段。
|
||
func (c *MaxKBClient) Search(ctx context.Context, query string, topK int) ([]string, error) {
|
||
if topK <= 0 {
|
||
topK = 5
|
||
}
|
||
|
||
body := map[string]any{
|
||
"query": query,
|
||
"top_k": topK,
|
||
}
|
||
buf, _ := json.Marshal(body)
|
||
url := fmt.Sprintf("%s/api/application/%s/search", c.BaseURL, c.AppID)
|
||
|
||
req, _ := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(buf))
|
||
req.Header.Set("Authorization", "Bearer "+c.APIKey)
|
||
req.Header.Set("Content-Type", "application/json")
|
||
|
||
resp, err := c.Client.Do(req)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("[MaxKB] 检索失败: %w", err)
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
var result struct {
|
||
Documents []struct {
|
||
Content string `json:"content"`
|
||
Score float64 `json:"score"`
|
||
} `json:"documents"`
|
||
}
|
||
json.NewDecoder(resp.Body).Decode(&result)
|
||
|
||
docs := make([]string, 0, len(result.Documents))
|
||
for _, d := range result.Documents {
|
||
docs = append(docs, d.Content)
|
||
}
|
||
return docs, nil
|
||
}
|