AI代码审查 - 关键点优化

1
2
3
4
5
6
7
8
9
10
11
12
response = client.chat.completions.create(
    model=MODEL_NAME,
    messages=[
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ],
    stream=False,
    max_tokens=MAX_TOKENS,
    extra_headers={
        "M-TraceId": "unique_id_for_troubleshooting"
    }
)

1
2
3
4
5
6
for chunk in streamResult:
    if chunk.choices[0].delta.content is not None:
        content = chunk.choices[0].delta.content
        sys.stdout.write(content)
        sys.stdout.flush()
        full_content += content

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
# 等待可用APP_ID
if app_id is None and wait:
    # 使用条件变量等待
    if timeout is not None:
        remaining = timeout - (time.time() - start_time)
        if remaining <= 0:
            break
        self.condition.wait(remaining)  # 线程休眠并释放锁
    else:
        self.condition.wait()  # 无限等待

# 释放APP_ID时通知等待的线程
def release_app_id(self, app_id, model_name, success=True):
    with self.lock:

# 释放资源的逻辑..

# 通知所有等待的线程
    self.condition.notify_all()

1
2
3
4
5
	# 分桶计数器 (每个桶记录10秒内的请求数)
	self.bucket_size = 6  # 6个桶，每个桶10秒，共60秒
	self.bucket_duration = 10  # 每个桶的时间跨度(秒)
	self.app_id_buckets = defaultdict(lambda: [0] * self.bucket_size)
	self.app_id_bucket_timestamps = defaultdict(lambda: time.time())

1
2
3
4
5
6
7
8
9
# 启动进程
processes = []
for i in range(num_processes):
    p = mp.Process(target=ProcessWorker(
        i, task_queue, result_queue, global_app_id_manager,
        check_model_list, review_model, mode, threads_per_process
    ).run)
    processes.append(p)
    p.start()

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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
import multiprocessing as mp
import concurrent.futures
import time
import os
import json
from collections import defaultdict, deque
import threading
import queue


class GlobalAppIDManager:
    """全局APP_ID管理器，使用共享内存实现进程间通信"""

    def __init__(self, valid_app_ids, rate_limit_per_minute=3):
        # 使用Manager创建进程间共享对象
        self.manager = mp.Manager()
        self.app_id_pool = self.manager.list(valid_app_ids)
        self.app_id_lock = self.manager.Lock()
        self.app_id_usage = self.manager.dict()  # 记录使用情况
        self.app_id_request_times = self.manager.dict()  # 记录请求时间

        # 初始化请求时间记录
        for app_id in valid_app_ids:
            self.app_id_request_times[app_id] = self.manager.list()
            self.app_id_usage[app_id] = 0

        self.rate_limit_per_minute = rate_limit_per_minute

        # 统计信息
        self.total_requests = self.manager.Value('i', 0)
        self.total_success = self.manager.Value('i', 0)
        self.total_failed = self.manager.Value('i', 0)

    def _clean_old_requests(self, app_id):
        """清理一分钟前的请求记录"""
        current_time = time.time()
        with self.app_id_lock:
            request_times = list(self.app_id_request_times[app_id])
            self.app_id_request_times[app_id] = self.manager.list(
                [t for t in request_times if current_time - t <= 60]
            )

    def _can_use_app_id(self, app_id):
        """检查APP_ID是否可以使用（未超过限流）"""
        self._clean_old_requests(app_id)
        return len(self.app_id_request_times[app_id]) < self.rate_limit_per_minute

    def get_available_app_id(self):
        """获取一个可用的APP_ID，考虑限流"""
        with self.app_id_lock:
            # 寻找可以立即使用的APP_ID
            available_app_ids = []
            for app_id in self.app_id_pool:
                if self._can_use_app_id(app_id):
                    available_app_ids.append(app_id)

            if available_app_ids:
                # 选择使用次数最少的APP_ID
                selected_app_id = min(available_app_ids,
                                      key=lambda x: self.app_id_usage[x])
                self.app_id_pool.remove(selected_app_id)

                # 记录使用时间
                self.app_id_request_times[selected_app_id].append(time.time())
                self.app_id_usage[selected_app_id] += 1
                self.total_requests.value += 1

                return selected_app_id
            else:
                return None

    def release_app_id(self, app_id, success=True):
        """释放APP_ID回池中"""
        with self.app_id_lock:
            self.app_id_pool.append(app_id)

            if success:
                self.total_success.value += 1
            else:
                self.total_failed.value += 1

    def get_stats(self):
        """获取统计信息"""
        with self.app_id_lock:
            stats = {
                "total_requests": self.total_requests.value,
                "total_success": self.total_success.value,
                "total_failed": self.total_failed.value,
                "app_id_usage": {k: v for k, v in self.app_id_usage.items()},
                "available_app_ids": len(self.app_id_pool)
            }
            return stats


