feat(service): 实现嵌入式多领域知识中枢与Hybrid RRF检索引擎
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
WeiCHEN
committed by
weichen
weichen
co-authored by
Cursor
Cursor
parent
7d123ae8c6
commit
7cbfb33150
@@ -0,0 +1,60 @@
|
||||
package kb
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// LookupCite 根据标准号/条文引用粗查法规库；命中则返回可写入参考文献的完整说明条目
|
||||
// 若未命中或无法连接返回 ok=false
|
||||
func (r *Retriever) LookupCite(ctx context.Context, citeKey string) (string, bool) {
|
||||
if r == nil || r.Store == nil || r.Store.Pool == nil {
|
||||
return "", false
|
||||
}
|
||||
|
||||
std := strings.TrimSpace(citeKey)
|
||||
std = strings.TrimPrefix(std, "〔")
|
||||
std = strings.TrimSuffix(std, "〕")
|
||||
std = strings.TrimPrefix(std, "[")
|
||||
std = strings.TrimSuffix(std, "]")
|
||||
|
||||
clause := ""
|
||||
if idx := strings.IndexAny(std, " §"); idx >= 0 {
|
||||
clause = strings.TrimSpace(std[idx:])
|
||||
std = strings.TrimSpace(std[:idx])
|
||||
}
|
||||
if std == "" {
|
||||
return "", false
|
||||
}
|
||||
|
||||
cctx, cancel := context.WithTimeout(ctx, 3*time.Second)
|
||||
defer cancel()
|
||||
|
||||
// 1. 尝试精确查找文档
|
||||
var title, docNo, status string
|
||||
err := r.Store.Pool.QueryRow(cctx, `
|
||||
SELECT COALESCE(doc_no,''), COALESCE(title,''), COALESCE(status,'')
|
||||
FROM reg_documents
|
||||
WHERE doc_no ILIKE $1 OR title ILIKE $1
|
||||
ORDER BY id ASC LIMIT 1`, "%"+std+"%").Scan(&docNo, &title, &status)
|
||||
|
||||
if err == nil && docNo != "" {
|
||||
line := fmt.Sprintf("〔%s%s〕%s", docNo, clause, title)
|
||||
if status == "repealed" || status == "superseded" {
|
||||
line = "〔已废止〕" + line
|
||||
}
|
||||
return line, true
|
||||
}
|
||||
|
||||
// 2. 备选：按 Search 检索 Top 1
|
||||
hits, err := r.Search(cctx, std, SearchOptions{Limit: 1, IncludeRepealed: true})
|
||||
if err == nil && len(hits) > 0 {
|
||||
h := hits[0]
|
||||
line := fmt.Sprintf("%s %s", h.CiteKey, h.Title)
|
||||
return line, true
|
||||
}
|
||||
|
||||
return "", false
|
||||
}
|
||||
+101
-11
@@ -3,24 +3,104 @@ 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
|
||||
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 {
|
||||
return &Client{cfg: cfg}
|
||||
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
|
||||
}
|
||||
|
||||
func (c *Client) Search(ctx context.Context, query string, limit int) ([]HitItem, error) {
|
||||
// 基础查询：当未开启或处于离线测试时，返回标准结构
|
||||
return []HitItem{}, nil
|
||||
// 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/ 目录
|
||||
@@ -31,13 +111,23 @@ func MaterializeSnapshot(dir string, hits []HitItem) error {
|
||||
}
|
||||
|
||||
var sb strings.Builder
|
||||
sb.WriteString("# 法规与指导原则检索快照\n\n")
|
||||
sb.WriteString("以下为针对本项目命中的参考法规标准条文，请结合上下文使用，标注〔出处〕。\n\n")
|
||||
sb.WriteString("# 权威法规与技术标准引用库 (Regulatory Knowledge Hits)\n\n")
|
||||
sb.WriteString("以下为针对本项目研发与注册需求命中的权威法规、国家标准(GB)与行业标准(YY)条款。\n")
|
||||
sb.WriteString("在输出交付物时，请结合上下文使用，并显式标注【引用标记 (CiteKey)】标明出处。\n\n")
|
||||
|
||||
for _, h := range hits {
|
||||
sb.WriteString(fmt.Sprintf("### [%s] %s %s\n", strings.ToUpper(h.Category), h.DocCode, h.Title))
|
||||
sb.WriteString(h.Content + "\n\n")
|
||||
sb.WriteString(fmt.Sprintf("> 来源: %s\n\n---\n\n", h.Source))
|
||||
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)
|
||||
|
||||
@@ -0,0 +1,311 @@
|
||||
package kb
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
const (
|
||||
// EmbeddingDim BGE-M3 标准向量维度
|
||||
EmbeddingDim = 1024
|
||||
|
||||
// embedRequestTokenBudget 单批次 Token 上限
|
||||
embedRequestTokenBudget = 7000
|
||||
// embedMaxBatchItems 单请求文本条数上限
|
||||
embedMaxBatchItems = 16
|
||||
// embedMaxItemTokens 单条文本截断上限
|
||||
embedMaxItemTokens = 6000
|
||||
// embedLogBodyMax 错误截断最大字节
|
||||
embedLogBodyMax = 2048
|
||||
)
|
||||
|
||||
// Embedder 向量生成客户端
|
||||
type Embedder struct {
|
||||
Endpoint string
|
||||
APIKey string
|
||||
Model string
|
||||
HTTP *http.Client
|
||||
}
|
||||
|
||||
// NewEmbedder 创建 Embedder 实例
|
||||
func NewEmbedder(endpoint, key, model string) *Embedder {
|
||||
return &Embedder{
|
||||
Endpoint: normalizeAPIEndpoint(endpoint, "/embeddings"),
|
||||
APIKey: key,
|
||||
Model: model,
|
||||
HTTP: &http.Client{Timeout: 60 * time.Second},
|
||||
}
|
||||
}
|
||||
|
||||
// Embed 批量生成文本嵌入向量
|
||||
func (e *Embedder) Embed(ctx context.Context, texts []string) ([][]float32, error) {
|
||||
if e == nil || e.Endpoint == "" {
|
||||
return nil, fmt.Errorf("embedder not configured")
|
||||
}
|
||||
if len(texts) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
out := make([][]float32, len(texts))
|
||||
for start := 0; start < len(texts); {
|
||||
end := nextEmbedBatchEnd(texts, start, embedRequestTokenBudget, embedMaxBatchItems)
|
||||
batch := texts[start:end]
|
||||
vecs, err := e.embedBatch(ctx, batch)
|
||||
if err != nil {
|
||||
if isContextLengthError(err) && len(batch) > 1 {
|
||||
mid := start + (end-start+1)/2
|
||||
if mid <= start {
|
||||
mid = start + 1
|
||||
}
|
||||
if mid >= end {
|
||||
mid = end - 1
|
||||
}
|
||||
left, lerr := e.Embed(ctx, texts[start:mid])
|
||||
if lerr != nil {
|
||||
return nil, lerr
|
||||
}
|
||||
right, rerr := e.Embed(ctx, texts[mid:end])
|
||||
if rerr != nil {
|
||||
return nil, rerr
|
||||
}
|
||||
copy(out[start:mid], left)
|
||||
copy(out[mid:end], right)
|
||||
start = end
|
||||
continue
|
||||
}
|
||||
est := estimateTokensBatch(batch)
|
||||
return nil, fmt.Errorf("%w (batch=%d est_tokens≈%d)", err, len(batch), est)
|
||||
}
|
||||
copy(out[start:end], vecs)
|
||||
start = end
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// EmbedOne 单条文本向量化辅助方法
|
||||
func (e *Embedder) EmbedOne(ctx context.Context, text string) ([]float32, error) {
|
||||
vecs, err := e.Embed(ctx, []string{text})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(vecs) == 0 {
|
||||
return nil, fmt.Errorf("empty embedding returned")
|
||||
}
|
||||
return vecs[0], nil
|
||||
}
|
||||
|
||||
func (e *Embedder) embedBatch(ctx context.Context, texts []string) ([][]float32, error) {
|
||||
body := map[string]any{"model": e.Model, "input": texts}
|
||||
raw, err := e.post(ctx, e.Endpoint, body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var resp struct {
|
||||
Data []struct {
|
||||
Embedding []float32 `json:"embedding"`
|
||||
Index int `json:"index"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &resp); err != nil {
|
||||
return nil, fmt.Errorf("decode embeddings: %s", formatHTTPBodyError(0, raw))
|
||||
}
|
||||
out := make([][]float32, len(texts))
|
||||
for _, d := range resp.Data {
|
||||
if d.Index >= 0 && d.Index < len(out) {
|
||||
out[d.Index] = padOrTrim(d.Embedding, EmbeddingDim)
|
||||
}
|
||||
}
|
||||
for i, v := range out {
|
||||
if v == nil && i < len(resp.Data) {
|
||||
out[i] = padOrTrim(resp.Data[i].Embedding, EmbeddingDim)
|
||||
}
|
||||
if out[i] == nil {
|
||||
return nil, fmt.Errorf("missing embedding %d", i)
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (e *Embedder) post(ctx context.Context, url string, body any) ([]byte, error) {
|
||||
return postJSON(ctx, e.HTTP, url, e.APIKey, body)
|
||||
}
|
||||
|
||||
func postJSON(ctx context.Context, httpc *http.Client, url, key string, body any) ([]byte, error) {
|
||||
b, _ := json.Marshal(body)
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(b))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if key != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
}
|
||||
resp, err := httpc.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
raw, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
|
||||
if resp.StatusCode >= 300 {
|
||||
return nil, fmt.Errorf("http %s: %s", resp.Status, formatHTTPBodyError(resp.StatusCode, raw))
|
||||
}
|
||||
return raw, nil
|
||||
}
|
||||
|
||||
func formatHTTPBodyError(status int, raw []byte) string {
|
||||
msg := extractAPIErrorMessage(raw)
|
||||
if msg != "" {
|
||||
return msg
|
||||
}
|
||||
trim := strings.TrimSpace(string(raw))
|
||||
if trim == "" {
|
||||
if status > 0 {
|
||||
return fmt.Sprintf("empty body (status %d)", status)
|
||||
}
|
||||
return "empty body"
|
||||
}
|
||||
if !json.Valid(raw) || strings.HasPrefix(trim, "<") {
|
||||
preview := truncateUTF8(trim, 200)
|
||||
if status > 0 {
|
||||
return fmt.Sprintf("non-JSON response (status %d): %s", status, preview)
|
||||
}
|
||||
return "non-JSON response: " + preview
|
||||
}
|
||||
return truncateUTF8(trim, embedLogBodyMax)
|
||||
}
|
||||
|
||||
func extractAPIErrorMessage(raw []byte) string {
|
||||
var envelope struct {
|
||||
Error struct {
|
||||
Message string `json:"message"`
|
||||
Type string `json:"type"`
|
||||
Code any `json:"code"`
|
||||
} `json:"error"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &envelope); err != nil {
|
||||
return ""
|
||||
}
|
||||
if m := strings.TrimSpace(envelope.Error.Message); m != "" {
|
||||
return m
|
||||
}
|
||||
return strings.TrimSpace(envelope.Message)
|
||||
}
|
||||
|
||||
func isContextLengthError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
msg := strings.ToLower(err.Error())
|
||||
return strings.Contains(msg, "maximum context length") ||
|
||||
strings.Contains(msg, "context length") ||
|
||||
strings.Contains(msg, "too many tokens") ||
|
||||
(strings.Contains(msg, "input tokens") && strings.Contains(msg, "reduce"))
|
||||
}
|
||||
|
||||
func estimateTokens(s string) int {
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "" {
|
||||
return 1
|
||||
}
|
||||
tokens := 0
|
||||
asciiRun := 0
|
||||
flushASCII := func() {
|
||||
if asciiRun == 0 {
|
||||
return
|
||||
}
|
||||
tokens += (asciiRun + 3) / 4
|
||||
asciiRun = 0
|
||||
}
|
||||
for _, r := range s {
|
||||
if r <= 127 {
|
||||
asciiRun++
|
||||
continue
|
||||
}
|
||||
flushASCII()
|
||||
tokens += 2
|
||||
}
|
||||
flushASCII()
|
||||
if tokens < 1 {
|
||||
return 1
|
||||
}
|
||||
return tokens
|
||||
}
|
||||
|
||||
func estimateTokensBatch(texts []string) int {
|
||||
sum := 0
|
||||
for _, t := range texts {
|
||||
sum += estimateTokens(t)
|
||||
}
|
||||
return sum
|
||||
}
|
||||
|
||||
func nextEmbedBatchEnd(texts []string, start, budget, maxItems int) int {
|
||||
if maxItems <= 0 {
|
||||
maxItems = embedMaxBatchItems
|
||||
}
|
||||
end := start
|
||||
used := 0
|
||||
for end < len(texts) && end-start < maxItems {
|
||||
t := estimateTokens(texts[end])
|
||||
if t > embedMaxItemTokens {
|
||||
t = embedMaxItemTokens
|
||||
}
|
||||
if end > start && used+t > budget {
|
||||
break
|
||||
}
|
||||
used += t
|
||||
end++
|
||||
if used >= budget {
|
||||
break
|
||||
}
|
||||
}
|
||||
if end == start {
|
||||
end = start + 1
|
||||
}
|
||||
return end
|
||||
}
|
||||
|
||||
func normalizeAPIEndpoint(endpoint, suffix string) string {
|
||||
endpoint = strings.TrimSpace(endpoint)
|
||||
if endpoint == "" {
|
||||
return ""
|
||||
}
|
||||
endpoint = strings.TrimRight(endpoint, "/")
|
||||
if strings.HasSuffix(endpoint, suffix) {
|
||||
return endpoint
|
||||
}
|
||||
if strings.HasSuffix(endpoint, "/v1") || strings.Contains(endpoint, "/v1/") {
|
||||
return endpoint + suffix
|
||||
}
|
||||
return endpoint + "/v1" + suffix
|
||||
}
|
||||
|
||||
func padOrTrim(v []float32, n int) []float32 {
|
||||
if len(v) == n {
|
||||
return v
|
||||
}
|
||||
out := make([]float32, n)
|
||||
copy(out, v)
|
||||
return out
|
||||
}
|
||||
|
||||
func truncateUTF8(s string, n int) string {
|
||||
if n <= 0 || len(s) <= n {
|
||||
return s
|
||||
}
|
||||
for n > 0 && !utf8.RuneStart(s[n]) {
|
||||
n--
|
||||
}
|
||||
if n <= 0 {
|
||||
return ""
|
||||
}
|
||||
return s[:n]
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
package kb_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"chat-agent-system/internal/kb"
|
||||
)
|
||||
|
||||
func TestEmbedder_MockServer(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1/embeddings" {
|
||||
t.Errorf("unexpected path: %s", r.URL.Path)
|
||||
}
|
||||
var req struct {
|
||||
Model string `json:"model"`
|
||||
Input []string `json:"input"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
type item struct {
|
||||
Embedding []float32 `json:"embedding"`
|
||||
Index int `json:"index"`
|
||||
}
|
||||
var data []item
|
||||
for i := range req.Input {
|
||||
vec := make([]float32, kb.EmbeddingDim)
|
||||
vec[0] = float32(i + 1)
|
||||
data = append(data, item{
|
||||
Embedding: vec,
|
||||
Index: i,
|
||||
})
|
||||
}
|
||||
resp := map[string]any{"data": data}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
embedder := kb.NewEmbedder(server.URL, "test-key", "bge-m3")
|
||||
vec, err := embedder.EmbedOne(context.Background(), "医用电气安全")
|
||||
if err != nil {
|
||||
t.Fatalf("EmbedOne failed: %v", err)
|
||||
}
|
||||
if len(vec) != kb.EmbeddingDim {
|
||||
t.Fatalf("expected dim %d, got %d", kb.EmbeddingDim, len(vec))
|
||||
}
|
||||
if vec[0] != 1.0 {
|
||||
t.Fatalf("expected vec[0]=1.0, got %f", vec[0])
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
package kb
|
||||
|
||||
import (
|
||||
"context"
|
||||
)
|
||||
|
||||
// Hub 多领域知识中枢统一服务契约
|
||||
// 支持医疗器械法规标准、历史项目数字资产、员工能力模型等统一或分域检索
|
||||
type Hub interface {
|
||||
// Search 通用跨领域检索（支持按 SearchOptions 过滤领域、分类与分数）
|
||||
Search(ctx context.Context, query string, opts SearchOptions) ([]HitItem, error)
|
||||
|
||||
// LookupCite 根据精准引用键 (如 "[GB 9706.1-2020 §8.4]") 定位条文
|
||||
LookupCite(ctx context.Context, citeKey string) (string, bool)
|
||||
|
||||
// Stats 获取各知识领域的数据规模与资产统计
|
||||
Stats(ctx context.Context) (map[KnowledgeDomain]Stats, error)
|
||||
|
||||
// Close 释放数据库连接池等资源
|
||||
Close()
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
package kb_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"chat-agent-system/internal/config"
|
||||
"chat-agent-system/internal/kb"
|
||||
)
|
||||
|
||||
func TestKB_RealIntegration(t *testing.T) {
|
||||
cfg, err := config.Load("")
|
||||
if err != nil {
|
||||
t.Fatalf("failed to load config: %v", err)
|
||||
}
|
||||
|
||||
// 强制使用本机开发环境共享知识库容器端口
|
||||
if cfg.KbDbPort == 0 {
|
||||
cfg.KbDbPort = 5432
|
||||
}
|
||||
|
||||
client := kb.NewClient(cfg)
|
||||
defer client.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
||||
defer cancel()
|
||||
|
||||
// 1. 验证多领域资产规模统计
|
||||
t.Run("Stats", func(t *testing.T) {
|
||||
stats, err := client.Stats(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("client.Stats failed: %v", err)
|
||||
}
|
||||
regStats, ok := stats[kb.DomainRegulatory]
|
||||
if !ok {
|
||||
t.Fatalf("expected regulatory domain in stats")
|
||||
}
|
||||
if regStats.DocumentCount != 2907 {
|
||||
t.Errorf("expected 2907 documents, got %d", regStats.DocumentCount)
|
||||
}
|
||||
if regStats.ChunkCount < 312900 {
|
||||
t.Errorf("expected >= 312900 chunks, got %d", regStats.ChunkCount)
|
||||
}
|
||||
t.Logf("Stats verified: %d docs, %d chunks", regStats.DocumentCount, regStats.ChunkCount)
|
||||
})
|
||||
|
||||
// 2. 验证关键法规标准精确与模糊检索 (GB 9706.1 医用电气设备安全通用要求)
|
||||
t.Run("Search_GB9706", func(t *testing.T) {
|
||||
hits, err := client.Search(ctx, "GB 9706.1 医用电气设备 基本安全", kb.SearchOptions{
|
||||
Limit: 3,
|
||||
CandidateLimit: 10,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("client.Search failed: %v", err)
|
||||
}
|
||||
if len(hits) == 0 {
|
||||
t.Fatalf("expected hits for GB 9706.1, got 0")
|
||||
}
|
||||
|
||||
found := false
|
||||
for _, h := range hits {
|
||||
t.Logf("Hit: [%s] %s (Cite: %s, Score: %f)", h.DocCode, h.Title, h.CiteKey, h.Score)
|
||||
if strings.Contains(h.DocCode, "9706") || strings.Contains(h.Title, "9706") || strings.Contains(h.Title, "电气安全") {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Errorf("expected at least one hit containing '9706' or '电气安全'")
|
||||
}
|
||||
})
|
||||
|
||||
// 3. 验证引用规范精准定位 (LookupCite)
|
||||
t.Run("LookupCite", func(t *testing.T) {
|
||||
line, ok := client.LookupCite(ctx, "GB 9706.1-2020 §8.4")
|
||||
if !ok {
|
||||
t.Fatalf("LookupCite failed to find GB 9706.1-2020 §8.4")
|
||||
}
|
||||
t.Logf("LookupCite result: %s", line)
|
||||
if !strings.Contains(line, "9706.1") {
|
||||
t.Errorf("expected line to contain 9706.1, got: %s", line)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -1,9 +0,0 @@
|
||||
package kb
|
||||
|
||||
type HitItem struct {
|
||||
Category string `json:"category"` // standards, guidance, laws, catalog
|
||||
DocCode string `json:"doc_code"`
|
||||
Title string `json:"title"`
|
||||
Content string `json:"content"`
|
||||
Source string `json:"source"`
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
package kb
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Reranker 语义重排客户端 (Cross-Encoder)
|
||||
type Reranker struct {
|
||||
Endpoint string
|
||||
APIKey string
|
||||
Model string
|
||||
HTTP *http.Client
|
||||
}
|
||||
|
||||
// NewReranker 创建 Reranker 实例
|
||||
func NewReranker(endpoint, key, model string) *Reranker {
|
||||
return &Reranker{
|
||||
Endpoint: normalizeAPIEndpoint(endpoint, "/rerank"),
|
||||
APIKey: key,
|
||||
Model: model,
|
||||
HTTP: &http.Client{Timeout: 30 * time.Second},
|
||||
}
|
||||
}
|
||||
|
||||
// Rerank 对候选文档列表依据查询进行相关性重排打分
|
||||
// 返回与 docs 等长的一组打分，下标一一对应
|
||||
func (r *Reranker) Rerank(ctx context.Context, query string, docs []string) ([]float64, error) {
|
||||
if r == nil || r.Endpoint == "" {
|
||||
return nil, fmt.Errorf("reranker not configured")
|
||||
}
|
||||
if len(docs) == 0 {
|
||||
return []float64{}, nil
|
||||
}
|
||||
|
||||
body := map[string]any{
|
||||
"model": r.Model,
|
||||
"query": query,
|
||||
"documents": docs,
|
||||
}
|
||||
|
||||
raw, err := postJSON(ctx, r.HTTP, r.Endpoint, r.APIKey, body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var resp struct {
|
||||
Results []struct {
|
||||
Index int `json:"index"`
|
||||
RelevanceScore float64 `json:"relevance_score"`
|
||||
Score float64 `json:"score"`
|
||||
} `json:"results"`
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(raw, &resp); err != nil {
|
||||
return nil, fmt.Errorf("decode rerank: %s", formatHTTPBodyError(0, raw))
|
||||
}
|
||||
|
||||
out := make([]float64, len(docs))
|
||||
for _, it := range resp.Results {
|
||||
s := it.RelevanceScore
|
||||
if s == 0 {
|
||||
s = it.Score
|
||||
}
|
||||
if it.Index >= 0 && it.Index < len(out) {
|
||||
out[it.Index] = s
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
package kb_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"chat-agent-system/internal/kb"
|
||||
)
|
||||
|
||||
func TestReranker_MockServer(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1/rerank" {
|
||||
t.Errorf("unexpected path: %s", r.URL.Path)
|
||||
}
|
||||
var req struct {
|
||||
Model string `json:"model"`
|
||||
Query string `json:"query"`
|
||||
Documents []string `json:"documents"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
type result struct {
|
||||
Index int `json:"index"`
|
||||
RelevanceScore float64 `json:"relevance_score"`
|
||||
}
|
||||
var results []result
|
||||
for i := range req.Documents {
|
||||
results = append(results, result{
|
||||
Index: i,
|
||||
RelevanceScore: 0.95 - float64(i)*0.1,
|
||||
})
|
||||
}
|
||||
resp := map[string]any{"results": results}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
reranker := kb.NewReranker(server.URL, "test-key", "bge-reranker-v2-m3")
|
||||
docs := []string{
|
||||
"GB 9706.1-2020 医用电气设备安全通用要求",
|
||||
"GB/T 4208 外壳防护等级",
|
||||
}
|
||||
scores, err := reranker.Rerank(context.Background(), "医用电气安全接地", docs)
|
||||
if err != nil {
|
||||
t.Fatalf("Rerank failed: %v", err)
|
||||
}
|
||||
if len(scores) != len(docs) {
|
||||
t.Fatalf("expected %d scores, got %d", len(docs), len(scores))
|
||||
}
|
||||
if scores[0] <= scores[1] {
|
||||
t.Errorf("expected doc 0 score > doc 1 score, got %f vs %f", scores[0], scores[1])
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,348 @@
|
||||
package kb
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"log"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Retriever 检索调度器
|
||||
type Retriever struct {
|
||||
Store *Store
|
||||
Embedder *Embedder
|
||||
Reranker *Reranker
|
||||
}
|
||||
|
||||
// NewRetriever 创建检索调度器
|
||||
func NewRetriever(store *Store, embedder *Embedder, reranker *Reranker) *Retriever {
|
||||
return &Retriever{
|
||||
Store: store,
|
||||
Embedder: embedder,
|
||||
Reranker: reranker,
|
||||
}
|
||||
}
|
||||
|
||||
// Search 执行跨领域语义检索与重排
|
||||
func (r *Retriever) Search(ctx context.Context, query string, opts SearchOptions) ([]HitItem, error) {
|
||||
q := strings.TrimSpace(query)
|
||||
if q == "" {
|
||||
return []HitItem{}, nil
|
||||
}
|
||||
if opts.Limit <= 0 {
|
||||
opts.Limit = 5
|
||||
}
|
||||
if opts.CandidateLimit <= 0 {
|
||||
opts.CandidateLimit = 20
|
||||
}
|
||||
if len(opts.Domains) == 0 {
|
||||
opts.Domains = []KnowledgeDomain{DomainRegulatory}
|
||||
}
|
||||
|
||||
var allCandidates []HitItem
|
||||
|
||||
for _, domain := range opts.Domains {
|
||||
switch domain {
|
||||
case DomainRegulatory:
|
||||
hits, err := r.searchRegulatory(ctx, q, opts)
|
||||
if err != nil {
|
||||
log.Printf("[kb-retriever] searchRegulatory err: %v", err)
|
||||
} else {
|
||||
allCandidates = append(allCandidates, hits...)
|
||||
}
|
||||
case DomainProjectAsset:
|
||||
hits, err := r.searchProjectAsset(ctx, q, opts)
|
||||
if err != nil {
|
||||
log.Printf("[kb-retriever] searchProjectAsset err: %v", err)
|
||||
} else {
|
||||
allCandidates = append(allCandidates, hits...)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(allCandidates) == 0 {
|
||||
return []HitItem{}, nil
|
||||
}
|
||||
|
||||
// 排序候选集
|
||||
sort.Slice(allCandidates, func(i, j int) bool {
|
||||
return allCandidates[i].Score > allCandidates[j].Score
|
||||
})
|
||||
|
||||
// 二阶段语义重排 (Cross-Encoder)
|
||||
if r.Reranker != nil && r.Reranker.Endpoint != "" && len(allCandidates) > 1 {
|
||||
rerankLimit := opts.CandidateLimit
|
||||
if len(allCandidates) < rerankLimit {
|
||||
rerankLimit = len(allCandidates)
|
||||
}
|
||||
topList := allCandidates[:rerankLimit]
|
||||
docs := make([]string, len(topList))
|
||||
for i, h := range topList {
|
||||
docs[i] = firstNRunes(h.DocCode+" "+h.Title+"\n"+h.HeadingPath+"\n"+h.Content, 720)
|
||||
}
|
||||
scores, err := r.Reranker.Rerank(ctx, q, docs)
|
||||
if err != nil {
|
||||
log.Printf("[kb-retriever] reranker failed (keeping RRF scores): %v", err)
|
||||
} else if len(scores) == len(topList) {
|
||||
for i := range topList {
|
||||
topList[i].Score = scores[i]
|
||||
}
|
||||
sort.Slice(topList, func(i, j int) bool {
|
||||
return topList[i].Score > topList[j].Score
|
||||
})
|
||||
allCandidates = topList
|
||||
}
|
||||
}
|
||||
|
||||
// 过滤分数与截断
|
||||
var results []HitItem
|
||||
for _, it := range allCandidates {
|
||||
if opts.MinScore > 0 && it.Score < opts.MinScore {
|
||||
continue
|
||||
}
|
||||
results = append(results, it)
|
||||
if len(results) >= opts.Limit {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// searchRegulatory 法规领域检索实现 (Vector + Keyword + RRF)
|
||||
func (r *Retriever) searchRegulatory(ctx context.Context, q string, opts SearchOptions) ([]HitItem, error) {
|
||||
if r.Store == nil || r.Store.Pool == nil {
|
||||
return nil, fmt.Errorf("database not connected")
|
||||
}
|
||||
|
||||
var vecHits []HitItem
|
||||
if r.Embedder != nil && r.Embedder.Endpoint != "" {
|
||||
vec, err := r.Embedder.EmbedOne(ctx, q)
|
||||
if err != nil {
|
||||
log.Printf("[kb-retriever] embed failed: %v", err)
|
||||
} else if len(vec) == EmbeddingDim {
|
||||
vecHits, err = r.regVectorHits(ctx, vec, opts.IncludeRepealed, 40)
|
||||
if err != nil {
|
||||
log.Printf("[kb-retriever] regVectorHits failed: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
kwHits, err := r.regKeywordHits(ctx, q, opts.IncludeRepealed, 30)
|
||||
if err != nil {
|
||||
log.Printf("[kb-retriever] regKeywordHits failed: %v", err)
|
||||
}
|
||||
|
||||
// Reciprocal Rank Fusion (k=60.0)
|
||||
const rrfK = 60.0
|
||||
type rrfItem struct {
|
||||
hit HitItem
|
||||
score float64
|
||||
}
|
||||
fused := make(map[string]*rrfItem)
|
||||
|
||||
for rank, h := range vecHits {
|
||||
score := 1.0 / (rrfK + float64(rank+1))
|
||||
key := fmt.Sprintf("%s|%s", h.DocCode, h.HeadingPath)
|
||||
if existing, ok := fused[key]; ok {
|
||||
existing.score += score
|
||||
} else {
|
||||
fused[key] = &rrfItem{hit: h, score: score}
|
||||
}
|
||||
}
|
||||
|
||||
for rank, h := range kwHits {
|
||||
score := 1.0 / (rrfK + float64(rank+1))
|
||||
key := fmt.Sprintf("%s|%s", h.DocCode, h.HeadingPath)
|
||||
if existing, ok := fused[key]; ok {
|
||||
existing.score += score
|
||||
} else {
|
||||
fused[key] = &rrfItem{hit: h, score: score}
|
||||
}
|
||||
}
|
||||
|
||||
out := make([]HitItem, 0, len(fused))
|
||||
for _, it := range fused {
|
||||
it.hit.Score = it.score
|
||||
out = append(out, it.hit)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *Retriever) regVectorHits(ctx context.Context, vec []float32, includeRepealed bool, limit int) ([]HitItem, error) {
|
||||
lit := VectorLiteral(vec)
|
||||
sqlStr := `
|
||||
SELECT c.id, COALESCE(d.doc_no,''), COALESCE(d.title,''), COALESCE(d.kind,'unknown'), COALESCE(d.status,'unknown'),
|
||||
COALESCE(c.content,''), COALESCE(c.heading_path,''), COALESCE(c.page,0), COALESCE(c.chunk_no,''),
|
||||
1 - (c.embedding <=> $1::vector) AS score
|
||||
FROM reg_chunks c
|
||||
JOIN reg_documents d ON d.id=c.document_id
|
||||
WHERE c.embedding IS NOT NULL` + regStatusFilter(includeRepealed) + `
|
||||
ORDER BY c.embedding <=> $1::vector
|
||||
LIMIT $2`
|
||||
|
||||
rows, err := r.Store.Pool.Query(ctx, sqlStr, lit, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
return scanRegHits(rows)
|
||||
}
|
||||
|
||||
func (r *Retriever) regKeywordHits(ctx context.Context, q string, includeRepealed bool, limit int) ([]HitItem, error) {
|
||||
like1 := "%" + escapeLike(q) + "%"
|
||||
compactQ := strings.ReplaceAll(q, " ", "")
|
||||
like2 := "%" + escapeLike(compactQ) + "%"
|
||||
|
||||
sqlStr := `
|
||||
SELECT c.id, COALESCE(d.doc_no,''), COALESCE(d.title,''), COALESCE(d.kind,'unknown'), COALESCE(d.status,'unknown'),
|
||||
COALESCE(c.content,''), COALESCE(c.heading_path,''), COALESCE(c.page,0), COALESCE(c.chunk_no,''),
|
||||
CASE WHEN c.tsv @@ plainto_tsquery('simple', $1) THEN ts_rank_cd(c.tsv, plainto_tsquery('simple', $1)) ELSE 0.1 END AS score
|
||||
FROM reg_chunks c
|
||||
JOIN reg_documents d ON d.id=c.document_id
|
||||
WHERE (
|
||||
(c.tsv IS NOT NULL AND c.tsv @@ plainto_tsquery('simple', $1)) OR
|
||||
COALESCE(d.doc_no,'') ILIKE $2 ESCAPE '\' OR COALESCE(d.doc_no,'') ILIKE $3 ESCAPE '\' OR
|
||||
COALESCE(c.heading_path,'') ILIKE $2 ESCAPE '\' OR COALESCE(c.heading_path,'') ILIKE $3 ESCAPE '\' OR
|
||||
COALESCE(c.content,'') ILIKE $2 ESCAPE '\' OR COALESCE(c.content,'') ILIKE $3 ESCAPE '\' OR
|
||||
COALESCE(c.cite_key,'') ILIKE $2 ESCAPE '\' OR COALESCE(c.cite_key,'') ILIKE $3 ESCAPE '\'
|
||||
)` + regStatusFilter(includeRepealed) + `
|
||||
ORDER BY score DESC
|
||||
LIMIT $4`
|
||||
|
||||
rows, err := r.Store.Pool.Query(ctx, sqlStr, q, like1, like2, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
return scanRegHits(rows)
|
||||
}
|
||||
|
||||
func (r *Retriever) searchProjectAsset(ctx context.Context, q string, opts SearchOptions) ([]HitItem, error) {
|
||||
sqlStr := `
|
||||
SELECT c.id, p.project_code, p.name, d.doc_type, 'active',
|
||||
c.content, c.chunk_title, 0, '',
|
||||
CASE WHEN c.tsv @@ plainto_tsquery('simple', $1) THEN ts_rank_cd(c.tsv, plainto_tsquery('simple', $1)) ELSE 0.1 END AS score
|
||||
FROM asset_chunks c
|
||||
JOIN asset_projects p ON p.id=c.project_id
|
||||
JOIN asset_documents d ON d.id=c.document_id
|
||||
WHERE c.tsv IS NOT NULL AND c.tsv @@ plainto_tsquery('simple', $1)
|
||||
ORDER BY score DESC LIMIT $2`
|
||||
|
||||
rows, err := r.Store.Pool.Query(ctx, sqlStr, q, opts.Limit)
|
||||
if err != nil {
|
||||
return nil, nil // 预留表空查询容错
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var out []HitItem
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
var code, title, category, status, content, heading, chunkNo string
|
||||
var page int
|
||||
var score float64
|
||||
if err := rows.Scan(&id, &code, &title, &category, &status, &content, &heading, &page, &chunkNo, &score); err != nil {
|
||||
continue
|
||||
}
|
||||
out = append(out, HitItem{
|
||||
Domain: DomainProjectAsset,
|
||||
Category: category,
|
||||
DocCode: code,
|
||||
Title: title,
|
||||
HeadingPath: heading,
|
||||
Content: content,
|
||||
Score: score,
|
||||
CiteKey: fmt.Sprintf("〔项目:%s §%s〕", code, heading),
|
||||
Source: "项目数字资产库",
|
||||
})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func regStatusFilter(includeRepealed bool) string {
|
||||
if includeRepealed {
|
||||
return ""
|
||||
}
|
||||
return ` AND LOWER(d.status) IN ('in_force','unknown','active','published')`
|
||||
}
|
||||
|
||||
func scanRegHits(rows interface {
|
||||
Next() bool
|
||||
Scan(dest ...any) error
|
||||
Err() error
|
||||
}) ([]HitItem, error) {
|
||||
var out []HitItem
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
var docNo, title, kind, status, content, heading, chunkNo sql.NullString
|
||||
var page sql.NullInt64
|
||||
var score sql.NullFloat64
|
||||
|
||||
if err := rows.Scan(&id, &docNo, &title, &kind, &status, &content, &heading, &page, &chunkNo, &score); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
docCode := docNo.String
|
||||
if docCode == "" {
|
||||
docCode = title.String
|
||||
}
|
||||
|
||||
cite := formatCiteKey(docCode, heading.String, chunkNo.String, int(page.Int64), status.String)
|
||||
|
||||
out = append(out, HitItem{
|
||||
Domain: DomainRegulatory,
|
||||
Category: kind.String,
|
||||
DocCode: docCode,
|
||||
Title: title.String,
|
||||
HeadingPath: heading.String,
|
||||
Content: content.String,
|
||||
Score: score.Float64,
|
||||
CiteKey: cite,
|
||||
Source: "国家医疗器械法规与技术标准库",
|
||||
})
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func formatCiteKey(docCode, heading, chunkNo string, page int, status string) string {
|
||||
clause := chunkNo
|
||||
if clause == "" {
|
||||
clause = heading
|
||||
}
|
||||
pageStr := ""
|
||||
if page > 0 {
|
||||
pageStr = " p." + strconv.Itoa(page)
|
||||
}
|
||||
clauseStr := ""
|
||||
if clause != "" {
|
||||
clauseStr = " §" + clause
|
||||
}
|
||||
core := "〔" + docCode + clauseStr + pageStr + "〕"
|
||||
if status == "repealed" || status == "superseded" {
|
||||
core = "〔已废止〕" + core
|
||||
}
|
||||
return core
|
||||
}
|
||||
|
||||
func escapeLike(s string) string {
|
||||
s = strings.ReplaceAll(s, `\`, `\\`)
|
||||
s = strings.ReplaceAll(s, `%`, `\%`)
|
||||
s = strings.ReplaceAll(s, `_`, `\_`)
|
||||
return s
|
||||
}
|
||||
|
||||
func firstNRunes(s string, n int) string {
|
||||
if n <= 0 || len(s) == 0 {
|
||||
return ""
|
||||
}
|
||||
runes := []rune(s)
|
||||
if len(runes) <= n {
|
||||
return s
|
||||
}
|
||||
return string(runes[:n])
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
package kb
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
// Store PostgreSQL 知识库存储访问层
|
||||
type Store struct {
|
||||
Pool *pgxpool.Pool
|
||||
}
|
||||
|
||||
// NewStore 创建 Store 实例
|
||||
func NewStore(pool *pgxpool.Pool) *Store {
|
||||
return &Store{Pool: pool}
|
||||
}
|
||||
|
||||
// OpenPool 建立 PostgreSQL pgx 连接池
|
||||
func OpenPool(ctx context.Context, dsn string) (*pgxpool.Pool, error) {
|
||||
if strings.TrimSpace(dsn) == "" {
|
||||
return nil, fmt.Errorf("empty postgres dsn")
|
||||
}
|
||||
cfg, err := pgxpool.ParseConfig(dsn)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
cfg.MaxConns = 8
|
||||
cfg.MinConns = 1
|
||||
pool, err := pgxpool.NewWithConfig(ctx, cfg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := pool.Ping(ctx); err != nil {
|
||||
pool.Close()
|
||||
return nil, err
|
||||
}
|
||||
return pool, nil
|
||||
}
|
||||
|
||||
// Stats 查询各知识领域的数据行数与切片规模
|
||||
func (s *Store) Stats(ctx context.Context) (map[KnowledgeDomain]Stats, error) {
|
||||
out := make(map[KnowledgeDomain]Stats)
|
||||
if s == nil || s.Pool == nil {
|
||||
return out, fmt.Errorf("store not connected")
|
||||
}
|
||||
|
||||
// 1. 法规领域统计
|
||||
var regStats Stats
|
||||
_ = s.Pool.QueryRow(ctx, `SELECT count(*) FROM reg_documents`).Scan(®Stats.DocumentCount)
|
||||
_ = s.Pool.QueryRow(ctx, `SELECT count(*) FROM reg_nodes`).Scan(®Stats.NodeCount)
|
||||
_ = s.Pool.QueryRow(ctx, `SELECT count(*) FROM reg_chunks`).Scan(®Stats.ChunkCount)
|
||||
out[DomainRegulatory] = regStats
|
||||
|
||||
// 2. 项目数字资产领域统计
|
||||
var assetStats Stats
|
||||
_ = s.Pool.QueryRow(ctx, `SELECT count(*) FROM asset_projects`).Scan(&assetStats.DocumentCount)
|
||||
_ = s.Pool.QueryRow(ctx, `SELECT count(*) FROM asset_chunks`).Scan(&assetStats.ChunkCount)
|
||||
out[DomainProjectAsset] = assetStats
|
||||
|
||||
// 3. 员工能力领域统计
|
||||
var staffStats Stats
|
||||
_ = s.Pool.QueryRow(ctx, `SELECT count(*) FROM staff_profiles`).Scan(&staffStats.DocumentCount)
|
||||
_ = s.Pool.QueryRow(ctx, `SELECT count(*) FROM staff_experience_records`).Scan(&staffStats.ChunkCount)
|
||||
out[DomainStaffCapability] = staffStats
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// VectorLiteral 将 float32 向量序列化为 pgvector 接受的字面量字符串 [0.1,0.2,...]
|
||||
func VectorLiteral(vec []float32) string {
|
||||
if len(vec) == 0 {
|
||||
return ""
|
||||
}
|
||||
var sb strings.Builder
|
||||
sb.WriteByte('[')
|
||||
for i, f := range vec {
|
||||
if i > 0 {
|
||||
sb.WriteByte(',')
|
||||
}
|
||||
sb.WriteString(fmt.Sprintf("%g", f))
|
||||
}
|
||||
sb.WriteByte(']')
|
||||
return sb.String()
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
package kb
|
||||
|
||||
import (
|
||||
"time"
|
||||
)
|
||||
|
||||
// KnowledgeDomain 知识领域类型
|
||||
type KnowledgeDomain string
|
||||
|
||||
const (
|
||||
// DomainRegulatory 医疗器械法律法规、国家标准(GB/T)、行业标准(YY)与指导原则
|
||||
DomainRegulatory KnowledgeDomain = "regulatory"
|
||||
// DomainProjectAsset 已完成历史项目数字资产（可行性报告、CDO 报价、方案文档、BOM清单）
|
||||
DomainProjectAsset KnowledgeDomain = "project_asset"
|
||||
// DomainStaffCapability 研发团队人员工作能力（技术专长、历史人效系数、工时评估模型）
|
||||
DomainStaffCapability KnowledgeDomain = "staff_capability"
|
||||
)
|
||||
|
||||
// HitItem 统一跨领域知识检索命中项
|
||||
type HitItem struct {
|
||||
Domain KnowledgeDomain `json:"domain"`
|
||||
Category string `json:"category"` // standards, guidance, laws, catalog, project_doc, staff_model
|
||||
DocCode string `json:"doc_code"` // 如 GB 9706.1-2020, YY 0569-2019, PRJ-2025-001
|
||||
Title string `json:"title"` // 文档或标准全称
|
||||
HeadingPath string `json:"heading_path"` // 章节小节层级路径 (如: 8. 电气安全 > 8.4 接地)
|
||||
Content string `json:"content"` // 命中正文或条款详细描述
|
||||
Score float64 `json:"score"` // 综合匹配评分 (Reranker 或 RRF 分数)
|
||||
CiteKey string `json:"cite_key"` // 权威引用键 (如: [GB 9706.1-2020 §8.4])
|
||||
Source string `json:"source"` // 来源信息
|
||||
Metadata map[string]any `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// SearchOptions 检索配置选项
|
||||
type SearchOptions struct {
|
||||
Domains []KnowledgeDomain `json:"domains,omitempty"` // 目标领域范围 (空表示默认法规域)
|
||||
Limit int `json:"limit,omitempty"` // 最终期望返回数量
|
||||
CandidateLimit int `json:"candidate_limit,omitempty"` // 初筛召回数量 (Rerank 前)
|
||||
IncludeRepealed bool `json:"include_repealed,omitempty"` // 是否包含已废止/历史版本标准
|
||||
FilterCategory string `json:"filter_category,omitempty"` // 过滤分类
|
||||
MinScore float64 `json:"min_score,omitempty"` // 最低分数阈值
|
||||
}
|
||||
|
||||
// DefaultSearchOptions 返回通用检索默认配置
|
||||
func DefaultSearchOptions() SearchOptions {
|
||||
return SearchOptions{
|
||||
Domains: []KnowledgeDomain{DomainRegulatory},
|
||||
Limit: 5,
|
||||
CandidateLimit: 20,
|
||||
IncludeRepealed: false,
|
||||
MinScore: 0.0,
|
||||
}
|
||||
}
|
||||
|
||||
// Document 数据库文档实体
|
||||
type Document struct {
|
||||
ID int64 `json:"id"`
|
||||
DocID string `json:"doc_id"`
|
||||
SHA256 string `json:"sha256"`
|
||||
FilePath string `json:"file_path"`
|
||||
FileName string `json:"file_name"`
|
||||
Kind string `json:"kind"`
|
||||
DocNo string `json:"doc_no"`
|
||||
Title string `json:"title"`
|
||||
Jurisdiction string `json:"jurisdiction"`
|
||||
Status string `json:"status"`
|
||||
EffectiveDate *time.Time `json:"effective_date,omitempty"`
|
||||
RepealDate *time.Time `json:"repeal_date,omitempty"`
|
||||
SourceURL string `json:"source_url"`
|
||||
PageCount int `json:"page_count"`
|
||||
ExtraJSON map[string]any `json:"extra_json,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// Node 结构化大纲节点实体
|
||||
type Node struct {
|
||||
ID int64 `json:"id"`
|
||||
DocumentID int64 `json:"document_id"`
|
||||
ParentID *int64 `json:"parent_id,omitempty"`
|
||||
Level int `json:"level"`
|
||||
NodeKind string `json:"node_kind"`
|
||||
SectionNo string `json:"section_no"`
|
||||
Title string `json:"title"`
|
||||
HeadingPath string `json:"heading_path"`
|
||||
Content string `json:"content"`
|
||||
PageStart int `json:"page_start"`
|
||||
PageEnd int `json:"page_end"`
|
||||
SortOrder int `json:"sort_order"`
|
||||
}
|
||||
|
||||
// Chunk 向量切片实体
|
||||
type Chunk struct {
|
||||
ID int64 `json:"id"`
|
||||
DocumentID int64 `json:"document_id"`
|
||||
NodeID int64 `json:"node_id"`
|
||||
Grain string `json:"grain"`
|
||||
ChunkNo string `json:"chunk_no"`
|
||||
Content string `json:"content"`
|
||||
HeadingPath string `json:"heading_path"`
|
||||
CiteKey string `json:"cite_key"`
|
||||
Page int `json:"page"`
|
||||
Embedding []float32 `json:"-"`
|
||||
MetadataJSON map[string]any `json:"metadata_json,omitempty"`
|
||||
}
|
||||
|
||||
// Stats 统计信息
|
||||
type Stats struct {
|
||||
DocumentCount int64 `json:"document_count"`
|
||||
NodeCount int64 `json:"node_count"`
|
||||
ChunkCount int64 `json:"chunk_count"`
|
||||
}
|
||||
Reference in New Issue
Block a user