All files / modules/book-generator/utils llm-concurrency.ts

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
/**
 * 全局 LLM 并发控制器
 *
 * 解决 issue：多个书籍同时生成时，各自的 AsyncPool 互不感知，
 * 导致总计并发数可能超过供应商 Rate Limit。
 *
 * 使用信号量模式，全局统一控制所有书籍生成任务的最大并发数。
 *
 * 并发数可通过环境变量 LLM_GLOBAL_CONCURRENCY 配置，默认 10。
 */
 
/** 默认全局最大并发数 */
const DEFAULT_GLOBAL_CONCURRENCY = 10;
 
/** 默认单书最大并发数 */
const DEFAULT_BOOK_CONCURRENCY = 8;
 
/**
 * 简易信号量实现（无外部依赖）
 */
class SimpleSemaphore {
  private current = 0;
  private waitQueue: Array<() => void> = [];
 
  constructor(private max: number) {}
 
  async acquire(): Promise<() => void> {
    if (this.current < this.max) {
      this.current++;
      return () => { this.current--; this.processQueue(); };
    }
    // 等待队列
    return new Promise<() => void>((resolve) => {
      this.waitQueue.push(() => {
        this.current++;
        resolve(() => { this.current--; this.processQueue(); });
      });
    });
  }
 
  private processQueue() {
    if (this.waitQueue.length > 0 && this.current < this.max) {
      const next = this.waitQueue.shift()!;
      next();
    }
  }
}
 
/** 全局信号量实例（单例） */
let globalSemaphore: SimpleSemaphore | null = null;
 
/**
 * 获取全局并发限制
 */
function getGlobalMaxConcurrency(): number {
  const env = process.env['LLM_GLOBAL_CONCURRENCY'];
  if (env) {
    const parsed = parseInt(env, 10);
    if (!isNaN(parsed) && parsed > 0) return parsed;
  }
  return DEFAULT_GLOBAL_CONCURRENCY;
}
 
/**
 * 获取单书最大并发数（可配置）
 */
export function getBookConcurrency(): number {
  const env = process.env['LLM_BOOK_CONCURRENCY'];
  if (env) {
    const parsed = parseInt(env, 10);
    if (!isNaN(parsed) && parsed > 0) return parsed;
  }
  return DEFAULT_BOOK_CONCURRENCY;
}
 
/**
 * 获取全局信号量（懒初始化单例）
 */
function getGlobalSemaphore(): SimpleSemaphore {
  if (!globalSemaphore) {
    globalSemaphore = new SimpleSemaphore(getGlobalMaxConcurrency());
    console.log(`[LLM-Concurrency] 全局信号量初始化: max=${getGlobalMaxConcurrency()}`);
  }
  return globalSemaphore;
}
 
/**
 * 在全局并发控制下执行异步任务。
 *
 * 使用方式：
 *   const result = await withGlobalLLMConcurrency(() => callLLMWithMessages(...));
 *
 * 如果全局槽位满了，调用方会排队等待，而非直接报错。
 *
 * @param fn 需要执行的 LLM 调用函数
 * @returns fn 的返回值
 */
export async function withGlobalLLMConcurrency<T>(fn: () => Promise<T>): Promise<T> {
  const semaphore = getGlobalSemaphore();
  const release = await semaphore.acquire();
  try {
    return await fn();
  } finally {
    release();
  }
}
 
/**
 * 获取当前全局信号量状态（用于监控/调试）
 */
export function getConcurrencyStatus(): { max: number; running: number; waiting: number } {
  const semaphore = getGlobalSemaphore();
  // async-mutex 的 Semaphore 没有公共 API 暴露等待数，这里做一个合理近似
  return {
    max: getGlobalMaxConcurrency(),
    running: -1, // 无法精确获取，需要自行维护计数器
    waiting: -1,
  };
}