class ProcessWorker:
    """进程工作器，每个进程处理一批数据"""

    def __init__(self, process_id, task_queue, result_queue, global_app_id_manager,
                 check_model_list, review_model, mode, max_threads=5):
        self.process_id = process_id
        self.task_queue = task_queue
        self.result_queue = result_queue
        self.global_app_id_manager = global_app_id_manager
        self.check_model_list = check_model_list
        self.review_model = review_model
        self.mode = mode
        self.max_threads = max_threads

        # 进程内连接池
        self.connection_pool = {}  # 简化版连接池
        self.pool_lock = threading.Lock()

    def get_connection(self, app_id, request_url):
        """从进程内连接池获取连接"""
        key = f"{app_id}:{request_url}"
        with self.pool_lock:
            if key not in self.connection_pool:
                # 创建新连接（简化版）
                self.connection_pool[key] = {
                    "client": f"Client-{key}",  # 实际应该创建真正的客户端
                    "created_time": time.time()
                }
            return self.connection_pool[key]

    def process_task(self, task):
        """处理单个任务"""
        target_data, check_model = task
        input_id = target_data["input_id"]

        # 获取APP_ID
        app_id = None
        while app_id is None:
            app_id = self.global_app_id_manager.get_available_app_id()
            if app_id is None:
                time.sleep(0.5)

        try:
            # 获取连接
            connection = self.get_connection(app_id, "REQUEST_URL")

            # 模拟处理过程
            print(f"Process {self.process_id}, Thread {threading.get_ident()}: "
                  f"处理 input_id={input_id}, model={check_model}, app_id={app_id[-6:]}")

            # 模拟API调用
            time.sleep(1)  # 模拟处理时间

            # 模拟结果
            result = {check_model: [{"comment": f"Comment for {input_id} using {check_model}"}]}

            # 释放APP_ID
            self.global_app_id_manager.release_app_id(app_id, success=True)

            return input_id, check_model, result

        except Exception as e:
            print(f"处理任务出错: {str(e)}")
            self.global_app_id_manager.release_app_id(app_id, success=False)
            return input_id, check_model, None

    def thread_worker(self):
        """线程工作函数"""
        while True:
            try:
                task = self.task_queue.get(timeout=1)
                if task is None:  # 结束信号
                    break

                result = self.process_task(task)
                self.result_queue.put(result)

            except queue.Empty:
                break
            except Exception as e:
                print(f"线程异常: {str(e)}")

    def run(self):
        """运行进程工作器"""
        print(f"进程 {self.process_id} 启动，最大线程数: {self.max_threads}")

        # 创建线程池
        threads = []
        for _ in range(self.max_threads):
            t = threading.Thread(target=self.thread_worker)
            threads.append(t)
            t.start()

        # 等待所有线程完成
        for t in threads:
            t.join()

        print(f"进程 {self.process_id} 完成")


def run_multi_process_multi_thread(origin_data, check_model_list, review_model, mode,
                                   valid_app_ids, num_processes=4, threads_per_process=5):
    """多进程+多线程处理函数"""
    start_time = time.time()

    # 创建全局APP_ID管理器
    global_app_id_manager = GlobalAppIDManager(valid_app_ids, rate_limit_per_minute=3)

    # 创建任务队列和结果队列
    task_queue = mp.Queue()
    result_queue = mp.Queue()

    # 填充任务队列
    task_count = 0
    for input_id, target_data in origin_data.items():
        for model_name in check_model_list:
            task_queue.put((target_data, model_name))
            task_count += 1

    # 添加结束信号
    for _ in range(num_processes * threads_per_process):
        task_queue.put(None)

    print(f"创建了 {task_count} 个任务，使用 {num_processes} 个进程，每个进程 {threads_per_process} 个线程")

    # 启动进程
    processes = []
    for i in range(num_processes):
        p = mp.Process(target=ProcessWorker(
            i, task_queue, result_queue, global_app_id_manager,
            check_model_list, review_model, mode, threads_per_process
        ).run)
        processes.append(p)
        p.start()

    # 收集结果
    results = {}
    completed = 0

    # 监控线程
    def monitor_progress():
        while completed < task_count:
            stats = global_app_id_manager.get_stats()
            print(f"\n进度: {completed}/{task_count} | "
                  f"成功: {stats['total_success']} | "
                  f"失败: {stats['total_failed']} | "
                  f"可用APP_ID: {stats['available_app_ids']}")
            time.sleep(10)

    monitor_thread = threading.Thread(target=monitor_progress)
    monitor_thread.daemon = True
    monitor_thread.start()

    # 收集结果
    while completed < task_count:
        try:
            input_id, model_name, result = result_queue.get(timeout=1)
            if result is not None:
                if input_id not in results:
                    results[input_id] = {}
                results[input_id][model_name] = result
            completed += 1

            if completed % 10 == 0:
                print(f"已完成: {completed}/{task_count}")

        except queue.Empty:
            if all(not p.is_alive() for p in processes):
                break

    # 等待所有进程完成
    for p in processes:
        p.join()

    # 处理结果
    final_results = {}
    for input_id, model_results in results.items():
        if input_id in origin_data:
            target_data = origin_data[input_id].copy()
            target_data["output_pred"] = {
                k: v for model, result in model_results.items()
                for k, v in result.items()
            }
            final_results[input_id] = target_data

    elapsed_time = time.time() - start_time

    # 打印统计信息
    stats = global_app_id_manager.get_stats()
    print(f"\n处理完成!")
    print(f"总任务数: {task_count}")
    print(f"成功: {stats['total_success']}")
    print(f"失败: {stats['total_failed']}")
    print(f"总耗时: {elapsed_time:.2f} 秒")
    print(f"平均任务处理时间: {elapsed_time / task_count:.2f} 秒")
    print(f"完整处理的数据项: {len(final_results)}")

    return final_results, elapsed_time


