aiRouter.py 7.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194
  1. import os, json, time, asyncio
  2. from base_config import ai_key, path
  3. from fastapi import APIRouter, BackgroundTasks
  4. from pathlib import Path
  5. from pydantic import BaseModel
  6. from git import Repo
  7. from http import HTTPStatus
  8. from dashscope import Application
  9. from models.aiModels import Scan_Tasks, Commit_Summary_Tasks
  10. from models.gitModels import Repos
  11. airouter = APIRouter()
  12. class RequestBody(BaseModel):
  13. uuid: str
  14. repo_url: str
  15. def generate_repo_path(uuid, repo_url):
  16. repo_name = repo_url.split("/")[-1].replace(".git", "")
  17. base_path = os.path.join(path, uuid)
  18. return os.path.join(base_path, repo_name), repo_name
  19. async def commit_summary(content):
  20. response = Application.call(
  21. # 若没有配置环境变量,可用百炼API Key将下行替换为:api_key="sk-xxx"。但不建议在生产环境中直接将API Key硬编码到代码中,以减少API Key泄露风险。
  22. api_key=ai_key,
  23. app_id='31a822df773248b19c3c9779f41a5507',
  24. prompt=content)
  25. if response.status_code == HTTPStatus.OK:
  26. try:
  27. json_data = json.loads(response.output.text)
  28. print(json_data)
  29. except json.JSONDecodeError:
  30. print("返回内容不是有效的 JSON 格式!")
  31. json_data = {"null": []}
  32. else:
  33. print(f"请求失败: {response.message}")
  34. json_data = {"null": []}
  35. return json_data
  36. def filter_code_files(prompt):
  37. response = Application.call(
  38. # 若没有配置环境变量,可用百炼API Key将下行替换为:api_key="sk-xxx"。但不建议在生产环境中直接将API Key硬编码到代码中,以减少API Key泄露风险。
  39. api_key=ai_key,
  40. app_id='c1a6dbb6d2314e469bfcbe44c2fe0a5f',
  41. prompt=prompt)
  42. if response.status_code == HTTPStatus.OK:
  43. try:
  44. json_data = json.loads(response.output.text)
  45. print(json_data)
  46. except json.JSONDecodeError:
  47. print("返回内容不是有效的 JSON 格式!")
  48. json_data={"files":[]}
  49. else:
  50. print(f"请求失败: {response.message}")
  51. json_data = {"files": []}
  52. return json_data
  53. def analysis_results(local_path,path):
  54. prompt=""
  55. file_path=os.path.join(local_path,path)
  56. with open(file_path, 'r',encoding="utf8") as f:
  57. for line_num, line in enumerate(f, start=1):
  58. prompt+=f"{line_num}\t{line}"
  59. response = Application.call(
  60. # 若没有配置环境变量,可用百炼API Key将下行替换为:api_key="sk-xxx"。但不建议在生产环境中直接将API Key硬编码到代码中,以减少API Key泄露风险。
  61. api_key=ai_key,
  62. app_id='2f288f146e2d492abb3fe22695e70635', # 替换为实际的应用 ID
  63. prompt=prompt)
  64. if response.status_code == HTTPStatus.OK:
  65. try:
  66. json_data = json.loads(response.output.text)
  67. except json.JSONDecodeError:
  68. print("返回内容不是有效的 JSON 格式!")
  69. print(response.output.text)
  70. json_data={"summary":None}
  71. else:
  72. print(f"请求失败: {response.message}")
  73. json_data = {"summary":None}
  74. json_data["path"]=file_path
  75. return json_data
  76. async def get_filtered_files(folder_path):
  77. base_path = Path(folder_path).resolve()
  78. if not base_path.is_dir():
  79. raise ValueError("无效的目录路径")
  80. file_list = []
  81. for root, dirs, files in os.walk(base_path):
  82. dirs[:] = [d for d in dirs if not d.startswith('.')]
  83. files = [f for f in files if not f.startswith('.')]
  84. for file in files:
  85. abs_path = Path(root) / file
  86. rel_path = abs_path.relative_to(base_path)
  87. file_list.append(str(rel_path))
  88. return file_list
  89. async def process_batch1(batch_files):
  90. """多线程处理单个文件批次的函数"""
  91. try:
  92. js = filter_code_files(str(batch_files))
  93. return js["files"]
  94. except Exception as e:
  95. print(f"处理批次时出错: {e}")
  96. return []
  97. async def get_code_files(path):
  98. file_list = []
  99. files = await get_filtered_files(path)
  100. print(files)
  101. print(f"找到 {len(files)} 个文件")
  102. # 将文件列表分块(每500个一组)
  103. chunks = [files[i * 500: (i + 1) * 500]
  104. for i in range(0, len(files) // 500 + 1)]
  105. # 提交所有批次任务
  106. # futures = [executor.submit(process_batch1, chunk) for chunk in chunks]
  107. tasks = [process_batch1(chunk) for chunk in chunks]
  108. futures = await asyncio.gather(*tasks, return_exceptions=True)
  109. # 实时获取已完成任务的结果
  110. for future in futures[0]:
  111. if isinstance(future, Exception):
  112. print(f"处理出错: {future}")
  113. else:
  114. file_list.append(future)
  115. return file_list
  116. async def process_batch2(local_path,path):
  117. """多线程处理单个文件批次的函数"""
  118. try:
  119. # print(local_path, path)
  120. js = analysis_results(local_path,path)
  121. return js
  122. except Exception as e:
  123. print(11111)
  124. print(f"处理批次时出错: {e}")
  125. return {"summary":None}
  126. async def analysis(local_path, repo_id):
  127. file_list = await get_code_files(local_path)
  128. print(file_list)
  129. results = []
  130. tasks = [process_batch2(local_path, file) for file in file_list] # 假设process_batch2已改为异步函数
  131. batch_results = await asyncio.gather(*tasks, return_exceptions=True)
  132. for result in batch_results:
  133. if isinstance(result, Exception):
  134. print(f"处理出错: {result}")
  135. await write_to_db({"results": results}, repo_id, 3)
  136. else:
  137. results.append(result)
  138. await write_to_db({"results": results}, repo_id, 2)
  139. print("扫描完成")
  140. async def write_to_db(results_dict,repo_id,state):
  141. await Scan_Tasks.filter(repo_id=repo_id).update(state=state, result=results_dict,scan_end_time=int(time.time()))
  142. async def commit_task(content,repo_id):
  143. commit_result=await asyncio.gather(commit_summary(content), return_exceptions=True)
  144. if isinstance(commit_result, Exception):
  145. print(f"提交出错: {commit_result}")
  146. else:
  147. print("提交成功")
  148. await Commit_Summary_Tasks.filter(repo_id=repo_id).update(result=commit_result,end_time=int(time.time()))
  149. @airouter.post("/scan")
  150. async def scan(request: RequestBody, background_tasks: BackgroundTasks):
  151. local_path, repo_name = generate_repo_path(request.uuid, request.repo_url)
  152. repo = await Repos.get(name=repo_name)
  153. repo_id = repo.id
  154. print(f"开始扫描仓库: {repo_name}")
  155. await Scan_Tasks.filter(repo_id=repo_id).update(state=1, scan_start_time=int(time.time()))
  156. background_tasks.add_task(analysis, local_path, repo_id)
  157. return {"code": 200, "meg": "添加扫描任务成功"}
  158. @airouter.post("/summary")
  159. async def summary(request: RequestBody, background_tasks: BackgroundTasks):
  160. local_path, repo_name = generate_repo_path(request.uuid, request.repo_url)
  161. repo=await Repos.get(name=repo_name)
  162. repo_id=repo.id
  163. await Commit_Summary_Tasks.filter(repo_id=repo_id).update(start_time=int(time.time()))
  164. commit_content = Repo(local_path).git.log('-1', '-p', '--pretty=format:%h %s')
  165. background_tasks.add_task(commit_task,commit_content,repo_id)
  166. return {"code": 200, "meg": "添加提交任务成功"}