jarvis-ai-assistant 0.7.16__py3-none-any.whl → 1.0.2__py3-none-any.whl

jarvis/jarvis_sec/review.py

@@ -1,18 +1,27 @@
 # -*- coding: utf-8 -*-
 """复核相关模块"""
 
-from typing import Dict, List, Optional
-import typer
+from typing import Dict
+from typing import List
+from typing import Optional
 
-from jarvis.jarvis_sec.prompts import get_review_system_prompt, get_review_summary_prompt, build_verification_summary_prompt
+from jarvis.jarvis_utils.output import PrettyOutput
+
+from jarvis.jarvis_sec.agents import create_review_agent
+from jarvis.jarvis_sec.agents import subscribe_summary_event
 from jarvis.jarvis_sec.parsers import try_parse_summary_report
-from jarvis.jarvis_sec.agents import create_review_agent, subscribe_summary_event
+from jarvis.jarvis_sec.prompts import build_verification_summary_prompt
+from jarvis.jarvis_sec.prompts import get_review_summary_prompt
+from jarvis.jarvis_sec.prompts import get_review_system_prompt
 from jarvis.jarvis_sec.utils import git_restore_if_dirty
 
 
-def build_review_task(review_batch: List[Dict], entry_path: str, langs: List[str]) -> str:
+def build_review_task(
+    review_batch: List[Dict], entry_path: str, langs: List[str]
+) -> str:
     """构建复核任务上下文"""
     import json as _json_review
