summary.py 3.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394
  1. import requests
  2. from config import conf
  3. from common.log import logger
  4. import os
  5. class LinkSummary:
  6. def __init__(self):
  7. pass
  8. def summary_file(self, file_path: str):
  9. file_body = {
  10. "file": open(file_path, "rb"),
  11. "name": file_path.split("/")[-1],
  12. }
  13. url = self.base_url() + "/v1/summary/file"
  14. res = requests.post(url, headers=self.headers(), files=file_body, timeout=(5, 300))
  15. return self._parse_summary_res(res)
  16. def summary_url(self, url: str):
  17. body = {
  18. "url": url
  19. }
  20. res = requests.post(url=self.base_url() + "/v1/summary/url", headers=self.headers(), json=body, timeout=(5, 180))
  21. return self._parse_summary_res(res)
  22. def summary_chat(self, summary_id: str):
  23. body = {
  24. "summary_id": summary_id
  25. }
  26. res = requests.post(url=self.base_url() + "/v1/summary/chat", headers=self.headers(), json=body, timeout=(5, 180))
  27. if res.status_code == 200:
  28. res = res.json()
  29. logger.debug(f"[LinkSum] chat open, res={res}")
  30. if res.get("code") == 200:
  31. data = res.get("data")
  32. return {
  33. "questions": data.get("questions"),
  34. "file_id": data.get("file_id")
  35. }
  36. else:
  37. res_json = res.json()
  38. logger.error(f"[LinkSum] summary error, status_code={res.status_code}, msg={res_json.get('message')}")
  39. return None
  40. def _parse_summary_res(self, res):
  41. if res.status_code == 200:
  42. res = res.json()
  43. logger.debug(f"[LinkSum] url summary, res={res}")
  44. if res.get("code") == 200:
  45. data = res.get("data")
  46. return {
  47. "summary": data.get("summary"),
  48. "summary_id": data.get("summary_id")
  49. }
  50. else:
  51. res_json = res.json()
  52. logger.error(f"[LinkSum] summary error, status_code={res.status_code}, msg={res_json.get('message')}")
  53. return None
  54. def base_url(self):
  55. return conf().get("linkai_api_base", "https://api.link-ai.chat")
  56. def headers(self):
  57. return {"Authorization": "Bearer " + conf().get("linkai_api_key")}
  58. def check_file(self, file_path: str, sum_config: dict) -> bool:
  59. file_size = os.path.getsize(file_path) // 1000
  60. if (sum_config.get("max_file_size") and file_size > sum_config.get("max_file_size")) or file_size > 15000:
  61. logger.warn(f"[LinkSum] file size exceeds limit, No processing, file_size={file_size}KB")
  62. return False
  63. suffix = file_path.split(".")[-1]
  64. support_list = ["txt", "csv", "docx", "pdf", "md", "jpg", "jpeg", "png"]
  65. if suffix not in support_list:
  66. logger.warn(f"[LinkSum] unsupported file, suffix={suffix}, support_list={support_list}")
  67. return False
  68. return True
  69. def check_url(self, url: str):
  70. if not url:
  71. return False
  72. support_list = ["http://mp.weixin.qq.com", "https://mp.weixin.qq.com"]
  73. black_support_list = ["https://mp.weixin.qq.com/mp/waerrpage"]
  74. for black_url_prefix in black_support_list:
  75. if url.strip().startswith(black_url_prefix):
  76. logger.warn(f"[LinkSum] unsupported url, no need to process, url={url}")
  77. return False
  78. for support_url in support_list:
  79. if url.strip().startswith(support_url):
  80. return True
  81. return False