AI代码审查

1
2
3
4
5
def inferenceOpenAI(APP_ID, REQUEST_URL, MODEL_NAME, user_prompt, system_prompt=""):
    http_client = httpx.Client(base_url=REQUEST_URL, timeout=600.0)
    client = OpenAI(api_key=APP_ID, base_url=REQUEST_URL, http_client=http_client)
    response = client.chat.completions.create(...)
    return packageResult(response)

1
2
3
4
5
async def inferenceOpenAIAsync(APP_ID, REQUEST_URL, MODEL_NAME, user_prompt, system_prompt=""):
    async with httpx.AsyncClient(base_url=REQUEST_URL, timeout=600.0) as http_client:
        client = AsyncOpenAI(api_key=APP_ID, base_url=REQUEST_URL, http_client=http_client)
        response = await client.chat.completions.create(...)
        return packageResult(response)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# 全局客户端池
client_pool = {}

def get_client(APP_ID, REQUEST_URL):
    key = f"{APP_ID}:{REQUEST_URL}"
    if key not in client_pool:
        http_client = httpx.Client(base_url=REQUEST_URL, timeout=600.0)
        client_pool[key] = OpenAI(api_key=APP_ID, base_url=REQUEST_URL, http_client=http_client)
    return client_pool[key]

def inferenceOpenAI(APP_ID, REQUEST_URL, MODEL_NAME, user_prompt, system_prompt=""):
    client = get_client(APP_ID, REQUEST_URL)
    response = client.chat.completions.create(...)
    return packageResult(response)

1
2
3
4
5
6
7
8
from tenacity import retry, stop_after_attempt, wait_exponential
from circuitbreaker import circuit

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=10))
@circuit(failure_threshold=5, recovery_timeout=30)
def inferenceOpenAI(APP_ID, REQUEST_URL, MODEL_NAME, user_prompt, system_prompt=""):
    # 实现与当前类似，但添加了重试和断路器机制
    ...

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import hashlib
from functools import lru_cache

def get_cache_key(MODEL_NAME, system_prompt, user_prompt):
    content = f"{MODEL_NAME}:{system_prompt}:{user_prompt}"
    return hashlib.md5(content.encode()).hexdigest()

@lru_cache(maxsize=1000)
def inferenceOpenAICached(cache_key, APP_ID, REQUEST_URL, MODEL_NAME, user_prompt, system_prompt=""):
    # 实际调用逻辑
    return inferenceOpenAI(APP_ID, REQUEST_URL, MODEL_NAME, user_prompt, system_prompt)

def inferenceOpenAI(APP_ID, REQUEST_URL, MODEL_NAME, user_prompt, system_prompt=""):
    cache_key = get_cache_key(MODEL_NAME, system_prompt, user_prompt)
    return inferenceOpenAICached(cache_key, APP_ID, REQUEST_URL, MODEL_NAME, user_prompt, system_prompt)

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
import httpx
from openai import OpenAI
import time
import threading
from functools import wraps
import logging

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

class OpenAIClientManager:
    """管理OpenAI客户端连接池"""
    
    def __init__(self):
        self.clients = {}
        self.lock = threading.RLock()
        self.stats = {
            "total_requests": 0,
            "successful_requests": 0,
            "failed_requests": 0,
            "total_latency": 0,
        }
        
    def get_client(self, app_id, base_url, timeout=600.0):
        """获取或创建OpenAI客户端"""
        key = f"{app_id}:{base_url}"
        
        with self.lock:
            if key not in self.clients:
                http_client = httpx.Client(
                    base_url=base_url,
                    timeout=timeout,
                    limits=httpx.Limits(max_connections=100, max_keepalive_connections=20)
                )
                self.clients[key] = OpenAI(
                    api_key=app_id,
                    base_url=base_url,
                    http_client=http_client
                )
                logger.info(f"Created new OpenAI client for APP_ID: {app_id[-6:]}")
            
            return self.clients[key]
            
    def update_stats(self, success, latency):
        """更新统计信息"""
        with self.lock:
            self.stats["total_requests"] += 1
            if success:
                self.stats["successful_requests"] += 1
            else:
                self.stats["failed_requests"] += 1
            self.stats["total_latency"] += latency
            
    def get_stats(self):
        """获取统计信息"""
        with self.lock:
            stats = self.stats.copy()
            avg_latency = stats["total_latency"] / max(1, stats["total_requests"])
            success_rate = stats["successful_requests"] / max(1, stats["total_requests"]) * 100
            
            return {
                **stats,
                "avg_latency": avg_latency,
                "success_rate": success_rate
            }
            
    def close_all(self):
        """关闭所有客户端连接"""
        with self.lock:
            for key, client in self.clients.items():
                try:
                    if hasattr(client, 'http_client') and client.http_client:
                        client.http_client.close()
                except Exception as e:
                    logger.error(f"Error closing client {key}: {str(e)}")
            self.clients.clear()

