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) // RequireOnline reports whether KB must be reachable. Default is fail-open (offline OK) // for short unit tests. Fail-closed when KB_REQUIRE_ONLINE=1 or RUN_PI_E2E=1 unless // KB_ALLOW_OFFLINE=1 explicitly opts into degradation tests. func RequireOnline() bool { if envTruthy("KB_ALLOW_OFFLINE") { return false } return envTruthy("KB_REQUIRE_ONLINE") || envTruthy("RUN_PI_E2E") } func envTruthy(key string) bool { v := strings.TrimSpace(os.Getenv(key)) return v == "1" || strings.EqualFold(v, "true") || strings.EqualFold(v, "yes") } // NewClient 创建知识库中枢客户端实例。连接失败时默认离线降级;RequireOnline 时 Fatal。 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) } if err := cfg.RequireKbPassword(); err != nil { reason := "empty_kb_db_password" log.Printf("[kb] offline=true reason=%s detail=%v", reason, err) if envTruthy("KB_ALLOW_OFFLINE") { log.Printf("[kb] NOTE: KB offline degradation test (KB_ALLOW_OFFLINE=1)") } else if RequireOnline() { log.Printf("[kb] KB_REQUIRE_ONLINE: offline after empty password; callers must fail closed") } return c } // 3. 尝试建立 PostgreSQL 连接池 dsn := cfg.KbDSN() redacted := config.RedactPostgresDSN(dsn) ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) defer cancel() pool, err := OpenPool(ctx, dsn) if err != nil { log.Printf("[kb] offline=true reason=connect_failed dsn=%s err=%v", redacted, err) if envTruthy("KB_ALLOW_OFFLINE") { log.Printf("[kb] NOTE: KB offline degradation test (KB_ALLOW_OFFLINE=1)") } else if RequireOnline() { log.Printf("[kb] KB_REQUIRE_ONLINE: offline after connect failure; callers must fail closed") } return c } c.pool = pool c.store = NewStore(pool) c.retriever = NewRetriever(c.store, c.embedder, c.reranker) log.Printf("[kb] online=true dsn=%s", redacted) return c } // Online reports whether the PostgreSQL pool is connected. func (c *Client) Online() bool { return c != nil && c.pool != nil } // EnsureOnline returns an error when RequireOnline is set and the client is offline. func EnsureOnline(c *Client) error { if !RequireOnline() { return nil } if c == nil || !c.Online() { return fmt.Errorf("kb: required online (KB_REQUIRE_ONLINE/RUN_PI_E2E) but client is offline; set KB_ALLOW_OFFLINE=1 only for explicit degradation tests") } return 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/ 目录 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) }