def run_serial(origin_data, check_model_list, review_model, mode, valid_app_ids):
    """串行处理函数"""
    start_time = time.time()

    # 简单的APP_ID轮询
    app_id_index = 0

    results = {}
    task_count = 0

    for input_id, target_data in origin_data.items():
        results[input_id] = {"output_pred": {}}

        for model_name in check_model_list:
            task_count += 1

            # 轮询获取APP_ID
            app_id = valid_app_ids[app_id_index]
            app_id_index = (app_id_index + 1) % len(valid_app_ids)

            try:
                # 模拟处理
                print(f"串行处理: input_id={input_id}, model={model_name}, app_id={app_id[-6:]}")
                time.sleep(1)  # 模拟处理时间

                # 模拟结果
                result = {model_name: [{"comment": f"Comment for {input_id} using {model_name}"}]}
                results[input_id]["output_pred"].update(result)

            except Exception as e:
                print(f"处理任务出错: {str(e)}")

    elapsed_time = time.time() - start_time

    print(f"\n串行处理完成!")
    print(f"总任务数: {task_count}")
    print(f"总耗时: {elapsed_time:.2f} 秒")
    print(f"平均任务处理时间: {elapsed_time / task_count:.2f} 秒")

    return results, elapsed_time


def performance_comparison():
    """性能对比测试"""
    # 模拟数据
    origin_data = {
        f"id_{i}": {"input_id": f"id_{i}", "input_code": f"code_{i}"}
        for i in range(50)
    }

    check_model_list = ["model_1", "model_2", "model_3"]
    review_model = "review_model"
    mode = "check_and_review"
    valid_app_ids = [f"app_id_{i}" for i in range(10)]

    # 运行多进程+多线程版本
    print("\n" + "=" * 50)
    print("运行多进程+多线程版本")
    print("=" * 50)
    mp_results, mp_time = run_multi_process_multi_thread(
        origin_data, check_model_list, review_model, mode, valid_app_ids,
        num_processes=4, threads_per_process=5
    )

    # 运行串行版本
    print("\n" + "=" * 50)
    print("运行串行版本")
    print("=" * 50)
    serial_results, serial_time = run_serial(
        origin_data, check_model_list, review_model, mode, valid_app_ids
    )

    # 性能对比
    task_count = len(origin_data) * len(check_model_list)
    speedup = serial_time / mp_time

    print("\n" + "=" * 50)
    print("性能对比")
    print("=" * 50)
    print(f"{'指标':<25} {'多进程+多线程':<20} {'串行':<20} {'提升比例':<15}")
    print("-" * 80)
    print(f"{'总任务数':<25} {task_count:<20} {task_count:<20} {'-':<15}")
    print(f"{'总耗时(秒)':<25} {mp_time:<20.2f} {serial_time:<20.2f} {speedup:>14.2f}x")
    print(f"{'平均任务时间(秒)':<25} {mp_time / task_count:<20.2f} {serial_time / task_count:<20.2f} {speedup:>14.2f}x")
    print(f"{'理论最大加速比':<25} {4 * 5:<20} {'-':<20} {'-':<15}")
    print(f"{'实际加速比':<25} {speedup:<20.2f} {'-':<20} {'-':<15}")
    print(f"{'并行效率':<25} {speedup / (4 * 5) * 100:<19.2f}% {'-':<20} {'-':<15}")


