text_search.py 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588
  1. from fastapi import APIRouter, HTTPException, Depends
  2. from pydantic import BaseModel, Field, validator
  3. from typing import List, Optional
  4. from service.trunks_service import TrunksService
  5. from utils.sentence_util import SentenceUtil
  6. from utils.vector_distance import VectorDistance
  7. from model.response import StandardResponse
  8. from utils.vectorizer import Vectorizer
  9. # from utils.find_text_in_pdf import find_text_in_pdf
  10. import os
  11. DISTANCE_THRESHOLD = 0.73
  12. import logging
  13. import time
  14. from db.session import get_db
  15. from sqlalchemy.orm import Session
  16. from service.kg_node_service import KGNodeService
  17. from service.kg_prop_service import KGPropService
  18. from cachetools import TTLCache
  19. logger = logging.getLogger(__name__)
  20. router = APIRouter(tags=["Text Search"])
  21. # 创建全局缓存实例
  22. cache = TTLCache(maxsize=1000, ttl=3600)
  23. class TextSearchRequest(BaseModel):
  24. text: str
  25. conversation_id: Optional[str] = None
  26. need_convert: Optional[bool] = False
  27. class TextCompareRequest(BaseModel):
  28. sentence: str
  29. text: str
  30. class TextMatchRequest(BaseModel):
  31. text: str = Field(..., min_length=1, max_length=10000, description="需要搜索的文本内容")
  32. @validator('text')
  33. def validate_text(cls, v):
  34. # 保留所有可打印字符、换行符和中文字符
  35. v = ''.join(char for char in v if char.isprintable() or char in '\n\r')
  36. # 转义JSON特殊字符
  37. # 先处理反斜杠,避免后续转义时出现问题
  38. v = v.replace('\\', '\\\\')
  39. # 处理引号和其他特殊字符
  40. v = v.replace('"', '\\"')
  41. v = v.replace('/', '\\/')
  42. # 处理控制字符
  43. v = v.replace('\n', '\\n')
  44. v = v.replace('\r', '\\r')
  45. v = v.replace('\t', '\\t')
  46. v = v.replace('\b', '\\b')
  47. v = v.replace('\f', '\\f')
  48. # 处理Unicode转义
  49. # v = v.replace('\u', '\\u')
  50. return v
  51. class TextCompareMultiRequest(BaseModel):
  52. origin: str
  53. similar: str
  54. class NodePropsSearchRequest(BaseModel):
  55. node_id: int
  56. props_ids: List[int]
  57. @router.post("/kgrt_api/text/clear_cache", response_model=StandardResponse)
  58. async def clear_cache():
  59. try:
  60. # 清除全局缓存
  61. cache.clear()
  62. return StandardResponse(success=True, data={"message": "缓存已清除"})
  63. except Exception as e:
  64. logger.error(f"清除缓存失败: {str(e)}")
  65. raise HTTPException(status_code=500, detail=str(e))
  66. @router.post("/kgrt_api/text/search", response_model=StandardResponse)
  67. @router.post("/knowledge/text/search", response_model=StandardResponse)
  68. async def search_text(request: TextSearchRequest):
  69. try:
  70. #判断request.text是否为json格式,如果是,使用JsonToText的convert方法转换为text
  71. if request.text.startswith('{') and request.text.endswith('}'):
  72. from utils.json_to_text import JsonToTextConverter
  73. converter = JsonToTextConverter()
  74. request.text = converter.convert(request.text)
  75. # 使用TextSplitter拆分文本
  76. sentences = SentenceUtil.split_text(request.text)
  77. if not sentences:
  78. return StandardResponse(success=True, data={"answer": "", "references": []})
  79. # 初始化服务和结果列表
  80. trunks_service = TrunksService()
  81. result_sentences = []
  82. all_references = []
  83. reference_index = 1
  84. # 根据conversation_id获取缓存结果
  85. cached_results = trunks_service.get_cached_result(request.conversation_id) if request.conversation_id else []
  86. for sentence in sentences:
  87. # if request.need_convert:
  88. sentence = sentence.replace("\n", "<br>")
  89. if len(sentence) < 10:
  90. result_sentences.append(sentence)
  91. continue
  92. if cached_results:
  93. # 如果有缓存结果,计算向量距离
  94. min_distance = float('inf')
  95. best_result = None
  96. sentence_vector = Vectorizer.get_embedding(sentence)
  97. for cached_result in cached_results:
  98. content_vector = cached_result['embedding']
  99. distance = VectorDistance.calculate_distance(sentence_vector, content_vector)
  100. if distance < min_distance:
  101. min_distance = distance
  102. best_result = {**cached_result, 'distance': distance}
  103. if best_result and best_result['distance'] < DISTANCE_THRESHOLD:
  104. search_results = [best_result]
  105. else:
  106. search_results = []
  107. else:
  108. # 如果没有缓存结果,进行向量搜索
  109. search_results = trunks_service.search_by_vector(
  110. text=sentence,
  111. limit=1,
  112. type='trunk'
  113. )
  114. # 处理搜索结果
  115. for search_result in search_results:
  116. distance = search_result.get("distance", DISTANCE_THRESHOLD)
  117. if distance >= DISTANCE_THRESHOLD:
  118. result_sentences.append(sentence)
  119. continue
  120. # 检查是否已存在相同引用
  121. existing_ref = next((ref for ref in all_references if ref["id"] == search_result["id"]), None)
  122. current_index = reference_index
  123. if existing_ref:
  124. current_index = int(existing_ref["index"])
  125. else:
  126. # 添加到引用列表
  127. # 从referrence中提取文件名
  128. file_name = ""
  129. referrence = search_result.get("referrence", "")
  130. if referrence and "/books/" in referrence:
  131. file_name = referrence.split("/books/")[-1]
  132. # 去除文件扩展名
  133. file_name = os.path.splitext(file_name)[0]
  134. reference = {
  135. "index": str(reference_index),
  136. "id": search_result["id"],
  137. "content": search_result["content"],
  138. "file_path": search_result.get("file_path", ""),
  139. "title": search_result.get("title", ""),
  140. "distance": distance,
  141. "file_name": file_name,
  142. "referrence": referrence
  143. }
  144. all_references.append(reference)
  145. reference_index += 1
  146. # 添加引用标记
  147. if sentence.endswith('<br>'):
  148. # 如果有多个<br>,在所有<br>前添加^[current_index]^
  149. result_sentence = sentence.replace('<br>', f'^[{current_index}]^<br>')
  150. else:
  151. # 直接在句子末尾添加^[current_index]^
  152. result_sentence = f'{sentence}^[{current_index}]^'
  153. result_sentences.append(result_sentence)
  154. # 组装返回数据
  155. response_data = {
  156. "answer": result_sentences,
  157. "references": all_references
  158. }
  159. return StandardResponse(success=True, data=response_data)
  160. except Exception as e:
  161. logger.error(f"Text search failed: {str(e)}")
  162. raise HTTPException(status_code=500, detail=str(e))
  163. @router.post("/kgrt_api/text/match", response_model=StandardResponse)
  164. @router.post("/knowledge/text/match", response_model=StandardResponse)
  165. async def match_text(request: TextCompareRequest):
  166. try:
  167. sentences = SentenceUtil.split_text(request.text)
  168. sentence_vector = Vectorizer.get_embedding(request.sentence)
  169. min_distance = float('inf')
  170. best_sentence = ""
  171. result_sentences = []
  172. for temp in sentences:
  173. result_sentences.append(temp)
  174. if len(temp) < 10:
  175. continue
  176. temp_vector = Vectorizer.get_embedding(temp)
  177. distance = VectorDistance.calculate_distance(sentence_vector, temp_vector)
  178. if distance < min_distance and distance < DISTANCE_THRESHOLD:
  179. min_distance = distance
  180. best_sentence = temp
  181. for i in range(len(result_sentences)):
  182. result_sentences[i] = {"sentence": result_sentences[i], "matched": False}
  183. if result_sentences[i]["sentence"] == best_sentence:
  184. result_sentences[i]["matched"] = True
  185. return StandardResponse(success=True, records=result_sentences)
  186. except Exception as e:
  187. logger.error(f"Text comparison failed: {str(e)}")
  188. raise HTTPException(status_code=500, detail=str(e))
  189. @router.post("/kgrt_api/text/mr_search", response_model=StandardResponse)
  190. @router.post("/knowledge/text/mr_search", response_model=StandardResponse)
  191. async def mr_search_text_content(request: TextMatchRequest):
  192. try:
  193. # 初始化服务
  194. trunks_service = TrunksService()
  195. # 获取文本向量并搜索相似内容
  196. search_results = trunks_service.search_by_vector(
  197. text=request.text,
  198. limit=10,
  199. type="mr"
  200. )
  201. # 处理搜索结果
  202. records = []
  203. for result in search_results:
  204. distance = result.get("distance", DISTANCE_THRESHOLD)
  205. if distance >= DISTANCE_THRESHOLD:
  206. continue
  207. # 添加到引用列表
  208. record = {
  209. "content": result["content"],
  210. "file_path": result.get("file_path", ""),
  211. "title": result.get("title", ""),
  212. "distance": distance,
  213. }
  214. records.append(record)
  215. # 组装返回数据
  216. response_data = {
  217. "records": records
  218. }
  219. return StandardResponse(success=True, data=response_data)
  220. except Exception as e:
  221. logger.error(f"Mr search failed: {str(e)}")
  222. raise HTTPException(status_code=500, detail=str(e))
  223. @router.post("/kgrt_api/text/mr_match", response_model=StandardResponse)
  224. @router.post("/knowledge/text/mr_match", response_model=StandardResponse)
  225. async def compare_text(request: TextCompareMultiRequest):
  226. start_time = time.time()
  227. try:
  228. # 拆分两段文本
  229. origin_sentences = SentenceUtil.split_text(request.origin)
  230. similar_sentences = SentenceUtil.split_text(request.similar)
  231. end_time = time.time()
  232. logger.info(f"mr_match接口处理文本耗时: {(end_time - start_time) * 1000:.2f}ms")
  233. # 初始化结果列表
  234. origin_results = []
  235. # 过滤短句并预计算向量
  236. valid_origin_sentences = [(sent, len(sent) >= 10) for sent in origin_sentences]
  237. valid_similar_sentences = [(sent, len(sent) >= 10) for sent in similar_sentences]
  238. # 初始化similar_results,所有matched设为False
  239. similar_results = [{"sentence": sent, "matched": False} for sent, _ in valid_similar_sentences]
  240. # 批量获取向量
  241. origin_vectors = {}
  242. similar_vectors = {}
  243. origin_batch = [sent for sent, is_valid in valid_origin_sentences if is_valid]
  244. similar_batch = [sent for sent, is_valid in valid_similar_sentences if is_valid]
  245. if origin_batch:
  246. origin_embeddings = [Vectorizer.get_embedding(sent) for sent in origin_batch]
  247. origin_vectors = dict(zip(origin_batch, origin_embeddings))
  248. if similar_batch:
  249. similar_embeddings = [Vectorizer.get_embedding(sent) for sent in similar_batch]
  250. similar_vectors = dict(zip(similar_batch, similar_embeddings))
  251. end_time = time.time()
  252. logger.info(f"mr_match接口处理向量耗时: {(end_time - start_time) * 1000:.2f}ms")
  253. # 处理origin文本
  254. for origin_sent, is_valid in valid_origin_sentences:
  255. if not is_valid:
  256. origin_results.append({"sentence": origin_sent, "matched": False})
  257. continue
  258. origin_vector = origin_vectors[origin_sent]
  259. matched = False
  260. # 优化的相似度计算
  261. for i, similar_result in enumerate(similar_results):
  262. if similar_result["matched"]:
  263. continue
  264. similar_sent = similar_result["sentence"]
  265. if len(similar_sent) < 10:
  266. continue
  267. similar_vector = similar_vectors.get(similar_sent)
  268. if not similar_vector:
  269. continue
  270. distance = VectorDistance.calculate_distance(origin_vector, similar_vector)
  271. if distance < DISTANCE_THRESHOLD:
  272. matched = True
  273. similar_results[i]["matched"] = True
  274. break
  275. origin_results.append({"sentence": origin_sent, "matched": matched})
  276. response_data = {
  277. "origin": origin_results,
  278. "similar": similar_results
  279. }
  280. end_time = time.time()
  281. logger.info(f"mr_match接口耗时: {(end_time - start_time) * 1000:.2f}ms")
  282. return StandardResponse(success=True, data=response_data)
  283. except Exception as e:
  284. end_time = time.time()
  285. logger.error(f"Text comparison failed: {str(e)}")
  286. logger.info(f"mr_match接口耗时: {(end_time - start_time) * 1000:.2f}ms")
  287. raise HTTPException(status_code=500, detail=str(e))
  288. def _check_cache(node_id: int) -> Optional[dict]:
  289. """检查并返回缓存结果"""
  290. cache_key = f"xunzheng_{node_id}"
  291. cached_result = cache.get(cache_key)
  292. if cached_result:
  293. logger.info(f"从缓存获取结果,node_id: {node_id}")
  294. return cached_result
  295. return None
  296. def _get_node_info(node_service: KGNodeService, node_id: int) -> dict:
  297. """获取并验证节点信息"""
  298. node = node_service.get_node(node_id)
  299. if not node:
  300. raise ValueError(f"节点不存在: {node_id}")
  301. return {
  302. "id": node_id,
  303. "name": node.get('name', ''),
  304. "category": node.get('category', ''),
  305. "props": [],
  306. "files": [],
  307. "distance": 0
  308. }
  309. def _process_search_result(search_result: dict, reference_index: int) -> tuple[dict, str]:
  310. """处理搜索结果,返回引用信息和文件名"""
  311. file_name = ""
  312. referrence = search_result.get("referrence", "")
  313. if referrence and "/books/" in referrence:
  314. file_name = referrence.split("/books/")[-1]
  315. file_name = os.path.splitext(file_name)[0]
  316. reference = {
  317. "index": str(reference_index),
  318. "id": search_result["id"],
  319. "content": search_result["content"],
  320. "file_path": search_result.get("file_path", ""),
  321. "title": search_result.get("title", ""),
  322. "distance": search_result.get("distance", DISTANCE_THRESHOLD),
  323. "page_no": search_result.get("page_no", ""),
  324. "file_name": file_name,
  325. "referrence": referrence
  326. }
  327. return reference, file_name
  328. def _get_file_type(file_name: str) -> str:
  329. """根据文件名确定文件类型"""
  330. file_name_lower = file_name.lower()
  331. if file_name_lower.endswith(".pdf"):
  332. return "pdf"
  333. elif file_name_lower.endswith((".doc", ".docx")):
  334. return "doc"
  335. elif file_name_lower.endswith((".xls", ".xlsx")):
  336. return "excel"
  337. elif file_name_lower.endswith((".ppt", ".pptx")):
  338. return "ppt"
  339. return "other"
  340. def _process_sentence_search(node_name: str, prop_title: str, sentences: list, trunks_service: TrunksService) -> tuple[list, list]:
  341. """处理句子搜索,返回结果句子和引用列表"""
  342. result_sentences = []
  343. all_references = []
  344. reference_index = 1
  345. i = 0
  346. while i < len(sentences):
  347. sentence = sentences[i]
  348. if len(sentence) < 10 and i + 1 < len(sentences):
  349. next_sentence = sentences[i + 1]
  350. result_sentences.append({"sentence": sentence, "flag": ""})
  351. search_text = f"{node_name}:{prop_title}:{sentence} {next_sentence}"
  352. i += 1
  353. elif len(sentence) < 10:
  354. result_sentences.append({"sentence": sentence, "flag": ""})
  355. i += 1
  356. continue
  357. else:
  358. search_text = f"{node_name}:{prop_title}:{sentence}"
  359. i += 1
  360. search_results = trunks_service.search_by_vector(text=search_text, limit=1, type='trunk')
  361. if not search_results:
  362. result_sentences.append({"sentence": sentence, "flag": ""})
  363. continue
  364. for search_result in search_results:
  365. if search_result.get("distance", DISTANCE_THRESHOLD) >= DISTANCE_THRESHOLD:
  366. result_sentences.append({"sentence": sentence, "flag": ""})
  367. continue
  368. existing_ref = next((ref for ref in all_references if ref["id"] == search_result["id"]), None)
  369. current_index = int(existing_ref["index"]) if existing_ref else reference_index
  370. if not existing_ref:
  371. reference, _ = _process_search_result(search_result, reference_index)
  372. all_references.append(reference)
  373. reference_index += 1
  374. result_sentences.append({"sentence": sentence, "flag": str(current_index)})
  375. return result_sentences, all_references
  376. @router.post("/kgrt_api/text/eb_search", response_model=StandardResponse)
  377. @router.post("/knowledge/text/eb_search", response_model=StandardResponse)
  378. async def node_props_search(request: NodePropsSearchRequest, db: Session = Depends(get_db)):
  379. try:
  380. start_time = time.time()
  381. # 检查缓存
  382. cached_result = _check_cache(request.node_id)
  383. if cached_result:
  384. return StandardResponse(success=True, data=cached_result)
  385. # 初始化服务
  386. trunks_service = TrunksService()
  387. node_service = KGNodeService(db)
  388. prop_service = KGPropService(db)
  389. # 获取节点信息
  390. result = _get_node_info(node_service, request.node_id)
  391. node_name = result["name"]
  392. # 遍历props_ids查询属性信息
  393. for prop_id in request.props_ids:
  394. prop = prop_service.get_props_by_id(prop_id)
  395. if not prop:
  396. logger.warning(f"属性不存在: {prop_id}")
  397. continue
  398. prop_title = prop.get('prop_title', '')
  399. prop_value = prop.get('prop_value', '')
  400. # 创建属性结果对象
  401. prop_result = {
  402. "id": prop_id,
  403. "category": prop.get('category', 0),
  404. "prop_name": prop.get('prop_name', ''),
  405. "prop_value": prop_value,
  406. "prop_title": prop_title,
  407. "type": prop.get('type', 1)
  408. }
  409. result["props"].append(prop_result)
  410. # 如果prop_value为'无',则跳过搜索
  411. if prop_value == '无':
  412. prop_result["answer"] = [{
  413. "sentence": prop_value,
  414. "flag": ""
  415. }]
  416. continue
  417. # 先用完整的prop_value进行搜索
  418. search_text = f"{node_name}:{prop_title}:{prop_value}"
  419. full_search_results = trunks_service.search_by_vector(
  420. text=search_text,
  421. limit=1,
  422. type='trunk'
  423. )
  424. # 处理搜索结果
  425. if full_search_results and full_search_results[0].get("distance", DISTANCE_THRESHOLD) < DISTANCE_THRESHOLD:
  426. search_result = full_search_results[0]
  427. reference, _ = _process_search_result(search_result, 1)
  428. prop_result["references"] = [reference]
  429. prop_result["answer"] = [{
  430. "sentence": prop_value,
  431. "flag": "1"
  432. }]
  433. else:
  434. # 如果整体搜索没有找到匹配结果,则进行句子拆分搜索
  435. sentences = SentenceUtil.split_text(prop_value)
  436. result_sentences, references = _process_sentence_search(
  437. node_name, prop_title, sentences, trunks_service
  438. )
  439. if references:
  440. prop_result["references"] = references
  441. if result_sentences:
  442. prop_result["answer"] = result_sentences
  443. # 处理文件信息
  444. all_files = set()
  445. file_index_map = {}
  446. file_index = 1
  447. # 收集文件信息
  448. for prop_result in result["props"]:
  449. if "references" not in prop_result:
  450. continue
  451. for ref in prop_result["references"]:
  452. referrence = ref.get("referrence", "")
  453. if not (referrence and "/books/" in referrence):
  454. continue
  455. file_name = referrence.split("/books/")[-1]
  456. if not file_name:
  457. continue
  458. file_type = _get_file_type(file_name)
  459. if file_name not in file_index_map:
  460. file_index_map[file_name] = file_index
  461. file_index += 1
  462. all_files.add((file_name, file_type))
  463. # 更新引用索引
  464. for prop_result in result["props"]:
  465. if "references" not in prop_result:
  466. continue
  467. for ref in prop_result["references"]:
  468. referrence = ref.get("referrence", "")
  469. if referrence and "/books/" in referrence:
  470. file_name = referrence.split("/books/")[-1]
  471. if file_name in file_index_map:
  472. ref["index"] = f"{file_index_map[file_name]}-{ref['index']}"
  473. # 更新answer中的flag
  474. if "answer" in prop_result:
  475. for sentence in prop_result["answer"]:
  476. if sentence["flag"]:
  477. for ref in prop_result["references"]:
  478. if ref["index"].endswith(f"-{sentence['flag']}"):
  479. sentence["flag"] = ref["index"]
  480. break
  481. # 添加文件信息到结果
  482. result["files"] = sorted([{
  483. "file_name": file_name,
  484. "file_type": file_type,
  485. "index": str(file_index_map[file_name])
  486. } for file_name, file_type in all_files], key=lambda x: int(x["index"]))
  487. end_time = time.time()
  488. logger.info(f"node_props_search接口耗时: {(end_time - start_time) * 1000:.2f}ms")
  489. # 缓存结果
  490. cache_key = f"xunzheng_{request.node_id}"
  491. cache[cache_key] = result
  492. return StandardResponse(success=True, data=result)
  493. except Exception as e:
  494. logger.error(f"Node props search failed: {str(e)}")
  495. raise HTTPException(status_code=500, detail=str(e))
  496. text_search_router = router