+
     return f"""
 # 复核无效聚类任务
 上下文参数：
@@ -56,9 +65,6 @@ def build_gid_to_review_mapping(review_results: List[Dict]) -> Dict[int, Dict]:
     """构建gid到复核结果的映射（支持 gid 和 gids 两种格式）"""
     gid_to_review: Dict[int, Dict] = {}
     for rr in review_results:
-        if not isinstance(rr, dict):
-            continue
-        
         # 支持 gid 和 gids 两种格式
         gids_to_process: List[int] = []
         if "gids" in rr and isinstance(rr.get("gids"), list):
@@ -78,14 +84,14 @@ def build_gid_to_review_mapping(review_results: List[Dict]) -> Dict[int, Dict]:
                     gids_to_process.append(gid_int)
             except Exception:
                 pass
-        
+
         # 为每个 gid 创建复核结果映射
         is_reason_sufficient = rr.get("is_reason_sufficient")
         review_notes = str(rr.get("review_notes", "")).strip()
         for gid in gids_to_process:
             gid_to_review[gid] = {
                 "is_reason_sufficient": is_reason_sufficient,
-                "review_notes": review_notes
+                "review_notes": review_notes,
             }
     return gid_to_review
 
@@ -101,12 +107,12 @@ def process_review_batch(
     if review_results:
         # 构建gid到复核结果的映射
         gid_to_review = build_gid_to_review_mapping(review_results)
-        
+
         # 处理每个无效聚类
         for invalid_cluster in review_batch:
             cluster_gids = invalid_cluster.get("gids", [])
             cluster_members = invalid_cluster.get("members", [])
-            
+
             # 检查该聚类中的所有gid的复核结果
             all_sufficient = True
             any_reviewed = False
@@ -120,41 +126,56 @@ def process_review_batch(
                         if not insufficient_review_result:
                             insufficient_review_result = review_result
                         break
-            
+
             if any_reviewed and not all_sufficient:
                 # 理由不充分，重新加入验证流程
-                typer.secho(f"[jarvis-sec] 复核结果：无效聚类（gids={cluster_gids}）理由不充分，重新加入验证流程", fg=typer.colors.BLUE)
+                PrettyOutput.auto_print(
+                    f"🔵 [jarvis-sec] 复核结果：无效聚类（gids={cluster_gids}）理由不充分，重新加入验证流程"
+                )
                 for member in cluster_members:
                     reinstated_candidates.append(member)
-                reviewed_clusters.append({
-                    **invalid_cluster,
-                    "review_result": "reinstated",
-                    "review_notes": insufficient_review_result.get("review_notes", "") if insufficient_review_result else "",
-                })
+                reviewed_clusters.append(
+                    {
+                        **invalid_cluster,
+                        "review_result": "reinstated",
+                        "review_notes": insufficient_review_result.get(
+                            "review_notes", ""
+                        )
+                        if insufficient_review_result
+                        else "",
+                    }
+                )
             else:
                 # 理由充分，确认无效
                 review_notes = ""
                 if cluster_gids and gid_to_review.get(cluster_gids[0]):
-                    review_notes = gid_to_review[cluster_gids[0]].get("review_notes", "")
-                typer.secho(f"[jarvis-sec] 复核结果：无效聚类（gids={cluster_gids}）理由充分，确认为无效", fg=typer.colors.GREEN)
-                reviewed_clusters.append({
-                    **invalid_cluster,
-                    "review_result": "confirmed_invalid",
-                    "review_notes": review_notes,
-                })
-                
+                    review_notes = gid_to_review[cluster_gids[0]].get(
+                        "review_notes", ""
+                    )
+                PrettyOutput.auto_print(
+                    f"✅ [jarvis-sec] 复核结果：无效聚类（gids={cluster_gids}）理由充分，确认为无效"
+                )
+                reviewed_clusters.append(
+                    {
+                        **invalid_cluster,
+                        "review_result": "confirmed_invalid",
+                        "review_notes": review_notes,
+                    }
+                )
+
                 # 将确认无效的gids保存到analysis.jsonl的false_positive_gids中，以便断点恢复时能正确识别已复核的无效聚类
                 if sec_dir:
                     try:
+                        from datetime import datetime
                         from pathlib import Path
+
                         from jarvis.jarvis_sec.file_manager import save_analysis_result
-                        from datetime import datetime
-                        
+
                         # 构建cluster_id
                         file_name = invalid_cluster.get("file", "")
                         batch_index = invalid_cluster.get("batch_index", 1)
                         cluster_id = f"{file_name}|{batch_index}|review"
-                        
+
                         # 将gids转换为整数列表
                         false_positive_gids = []
                         for gid_val in cluster_gids:
@@ -164,7 +185,7 @@ def process_review_batch(
                                     false_positive_gids.append(gid_int)
                             except Exception:
                                 pass
-                        
+
                         # 保存分析结果
                         if false_positive_gids:
                             analysis_result = {
@@ -181,21 +202,27 @@ def process_review_batch(
                     except Exception as e:
                         # 保存失败不影响主流程，只记录警告
                         try:
-                            typer.secho(f"[jarvis-sec] 警告：保存复核结果到analysis.jsonl失败: {str(e)}", fg=typer.colors.YELLOW)
+                            PrettyOutput.auto_print(
+                                f"⚠️ [jarvis-sec] 警告：保存复核结果到analysis.jsonl失败: {str(e)}"
+                            )
                         except Exception:
                             pass
     else:
         # 复核结果解析失败，保守策略：重新加入验证流程
-        typer.secho("[jarvis-sec] 警告：复核结果解析失败，保守策略：将批次中的所有候选重新加入验证流程", fg=typer.colors.YELLOW)
+        PrettyOutput.auto_print(
+            "⚠️ [jarvis-sec] 警告：复核结果解析失败，保守策略：将批次中的所有候选重新加入验证流程"
+        )
         for invalid_cluster in review_batch:
             cluster_members = invalid_cluster.get("members", [])
             for member in cluster_members:
                 reinstated_candidates.append(member)
-            reviewed_clusters.append({
-                **invalid_cluster,
-                "review_result": "reinstated",
-                "review_notes": "复核结果解析失败，保守策略重新加入验证",
-            })
+            reviewed_clusters.append(
+                {
+                    **invalid_cluster,
+                    "review_result": "reinstated",
+                    "review_notes": "复核结果解析失败，保守策略重新加入验证",
+                }
+            )
 
 
 def process_review_batch_items(
@@ -222,27 +249,31 @@ def reinstated_candidates_to_cluster_batches(
 ) -> None:
     """将重新加入的候选添加到cluster_batches"""
     from collections import defaultdict as _dd2
-    
+
     if not reinstated_candidates:
         return
-    
-    typer.secho(f"[jarvis-sec] 复核完成：{len(reinstated_candidates)} 个候选重新加入验证流程", fg=typer.colors.GREEN)
+
+    PrettyOutput.auto_print(
+        f"🔄 [jarvis-sec] 复核完成：{len(reinstated_candidates)} 个候选重新加入验证流程"
+    )
     # 按文件分组重新加入的候选
     reinstated_by_file: Dict[str, List[Dict]] = _dd2(list)
     for cand in reinstated_candidates:
         file_key = str(cand.get("file") or "")
         reinstated_by_file[file_key].append(cand)
-    
+
     # 为每个文件的重新加入候选创建批次
     for file_key, cands in reinstated_by_file.items():
         if cands:
             cluster_batches.append(cands)
-            _progress_append({
-                "event": "review_reinstated",
-                "file": file_key,
-                "gids": [c.get("gid") for c in cands],
-                "count": len(cands),
-            })
+            _progress_append(
+                {
+                    "event": "review_reinstated",
+                    "file": file_key,
+                    "gids": [c.get("gid") for c in cands],
+                    "count": len(cands),
+                }
+            )
 
 
 def run_review_agent_with_retry(
@@ -256,70 +287,93 @@ def run_review_agent_with_retry(
     use_direct_model_review = False
     prev_parse_error_review: Optional[str] = None
     review_attempt = 0
-    
+
     while True:
         review_attempt += 1
         review_summary_container["text"] = ""
-        
+
         if use_direct_model_review:
             # 格式校验失败后，直接调用模型接口
             review_summary_prompt_text = build_verification_summary_prompt()
             error_guidance = ""
             if prev_parse_error_review:
                 error_guidance = f"\n\n**格式错误详情（请根据以下错误修复输出格式）：**\n- JSON解析失败: {prev_parse_error_review}\n\n请确保输出的JSON格式正确，包括正确的引号、逗号、大括号等。仅输出一个 <REPORT> 块，块内直接包含 JSON 数组（不需要额外的标签）。支持jsonnet语法（如尾随逗号、注释、||| 或 ``` 分隔符多行字符串等）。"
-            
-            full_review_prompt = f"{review_task}{error_guidance}\n\n{review_summary_prompt_text}"
+
+            full_review_prompt = (
+                f"{review_task}{error_guidance}\n\n{review_summary_prompt_text}"
+            )
             try:
-                review_response = review_agent.model.chat_until_success(full_review_prompt)  # type: ignore
+                review_response = review_agent.model.chat_until_success(
+                    full_review_prompt
+                )
                 review_summary_container["text"] = review_response
             except Exception as e:
                 try:
-                    typer.secho(f"[jarvis-sec] 复核阶段直接模型调用失败: {e}，回退到 run()", fg=typer.colors.YELLOW)
+                    PrettyOutput.auto_print(
+                        f"✨ [jarvis-sec] 复核阶段直接模型调用失败: {e}，回退到 run()",
+                        timestamp=True,
+                    )
                 except Exception:
                     pass
                 review_agent.run(review_task)
         else:
             review_agent.run(review_task)
-        
+
         # 工作区保护
         try:
             _changed_review = git_restore_if_dirty(entry_path)
             if _changed_review:
                 try:
-                    typer.secho(f"[jarvis-sec] 复核 Agent 工作区已恢复 ({_changed_review} 个文件）", fg=typer.colors.BLUE)
+                    PrettyOutput.auto_print(
+                        f"✨ [jarvis-sec] 复核 Agent 工作区已恢复 ({_changed_review} 个文件）",
+                        timestamp=True,
+                    )
                 except Exception:
                     pass
         except Exception:
             pass
-        
+
         # 解析复核结果
         review_summary_text = review_summary_container.get("text", "")
         parse_error_review = None
         if review_summary_text:
-            review_parsed, parse_error_review = try_parse_summary_report(review_summary_text)
+            review_parsed, parse_error_review = try_parse_summary_report(
+                review_summary_text
+            )
             if parse_error_review:
                 prev_parse_error_review = parse_error_review
                 try:
-                    typer.secho(f"[jarvis-sec] 复核结果JSON解析失败: {parse_error_review}", fg=typer.colors.YELLOW)
+                    PrettyOutput.auto_print(
+                        f"✨ [jarvis-sec] 复核结果JSON解析失败: {parse_error_review}",
+                        timestamp=True,
+                    )
                 except Exception:
                     pass
             else:
                 prev_parse_error_review = None
                 if isinstance(review_parsed, list):
                     # 验证复核结果格式
-                    if review_parsed and all(is_valid_review_item(item) for item in review_parsed):
+                    if review_parsed and all(
+                        is_valid_review_item(item) for item in review_parsed
+                    ):
                         return review_parsed, None
-        
+
         # 格式校验失败，后续重试使用直接模型调用
         use_direct_model_review = True
         if parse_error_review:
             try:
-                typer.secho(f"[jarvis-sec] 复核结果JSON解析失败 -> 重试第 {review_attempt} 次 (使用直接模型调用，将反馈解析错误)", fg=typer.colors.YELLOW)
+                PrettyOutput.auto_print(
+                    f"✨ [jarvis-sec] 复核结果JSON解析失败 -> 重试第 {review_attempt} 次 (使用直接模型调用，将反馈解析错误)",
+                    timestamp=True,
+                )
             except Exception:
                 pass
         else:
             try:
-                typer.secho(f"[jarvis-sec] 复核结果格式无效 -> 重试第 {review_attempt} 次 (使用直接模型调用)", fg=typer.colors.YELLOW)
+                PrettyOutput.auto_print(
+                    f"✨ [jarvis-sec] 复核结果格式无效 -> 重试第 {review_attempt} 次 (使用直接模型调用)",
+                    timestamp=True,
+                )
             except Exception:
                 pass
 
@@ -336,49 +390,59 @@ def process_review_phase(
 ) -> List[List[Dict]]:
     """
     处理复核阶段：验证所有标记为无效的聚类。
