Agent-System/internal/kb/client.go

package kb

import (
	"context"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"chat-agent-system/internal/config"
	"github.com/jackc/pgx/v5/pgxpool"
)

// Client 知识库中枢客户端，实现 Hub 接口
type Client struct {
	cfg       *config.Config
	pool      *pgxpool.Pool
	store     *Store
	embedder  *Embedder
	reranker  *Reranker
	retriever *Retriever
}

// 确保 Client 实现了 Hub 接口
var _ Hub = (*Client)(nil)

// 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)
}