llm_service.py 3.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788
  1. import requests
  2. import json
  3. import re
  4. from app.core.config import settings
  5. from app.core.logger import get_logger
  6. from app.schemas.models import CardInfoOutput
  7. logger = get_logger("LLMService")
  8. class LLMService:
  9. def split_text_with_overlap(self, text, chunk_size, overlap):
  10. start = 0
  11. text_len = len(text)
  12. while start < text_len:
  13. end = min(start + chunk_size, text_len)
  14. yield text[start:end]
  15. if end == text_len:
  16. break
  17. start += (chunk_size - overlap)
  18. def parse_ai_json(self, json_str):
  19. if not json_str: return []
  20. cleaned_str = re.sub(r"```json\s*", "", json_str)
  21. cleaned_str = re.sub(r"```\s*", "", cleaned_str).strip()
  22. try:
  23. return json.loads(cleaned_str)
  24. except json.JSONDecodeError:
  25. logger.warning(f"⚠️ 解析 JSON 失败,AI 返回原始内容: \n{cleaned_str}\n")
  26. return []
  27. def send_to_ai(self, text_chunk):
  28. headers = {
  29. "Authorization": settings.API_KEY,
  30. "Content-Type": "application/json"
  31. }
  32. payload = {
  33. "inputs": {"question": text_chunk},
  34. "response_mode": "blocking",
  35. "user": settings.USER_ID
  36. }
  37. try:
  38. response = requests.post(settings.API_URL, headers=headers, json=payload)
  39. response.raise_for_status()
  40. return response.json().get('data', {}).get('outputs', {}).get('result', "[]")
  41. except Exception as e:
  42. logger.error(f"❌ 请求发生错误: {e}")
  43. return "[]"
  44. def process_text(self, full_text: str) -> list[CardInfoOutput]:
  45. logger.info(f"📄 开始处理文本,原文总长度: {len(full_text)} 字符")
  46. chunks = list(self.split_text_with_overlap(full_text, settings.CHUNK_SIZE, settings.OVERLAP_SIZE))
  47. total_chunks = len(chunks)
  48. logger.info(
  49. f"✂️ 将分割为 {total_chunks} 个片段进行处理 (Size: {settings.CHUNK_SIZE}, Overlap: {settings.OVERLAP_SIZE})")
  50. all_hits = []
  51. for i, chunk in enumerate(chunks):
  52. logger.info(f"⏳ [进度 {i + 1}/{total_chunks}] 正在发送请求,片段长度 {len(chunk)}...")
  53. ai_resp = self.send_to_ai(chunk)
  54. chunk_hits = self.parse_ai_json(ai_resp)
  55. if chunk_hits:
  56. logger.info(f" ✅ 第 {i + 1} 段识别到 {len(chunk_hits)} 张卡片")
  57. all_hits.extend(chunk_hits)
  58. else:
  59. logger.info(f" ⚪ 第 {i + 1} 段未发现目标")
  60. # 去重
  61. seen = set()
  62. unique_results = []
  63. for hit in all_hits:
  64. try:
  65. # 转换为输出模型
  66. card = CardInfoOutput(**hit)
  67. # key: 时间 + 英文名
  68. key = (card.time, card.card_name_en)
  69. if key not in seen:
  70. seen.add(key)
  71. unique_results.append(card)
  72. except Exception as e:
  73. logger.error(f"⚠️ 数据校验失败: {e}")
  74. logger.info(f"🎉 处理完成! 原始识别 {len(all_hits)} 条,去重后 {len(unique_results)} 条")
  75. return unique_results