# 创建全局客户端管理器
client_manager = OpenAIClientManager()

def with_retry(max_retries=3, backoff_factor=0.5):
    """重试装饰器"""
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            last_exception = None
            for attempt in range(max_retries):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    last_exception = e
                    wait_time = backoff_factor * (2 ** attempt)
                    logger.warning(f"Attempt {attempt+1}/{max_retries} failed: {str(e)}. Retrying in {wait_time:.2f}s")
                    time.sleep(wait_time)
            
            # 所有重试都失败
            logger.error(f"All {max_retries} attempts failed. Last error: {str(last_exception)}")
            raise last_exception
        return wrapper
    return decorator

@with_retry(max_retries=3)
def inferenceOpenAI(APP_ID, REQUEST_URL, MODEL_NAME, user_prompt, system_prompt="", MAX_TOKENS=30*1024):
    """改进的AI调用函数，使用连接池和重试机制"""
    start_time = time.time()
    success = False
    
    try:
        # 获取客户端（从连接池）
        client = client_manager.get_client(APP_ID, REQUEST_URL)
        
        # 记录请求信息
        token_count = get_token_count(system_prompt + user_prompt)
        logger.info(f"Request: model={MODEL_NAME}, tokens={token_count}, max_tokens={MAX_TOKENS}")
        
        # 发送请求
        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": f"trace_{int(time.time())}"
            }
        )
        
        # 处理响应
        result = packageResult(response)
        success = True
        return result
        
    except Exception as e:
        logger.error(f"Error calling {MODEL_NAME}: {str(e)}")
        raise
    finally:
        # 更新统计信息
        latency = time.time() - start_time
        client_manager.update_stats(success, latency)
        logger.debug(f"Request completed in {latency:.2f}s, success={success}")

# 同样方式改进流式调用函数
@with_retry(max_retries=2)  # 流式调用重试次数少一些
def inferenceOpenAIStream(APP_ID, REQUEST_URL, MODEL_NAME, user_prompt, system_prompt="", MAX_TOKENS=30*1024):
    """改进的流式AI调用函数"""
    # 实现类似，但处理流式响应
    ...

