464 lines
17 KiB
Python
464 lines
17 KiB
Python
"""
|
||
舷号搜索策略模块
|
||
提供三级搜索策略：
|
||
1. 指定舷号搜索
|
||
2. 同型号所有舷号搜索
|
||
3. 全局搜索
|
||
|
||
新增功能：
|
||
- 智能舷号映射：支持舷号、舰名的智能识别和映射
|
||
"""
|
||
from typing import List, Dict, Any, Optional, Tuple
|
||
import random
|
||
import re
|
||
import httpx
|
||
from utils.function_tracker import get_all_callbacks
|
||
from tools.ship_model_db import load_ship_model_mapping
|
||
from tools.rag_tools import rag_search
|
||
from modelsAPI.model_api import OpenaiAPI
|
||
from config import KB_TREE_CONFIG
|
||
|
||
|
||
_kb_name_to_id_cache: Optional[Dict[str, str]] = None
|
||
|
||
|
||
async def _load_kb_tree() -> Dict[str, str]:
|
||
"""
|
||
从知识库树 API 加载所有知识库，返回 {名称: kb_id} 的映射
|
||
递归遍历树结构，带缓存，只请求一次
|
||
"""
|
||
global _kb_name_to_id_cache
|
||
if _kb_name_to_id_cache is not None:
|
||
return _kb_name_to_id_cache
|
||
|
||
url = KB_TREE_CONFIG.get("url", "")
|
||
if not url:
|
||
print("[ship_number_search] KB_TREE_CONFIG.url 未配置")
|
||
return {}
|
||
|
||
headers = {
|
||
"accept": "application/json",
|
||
"x-user-id": "1",
|
||
"x-user-name": "testuser",
|
||
"x-role": "admin",
|
||
}
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=30.0, verify=False) as client:
|
||
response = await client.get(url, headers=headers)
|
||
if response.status_code != 200:
|
||
print(f"[ship_number_search] 知识库树 API 请求失败: HTTP {response.status_code}")
|
||
return {}
|
||
|
||
result = response.json()
|
||
|
||
if isinstance(result, list):
|
||
top_nodes = result
|
||
elif isinstance(result, dict):
|
||
kb_data = result.get("data") or result.get("list") or {}
|
||
top_nodes = kb_data.get("children", []) if isinstance(kb_data, dict) else []
|
||
else:
|
||
top_nodes = []
|
||
|
||
name_to_id = {}
|
||
|
||
def _collect(nodes):
|
||
for node in nodes:
|
||
kb_name = node.get("name") or ""
|
||
kb_id = str(node.get("id", "")) if node.get("id") is not None else ""
|
||
if kb_name and kb_id:
|
||
name_to_id[kb_name] = kb_id
|
||
children = node.get("children", [])
|
||
if children:
|
||
_collect(children)
|
||
|
||
_collect(top_nodes)
|
||
|
||
_kb_name_to_id_cache = name_to_id
|
||
return name_to_id
|
||
except Exception as e:
|
||
print(f"[ship_number_search] 加载知识库树失败: {str(e)}")
|
||
return {}
|
||
|
||
|
||
async def resolve_kb_id(ship_number: str) -> Optional[str]:
|
||
"""
|
||
根据舷号从知识库树中查找对应的真实 kb_id
|
||
|
||
匹配策略：在知识库名称中查找包含该舷号的条目
|
||
例如 ship_number="163"，知识库名称为 "163舰资料"，则匹配成功
|
||
|
||
Args:
|
||
ship_number: 舷号（如 "163"）
|
||
|
||
Returns:
|
||
匹配到的 kb_id，未找到返回 None
|
||
"""
|
||
if not ship_number:
|
||
return None
|
||
|
||
name_to_id = await _load_kb_tree()
|
||
|
||
if ship_number in name_to_id:
|
||
return name_to_id[ship_number]
|
||
|
||
for kb_name, kb_id in name_to_id.items():
|
||
if re.search(rf'(?<!\d){re.escape(ship_number)}(?!\d)', kb_name):
|
||
return kb_id
|
||
|
||
return None
|
||
|
||
|
||
async def get_ship_numbers_by_model(ship_number: str) -> list:
|
||
"""
|
||
根据舷号获取同型号的所有舷号
|
||
|
||
Args:
|
||
ship_number: 舷号
|
||
|
||
Returns:
|
||
list: 同型号的所有舷号列表，如果未找到则返回只包含输入舷号的列表
|
||
"""
|
||
ship_number = str(ship_number).strip()
|
||
ships = await load_ship_model_mapping()
|
||
target_model = None
|
||
for ship in ships:
|
||
if ship["hull_number"] == ship_number:
|
||
target_model = ship["model_name"]
|
||
break
|
||
if target_model:
|
||
return [s["hull_number"] for s in ships if s["model_name"] == target_model]
|
||
return [ship_number]
|
||
|
||
|
||
async def _build_ship_mapping_candidates() -> Dict[str, Dict[str, Any]]:
|
||
"""
|
||
构建舷号映射候选词库
|
||
从每艘舰的记录中构建候选词，支持舷号和舰名匹配
|
||
"""
|
||
candidates = {}
|
||
|
||
ships = await load_ship_model_mapping()
|
||
|
||
model_numbers: Dict[str, List[str]] = {}
|
||
for ship in ships:
|
||
model = ship["model_name"]
|
||
if model not in model_numbers:
|
||
model_numbers[model] = []
|
||
model_numbers[model].append(ship["hull_number"])
|
||
|
||
for ship in ships:
|
||
model = ship["model_name"]
|
||
hull_number = ship["hull_number"]
|
||
ship_name = ship.get("ship_name", "")
|
||
numbers = model_numbers.get(model, [])
|
||
|
||
mapping_info = {
|
||
"model": model,
|
||
"numbers": numbers.copy()
|
||
}
|
||
|
||
candidates[hull_number] = {
|
||
"type": "number",
|
||
"ship_number": hull_number,
|
||
**mapping_info
|
||
}
|
||
|
||
if ship_name:
|
||
candidates[ship_name] = {
|
||
"type": "name",
|
||
"ship_number": hull_number,
|
||
**mapping_info
|
||
}
|
||
|
||
return candidates
|
||
|
||
|
||
_SHIP_MAPPING_CANDIDATES = None
|
||
|
||
async def get_ship_mapping_candidates() -> Dict[str, Dict[str, Any]]:
|
||
"""
|
||
获取舷号映射候选词库（带缓存）
|
||
"""
|
||
global _SHIP_MAPPING_CANDIDATES
|
||
if _SHIP_MAPPING_CANDIDATES is None:
|
||
_SHIP_MAPPING_CANDIDATES = await _build_ship_mapping_candidates()
|
||
return _SHIP_MAPPING_CANDIDATES
|
||
|
||
async def smart_ship_number_mapping(
|
||
user_input: str,
|
||
similarity_threshold: float = 0.8
|
||
) -> Tuple[Optional[str], Optional[str], Optional[List[str]], str]:
|
||
"""
|
||
智能舷号映射函数
|
||
将用户输入的舷号/舰名映射到准确的舷号或型号
|
||
|
||
Returns:
|
||
Tuple: (matched_ship_number, matched_model, all_numbers, match_type)
|
||
- matched_ship_number: 匹配到的具体舷号（舰名也会映射到对应舷号）
|
||
- matched_model: 匹配到的型号名称
|
||
- all_numbers: 该型号下的所有舷号列表
|
||
- match_type: 匹配类型 ("exact_number", "exact_name", "embedding", "none")
|
||
"""
|
||
if not user_input or not user_input.strip():
|
||
return None, None, None, "none"
|
||
|
||
user_input = user_input.strip()
|
||
candidates = await get_ship_mapping_candidates()
|
||
|
||
if user_input in candidates:
|
||
info = candidates[user_input]
|
||
match_type = f"exact_{info['type']}"
|
||
model = info["model"]
|
||
numbers = info["numbers"]
|
||
ship_number = info.get("ship_number")
|
||
|
||
return ship_number, model, numbers, match_type
|
||
|
||
try:
|
||
candidate_texts = list(candidates.keys())
|
||
query_embedding = await OpenaiAPI.get_embeddings_async(user_input)
|
||
candidate_embeddings = await OpenaiAPI.get_embeddings_batch_async(candidate_texts)
|
||
|
||
best_score = -1
|
||
best_idx = -1
|
||
|
||
for idx, emb in enumerate(candidate_embeddings):
|
||
score = OpenaiAPI.cosine_similarity(query_embedding, emb)
|
||
if score > best_score:
|
||
best_score = score
|
||
best_idx = idx
|
||
|
||
if best_score >= similarity_threshold and best_idx >= 0:
|
||
matched_text = candidate_texts[best_idx]
|
||
info = candidates[matched_text]
|
||
model = info["model"]
|
||
numbers = info["numbers"]
|
||
ship_number = info.get("ship_number")
|
||
|
||
print(f"[智能映射] 用户输入 '{user_input}' -> 匹配到 '{matched_text}' (相似度: {best_score:.3f}, 类型: {info['type']})")
|
||
|
||
return ship_number, model, numbers, "embedding"
|
||
except Exception as e:
|
||
print(f"[智能映射] Embedding 匹配失败: {str(e)}")
|
||
|
||
return None, None, None, "none"
|
||
|
||
def send_search_status(title: str, details: str):
|
||
"""发送搜索状态事件"""
|
||
title_options = [f"🔍 {title}", f"🔎 {title}", f"📚 {title}"]
|
||
start_event = {
|
||
"type": "function_execution",
|
||
"title": random.choice(title_options),
|
||
"details": details
|
||
}
|
||
for callback in get_all_callbacks():
|
||
try:
|
||
callback(start_event)
|
||
except Exception:
|
||
pass
|
||
|
||
def convert_rag_result(rag_result):
|
||
"""将 rag_search 的返回结果转换为原来的格式"""
|
||
if not rag_result or not rag_result.get("success", False):
|
||
return []
|
||
source_citation = rag_result.get("sourceCitation", {})
|
||
flat_results = []
|
||
for resource, items in source_citation.items():
|
||
for item in items:
|
||
flat_results.append(item)
|
||
return flat_results
|
||
|
||
async def search_with_ship_number_strategy(
|
||
base_query: str,
|
||
ship_number: str,
|
||
top_k: int = 5,
|
||
search_type: str = "fault"
|
||
) -> Tuple[List[Dict[str, Any]], str, str, str]:
|
||
"""
|
||
使用舷号三级搜索策略进行RAG检索
|
||
返回: (rag_results, matched_kb_name, matched_kb_id, matched_ship_number)
|
||
matched_ship_number: 映射后的舷号，未映射到具体舷号时为空字符串
|
||
"""
|
||
rag_results = []
|
||
matched_kb_name = ""
|
||
matched_kb_id = ""
|
||
|
||
matched_ship_number = None
|
||
matched_model = None
|
||
model_all_numbers = None
|
||
match_type = "none"
|
||
|
||
if ship_number and ship_number.strip():
|
||
send_search_status(
|
||
"正在智能识别",
|
||
f"识别 '{ship_number}' 的含义（舷号/舰名）..."
|
||
)
|
||
|
||
matched_ship_number, matched_model, model_all_numbers, match_type = await smart_ship_number_mapping(ship_number)
|
||
|
||
if match_type != "none":
|
||
print(f"[智能映射结果] 输入='{ship_number}' -> 舷号={matched_ship_number}, 型号={matched_model}, 所有舷号={model_all_numbers}, 匹配类型={match_type}")
|
||
else:
|
||
print(f"[智能映射结果] 输入='{ship_number}' -> 未找到匹配，将使用原始值进行检索")
|
||
|
||
if match_type in ["exact_number", "exact_name", "embedding"] and matched_ship_number:
|
||
real_kb_id = await resolve_kb_id(matched_ship_number)
|
||
if real_kb_id:
|
||
send_search_status(
|
||
"正在指定舷号检索",
|
||
f"使用舷号 {matched_ship_number} 进行检索..."
|
||
)
|
||
|
||
rag_query = f"{base_query}"
|
||
|
||
try:
|
||
rag_result_raw = await rag_search(rag_query, top_k=top_k, kb_id=real_kb_id)
|
||
rag_results = convert_rag_result(rag_result_raw)
|
||
|
||
if rag_results:
|
||
print(f"步骤1 - 舷号 {matched_ship_number} 检索完成，返回 {len(rag_results)} 条结果")
|
||
for item in rag_results:
|
||
if isinstance(item, dict):
|
||
found_kb_name = item.get("kb_name", "")
|
||
found_kb_id = item.get("kb_id", "")
|
||
if found_kb_name and not matched_kb_name:
|
||
matched_kb_name = str(found_kb_name).strip()
|
||
if found_kb_id and not matched_kb_id:
|
||
matched_kb_id = str(found_kb_id).strip()
|
||
return rag_results, matched_kb_name, matched_kb_id, matched_ship_number or ""
|
||
except Exception as e:
|
||
print(f"步骤1 - 舷号 {matched_ship_number} 检索失败: {str(e)}")
|
||
else:
|
||
print(f"步骤1 - 舷号 {matched_ship_number} 无对应知识库，跳过指定舷号检索")
|
||
|
||
if match_type in ["exact_number", "exact_name", "embedding"] and model_all_numbers:
|
||
send_search_status(
|
||
"正在型号检索",
|
||
f"识别为型号 '{matched_model}'，使用该型号所有舷号 {', '.join(model_all_numbers)} 进行检索..."
|
||
)
|
||
|
||
for num in model_all_numbers:
|
||
rag_query = f"{base_query}"
|
||
|
||
try:
|
||
real_kb_id = await resolve_kb_id(num)
|
||
if not real_kb_id:
|
||
continue
|
||
rag_result_raw = await rag_search(rag_query, top_k=top_k, kb_id=real_kb_id)
|
||
current_results = convert_rag_result(rag_result_raw)
|
||
|
||
if current_results:
|
||
rag_results.extend(current_results)
|
||
print(f"步骤2 - 舷号 {num} 检索完成，返回 {len(current_results)} 条结果")
|
||
|
||
if not matched_kb_name:
|
||
for item in current_results:
|
||
if isinstance(item, dict):
|
||
found_kb_name = item.get("kb_name", "")
|
||
found_kb_id = item.get("kb_id", "")
|
||
if found_kb_name and not matched_kb_name:
|
||
matched_kb_name = str(found_kb_name).strip()
|
||
if found_kb_id and not matched_kb_id:
|
||
matched_kb_id = str(found_kb_id).strip()
|
||
except Exception as e:
|
||
print(f"步骤2 - 舷号 {num} 检索失败: {str(e)}")
|
||
|
||
if rag_results:
|
||
return rag_results, matched_kb_name, matched_kb_id, matched_ship_number or ""
|
||
|
||
if ship_number and ship_number.strip() and match_type == "none":
|
||
real_kb_id = await resolve_kb_id(ship_number)
|
||
if real_kb_id:
|
||
send_search_status(
|
||
"正在指定舰名检索",
|
||
f"使用舰名 {ship_number} 进行检索..."
|
||
)
|
||
|
||
rag_query = f"{base_query}"
|
||
|
||
try:
|
||
rag_result_raw = await rag_search(rag_query, top_k=top_k, kb_id=real_kb_id)
|
||
rag_results = convert_rag_result(rag_result_raw)
|
||
|
||
if rag_results:
|
||
print(f"步骤1 - 舷号 {ship_number} 检索完成，返回 {len(rag_results)} 条结果")
|
||
for item in rag_results:
|
||
if isinstance(item, dict):
|
||
found_kb_name = item.get("kb_name", "")
|
||
found_kb_id = item.get("kb_id", "")
|
||
if found_kb_name and not matched_kb_name:
|
||
matched_kb_name = str(found_kb_name).strip()
|
||
if found_kb_id and not matched_kb_id:
|
||
matched_kb_id = str(found_kb_id).strip()
|
||
return rag_results, matched_kb_name, matched_kb_id, matched_ship_number or ""
|
||
except Exception as e:
|
||
print(f"步骤1 - 舷号 {ship_number} 检索失败: {str(e)}")
|
||
else:
|
||
print(f"步骤1 - 舷号 {ship_number} 无对应知识库，跳过指定舰名检索")
|
||
|
||
model_ship_numbers = await get_ship_numbers_by_model(ship_number)
|
||
if len(model_ship_numbers) > 1 or (len(model_ship_numbers) == 1 and ship_number not in model_ship_numbers):
|
||
send_search_status(
|
||
"正在同型号检索",
|
||
f"使用同型号舷号 {', '.join(model_ship_numbers)} 进行检索..."
|
||
)
|
||
|
||
for num in model_ship_numbers:
|
||
if num == ship_number:
|
||
continue
|
||
|
||
rag_query = f"{base_query}"
|
||
|
||
try:
|
||
real_kb_id = await resolve_kb_id(num)
|
||
if not real_kb_id:
|
||
continue
|
||
rag_result_raw = await rag_search(rag_query, top_k=top_k, kb_id=real_kb_id)
|
||
current_results = convert_rag_result(rag_result_raw)
|
||
|
||
if current_results:
|
||
rag_results.extend(current_results)
|
||
print(f"步骤2 - 舷号 {num} 检索完成，返回 {len(current_results)} 条结果")
|
||
|
||
if not matched_kb_name:
|
||
for item in current_results:
|
||
if isinstance(item, dict):
|
||
found_kb_name = item.get("kb_name", "")
|
||
found_kb_id = item.get("kb_id", "")
|
||
if found_kb_name and not matched_kb_name:
|
||
matched_kb_name = str(found_kb_name).strip()
|
||
if found_kb_id and not matched_kb_id:
|
||
matched_kb_id = str(found_kb_id).strip()
|
||
except Exception as e:
|
||
print(f"步骤2 - 舷号 {num} 检索失败: {str(e)}")
|
||
|
||
if rag_results:
|
||
return rag_results, matched_kb_name, matched_kb_id, matched_ship_number or ""
|
||
|
||
send_search_status(
|
||
"正在全局检索",
|
||
"未找到舷号相关资料，进行全局检索..."
|
||
)
|
||
|
||
try:
|
||
rag_result_raw = await rag_search(base_query, top_k=top_k)
|
||
rag_results = convert_rag_result(rag_result_raw)
|
||
|
||
if rag_results:
|
||
print(f"步骤3 - 全局检索完成，返回 {len(rag_results)} 条结果")
|
||
for item in rag_results:
|
||
if isinstance(item, dict):
|
||
found_kb_name = item.get("kb_name", "")
|
||
found_kb_id = item.get("kb_id", "")
|
||
if found_kb_name and not matched_kb_name:
|
||
matched_kb_name = str(found_kb_name).strip()
|
||
if found_kb_id and not matched_kb_id:
|
||
matched_kb_id = str(found_kb_id).strip()
|
||
except Exception as e:
|
||
print(f"步骤3 - 全局检索失败: {str(e)}")
|
||
rag_results = []
|
||
|
||
return rag_results, matched_kb_name, matched_kb_id, matched_ship_number or ""
|
||
|