-    
+
     返回: 更新后的 cluster_batches（包含重新加入验证的候选）
     """
     if not invalid_clusters_for_review:
-        typer.secho("[jarvis-sec] 无无效聚类需要复核", fg=typer.colors.BLUE)
+        PrettyOutput.auto_print("🔵 [jarvis-sec] 无无效聚类需要复核")
         return cluster_batches
-    
-    typer.secho(f"\n[jarvis-sec] 开始复核 {len(invalid_clusters_for_review)} 个无效聚类...", fg=typer.colors.MAGENTA)
+
+    PrettyOutput.auto_print(
+        f"✨ \n[jarvis-sec] 开始复核 {len(invalid_clusters_for_review)} 个无效聚类...",
+        timestamp=True,
+    )
     status_mgr.update_review(
         current_review=0,
         total_reviews=len(invalid_clusters_for_review),
-        message="开始复核无效聚类..."
+        message="开始复核无效聚类...",
     )
-    
+
     # 按批次复核（每批最多10个无效聚类，避免上下文过长）
     review_batch_size = 10
     reviewed_clusters: List[Dict] = []
     reinstated_candidates: List[Dict] = []  # 重新加入验证的候选
-    
+
     get_review_system_prompt()
     review_summary_prompt = get_review_summary_prompt()
-    
+
     for review_idx in range(0, len(invalid_clusters_for_review), review_batch_size):
-        review_batch = invalid_clusters_for_review[review_idx:review_idx + review_batch_size]
+        review_batch = invalid_clusters_for_review[
+            review_idx : review_idx + review_batch_size
+        ]
         current_review_num = review_idx // review_batch_size + 1
-        total_review_batches = (len(invalid_clusters_for_review) + review_batch_size - 1) // review_batch_size
-        
-        typer.secho(f"[jarvis-sec] 复核批次 {current_review_num}/{total_review_batches}: {len(review_batch)} 个无效聚类", fg=typer.colors.CYAN)
+        total_review_batches = (
+            len(invalid_clusters_for_review) + review_batch_size - 1
+        ) // review_batch_size
+
+        PrettyOutput.auto_print(
+            f"✨ [jarvis-sec] 复核批次 {current_review_num}/{total_review_batches}: {len(review_batch)} 个无效聚类",
+            timestamp=True,
+        )
         status_mgr.update_review(
             current_review=current_review_num,
             total_reviews=total_review_batches,
-            message=f"正在复核批次 {current_review_num}/{total_review_batches}"
+            message=f"正在复核批次 {current_review_num}/{total_review_batches}",
         )
-        
+
         # 构建复核任务
         review_task = build_review_task(review_batch, entry_path, langs)
-        
+
         # 创建复核Agent
         review_agent = create_review_agent(current_review_num, llm_group)
-        
+
         # 订阅复核Agent的摘要
         review_summary_container = subscribe_summary_event(review_agent)
-        
+
         # 运行复核Agent（永久重试直到格式正确）
         review_results, parse_error = run_review_agent_with_retry(
             review_agent,
@@ -387,7 +451,7 @@ def process_review_phase(
             entry_path,
             review_summary_container,
         )
-        
+
         # 处理复核结果
         process_review_batch_items(
             review_batch,
@@ -396,36 +460,43 @@ def process_review_phase(
             reinstated_candidates,
             sec_dir,
         )
-        
+
         # 记录每个已复核的无效聚类的 gids（包括确认无效的和重新加入验证的）
         for invalid_cluster in review_batch:
             cluster_gids = invalid_cluster.get("gids", [])
             if cluster_gids:
-                _progress_append({
-                    "event": "review_invalid_cluster",
-                    "gids": cluster_gids,
-                    "file": invalid_cluster.get("file"),
-                    "batch_index": invalid_cluster.get("batch_index"),
-                })
-    
+                _progress_append(
+                    {
+                        "event": "review_invalid_cluster",
+                        "gids": cluster_gids,
+                        "file": invalid_cluster.get("file"),
+                        "batch_index": invalid_cluster.get("batch_index"),
+                    }
+                )
+
     # 将重新加入验证的候选添加到cluster_batches
     reinstated_candidates_to_cluster_batches(
         reinstated_candidates,
         cluster_batches,
         _progress_append,
     )
-    
+
     if not reinstated_candidates:
-        typer.secho("[jarvis-sec] 复核完成：所有无效聚类理由充分，确认为无效", fg=typer.colors.GREEN)
-    
+        PrettyOutput.auto_print(
+            "✨ [jarvis-sec] 复核完成：所有无效聚类理由充分，确认为无效", timestamp=True
+        )
+
     # 记录复核结果（汇总）
-    _progress_append({
-        "event": "review_completed",
-        "total_reviewed": len(invalid_clusters_for_review),
-        "reinstated": len(reinstated_candidates),
-        "confirmed_invalid": len(invalid_clusters_for_review) - len(reinstated_candidates),
-    })
-    
+    _progress_append(
+        {
+            "event": "review_completed",
+            "total_reviewed": len(invalid_clusters_for_review),
+            "reinstated": len(reinstated_candidates),
+            "confirmed_invalid": len(invalid_clusters_for_review)
+            - len(reinstated_candidates),
+        }
+    )
+
     # 记录所有已复核的无效聚类的 gids（用于断点恢复时跳过已复核的聚类）
     all_reviewed_gids = set()
     for invalid_cluster in invalid_clusters_for_review:
@@ -437,17 +508,19 @@ def process_review_phase(
                     all_reviewed_gids.add(gid_int)
             except Exception:
                 pass
-    
+
     if all_reviewed_gids:
-        _progress_append({
-            "event": "review_all_gids",
-            "gids": sorted(list(all_reviewed_gids)),
-            "total": len(all_reviewed_gids),
-        })
+        _progress_append(
+            {
+                "event": "review_all_gids",
+                "gids": sorted(list(all_reviewed_gids)),
+                "total": len(all_reviewed_gids),
+            }
+        )
     status_mgr.update_review(
         current_review=len(invalid_clusters_for_review),
         total_reviews=len(invalid_clusters_for_review),
-        message=f"复核完成：{len(reinstated_candidates)} 个候选重新加入验证"
+        message=f"复核完成：{len(reinstated_candidates)} 个候选重新加入验证",
     )
-    
+
     return cluster_batches