if __name__ == "__main__":
    performance_comparison()

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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
"""
OpenAI 调用 LLM 接口 - 连接池优化版 vs 原始版本性能对比
"""

from openai import OpenAI
import sys
import tiktoken
import httpx
import threading
import time
import logging
import atexit

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("inferenceAI")

def get_token_count(text):
    """计算文本的token数量"""
    try:
        encoding = tiktoken.encoding_for_model("gpt-4")
        tokens = encoding.encode(text)
        return len(tokens)
    except Exception as e:
        logger.warning(f"计算token数量失败: {e}")
        # 估算token数量（每个字符约0.25个token）
        return len(text) // 4


class OpenAIClientPool:
    """OpenAI 客户端连接池，用于管理和复用 HTTP 连接"""

    def __init__(self, max_clients=50):
        """
        初始化连接池

        Args:
            max_clients: 连接池中最大客户端数量
        """
        self.clients = {}  # 存储客户端的字典，键为 "app_id:request_url"
        self.lock = threading.RLock()  # 线程安全锁
        self.max_clients = max_clients  # 最大客户端数量
        self.stats = {  # 统计信息
            "created_clients": 0,
            "reused_clients": 0,
            "total_requests": 0,
            "total_time": 0.0
        }
        logger.info(f"初始化 OpenAI 客户端连接池，最大客户端数: {max_clients}")

    def get_client(self, app_id, request_url, timeout=600.0):
        """
        获取或创建 OpenAI 客户端

        Args:
            app_id: OpenAI API 密钥
            request_url: API 请求地址
            timeout: 请求超时时间（秒）

        Returns:
            OpenAI 客户端实例
        """
        key = f"{app_id}:{request_url}"  # 客户端的唯一标识

        with self.lock:
            # 检查是否已有客户端
            if key in self.clients:
                self.stats["reused_clients"] += 1
                self.stats["total_requests"] += 1
                logger.debug(f"复用现有客户端: {app_id[-6:]}")
                return self.clients[key]

            # 如果连接池已满，关闭最早创建的客户端
            if len(self.clients) >= self.max_clients:
                oldest_key = next(iter(self.clients))
                logger.info(f"连接池已满，关闭最早的客户端: {oldest_key.split(':')[0][-6:]}")
                self._close_client(oldest_key)

            # 创建新客户端
            http_client = httpx.Client(
                base_url=request_url,
                timeout=timeout,
                limits=httpx.Limits(max_connections=20, max_keepalive_connections=10)
            )

            client = OpenAI(
                api_key=app_id,
                base_url=request_url,
                http_client=http_client
            )

            # 存储客户端
            self.clients[key] = client
            self.stats["created_clients"] += 1
            self.stats["total_requests"] += 1
            logger.info(f"创建新客户端: {app_id[-6:]}，当前连接池大小: {len(self.clients)}")

            return client

    def _close_client(self, key):
        """
        关闭并移除指定的客户端

        Args:
            key: 客户端的唯一标识
        """
        if key in self.clients:
            try:
                if hasattr(self.clients[key], 'http_client'):
                    self.clients[key].http_client.close()
                    logger.debug(f"关闭客户端: {key.split(':')[0][-6:]}")
            except Exception as e:
                logger.error(f"关闭客户端 {key.split(':')[0][-6:]} 时出错: {e}")
            del self.clients[key]

    def close_all(self):
        """关闭所有客户端连接并清理资源"""
        with self.lock:
            logger.info(f"关闭所有客户端连接，共 {len(self.clients)} 个")
            for key in list(self.clients.keys()):
                self._close_client(key)
            logger.info(f"连接池统计: 创建 {self.stats['created_clients']} 个客户端，"
                        f"复用 {self.stats['reused_clients']} 次，"
                        f"总请求 {self.stats['total_requests']} 次")

    def get_stats(self):
        """获取连接池统计信息"""
        with self.lock:
            stats = self.stats.copy()
            if stats['total_requests'] > 0:
                stats['avg_time'] = stats['total_time'] / stats['total_requests']
            else:
                stats['avg_time'] = 0
            return stats

    def update_stats(self, elapsed_time):
        """更新统计信息"""
        with self.lock:
            self.stats['total_time'] += elapsed_time


# 创建全局客户端池
client_pool = OpenAIClientPool(max_clients=50)