# 应用关闭时清理资源
def cleanup():
    """关闭所有连接并清理资源"""
    logger.info("Cleaning up resources...")
    client_manager.close_all()

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

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
import threading
import time
import heapq
from collections import defaultdict, deque
from queue import PriorityQueue, Empty
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("appid_manager.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("AppIDManager")

class TokenBucket:
    """令牌桶实现，用于限流控制"""
    def __init__(self, capacity, refill_rate):
        """
        初始化令牌桶
        
        Args:
            capacity: 桶容量（最大令牌数）
            refill_rate: 令牌补充速率（每秒补充的令牌数）
        """
        self.capacity = capacity
        self.refill_rate = refill_rate
        self.tokens = capacity  # 初始满桶
        self.last_refill_time = time.time()
        self.lock = threading.RLock()
        
    def _refill(self):
        """补充令牌"""
        now = time.time()
        elapsed = now - self.last_refill_time
        # 计算需要补充的令牌数
        new_tokens = elapsed * self.refill_rate
        
        if new_tokens > 0:
            self.tokens = min(self.capacity, self.tokens + new_tokens)
            self.last_refill_time = now
            
    def try_acquire(self, tokens=1):
        """
        尝试获取令牌，非阻塞
        
        Args:
            tokens: 需要获取的令牌数
            
        Returns:
            bool: 是否成功获取令牌
        """
        with self.lock:
            self._refill()
            if self.tokens >= tokens:
                self.tokens -= tokens
                return True
            return False
            
    def acquire(self, tokens=1, timeout=None):
        """
        获取令牌，支持超时
        
        Args:
            tokens: 需要获取的令牌数
            timeout: 超时时间（秒），None表示无限等待
            
        Returns:
            bool: 是否成功获取令牌
        """
        start_time = time.time()
        
        while True:
            if self.try_acquire(tokens):
                return True
                
            # 检查是否超时
            if timeout is not None and time.time() - start_time >= timeout:
                return False
                
            # 计算下一个令牌可用的时间
            with self.lock:
                self._refill()
                if self.tokens >= tokens:
                    self.tokens -= tokens
                    return True
                
                # 计算需要等待的时间
                required_tokens = tokens - self.tokens
                wait_time = required_tokens / self.refill_rate
                
                # 考虑超时
                if timeout is not None:
                    remaining = timeout - (time.time() - start_time)
                    if remaining <= 0:
                        return False
                    wait_time = min(wait_time, remaining)
                    
            # 等待一段时间再尝试
            time.sleep(min(wait_time, 0.1))  # 最多等待100ms，避免过长等待

class Task:
    """表示一个需要APP_ID的任务"""
    def __init__(self, model_name, priority=0, timestamp=None):
        self.model_name = model_name
        self.priority = priority  # 优先级，数字越小优先级越高
        self.timestamp = timestamp or time.time()  # 创建时间戳
        self.event = threading.Event()  # 用于通知任务已分配APP_ID
        self.app_id = None  # 分配的APP_ID
        
    def __lt__(self, other):
        """比较运算符，用于优先级队列排序"""
        # 先按优先级排序，再按时间戳排序（FIFO）
        if self.priority != other.priority:
            return self.priority < other.priority
        return self.timestamp < other.timestamp

class EnhancedAppIDManager:
    """增强版APP_ID管理器，使用条件变量、令牌桶和优先级队列"""
    def __init__(self, valid_app_ids, review_model, rate_limit_per_minute=3):
        """
        初始化APP_ID管理器
        
        Args:
            valid_app_ids: 有效的APP_ID列表
            review_model: 审核模型名称
            rate_limit_per_minute: 每分钟每个APP_ID的请求限制
        """
        # 基本属性
        self.review_model = review_model
        self.rate_limit_per_minute = rate_limit_per_minute
        
        # 资源池管理
        self.app_id_pool = valid_app_ids.copy()
        self.app_id_in_use = set()  # 正在使用的APP_ID
        
        # 线程同步
        self.lock = threading.RLock()
        self.condition = threading.Condition(self.lock)
        
        # 任务队列（按优先级）
        self.waiting_tasks = PriorityQueue()
        
        # 令牌桶（限流）
        refill_rate = rate_limit_per_minute / 60.0  # 每秒补充的令牌数
        self.token_buckets = {
            app_id: TokenBucket(capacity=rate_limit_per_minute, refill_rate=refill_rate)
            for app_id in valid_app_ids
        }
        
        # 使用统计
        self.app_id_usage_count = defaultdict(int)  # 总使用次数
        self.app_id_success_count = defaultdict(int)  # 成功次数
        self.app_id_failure_count = defaultdict(int)  # 失败次数
        self.app_id_current_tasks = defaultdict(set)  # 当前正在处理的任务
        
        # 性能监控
        self.last_status_print_time = time.time()
        self.total_requests = 0
        self.total_success = 0
        self.total_failed = 0
        self.total_wait_time = 0
        self.wait_count = 0
        
        # 启动资源分配线程
        self.running = True
        self.allocator_thread = threading.Thread(target=self._resource_allocator, daemon=True)
        self.allocator_thread.start()
        
        logger.info(f"初始化AppIDManager，APP_ID数量: {len(valid_app_ids)}，限流: {rate_limit_per_minute}/分钟")
        
    def _resource_allocator(self):
        """后台线程，负责分配资源给等待的任务"""
        while self.running:
            allocated = False
            
            with self.lock:
                # 检查是否有可用的APP_ID
                available_app_ids = [
                    app_id for app_id in self.app_id_pool 
                    if app_id not in self.app_id_in_use and self.token_buckets[app_id].try_acquire()
                ]
                
                if available_app_ids and not self.waiting_tasks.empty():
                    try:
                        # 获取优先级最高的任务
                        task = self.waiting_tasks.get_nowait()
                        
                        # 选择使用次数最少的APP_ID（负载均衡）
                        selected_app_id = min(
                            available_app_ids, 
                            key=lambda x: self.app_id_usage_count[x]
                        )
                        
                        # 分配APP_ID给任务
                        self.app_id_in_use.add(selected_app_id)
                        task.app_id = selected_app_id
                        self.app_id_usage_count[selected_app_id] += 1
                        self.app_id_current_tasks[selected_app_id].add(task.model_name)
                        self.total_requests += 1
                        
                        # 通知等待的线程
                        task.event.set()
                        allocated = True
                        
                        logger.debug(f"分配APP_ID: {selected_app_id[-6:]} 给模型: {task.model_name}")
                    except Empty:
                        pass
            
            # 如果没有分配资源，短暂休眠避免CPU占用
            if not allocated:
                time.sleep(0.01)
                
    def get_app_id(self, model_name, priority=0, timeout=None):
        """
        获取一个可用的APP_ID
        
        Args:
            model_name: 模型名称
            priority: 优先级（数字越小优先级越高）
            timeout: 超时时间（秒），None表示无限等待
            
        Returns:
            str: 分配的APP_ID，如果超时返回None
        """
        start_time = time.time()
        
        # 创建任务
        task = Task(model_name, priority)
        
        # 将任务加入等待队列
        self.waiting_tasks.put(task)
        logger.debug(f"模型 {model_name} 请求APP_ID，优先级: {priority}")
        
        # 等待分配APP_ID或超时
        if timeout is not None:
            success = task.event.wait(timeout)
            if not success:
                # 超时，从队列中移除任务（实际无法直接从优先级队列中移除）
                logger.warning(f"获取APP_ID超时: 模型={model_name}, 优先级={priority}")
                return None
        else:
            task.event.wait()
        
        # 计算等待时间
        wait_time = time.time() - start_time
        with self.lock:
            self.total_wait_time += wait_time
            self.wait_count += 1
            
        logger.debug(f"模型 {model_name} 获得APP_ID: {task.app_id[-6:]}，等待时间: {wait_time:.2f}秒")
        return task.app_id
        
    def release_app_id(self, app_id, model_name, success=True):
        """
        释放APP_ID回池中
        
        Args:
            app_id: 要释放的APP_ID
            model_name: 使用的模型名称
            success: 是否成功完成任务
        """
        with self.lock:
            if app_id in self.app_id_in_use:
                self.app_id_in_use.remove(app_id)
                self.app_id_pool.append(app_id)
                
                # 更新统计信息
                if success:
                    self.app_id_success_count[app_id] += 1
                    self.total_success += 1
                else:
                    self.app_id_failure_count[app_id] += 1
                    self.total_failed += 1
                    
                # 移除当前任务记录
                if model_name in self.app_id_current_tasks[app_id]:
                    self.app_id_current_tasks[app_id].remove(model_name)
                
                # 通知等待的线程
                self.condition.notify_all()
                
                logger.debug(f"释放APP_ID: {app_id[-6:]}，模型: {model_name}，成功: {success}")
            else:
                logger.warning(f"尝试释放未使用的APP_ID: {app_id[-6:]}，模型: {model_name}")
                
    def print_status(self, force=False):
        """
        打印当前状态
        
        Args:
            force: 是否强制打印，不考虑时间间隔
        """
        current_time = time.time()
        if force or (current_time - self.last_status_print_time >= 30):  # 每30秒打印一次
            with self.lock:
                # 计算平均等待时间
                avg_wait_time = self.total_wait_time / max(1, self.wait_count)
                
                status_lines = []
                status_lines.append(f"┌{'─' * 90}┐")
                status_lines.append(f"│ 当前状态: 使用中APP_ID: {len(self.app_id_in_use)}/{len(self.app_id_pool) + len(self.app_id_in_use)} | 等待任务: {self.waiting_tasks.qsize()} │")
                
                # APP_ID使用统计
                app_id_stats = []
                for app_id in sorted(self.token_buckets.keys(), 
                                    key=lambda x: self.app_id_usage_count[x], 
                                    reverse=True)[:5]:  # 只显示使用最多的5个
                    short_id = app_id[-6:] if len(app_id) > 6 else app_id
                    usage = self.app_id_usage_count[app_id]
                    success = self.app_id_success_count[app_id]
                    failure = self.app_id_failure_count[app_id]
                    current = len(self.app_id_current_tasks[app_id])
                    app_id_stats.append(f"{short_id}({usage}次:{success}/{failure}, 当前:{current})")
                
                if app_id_stats:
                    stats_str = ', '.join(app_id_stats)
                    status_lines.append(f"│ APP_ID统计(总次数:成功/失败, 当前): {stats_str:<50}│")
                
                status_lines.append(f"│ 总请求: {self.total_requests} | 成功: {self.total_success} | 失败: {self.total_failed} | 平均等待: {avg_wait_time:.2f}秒 │")
                status_lines.append(f"│ 限流: {self.rate_limit_per_minute}/分钟/APP_ID │")
                status_lines.append(f"└{'─' * 90}┘")
                
                print('\n'.join(status_lines))
                self.last_status_print_time = current_time
                
    def shutdown(self):
        """关闭管理器，停止后台线程"""
        self.running = False
        if hasattr(self, 'allocator_thread') and self.allocator_thread.is_alive():
            self.allocator_thread.join(timeout=1)
        logger.info("AppIDManager已关闭")

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
# 修改 process_with_dynamic_app_id_v2 函数使用新的 AppIDManager
def process_with_dynamic_app_id_v2(task_info, app_id_manager, mode="check_and_review", max_retries=3):
    """使用动态分配的APP_ID处理单个任务"""
    target_data, check_model, review_model = task_info
    
    # 确定任务优先级（可以根据业务需求自定义）
    priority = 0  # 默认优先级
    if "priority" in target_data:
        priority = target_data["priority"]
                
    try:
        print(f"开始处理任务 - input_id: {target_data.get('input_id', 'Unknown')}, model: {check_model}")
    except Exception as e:
        print(f"获取任务信息时出错: {str(e)}")
        return None, None, None

    retry_count = 0
    while retry_count < max_retries:
        # 使用优先级获取APP_ID，最多等待30秒
        app_id = app_id_manager.get_app_id(check_model, priority=priority, timeout=30)
        
        if app_id is None:
            # 如果超时未获取到APP_ID，增加重试次数
            retry_count += 1
            print(f"获取APP_ID超时，重试 {retry_count}/{max_retries}")
            continue
                
        try:
            print(f"使用 APP_ID: {app_id[-6:]}... 处理 input_id: {target_data.get('input_id', 'Unknown')}")

            result = process_single_item_with_app_id(
                target_data=target_data,
                check_model=check_model,
                review_model=review_model,
                app_id=app_id,
                mode=mode
            )

            if result is not None:
                app_id_manager.release_app_id(app_id, check_model, success=True)
                return target_data["input_id"], check_model, result
            else:
                retry_count += 1
                app_id_manager.release_app_id(app_id, check_model, success=False)
                if retry_count < max_retries:
                    print(f"处理结果为空，重试 {retry_count}/{max_retries} for input_id: {target_data.get('input_id', 'Unknown')}, model: {check_model}")
                    time.sleep(1)

        except Exception as e:
            retry_count += 1
            app_id_manager.release_app_id(app_id, check_model, success=False)
            # 错误处理代码...

    return target_data.get("input_id", "Unknown"), check_model, None

# 修改 run_comments_and_verify_parallel_v2 函数使用新的 AppIDManager
def run_comments_and_verify_parallel_v2(origin_data, model_list, valid_app_id, save_file,
                                       mode="check_and_review", review_model="gpt-4.1-mini",
                                       max_workers=None, rate_limit_per_minute=3):
    # ... 现有代码 ...
    
    # 创建增强版APP_ID管理器
    app_id_manager = EnhancedAppIDManager(valid_app_id, review_model, rate_limit_per_minute)
    
    try:
        # ... 现有处理代码 ...
    finally:
        # 确保关闭AppIDManager
        app_id_manager.shutdown()

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
# ... 保留现有代码 ...

import multiprocessing
from functools import partial

def worker_process(data_chunk, model_list, valid_app_id, review_model, mode, rate_limit_per_minute):
    """每个进程内的工作函数，使用线程池处理一部分数据"""
    # 每个进程创建自己的AppIDManager和线程池
    app_id_manager = AppIDManager(valid_app_id, review_model, rate_limit_per_minute)
    results = {}
    
    # 创建任务列表
    tasks = []
    for data_id, data in data_chunk.items():
        for model_name in model_list:
            tasks.append((data, model_name, review_model))
    
    # 使用线程池处理任务
    with concurrent.futures.ThreadPoolExecutor(max_workers=len(valid_app_id)*2) as executor:
        futures = {executor.submit(process_with_dynamic_app_id_v2, task, app_id_manager, mode): task for task in tasks}
        
        for future in concurrent.futures.as_completed(futures):
            task = futures[future]
            try:
                input_id, model_name, result = future.result()
                if input_id not in results:
                    results[input_id] = {}
                if result is not None:
                    results[input_id][model_name] = result
            except Exception as e:
                print(f"任务处理异常: {str(e)}")
    
    return results

def run_comments_and_verify_parallel_hybrid(origin_data, model_list, valid_app_id, save_file,
                                           mode="check_and_review", review_model="gpt-4.1-mini",
                                           rate_limit_per_minute=3):
    """混合进程池和线程池的并行处理函数"""
    # 确定进程数量 - 通常设置为CPU核心数
    num_processes = min(multiprocessing.cpu_count(), 8)  # 限制最大进程数
    print(f"使用 {num_processes} 个进程进行并行处理")
    
    # 读取已保存的数据
    try:
        saved_data = read_json_file(save_file)
        if saved_data is None:
            saved_data = {}
    except Exception as e:
        print(f"读取保存文件出错: {str(e)}")
        saved_data = {}
    
    # 过滤掉已处理的数据
    filtered_data = {}
    for input_id, data in origin_data.items():
        if str(input_id) not in saved_data:
            filtered_data[str(input_id)] = data
    
    if not filtered_data:
        print("所有数据都已处理完成")
        return
    
    print(f"需要处理的数据项: {len(filtered_data)}")
    
    # 将数据分成大致相等的块
    data_chunks = split_data_into_chunks(filtered_data, num_processes)
    print(f"数据已分成 {len(data_chunks)} 个块")
    
    # 准备进程池
    start_time = time.time()
    
    # 使用进程池处理数据块
    with multiprocessing.Pool(processes=num_processes) as pool:
        # 创建偏函数，固定除数据块外的其他参数
        process_func = partial(
            worker_process,
            model_list=model_list,
            valid_app_id=valid_app_id,
            review_model=review_model,
            mode=mode,
            rate_limit_per_minute=rate_limit_per_minute
        )
        
        # 并行处理所有数据块
        results = pool.map(process_func, data_chunks)
    
    # 合并所有进程的结果
    all_results = {}
    for result_dict in results:
        for input_id, model_results in result_dict.items():
            if input_id not in all_results:
                all_results[input_id] = {}
            all_results[input_id].update(model_results)
    
    # 构建最终结果并保存
    final_results = {}
    for input_id, model_results in all_results.items():
        if input_id in filtered_data:
            data_copy = filtered_data[input_id].copy()
            data_copy["output_pred"] = {k:v for k,v in model_results.items()}
            final_results[input_id] = data_copy
    
    # 合并已有结果和新结果
    final_data = saved_data.copy()
    final_data.update(final_results)
    
    # 保存结果
    try:
        with open(save_file, 'w', encoding='utf-8') as f:
            json.dump(final_data, f, ensure_ascii=False, indent=4)
        print(f"成功保存结果到: {save_file}")
    except Exception as e:
        print(f"保存结果时出错: {str(e)}")
        # 尝试保存到备份文件
        try:
            with open(save_file + ".backup", 'w', encoding='utf-8') as f:
                json.dump(final_results, f, ensure_ascii=False, indent=4)
            print(f"结果已保存到备份文件: {save_file}.backup")
        except Exception as backup_e:
            print(f"备份保存也失败: {str(backup_e)}")
    
    # 打印统计信息
    elapsed_time = time.time() - start_time
    print(f"\n处理完成!")
    print(f"处理的数据项: {len(filtered_data)}")
    print(f"成功处理的数据项: {len(final_results)}")
    print(f"总耗时: {elapsed_time/60:.2f} 分钟")

def split_data_into_chunks(data_dict, num_chunks):
    """将字典数据分成指定数量的块"""
    items = list(data_dict.items())
    chunk_size = max(1, len(items) // num_chunks)
    return [dict(items[i:i+chunk_size]) for i in range(0, len(items), chunk_size)]

# 修改主函数调用
def run_check_and_review(origin_data_json, save_file, check_model_list, review_model=None,
                        mode=None, rate_limit_per_minute=3):
    print("check_model_list: {}".format(check_model_list))
    print("review_model: {}".format(review_model))
    print("mode: {}".format(mode))
    print("save_file: {}".format(save_file))
    print("rate_limit_per_minute: {}".format(rate_limit_per_minute))

    if origin_data_json is not None:
        valid_app_id = get_available_app_ids(APP_ID_LIST)
        
        print(f"Using {multiprocessing.cpu_count()} processes with {len(valid_app_id)} available APP_IDs")
        print("Starting hybrid parallel processing...")

        try:
            run_comments_and_verify_parallel_hybrid(
                origin_data=origin_data_json,
                model_list=check_model_list,
                valid_app_id=valid_app_id,
                save_file=save_file,
                mode=mode,
                review_model=review_model,
                rate_limit_per_minute=rate_limit_per_minute
            )
            print("Processing completed successfully!")
        except Exception as e:
            print(f"Processing failed with error: {str(e)}")