135 lines
3.6 KiB
Go
135 lines
3.6 KiB
Go
package kb
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"time"
|
|
|
|
"chat-agent-system/internal/config"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
// Client 知识库中枢客户端，实现 Hub 接口
|
|
type Client struct {
|
|
cfg *config.Config
|
|
pool *pgxpool.Pool
|
|
store *Store
|
|
embedder *Embedder
|
|
reranker *Reranker
|
|
retriever *Retriever
|
|
}
|
|
|
|
// 确保 Client 实现了 Hub 接口
|
|
var _ Hub = (*Client)(nil)
|
|
|
|
// NewClient 创建知识库中枢客户端实例
|
|
func NewClient(cfg *config.Config) *Client {
|
|
c := &Client{
|
|
cfg: cfg,
|
|
}
|
|
|
|
if cfg == nil {
|
|
return c
|
|
}
|
|
|
|
// 1. 初始化 Embedder
|
|
if cfg.EmbeddingEndpointURL != "" {
|
|
c.embedder = NewEmbedder(cfg.EmbeddingEndpointURL, cfg.EmbeddingAPIKey, cfg.EmbeddingModel)
|
|
}
|
|
|
|
// 2. 初始化 Reranker
|
|
if cfg.RerankerEndpointURL != "" {
|
|
c.reranker = NewReranker(cfg.RerankerEndpointURL, cfg.RerankerAPIKey, cfg.RerankerModel)
|
|
}
|
|
|
|
// 3. 尝试建立 PostgreSQL 连接池
|
|
dsn := cfg.KbDSN()
|
|
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
|
defer cancel()
|
|
|
|
pool, err := OpenPool(ctx, dsn)
|
|
if err != nil {
|
|
log.Printf("[kb] 无法连接到 PostgreSQL 知识库 (%s): %v, 将降级运行在离线模式", dsn, err)
|
|
return c
|
|
}
|
|
|
|
c.pool = pool
|
|
c.store = NewStore(pool)
|
|
c.retriever = NewRetriever(c.store, c.embedder, c.reranker)
|
|
return c
|
|
}
|
|
|
|
// Search 通用跨领域检索
|
|
func (c *Client) Search(ctx context.Context, query string, opts SearchOptions) ([]HitItem, error) {
|
|
if c.retriever == nil {
|
|
return []HitItem{}, nil
|
|
}
|
|
return c.retriever.Search(ctx, query, opts)
|
|
}
|
|
|
|
// SearchSimple 兼容简易调用的检索方法
|
|
func (c *Client) SearchSimple(ctx context.Context, query string, limit int) ([]HitItem, error) {
|
|
opts := DefaultSearchOptions()
|
|
if limit > 0 {
|
|
opts.Limit = limit
|
|
}
|
|
return c.Search(ctx, query, opts)
|
|
}
|
|
|
|
// LookupCite 根据标准号/条文引用粗查知识库
|
|
func (c *Client) LookupCite(ctx context.Context, citeKey string) (string, bool) {
|
|
if c.retriever == nil {
|
|
return "", false
|
|
}
|
|
return c.retriever.LookupCite(ctx, citeKey)
|
|
}
|
|
|
|
// Stats 查询各知识领域的统计数据
|
|
func (c *Client) Stats(ctx context.Context) (map[KnowledgeDomain]Stats, error) {
|
|
if c.store == nil {
|
|
return make(map[KnowledgeDomain]Stats), nil
|
|
}
|
|
return c.store.Stats(ctx)
|
|
}
|
|
|
|
// Close 释放连接池资源
|
|
func (c *Client) Close() {
|
|
if c.pool != nil {
|
|
c.pool.Close()
|
|
}
|
|
}
|
|
|
|
// MaterializeSnapshot 将检索命中的法规写入任务工作空间的 knowledge/ 目录
|
|
func MaterializeSnapshot(dir string, hits []HitItem) error {
|
|
kbDir := filepath.Join(dir, "knowledge")
|
|
if err := os.MkdirAll(kbDir, 0o755); err != nil {
|
|
return err
|
|
}
|
|
|
|
var sb strings.Builder
|
|
sb.WriteString("# 权威法规与技术标准引用库 (Regulatory Knowledge Hits)\n\n")
|
|
sb.WriteString("以下为针对本项目研发与注册需求命中的权威法规、国家标准(GB)与行业标准(YY)条款。\n")
|
|
sb.WriteString("在输出交付物时，请结合上下文使用，并显式标注【引用标记 (CiteKey)】标明出处。\n\n")
|
|
|
|
for i, h := range hits {
|
|
sb.WriteString(fmt.Sprintf("### %d. [%s] %s\n", i+1, h.DocCode, h.Title))
|
|
if h.Category != "" {
|
|
sb.WriteString(fmt.Sprintf("- **类别**: %s\n", h.Category))
|
|
}
|
|
if h.HeadingPath != "" {
|
|
sb.WriteString(fmt.Sprintf("- **章节路径**: %s\n", h.HeadingPath))
|
|
}
|
|
if h.CiteKey != "" {
|
|
sb.WriteString(fmt.Sprintf("- **引用标记 (CiteKey)**: `%s`\n", h.CiteKey))
|
|
}
|
|
sb.WriteString(fmt.Sprintf("- **来源**: %s\n\n", h.Source))
|
|
sb.WriteString(h.Content + "\n\n---\n\n")
|
|
}
|
|
|
|
return os.WriteFile(filepath.Join(kbDir, "hits.md"), []byte(sb.String()), 0o644)
|
|
}
|