diff --git a/go.mod b/go.mod index 5b7456b8..a7855513 100644 --- a/go.mod +++ b/go.mod @@ -2,9 +2,15 @@ module chat-agent-system go 1.27.1 -require github.com/xuri/excelize/v2 v2.11.0 +require ( + github.com/jackc/pgx/v5 v5.10.0 + github.com/xuri/excelize/v2 v2.11.0 +) require ( + github.com/jackc/pgpassfile v1.0.0 // indirect + github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect + github.com/jackc/puddle/v2 v2.2.2 // indirect github.com/richardlehane/mscfb v1.0.7 // indirect github.com/richardlehane/msoleps v1.0.6 // indirect github.com/tiendc/go-deepcopy v1.7.2 // indirect @@ -12,5 +18,6 @@ require ( github.com/xuri/nfp v0.0.2-0.20250530014748-2ddeb826f9a9 // indirect golang.org/x/crypto v0.53.0 // indirect golang.org/x/net v0.56.0 // indirect + golang.org/x/sync v0.21.0 // indirect golang.org/x/text v0.38.0 // indirect ) diff --git a/go.sum b/go.sum index 3ea667a0..f6dba8e2 100644 --- a/go.sum +++ b/go.sum @@ -1,11 +1,23 @@ +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= +github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= +github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= +github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= +github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= +github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/richardlehane/mscfb v1.0.7 h1:oeoiM0WE79vHwE8RpIYYvIAc8ajTH2mb6UZm55/+EB0= github.com/richardlehane/mscfb v1.0.7/go.mod h1:pe0+IUIc0AHh0+teNzBlJCtSyZdFOGgV4ZK9bsoV+Jo= github.com/richardlehane/msoleps v1.0.6 h1:9BvkpjvD+iUBalUY4esMwv6uBkfOip/Lzvd93jvR9gg= github.com/richardlehane/msoleps v1.0.6/go.mod h1:BWev5JBpU9Ko2WAgmZEuiz4/u3ZYTKbjLycmwiWUfWg= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= github.com/tiendc/go-deepcopy v1.7.2 h1:Ut2yYR7W9tWjTQitganoIue4UGxZwCcJy3orjrrIj44= @@ -22,7 +34,11 @@ golang.org/x/image v0.38.0 h1:5l+q+Y9JDC7mBOMjo4/aPhMDcxEptsX+Tt3GgRQRPuE= golang.org/x/image v0.38.0/go.mod h1:/3f6vaXC+6CEanU4KJxbcUZyEePbyKbaLoDOe4ehFYY= golang.org/x/net v0.56.0 h1:Rw8j/hFzGvJUZwNBXnAtf5sVDVt+65SK2C7IxCxZt5o= golang.org/x/net v0.56.0/go.mod h1:D3Ku6r+V6JROoZK144D2XfMHFcMq/0zSfLelVTCFKec= +golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM= +golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/text v0.38.0 h1:sXmwo9DwP3OK9EZ7PqAdaooSGozfl/3a6/xJcbzPRhE= golang.org/x/text v0.38.0/go.mod h1:YXZt3QhHUKYT53r2lLKFIVi6Ao1jdzrTR/KQ09qyxF4= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/internal/config/config.go b/internal/config/config.go index 14b59ce4..2f387c45 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -27,6 +27,18 @@ type Config struct { KbDbUser string `json:"kb_db_user"` KbDbPassword string `json:"kb_db_password"` KbDbName string `json:"kb_db_name"` + + // 向量化 Embedding 服务配置 + EmbeddingEndpointURL string `json:"embedding_endpoint_url"` + EmbeddingProvider string `json:"embedding_provider"` + EmbeddingModel string `json:"embedding_model"` + EmbeddingAPIKey string `json:"embedding_api_key"` + + // 语义重排 Reranker 服务配置 + RerankerEndpointURL string `json:"reranker_endpoint_url"` + RerankerProvider string `json:"reranker_provider"` + RerankerModel string `json:"reranker_model"` + RerankerAPIKey string `json:"reranker_api_key"` } func Load(path string) (*Config, error) { @@ -43,10 +55,20 @@ func Load(path string) (*Config, error) { PiAgentProvider: "cc-switch-kimi-for-coding", PiAgentModel: "kimi-for-coding", KbDbHost: "127.0.0.1", - KbDbPort: 55432, + KbDbPort: 5432, KbDbUser: "regkb", KbDbPassword: "shared_kb_secret_pass", KbDbName: "regulations", + + EmbeddingEndpointURL: "https://www.dmxapi.cn/v1/embeddings", + EmbeddingProvider: "openai_compatible", + EmbeddingModel: "bge-m3", + EmbeddingAPIKey: "", + + RerankerEndpointURL: "https://www.dmxapi.cn/v1/rerank", + RerankerProvider: "openai_compatible", + RerankerModel: "bge-reranker-v2-m3-free", + RerankerAPIKey: "", } if path == "" { @@ -116,6 +138,22 @@ func Load(path string) (*Config, error) { cfg.KbDbPassword = val case "KB_DB_NAME": cfg.KbDbName = val + case "EMBEDDING_ENDPOINT_URL": + cfg.EmbeddingEndpointURL = val + case "EMBEDDING_PROVIDER": + cfg.EmbeddingProvider = val + case "EMBEDDING_MODEL": + cfg.EmbeddingModel = val + case "EMBEDDING_API_KEY": + cfg.EmbeddingAPIKey = val + case "RERANKER_ENDPOINT_URL": + cfg.RerankerEndpointURL = val + case "RERANKER_PROVIDER": + cfg.RerankerProvider = val + case "RERANKER_MODEL": + cfg.RerankerModel = val + case "RERANKER_API_KEY": + cfg.RerankerAPIKey = val } } } @@ -139,6 +177,46 @@ func Load(path string) (*Config, error) { if v := os.Getenv("PI_AGENT_MODEL"); v != "" { cfg.PiAgentModel = v } + if v := os.Getenv("KB_DB_HOST"); v != "" { + cfg.KbDbHost = v + } + if v := os.Getenv("KB_DB_PORT"); v != "" { + if p, err := strconv.Atoi(v); err == nil { + cfg.KbDbPort = p + } + } + if v := os.Getenv("KB_DB_USER"); v != "" { + cfg.KbDbUser = v + } + if v := os.Getenv("KB_DB_PASSWORD"); v != "" { + cfg.KbDbPassword = v + } + if v := os.Getenv("KB_DB_NAME"); v != "" { + cfg.KbDbName = v + } + if v := os.Getenv("EMBEDDING_ENDPOINT_URL"); v != "" { + cfg.EmbeddingEndpointURL = v + } + if v := os.Getenv("EMBEDDING_API_KEY"); v != "" { + cfg.EmbeddingAPIKey = v + } + if v := os.Getenv("EMBEDDING_MODEL"); v != "" { + cfg.EmbeddingModel = v + } + if v := os.Getenv("RERANKER_ENDPOINT_URL"); v != "" { + cfg.RerankerEndpointURL = v + } + if v := os.Getenv("RERANKER_API_KEY"); v != "" { + cfg.RerankerAPIKey = v + } + if v := os.Getenv("RERANKER_MODEL"); v != "" { + cfg.RerankerModel = v + } return cfg, nil } + +// KbDSN 返回 PostgreSQL 连接字符串 +func (c *Config) KbDSN() string { + return "postgres://" + c.KbDbUser + ":" + c.KbDbPassword + "@" + c.KbDbHost + ":" + strconv.Itoa(c.KbDbPort) + "/" + c.KbDbName + "?sslmode=disable" +} diff --git a/internal/kb/cite_lookup.go b/internal/kb/cite_lookup.go new file mode 100644 index 00000000..d6d62de9 --- /dev/null +++ b/internal/kb/cite_lookup.go @@ -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 +} diff --git a/internal/kb/client.go b/internal/kb/client.go index f6473251..cda59aed 100644 --- a/internal/kb/client.go +++ b/internal/kb/client.go @@ -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) diff --git a/internal/kb/embed.go b/internal/kb/embed.go new file mode 100644 index 00000000..06d1b289 --- /dev/null +++ b/internal/kb/embed.go @@ -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] +} diff --git a/internal/kb/embed_test.go b/internal/kb/embed_test.go new file mode 100644 index 00000000..31e1031e --- /dev/null +++ b/internal/kb/embed_test.go @@ -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]) + } +} diff --git a/internal/kb/hub.go b/internal/kb/hub.go new file mode 100644 index 00000000..5456bcfd --- /dev/null +++ b/internal/kb/hub.go @@ -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() +} diff --git a/internal/kb/integration_test.go b/internal/kb/integration_test.go new file mode 100644 index 00000000..8ad281de --- /dev/null +++ b/internal/kb/integration_test.go @@ -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) + } + }) +} diff --git a/internal/kb/models.go b/internal/kb/models.go deleted file mode 100644 index adb0ff9a..00000000 --- a/internal/kb/models.go +++ /dev/null @@ -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"` -} diff --git a/internal/kb/rerank.go b/internal/kb/rerank.go new file mode 100644 index 00000000..6da002a8 --- /dev/null +++ b/internal/kb/rerank.go @@ -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 +} diff --git a/internal/kb/rerank_test.go b/internal/kb/rerank_test.go new file mode 100644 index 00000000..82981aa0 --- /dev/null +++ b/internal/kb/rerank_test.go @@ -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]) + } +} diff --git a/internal/kb/retrieve.go b/internal/kb/retrieve.go new file mode 100644 index 00000000..c17e1681 --- /dev/null +++ b/internal/kb/retrieve.go @@ -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]) +} diff --git a/internal/kb/store.go b/internal/kb/store.go new file mode 100644 index 00000000..a421d5c2 --- /dev/null +++ b/internal/kb/store.go @@ -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() +} diff --git a/internal/kb/types.go b/internal/kb/types.go new file mode 100644 index 00000000..4cf184ad --- /dev/null +++ b/internal/kb/types.go @@ -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"` +}