# 性能统计类
class PerformanceStats:
    def __init__(self, name):
        self.name = name
        self.requests = 0
        self.total_time = 0.0
        self.min_time = float('inf')
        self.max_time = 0.0
        self.lock = threading.Lock()

    def add_request(self, elapsed_time):
        with self.lock:
            self.requests += 1
            self.total_time += elapsed_time
            self.min_time = min(self.min_time, elapsed_time)
            self.max_time = max(self.max_time, elapsed_time)

    def get_stats(self):
        with self.lock:
            if self.requests == 0:
                return {
                    'name': self.name,
                    'requests': 0,
                    'avg_time': 0,
                    'min_time': 0,
                    'max_time': 0,
                    'total_time': 0
                }
            return {
                'name': self.name,
                'requests': self.requests,
                'avg_time': self.total_time / self.requests,
                'min_time': self.min_time,
                'max_time': self.max_time,
                'total_time': self.total_time
            }

# 创建性能统计实例
pool_stats = PerformanceStats("连接池版本")
original_stats = PerformanceStats("原始版本")


def inferenceOpenAI(APP_ID, REQUEST_URL, MODEL_NAME, user_prompt, system_prompt="", MAX_TOKENS=30*1024):
    """
    普通方式调用 AI 服务（使用连接池优化）

    Args:
        APP_ID (str): AI 服务的应用 ID
        REQUEST_URL (str): API 请求地址
        MODEL_NAME (str): 使用的模型名称
        user_prompt (str): 用户输入的提示词
        system_prompt (str): 系统提示词
        MAX_TOKENS (int, optional): 最大生成 token 数，默认 30*1024

    Returns:
        dict: 包含 AI 响应结果的字典
    """
    start_time = time.time()
    token_count = get_token_count(system_prompt + user_prompt)

    print(f"[连接池版本] request, model name: {MODEL_NAME} , MAX_TOKENS: {MAX_TOKENS} ,token_used: {token_count}")
    logger.info(f"[连接池版本] 请求: 模型={MODEL_NAME}, tokens={token_count}, max_tokens={MAX_TOKENS}")

    try:
        # 从连接池获取客户端，而不是每次创建新的
        client = client_pool.get_client(APP_ID, REQUEST_URL)

        # 生成唯一的跟踪ID
        trace_id = f"pool_trace_{int(time.time())}_{threading.get_ident()}"

        # 发送请求
        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            stream=False,
            max_tokens=MAX_TOKENS,
            extra_headers={
                "M-TraceId": trace_id
            }
        )

        # 处理响应
        result = packageResult(response)

        # 记录请求完成信息
        elapsed = time.time() - start_time
        pool_stats.add_request(elapsed)
        client_pool.update_stats(elapsed)
        logger.info(f"[连接池版本] 请求完成: 模型={MODEL_NAME}, 耗时={elapsed:.2f}秒, tokens={result['total_tokens']}")

        return result

    except Exception as e:
        # 记录错误信息
        elapsed = time.time() - start_time
        pool_stats.add_request(elapsed)
        logger.error(f"[连接池版本] 请求失败: 模型={MODEL_NAME}, 耗时={elapsed:.2f}秒, 错误={str(e)}")
        raise


def inferenceOpenAI_Original(APP_ID, REQUEST_URL, MODEL_NAME, user_prompt, system_prompt="", MAX_TOKENS=30*1024):
    """
    原始方式调用 AI 服务（每次创建新连接）

    Args:
        APP_ID (str): AI 服务的应用 ID
        REQUEST_URL (str): API 请求地址
        MODEL_NAME (str): 使用的模型名称
        user_prompt (str): 用户输入的提示词
        system_prompt (str): 系统提示词
        MAX_TOKENS (int, optional): 最大生成 token 数，默认 30*1024

    Returns:
        dict: 包含 AI 响应结果的字典
    """
    start_time = time.time()
    token_count = get_token_count(system_prompt + user_prompt)

    print(f"[原始版本] request, model name: {MODEL_NAME} , MAX_TOKENS: {MAX_TOKENS} ,token_used: {token_count}")
    logger.info(f"[原始版本] 请求: 模型={MODEL_NAME}, tokens={token_count}, max_tokens={MAX_TOKENS}")

    try:
        # 每次创建新的 HTTP 客户端（原始方式）
        http_client = httpx.Client(
            base_url=REQUEST_URL,
            timeout=600.0  # 10分钟超时
        )

        client = OpenAI(
            api_key=APP_ID,
            base_url=REQUEST_URL,
            http_client=http_client
        )

        # 生成唯一的跟踪ID
        trace_id = f"orig_trace_{int(time.time())}_{threading.get_ident()}"

        # 发送请求
        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            stream=False,
            max_tokens=MAX_TOKENS,
            extra_headers={
                "M-TraceId": trace_id
            }
        )

        # 处理响应
        result = packageResult(response)

        # 记录请求完成信息
        elapsed = time.time() - start_time
        original_stats.add_request(elapsed)
        logger.info(f"[原始版本] 请求完成: 模型={MODEL_NAME}, 耗时={elapsed:.2f}秒, tokens={result['total_tokens']}")

        # 关闭HTTP客户端
        http_client.close()

        return result

    except Exception as e:
        # 记录错误信息
        elapsed = time.time() - start_time
        original_stats.add_request(elapsed)
        logger.error(f"[原始版本] 请求失败: 模型={MODEL_NAME}, 耗时={elapsed:.2f}秒, 错误={str(e)}")
        # 确保在异常情况下也关闭客户端
        try:
            http_client.close()
        except:
            pass
        raise

