88 lines
2.3 KiB
Go
88 lines
2.3 KiB
Go
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()
|
|
}
|