def packageResult(response):
    """
    解析 API 响应结果，支持非流式响应
    """
    try:
        # 处理普通响应
        content = response.choices[0].message.content
        prompt_tokens = response.usage.prompt_tokens
        completion_tokens = response.usage.completion_tokens
        total_tokens = response.usage.total_tokens
        model = response.model
        request_id = response.id
        created_time = response.created
        result = {
            "content": content,
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": total_tokens,
            "model": model,
            "request_id": request_id,
            "created_time": created_time
        }
        return result
    except Exception as e:
        logger.error(f"解析响应结果时出错: {e}")
        # 返回基本结果
        return {
            "content": response.choices[0].message.content if hasattr(response, 'choices') else "",
            "prompt_tokens": None,
            "completion_tokens": None,
            "total_tokens": None,
            "model": getattr(response, 'model', "unknown"),
            "request_id": getattr(response, 'id', "error"),
            "created_time": getattr(response, 'created', int(time.time()))
        }

def print_performance_comparison():
    """打印性能对比结果"""
    pool_data = pool_stats.get_stats()
    original_data = original_stats.get_stats()

    print("\n" + "="*80)
    print("性能对比结果")
    print("="*80)

    print(f"{'指标':<20} {'连接池版本':<20} {'原始版本':<20} {'性能提升':<20}")
    print("-" * 80)

    if pool_data['requests'] > 0 and original_data['requests'] > 0:
        improvement_avg = ((original_data['avg_time'] - pool_data['avg_time']) / original_data['avg_time']) * 100
        improvement_total = ((original_data['total_time'] - pool_data['total_time']) / original_data['total_time']) * 100

        print(f"{'请求次数':<20} {pool_data['requests']:<20} {original_data['requests']:<20} {'-':<20}")
        print(f"{'平均响应时间(s)':<20} {pool_data['avg_time']:<20.3f} {original_data['avg_time']:<20.3f} {improvement_avg:>18.1f}%")
        print(f"{'最小响应时间(s)':<20} {pool_data['min_time']:<20.3f} {original_data['min_time']:<20.3f} {'-':<20}")
        print(f"{'最大响应时间(s)':<20} {pool_data['max_time']:<20.3f} {original_data['max_time']:<20.3f} {'-':<20}")
        print(f"{'总耗时(s)':<20} {pool_data['total_time']:<20.3f} {original_data['total_time']:<20.3f} {improvement_total:>18.1f}%")
    else:
        print("暂无足够数据进行对比")

    # 显示连接池统计
    pool_client_stats = client_pool.get_stats()
    print("\n连接池详细统计:")
    print(f"  创建的客户端: {pool_client_stats['created_clients']}")
    print(f"  复用的客户端: {pool_client_stats['reused_clients']}")
    print(f"  总请求数: {pool_client_stats['total_requests']}")
    if pool_client_stats['total_requests'] > 0:
        reuse_rate = (pool_client_stats['reused_clients'] / pool_client_stats['total_requests']) * 100
        print(f"  连接复用率: {reuse_rate:.1f}%")


# 确保程序退出时关闭所有连接
def cleanup():
    """程序退出时的清理函数"""
    logger.info("程序退出，清理资源...")
    client_pool.close_all()

# 注册退出处理函数
atexit.register(cleanup)

if __name__ == "__main__":
    import concurrent.futures
    import random

    # 测试配置
    APP_IDS = [
        "21923277120655253527",
        "21921770414276694061",
        "21923297867650007079",
        "21923298145342619704"
    ]  # 多个APP_ID模拟真实场景
    REQUEST_URL = 'https://aigc.sankuai.com/v1/openai/native'
    MODEL_NAME = "gpt-4.1-mini"

    # 测试参数
    TOTAL_REQUESTS = 100  # 总请求数
    MAX_WORKERS = 20  # 最大并发线程数

    # 生成测试提示词
    test_prompts = [
        "Hello, how are you today?",
        "What is Python programming?",
        "Explain machine learning briefly",
        "Tell me a short joke",
        "What is the weather like?",
        "How does AI work?",
        "What is cloud computing?",
        "Explain blockchain technology",
        "What is data science?",
        "How to learn programming?",
        "What is artificial intelligence?",
        "Explain neural networks",
        "What is deep learning?",
        "How does the internet work?",
        "What is software engineering?",
        "Explain database concepts",
        "What is cybersecurity?",
        "How to build a website?",
        "What is mobile development?",
        "Explain API concepts"
    ]


    def worker_task(method, task_id):
        """工作线程任务"""
        app_id = random.choice(APP_IDS)
        prompt = random.choice(test_prompts) + f" (Request #{task_id})"

        try:
            start_time = time.time()
            result = method(app_id, REQUEST_URL, MODEL_NAME, prompt)
            end_time = time.time()

            return {
                'success': True,
                'task_id': task_id,
                'app_id': app_id[-6:],
                'elapsed_time': end_time - start_time,
                'content_length': len(result['content']) if result and 'content' in result else 0,
                'tokens': result.get('total_tokens', 0) if result else 0
            }
        except Exception as e:
            end_time = time.time()
            return {
                'success': False,
                'task_id': task_id,
                'app_id': app_id[-6:],
                'elapsed_time': end_time - start_time,
                'error': str(e),
                'content_length': 0,
                'tokens': 0
            }


    def run_performance_test(method, method_name, total_requests, max_workers):
        """运行性能测试"""
        print(f"\n{'=' * 60}")
        print(f"测试 {method_name}")
        print(f"总请求数: {total_requests}, 并发线程数: {max_workers}")
        print(f"{'=' * 60}")

        results = []
        start_time = time.time()

        # 使用线程池执行并发请求
        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
            # 提交所有任务
            futures = []
            for i in range(total_requests):
                future = executor.submit(worker_task, method, i + 1)
                futures.append(future)

            # 收集结果，显示进度
            completed = 0
            for future in concurrent.futures.as_completed(futures):
                result = future.result()
                results.append(result)
                completed += 1

                # 每完成10个请求显示一次进度
                if completed % 10 == 0 or completed == total_requests:
                    print(f"进度: {completed}/{total_requests} ({completed / total_requests * 100:.1f}%)")

        total_time = time.time() - start_time

        # 分析结果
        successful_results = [r for r in results if r['success']]
        failed_results = [r for r in results if not r['success']]

        if successful_results:
            response_times = [r['elapsed_time'] for r in successful_results]
            avg_response_time = sum(response_times) / len(response_times)
            min_response_time = min(response_times)
            max_response_time = max(response_times)

            total_tokens = sum(r['tokens'] for r in successful_results)
            total_content_length = sum(r['content_length'] for r in successful_results)

            # 按APP_ID统计
            app_id_stats = {}
            for r in successful_results:
                app_id = r['app_id']
                if app_id not in app_id_stats:
                    app_id_stats[app_id] = {'count': 0, 'total_time': 0}
                app_id_stats[app_id]['count'] += 1
                app_id_stats[app_id]['total_time'] += r['elapsed_time']
        else:
            avg_response_time = 0
            min_response_time = 0
            max_response_time = 0
            total_tokens = 0
            total_content_length = 0
            app_id_stats = {}

        # 打印详细统计
        print(f"\n{method_name} 测试结果:")
        print(f"{'=' * 50}")
        print(f"总耗时: {total_time:.2f}秒")
        print(f"成功请求: {len(successful_results)}/{total_requests} ({len(successful_results) / total_requests * 100:.1f}%)")
        print(f"失败请求: {len(failed_results)}")
        print(f"平均响应时间: {avg_response_time:.3f}秒")
        print(f"最小响应时间: {min_response_time:.3f}秒")
        print(f"最大响应时间: {max_response_time:.3f}秒")
        print(f"请求吞吐量: {len(successful_results) / total_time:.2f} 请求/秒")
        print(f"总生成tokens: {total_tokens}")
        print(f"总内容长度: {total_content_length} 字符")

        if app_id_stats:
            print(f"\nAPP_ID使用统计:")
            for app_id, stats in app_id_stats.items():
                avg_time = stats['total_time'] / stats['count']
                print(f"  {app_id}: {stats['count']}次请求, 平均耗时: {avg_time:.3f}秒")

        if failed_results:
            print(f"\n失败请求详情:")
            error_counts = {}
            for r in failed_results:
                error = r.get('error', 'Unknown error')
                error_counts[error] = error_counts.get(error, 0) + 1

            for error, count in error_counts.items():
                print(f"  {error}: {count}次")

        return {
            'method_name': method_name,
            'total_time': total_time,
            'successful_requests': len(successful_results),
            'failed_requests': len(failed_results),
            'avg_response_time': avg_response_time,
            'min_response_time': min_response_time,
            'max_response_time': max_response_time,
            'throughput': len(successful_results) / total_time if total_time > 0 else 0,
            'total_tokens': total_tokens,
            'app_id_stats': app_id_stats
        }


    def print_comparison(pool_result, original_result):
        """打印对比结果"""
        print(f"\n{'=' * 80}")
        print("高并发场景性能对比结果")
        print(f"{'=' * 80}")

        print(f"{'指标':<25} {'连接池版本':<20} {'原始版本':<20} {'性能提升':<15}")
        print("-" * 80)

        # 成功率对比
        pool_success_rate = pool_result['successful_requests'] / (
                    pool_result['successful_requests'] + pool_result['failed_requests']) * 100
        orig_success_rate = original_result['successful_requests'] / (
                    original_result['successful_requests'] + original_result['failed_requests']) * 100

        print(f"{'成功率(%)':<25} {pool_success_rate:<20.1f} {orig_success_rate:<20.1f} {pool_success_rate - orig_success_rate:>13.1f}%")

        # 响应时间对比
        if pool_result['avg_response_time'] > 0 and original_result['avg_response_time'] > 0:
            time_improvement = ((original_result['avg_response_time'] - pool_result['avg_response_time']) /
                                original_result['avg_response_time']) * 100
            print(f"{'平均响应时间(s)':<25} {pool_result['avg_response_time']:<20.3f} {original_result['avg_response_time']:<20.3f} {time_improvement:>13.1f}%")

        # 吞吐量对比
        if pool_result['throughput'] > 0 and original_result['throughput'] > 0:
            throughput_improvement = ((pool_result['throughput'] - original_result['throughput']) / original_result['throughput']) * 100
            print(f"{'吞吐量(请求/秒)':<25} {pool_result['throughput']:<20.2f} {original_result['throughput']:<20.2f} {throughput_improvement:>13.1f}%")

        # 总耗时对比
        if pool_result['total_time'] > 0 and original_result['total_time'] > 0:
            total_time_improvement = ((original_result['total_time'] - pool_result['total_time']) / original_result['total_time']) * 100
            print(f"{'总耗时(s)':<25} {pool_result['total_time']:<20.2f} {original_result['total_time']:<20.2f} {total_time_improvement:>13.1f}%")

        print(f"{'最小响应时间(s)':<25} {pool_result['min_response_time']:<20.3f} {original_result['min_response_time']:<20.3f} {'-':<15}")
        print(f"{'最大响应时间(s)':<25} {pool_result['max_response_time']:<20.3f} {original_result['max_response_time']:<20.3f} {'-':<15}")

        # 连接池统计
        pool_client_stats = client_pool.get_stats()
        print(f"\n连接池详细统计:")
        print(f"  创建的客户端: {pool_client_stats['created_clients']}")
        print(f"  复用的客户端: {pool_client_stats['reused_clients']}")
        print(f"  总请求数: {pool_client_stats['total_requests']}")
        if pool_client_stats['total_requests'] > 0:
            reuse_rate = (pool_client_stats['reused_clients'] / pool_client_stats['total_requests']) * 100
            print(f"  连接复用率: {reuse_rate:.1f}%")

    # 开始性能测试
    print("=== 高并发大量请求场景性能测试 ===")
    print(f"测试配置:")
    print(f"  总请求数: {TOTAL_REQUESTS}")
    print(f"  并发线程数: {MAX_WORKERS}")
    print(f"  APP_ID数量: {len(APP_IDS)}")
    print(f"  测试提示词数量: {len(test_prompts)}")

    # 重置统计数据
    pool_stats = PerformanceStats("连接池版本")
    original_stats = PerformanceStats("原始版本")

    # 测试连接池版本
    pool_result = run_performance_test(
        inferenceOpenAI,
        "连接池版本",
        TOTAL_REQUESTS,
        MAX_WORKERS
    )

    # 等待一段时间，让系统稳定
    print("\n等待5秒让系统稳定...")
    time.sleep(5)

    # 测试原始版本
    original_result = run_performance_test(
        inferenceOpenAI_Original,
        "原始版本",
        TOTAL_REQUESTS,
        MAX_WORKERS
    )

    # 打印对比结果
    print_comparison(pool_result, original_result)