1684 lines
58 KiB
Python
Executable File
1684 lines
58 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
RAG搜索路由 - 支持新的索引格式
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import pickle
|
||
import logging
|
||
from pathlib import Path
|
||
from typing import Optional, List, Dict, Any
|
||
from fastapi import APIRouter, HTTPException
|
||
from pydantic import BaseModel
|
||
from src.rag.project_config import PROJECT_ROOT, conda_env_path
|
||
|
||
# 配置日志
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 添加项目路径
|
||
sys.path.insert(0, str(PROJECT_ROOT))
|
||
gf_site_packages = conda_env_path(os.environ.get("GF_PYTHON_ENV", "gfpy312")) / "lib/python3.12/site-packages"
|
||
if gf_site_packages.exists():
|
||
sys.path.insert(0, str(gf_site_packages))
|
||
|
||
from src.rag.base import SearchResult
|
||
from src.rag.vector_store import InMemoryStore
|
||
from src.rag.embeddings import BGEEmbedding
|
||
from src.rag.family_facts import ensure_family_fact
|
||
|
||
router = APIRouter()
|
||
|
||
# agentic模块已移除，使用纯RAG功能
|
||
|
||
# 初始化LLM客户端
|
||
_llm_manager = None
|
||
def get_llm_manager():
|
||
global _llm_manager
|
||
if _llm_manager is None:
|
||
try:
|
||
from src.llm.manager import LLMManager
|
||
import yaml
|
||
|
||
# 读取配置文件
|
||
config_path = PROJECT_ROOT / "config" / "llm.yaml"
|
||
if config_path.exists():
|
||
with open(config_path, 'r') as f:
|
||
config = yaml.safe_load(f)
|
||
_llm_manager = LLMManager(config)
|
||
provider = config.get("default", {}).get("provider", "unknown")
|
||
print(f"LLMManager初始化成功（使用{provider}）")
|
||
else:
|
||
# 配置文件不存在，使用Ollama作为fallback
|
||
_llm_manager = LLMManager({
|
||
"default": {"provider": "ollama", "model": "qwen3:0.6b"}
|
||
})
|
||
print("LLMManager初始化成功（使用Ollama fallback）")
|
||
except Exception as e:
|
||
print(f"LLM客户端初始化失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
return _llm_manager
|
||
|
||
# 全局变量存储vector store
|
||
_vector_store = None
|
||
_embedding_model = None
|
||
|
||
def get_vector_store():
|
||
"""懒加载vector store"""
|
||
global _vector_store, _embedding_model
|
||
|
||
if _vector_store is None:
|
||
configured_index = os.environ.get("RAG_VECTOR_INDEX")
|
||
if configured_index:
|
||
index_path = Path(configured_index).expanduser()
|
||
if index_path.is_dir():
|
||
index_path = index_path / "vector_index.pkl"
|
||
else:
|
||
# 默认仍使用当前生产中的SPL索引；通用构建脚本可生成同格式索引。
|
||
index_path = PROJECT_ROOT / "data" / "rag_index_final_v15_spl" / "vector_index.pkl"
|
||
|
||
# 如果默认/配置索引不存在，回退到旧索引
|
||
if not index_path.exists():
|
||
index_path = PROJECT_ROOT / "data" / "rag_index_final_v15_1" / "vector_index.pkl"
|
||
|
||
if not index_path.exists():
|
||
raise HTTPException(status_code=500, detail="索引文件不存在，请先运行索引构建")
|
||
|
||
# 创建embedding model
|
||
_embedding_model = BGEEmbedding()
|
||
|
||
# 创建store
|
||
_vector_store = InMemoryStore(embedding_model=_embedding_model)
|
||
|
||
# 加载索引
|
||
with open(index_path, 'rb') as f:
|
||
data = pickle.load(f)
|
||
|
||
_vector_store._chunks = data['chunks']
|
||
|
||
# 将 embeddings 列表转换为字典（chunk_id -> embedding）
|
||
chunks_list = list(_vector_store._chunks.values())
|
||
embeddings_list = data['embeddings']
|
||
_vector_store._embeddings = {}
|
||
for i, chunk in enumerate(chunks_list):
|
||
if i < len(embeddings_list):
|
||
_vector_store._embeddings[chunk.chunk_id] = embeddings_list[i]
|
||
|
||
print(f"RAG索引加载完成: {len(_vector_store._chunks)} chunks, {len(_vector_store._embeddings)} embeddings")
|
||
|
||
return _vector_store
|
||
|
||
|
||
class RAGQueryRequest(BaseModel):
|
||
query: str
|
||
mode: str = "hybrid" # hybrid, normal, exact, alias, substring, bm25
|
||
top_k: int = 10
|
||
include_sources: bool = True
|
||
use_llm: bool = True # 是否使用LLM生成总结（默认启用）
|
||
strict_mode: bool = False # 严格模式：只返回文献原文，不生成总结（默认关闭）
|
||
|
||
|
||
class PhylogenyExtractionRequest(BaseModel):
|
||
"""系统发育提取请求"""
|
||
pdf_path: str # PDF文件路径
|
||
doi: Optional[str] = None # DOI（可选，用于获取补充材料）
|
||
pmcid: Optional[str] = None # PMCID（可选）
|
||
auto_fetch_supplementary: bool = True # 是否自动获取补充材料
|
||
|
||
|
||
class RAGSource(BaseModel):
|
||
doc_id: str
|
||
content: str
|
||
score: float
|
||
metadata: dict
|
||
pmid: Optional[str] = None # V15.1: 添加 PMID 字段
|
||
|
||
|
||
class KnowledgeGraphData(BaseModel):
|
||
nodes: List[dict] = []
|
||
edges: List[dict] = []
|
||
|
||
|
||
class KGEntityInfo(BaseModel):
|
||
"""知识图谱实体信息（用于前端可视化）"""
|
||
name: str
|
||
type: str # family, gene, species
|
||
display_name: str = "" # 用于展示的友好名称
|
||
|
||
|
||
class OmicsDataset(BaseModel):
|
||
"""多组学数据集信息"""
|
||
dataset_id: str
|
||
dataset_type: str # transcriptome, genome, etc.
|
||
species: str
|
||
source_paper: str
|
||
|
||
|
||
class RAGQueryResponse(BaseModel):
|
||
response: str
|
||
sources: List[RAGSource] = []
|
||
mode: str
|
||
total_chunks: int
|
||
knowledge_graph: Optional[KnowledgeGraphData] = None # 动态生成的知识图谱数据
|
||
kg_entities: List[KGEntityInfo] = [] # 可用于Neo4j可视化的实体列表
|
||
omics_datasets: Optional[List[OmicsDataset]] = None # 多组学数据集信息
|
||
|
||
|
||
class FamilyFactResponse(BaseModel):
|
||
family: str
|
||
source: str = ""
|
||
data: Dict[str, Any]
|
||
|
||
|
||
@router.post("/chat")
|
||
async def rag_chat(request: RAGQueryRequest):
|
||
"""RAG搜索端点"""
|
||
try:
|
||
logger.info(f"[RAG] 查询: '{request.query}' | 模式: {request.mode} | top_k: {request.top_k}")
|
||
store = get_vector_store()
|
||
|
||
# agentic模式已移除，使用hybrid
|
||
if request.mode == "agentic":
|
||
# 回退到hybrid search
|
||
results = await _hybrid_search(store, request.query, request.top_k)
|
||
elif request.mode == "exact":
|
||
results = await _exact_search(store, request.query, request.top_k)
|
||
elif request.mode == "substring":
|
||
results = await _substring_search(store, request.query, request.top_k)
|
||
elif request.mode == "bm25":
|
||
results = await _bm25_search(store, request.query, request.top_k)
|
||
else:
|
||
# hybrid search
|
||
results = await _hybrid_search(store, request.query, request.top_k)
|
||
|
||
logger.info(f"[RAG] 检索完成: {len(results)}条结果")
|
||
if results:
|
||
logger.info(f"[RAG] Top3: {[(r.chunk.doc_id, round(r.score,3)) for r in results[:3]]}")
|
||
|
||
# 格式化响应
|
||
sources = []
|
||
for result in results:
|
||
sources.append(RAGSource(
|
||
doc_id=result.chunk.doc_id,
|
||
content=result.chunk.content[:500],
|
||
score=result.score,
|
||
metadata=result.chunk.metadata,
|
||
pmid=result.chunk.metadata.get('pmid') # V15.1: 添加 PMID
|
||
))
|
||
|
||
# 生成回答（支持LLM总结或严格模式）
|
||
if request.strict_mode:
|
||
# 严格模式：只返回文献原文，不调用LLM
|
||
response_text = _generate_strict_response(request.query, sources)
|
||
logger.info(f"[RAG] 严格模式生成完成: {len(response_text)}字符")
|
||
else:
|
||
response_text = await _generate_response_async(request.query, sources, request.use_llm)
|
||
logger.info(f"[RAG] 生成完成: {len(response_text)}字符 | use_llm: {request.use_llm}")
|
||
|
||
# 从检索结果动态提取知识图谱
|
||
kg_data = _extract_knowledge_graph(sources, request.query)
|
||
|
||
# 提取可用于Neo4j可视化的实体
|
||
kg_entities = _extract_kg_entities(sources)
|
||
|
||
# 检测数据集查询意图，从知识图谱中查找多组学数据
|
||
omics_datasets = _query_omics_datasets(request.query)
|
||
if omics_datasets:
|
||
# 在回答中添加数据集信息
|
||
response_text = _add_omics_info_to_response(response_text, omics_datasets)
|
||
logger.info(f"[RAG] 找到 {len(omics_datasets)} 个多组学数据集")
|
||
|
||
return RAGQueryResponse(
|
||
response=response_text,
|
||
sources=sources if request.include_sources else [],
|
||
mode=request.mode,
|
||
total_chunks=len(store._chunks),
|
||
knowledge_graph=kg_data,
|
||
kg_entities=kg_entities,
|
||
omics_datasets=omics_datasets if omics_datasets else None
|
||
)
|
||
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=f"搜索失败: {str(e)}")
|
||
|
||
|
||
@router.get("/family/pfam")
|
||
async def family_pfam(family: str = "SPL"):
|
||
"""Return literature-backed Pfam IDs for a gene family."""
|
||
fact = ensure_family_fact(family)
|
||
pfam_items = fact.get("pfam") or []
|
||
pfam_ids = [item.get("pfam_id") for item in pfam_items if item.get("pfam_id")]
|
||
return {
|
||
"family": fact.get("family", family.upper()),
|
||
"pfam_ids": pfam_ids,
|
||
"pfams": pfam_items,
|
||
"confidence": "high" if any((item.get("support_count") or 0) >= 2 for item in pfam_items) else ("medium" if pfam_items else "none"),
|
||
"source": fact.get("source", ""),
|
||
"source_documents": fact.get("source_documents", []),
|
||
}
|
||
|
||
|
||
@router.get("/family/baseline")
|
||
async def family_baseline(family: str = "SPL"):
|
||
"""Return member-count baseline facts extracted from v15 literature outputs."""
|
||
fact = ensure_family_fact(family)
|
||
baseline = fact.get("baseline") or {}
|
||
return {
|
||
"family": fact.get("family", family.upper()),
|
||
"baseline": baseline,
|
||
"member_counts": fact.get("member_counts", []),
|
||
"source": fact.get("source", ""),
|
||
}
|
||
|
||
|
||
@router.get("/family/stage-context")
|
||
async def family_stage_context(family: str = "SPL", stage: str = "identify", limit: int = 12):
|
||
"""Return stage-specific methods/evidence context for analysis agents."""
|
||
fact = ensure_family_fact(family)
|
||
stage_key = (stage or "").strip().lower() or "identify"
|
||
contexts = (fact.get("stage_context") or {}).get(stage_key, [])
|
||
limit = max(1, min(int(limit or 12), 50))
|
||
return {
|
||
"family": fact.get("family", family.upper()),
|
||
"stage": stage_key,
|
||
"contexts": contexts[:limit],
|
||
"available_stages": sorted((fact.get("stage_context") or {}).keys()),
|
||
"source": fact.get("source", ""),
|
||
}
|
||
|
||
|
||
async def _exact_search(store, query: str, top_k: int):
|
||
"""精确匹配搜索"""
|
||
from src.rag.base import SearchResult
|
||
|
||
query_lower = query.lower().strip()
|
||
results = []
|
||
|
||
for chunk_id, chunk in store._chunks.items():
|
||
content = chunk.content.lower()
|
||
count = content.count(query_lower)
|
||
if count > 0:
|
||
score = 0.8 + 0.19 * min(count, 10) / 10
|
||
if content == query_lower:
|
||
score = 1.0
|
||
results.append(SearchResult(chunk=chunk, score=score))
|
||
|
||
results.sort(key=lambda x: x.score, reverse=True)
|
||
return results[:top_k]
|
||
|
||
|
||
def _expand_query(query_words: List[str]) -> List[str]:
|
||
"""查询扩展：添加中英文同义词，并拆分组合词"""
|
||
# 同义词映射
|
||
synonyms = {
|
||
# 中文 -> 英文
|
||
'转录组': ['transcriptome', 'rna-seq', 'rnaseq', 'rna seq'],
|
||
'转录': ['transcript', 'transcription'],
|
||
'基因组': ['genome', 'genomic'],
|
||
'蛋白': ['protein', 'peptide'],
|
||
'基因': ['gene', 'genetic'],
|
||
'表达': ['expression', 'express'],
|
||
'序列': ['sequence', 'seq'],
|
||
'数据': ['data', 'dataset'],
|
||
'编号': ['accession', 'id', 'number'],
|
||
'山核桃': ['pecan', 'carya', 'carya illinoinensis'],
|
||
'谷子': ['foxtail millet', 'setaria', 'setaria italica'],
|
||
'向日葵': ['sunflower', 'helianthus', 'helianthus annuus'],
|
||
'柑橘': ['citrus', 'orange', 'mandarin'],
|
||
# 英文 -> 相关词
|
||
'rna-seq': ['rnaseq', 'rna seq', 'transcriptome', '转录组'],
|
||
'transcriptome': ['rna-seq', 'rnaseq', '转录组'],
|
||
'pecan': ['carya', 'carya illinoinensis', '山核桃'],
|
||
}
|
||
|
||
# 先拆分组合词（如 "转录组数据" -> "转录组", "数据"）
|
||
split_words = []
|
||
for word in query_words:
|
||
split_words.append(word)
|
||
# 尝试拆分中文组合词
|
||
for key in synonyms.keys():
|
||
if key in word and word != key:
|
||
split_words.append(key)
|
||
|
||
expanded = set(split_words)
|
||
for word in split_words:
|
||
word_lower = word.lower()
|
||
if word_lower in synonyms:
|
||
expanded.update(synonyms[word_lower])
|
||
|
||
return list(expanded)
|
||
|
||
|
||
async def _substring_search(store, query: str, top_k: int):
|
||
"""子串搜索（支持查询扩展）"""
|
||
import re
|
||
from src.rag.base import SearchResult
|
||
|
||
query_lower = query.lower().strip()
|
||
|
||
# 提取查询词
|
||
if ' ' in query_lower:
|
||
query_words = [w.strip() for w in query_lower.split() if len(w.strip()) >= 2]
|
||
else:
|
||
query_words = re.findall(r'[\u4e00-\u9fa5]{2,}|[a-zA-Z]+', query_lower)
|
||
if not query_words:
|
||
query_words = [query_lower] if len(query_lower) >= 2 else []
|
||
|
||
if not query_words:
|
||
query_words = [query_lower] if query_lower else []
|
||
|
||
# 查询扩展
|
||
original_words = query_words.copy()
|
||
query_words = _expand_query(query_words)
|
||
|
||
if set(original_words) != set(query_words):
|
||
logger.info(f"[RAG] 查询扩展: {original_words} -> {query_words}")
|
||
|
||
results = []
|
||
|
||
# 物种关键词（高权重）
|
||
species_keywords = ['pecan', 'carya', 'carya illinoinensis', '山核桃', '谷子', 'setaria', '向日葵', 'helianthus', '柑橘', 'citrus']
|
||
# 数据集关键词（高权重）
|
||
dataset_keywords = ['transcriptome', 'rna-seq', 'rnaseq', 'rna seq', '转录组', 'prjna', 'srp', 'gse', 'gsm', 'srr', 'srx']
|
||
|
||
for chunk_id, chunk in store._chunks.items():
|
||
content_lower = chunk.content.lower()
|
||
|
||
total_score = 0
|
||
matched_words = 0
|
||
has_species = False
|
||
has_dataset = False
|
||
|
||
for word in query_words:
|
||
if word in content_lower:
|
||
count = content_lower.count(word)
|
||
# 基础权重
|
||
word_weight = min(len(word) / 10, 1.0) + 0.5
|
||
# 物种关键词额外加权
|
||
if word in species_keywords:
|
||
word_weight *= 2.0
|
||
has_species = True
|
||
# 数据集关键词额外加权
|
||
if word in dataset_keywords:
|
||
word_weight *= 1.5
|
||
has_dataset = True
|
||
total_score += 0.1 * count * word_weight
|
||
matched_words += 1
|
||
|
||
if matched_words > 0:
|
||
coverage = matched_words / len(query_words) if query_words else 0
|
||
base_score = 0.5
|
||
# 如果同时匹配物种和数据集，额外加分
|
||
if has_species and has_dataset:
|
||
base_score += 0.2
|
||
coverage_bonus = 0.3 * coverage
|
||
density_score = min(total_score, 0.2)
|
||
score = base_score + coverage_bonus + density_score
|
||
results.append(SearchResult(chunk=chunk, score=min(score, 1.0)))
|
||
|
||
results.sort(key=lambda x: x.score, reverse=True)
|
||
return results[:top_k]
|
||
|
||
|
||
async def _bm25_search(store, query: str, top_k: int):
|
||
"""BM25搜索"""
|
||
from src.rag.base import SearchResult
|
||
|
||
try:
|
||
from rank_bm25 import BM25Okapi
|
||
except ImportError:
|
||
# 如果没有BM25，使用简单关键词匹配
|
||
return await _substring_search(store, query, top_k)
|
||
|
||
# 构建corpus
|
||
corpus = []
|
||
chunk_ids = []
|
||
|
||
for chunk_id, chunk in store._chunks.items():
|
||
corpus.append(chunk.content)
|
||
chunk_ids.append(chunk_id)
|
||
|
||
if not corpus:
|
||
return []
|
||
|
||
bm25 = BM25Okapi(corpus)
|
||
tokenized_query = query.lower().split()
|
||
scores = bm25.get_scores(tokenized_query)
|
||
|
||
# 获取top_k
|
||
top_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:top_k]
|
||
|
||
results = []
|
||
for idx in top_indices:
|
||
if scores[idx] > 0:
|
||
chunk_id = chunk_ids[idx]
|
||
chunk = store._chunks[chunk_id]
|
||
results.append(SearchResult(chunk=chunk, score=min(scores[idx] / 10, 1.0)))
|
||
|
||
return results
|
||
|
||
|
||
async def _hybrid_search(store, query: str, top_k: int, alpha: float = 0.7):
|
||
"""混合搜索"""
|
||
# 向量搜索
|
||
try:
|
||
query_embedding = await store.embedding_model.aembed(query)
|
||
vector_results = store.search(query_embedding, top_k * 2, threshold=0.0)
|
||
except:
|
||
vector_results = []
|
||
|
||
# BM25搜索
|
||
bm25_results = await _bm25_search(store, query, top_k * 2)
|
||
|
||
# 子串搜索
|
||
substring_results = await _substring_search(store, query, top_k * 2)
|
||
|
||
# 合并结果
|
||
combined = {}
|
||
|
||
for r in vector_results:
|
||
chunk_id = r.chunk.chunk_id
|
||
combined[chunk_id] = (alpha * r.score, r.chunk)
|
||
|
||
for r in bm25_results:
|
||
chunk_id = r.chunk.chunk_id
|
||
if chunk_id in combined:
|
||
old_score, chunk = combined[chunk_id]
|
||
combined[chunk_id] = (old_score + (1 - alpha) * r.score * 0.5, chunk)
|
||
else:
|
||
combined[chunk_id] = ((1 - alpha) * r.score * 0.5, r.chunk)
|
||
|
||
for r in substring_results:
|
||
chunk_id = r.chunk.chunk_id
|
||
if chunk_id in combined:
|
||
old_score, chunk = combined[chunk_id]
|
||
combined[chunk_id] = (old_score + (1 - alpha) * r.score * 0.3, chunk)
|
||
else:
|
||
combined[chunk_id] = ((1 - alpha) * r.score * 0.3, r.chunk)
|
||
|
||
sorted_results = sorted(combined.items(), key=lambda x: x[1][0], reverse=True)
|
||
|
||
from src.rag.base import SearchResult
|
||
return [SearchResult(chunk=chunk, score=score) for _, (score, chunk) in sorted_results[:top_k]]
|
||
|
||
|
||
async def _generate_response_async(query: str, sources: List[RAGSource], use_llm: bool = False) -> str:
|
||
"""生成回答 - 支持LLM总结（异步版本）
|
||
|
||
优化：对于"有哪些物种"类问题，直接查询物种API返回完整列表
|
||
"""
|
||
import re
|
||
|
||
# 检测是否是"有哪些物种"类查询
|
||
is_species_list_query = any(re.search(p, query.lower()) for p in [
|
||
r"有哪些物种", r"哪些物种", r"什么物种", r"哪些.*做过.*分析", r"哪些.*研究过", r".*物种.*列表"
|
||
])
|
||
|
||
# 如果是物种列表查询，直接返回物种API的结果
|
||
if is_species_list_query:
|
||
gene_family = _extract_gene_family(query)
|
||
species_list_data = await _fetch_species_list(gene_family)
|
||
if species_list_data and species_list_data.get("species"):
|
||
return _format_species_list_answer(species_list_data, query)
|
||
|
||
if not sources:
|
||
return "未找到相关文献。"
|
||
|
||
# Pfam/domain accession questions need a deterministic answer. When the
|
||
# LLM is unavailable, falling through to the generic literature-count
|
||
# summary hides the actual accession users asked for.
|
||
pfam_answer = _try_generate_pfam_answer(query, sources)
|
||
if pfam_answer:
|
||
return pfam_answer
|
||
|
||
# 如果使用LLM生成总结
|
||
if use_llm:
|
||
llm_mgr = get_llm_manager()
|
||
if llm_mgr:
|
||
# 构造上下文
|
||
context = "\n\n".join([
|
||
f"【文献{i+1}】{s.content[:500]}"
|
||
for i, s in enumerate(sources[:5])
|
||
])
|
||
|
||
prompt = f"""你是一个基因家族分析专家。请**严格**根据以下检索到的文献内容回答用户问题。
|
||
|
||
【重要规则】
|
||
1. **只使用文献中明确提到的信息**，禁止补充任何训练知识
|
||
2. 如果文献中没有相关信息，必须明确说明："根据提供的文献，未找到相关信息"
|
||
3. 禁止推测、扩展或解释文献中未提及的内容
|
||
4. 每个关键事实必须标注来源（如【文献1】）
|
||
|
||
用户问题: {query}
|
||
|
||
文献内容:
|
||
{context}
|
||
|
||
请生成回答：
|
||
- 只总结文献中明确提到的内容
|
||
- 不添加背景知识、不解释概念、不做推测
|
||
- 如果文献内容不完整，如实说明"文献信息有限"
|
||
- 使用 Markdown 格式，关键信息加粗
|
||
|
||
如果文献内容与问题不相关，请直接说明："提供的文献内容与问题不相关，无法回答"。
|
||
"""
|
||
|
||
try:
|
||
from src.llm.base import LLMMessage
|
||
messages = [LLMMessage(role="user", content=prompt)]
|
||
# 使用异步调用
|
||
response = await llm_mgr.acall(messages)
|
||
return response.content if hasattr(response, 'content') else str(response)
|
||
except Exception as e:
|
||
print(f"LLM生成失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
# 默认统计信息
|
||
species_list = set()
|
||
gf_list = set()
|
||
|
||
for s in sources:
|
||
if 'species' in s.metadata:
|
||
species_list.add(s.metadata['species'])
|
||
if 'gene_family' in s.metadata:
|
||
gf_list.add(s.metadata['gene_family'])
|
||
|
||
response = f"找到 {len(sources)} 条相关文献。"
|
||
|
||
if species_list:
|
||
response += f"\n涉及物种: {', '.join(species_list)}。"
|
||
|
||
if gf_list:
|
||
response += f"\n涉及基因家族: {', '.join(gf_list)}。"
|
||
|
||
return response
|
||
|
||
|
||
def _try_generate_pfam_answer(query: str, sources: List[RAGSource]) -> str:
|
||
"""Answer Pfam accession questions from retrieved evidence when possible."""
|
||
import re
|
||
|
||
query_lower = query.lower()
|
||
is_pfam_query = any(token in query_lower for token in ["pfam", "pfam号", "pfam编号", "pfam id", "结构域编号"])
|
||
if not is_pfam_query:
|
||
return ""
|
||
|
||
gene_family = _extract_gene_family(query).upper()
|
||
try:
|
||
fact = ensure_family_fact(gene_family)
|
||
pfam_items = fact.get("pfam") or []
|
||
if pfam_items:
|
||
pfam_ids = [item.get("pfam_id") for item in pfam_items if item.get("pfam_id")]
|
||
lines = [
|
||
f"{gene_family} 家族的 Pfam 号是：**{', '.join(pfam_ids)}**。",
|
||
"",
|
||
"证据：",
|
||
]
|
||
for item in pfam_items:
|
||
pfam_id = item.get("pfam_id", "")
|
||
for ev in (item.get("evidence") or [])[:3]:
|
||
source_label = ev.get("paper_id") or ev.get("source_path") or "family_fact"
|
||
if ev.get("pmid"):
|
||
source_label += f", PMID: {ev['pmid']}"
|
||
lines.append(f"- {pfam_id}: {ev.get('evidence', '')}（来源: {source_label}）")
|
||
return "\n".join(lines)
|
||
except Exception as exc:
|
||
logger.warning(f"[RAG] family fact Pfam lookup failed: {exc}")
|
||
|
||
evidence_by_pfam: dict[str, dict] = {}
|
||
pfam_pattern = re.compile(r"\bPF\d{5}\b", re.I)
|
||
|
||
for source in sources:
|
||
text = source.content or ""
|
||
metadata_text = " ".join(str(v) for v in (source.metadata or {}).values())
|
||
combined = f"{text} {metadata_text}"
|
||
for match in pfam_pattern.finditer(combined):
|
||
pfam_id = match.group(0).upper()
|
||
if pfam_id in evidence_by_pfam:
|
||
continue
|
||
evidence = _extract_pfam_evidence_sentence(text, pfam_id)
|
||
evidence_by_pfam[pfam_id] = {
|
||
"evidence": evidence or text[:220],
|
||
"doc_id": source.doc_id,
|
||
"pmid": source.pmid or (source.metadata or {}).get("pmid", ""),
|
||
}
|
||
|
||
if not evidence_by_pfam:
|
||
return "根据检索到的文献片段，未找到明确的 Pfam 编号。"
|
||
|
||
pfam_ids = sorted(evidence_by_pfam)
|
||
lines = [
|
||
f"{gene_family} 家族的 Pfam 号是：**{', '.join(pfam_ids)}**。",
|
||
"",
|
||
"证据：",
|
||
]
|
||
for pfam_id in pfam_ids:
|
||
item = evidence_by_pfam[pfam_id]
|
||
source_label = item["doc_id"]
|
||
if item.get("pmid"):
|
||
source_label += f", PMID: {item['pmid']}"
|
||
lines.append(f"- {pfam_id}: {item['evidence']}（来源: {source_label}）")
|
||
|
||
return "\n".join(lines)
|
||
|
||
|
||
def _extract_pfam_evidence_sentence(text: str, pfam_id: str) -> str:
|
||
"""Return a compact sentence-like evidence snippet containing a Pfam ID."""
|
||
import re
|
||
|
||
if not text:
|
||
return ""
|
||
|
||
normalized = re.sub(r"\s+", " ", text).strip()
|
||
sentences = re.split(r"(?<=[。！？.!?])\s+", normalized)
|
||
for sentence in sentences:
|
||
if pfam_id.upper() in sentence.upper():
|
||
return sentence[:260]
|
||
|
||
idx = normalized.upper().find(pfam_id.upper())
|
||
if idx < 0:
|
||
return ""
|
||
start = max(0, idx - 110)
|
||
end = min(len(normalized), idx + 150)
|
||
return normalized[start:end].strip()
|
||
|
||
|
||
def _generate_strict_response(query: str, sources: List[RAGSource]) -> str:
|
||
"""Strict mode: only return original document chunks without summarization.
|
||
|
||
Features:
|
||
1. Only display retrieved document chunks
|
||
2. No summarization, inference, or knowledge supplementation
|
||
3. Clearly label source of each chunk
|
||
4. Honestly report if content is irrelevant
|
||
"""
|
||
import re
|
||
|
||
if not sources:
|
||
return "【严格模式】未检索到相关文献片段。"
|
||
|
||
lines = [
|
||
f"【严格模式 - 仅基于检索到的文献片段】",
|
||
f"查询: {query}",
|
||
f"检索到 {len(sources)} 条相关片段",
|
||
"",
|
||
"=" * 70,
|
||
""
|
||
]
|
||
|
||
# 检查相关性
|
||
query_keywords = set(re.findall(r'[\u4e00-\u9fa5]{2,}|[a-zA-Z]{3,}', query.lower()))
|
||
relevant_count = 0
|
||
|
||
for i, source in enumerate(sources, 1):
|
||
content_lower = source.content.lower()
|
||
matched_keywords = [k for k in query_keywords if k in content_lower]
|
||
relevance = len(matched_keywords) / len(query_keywords) if query_keywords else 0
|
||
|
||
if relevance > 0.3:
|
||
relevant_count += 1
|
||
|
||
lines.extend([
|
||
f"【文献片段 {i}】",
|
||
f"来源ID: {source.doc_id}",
|
||
f"相关度: {source.score:.3f}",
|
||
f"匹配关键词: {', '.join(matched_keywords) if matched_keywords else '无直接匹配'}",
|
||
f"内容类型: {source.metadata.get('type', 'unknown')}",
|
||
"-" * 50,
|
||
source.content,
|
||
"",
|
||
"=" * 70,
|
||
""
|
||
])
|
||
|
||
# 添加相关性评估
|
||
lines.extend([
|
||
"",
|
||
"【相关性评估】",
|
||
f"- 总检索片段: {len(sources)}",
|
||
f"- 高度相关片段: {relevant_count}",
|
||
f"- 相关度比例: {relevant_count/len(sources)*100:.1f}%" if sources else "N/A",
|
||
])
|
||
|
||
if relevant_count == 0:
|
||
lines.extend([
|
||
"",
|
||
"⚠️ 警告：检索到的片段与查询问题相关性较低，",
|
||
" 可能无法准确回答问题。"
|
||
])
|
||
|
||
return "\n".join(lines)
|
||
|
||
|
||
def _extract_gene_family(query: str) -> str:
|
||
"""从查询中提取基因家族名称"""
|
||
import re
|
||
gf_patterns = [
|
||
r"(SPL|NAC|WRKY|MYB|bHLH|AP2|ERF|bZIP|Dof|SBP)[^a-zA-Z]",
|
||
r"(\w+).*基因家族",
|
||
r"(\w+).*gene family",
|
||
]
|
||
query_upper = query.upper()
|
||
for pattern in gf_patterns:
|
||
match = re.search(pattern, query_upper)
|
||
if match:
|
||
return match.group(1)
|
||
return "SPL"
|
||
|
||
|
||
async def _fetch_species_list(gene_family: str) -> dict:
|
||
"""从物种索引API获取物种列表"""
|
||
try:
|
||
import aiohttp
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.get(
|
||
f"http://localhost:8000/api/v1/species/list?gene_family={gene_family}"
|
||
) as response:
|
||
if response.status == 200:
|
||
return await response.json()
|
||
except Exception as e:
|
||
logger.warning(f"[RAG] 获取物种列表失败: {e}")
|
||
return None
|
||
|
||
|
||
def _format_species_list_answer(species_data: dict, query: str) -> str:
|
||
"""格式化物种列表回答"""
|
||
if not species_data or "species" not in species_data:
|
||
return "未找到相关物种信息。"
|
||
|
||
species_list = species_data["species"]
|
||
gene_family = species_data.get("gene_family", "SPL")
|
||
total_species = species_data.get("total_species", len(species_list))
|
||
|
||
lines = [
|
||
f"## {gene_family}基因家族已分析的物种（共{total_species}种）",
|
||
"",
|
||
"根据基因家族-物种索引，已对以下物种开展全基因组鉴定与分析：",
|
||
"",
|
||
"| 物种 | 学名 | 比较物种 | 关系数 | 实体数 |",
|
||
"|:---|:---|:---|:---:|:---:|"
|
||
]
|
||
|
||
for sp in species_list:
|
||
name = sp.get('common_name_cn', 'N/A')
|
||
scientific = sp.get('scientific_name', 'N/A')
|
||
# 显示完整的比较物种列表
|
||
comparison = ', '.join(sp.get('comparison_species', []))
|
||
relations = sp.get('relations_count', 0)
|
||
entities = sp.get('entities_count', 0)
|
||
lines.append(f"| {name} | {scientific} | {comparison} | {relations} | {entities} |")
|
||
|
||
lines.extend([
|
||
"",
|
||
"### 说明",
|
||
"",
|
||
f"- **{gene_family}基因家族**：SQUAMOSA启动子结合蛋白样基因家族",
|
||
"- **研究覆盖**：木本经济作物、油料作物、蔬菜作物、药用植物及模式植物",
|
||
"- **比较分析**：所有研究均以拟南芥和水稻作为核心参照物种",
|
||
"",
|
||
"如需了解某个物种的详细分析结果，请直接提问，如'苦荞SPL基因家族功能'。"
|
||
])
|
||
|
||
return '\n'.join(lines)
|
||
|
||
|
||
@router.get("/rag/status")
|
||
async def rag_status():
|
||
"""RAG系统状态"""
|
||
try:
|
||
store = get_vector_store()
|
||
return {
|
||
"status": "ready",
|
||
"total_chunks": len(store._chunks),
|
||
"total_documents": len(set(c.doc_id for c in store._chunks.values())),
|
||
"agent_enabled": True
|
||
}
|
||
except Exception as e:
|
||
return {
|
||
"status": "error",
|
||
"error": str(e)
|
||
}
|
||
|
||
|
||
# ===== 用户反馈API =====
|
||
|
||
from datetime import datetime
|
||
import json
|
||
import sqlite3
|
||
|
||
# 反馈数据库路径
|
||
FEEDBACK_DB = PROJECT_ROOT / "data" / "feedback.db"
|
||
|
||
def init_feedback_db():
|
||
"""初始化反馈数据库"""
|
||
FEEDBACK_DB.parent.mkdir(parents=True, exist_ok=True)
|
||
conn = sqlite3.connect(FEEDBACK_DB)
|
||
conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS feedback (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
query TEXT NOT NULL,
|
||
answer TEXT,
|
||
rating INTEGER, -- 1-5 评分
|
||
feedback_type TEXT, -- thumbs_up/thumbs_down/correction
|
||
feedback_text TEXT, -- 用户反馈内容
|
||
session_id TEXT,
|
||
timestamp TEXT NOT NULL
|
||
)
|
||
""")
|
||
conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS evaluation_history (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
question_id TEXT,
|
||
query TEXT NOT NULL,
|
||
query_type TEXT,
|
||
ndcg REAL,
|
||
answer_quality REAL,
|
||
timestamp TEXT NOT NULL
|
||
)
|
||
""")
|
||
conn.commit()
|
||
conn.close()
|
||
|
||
|
||
class FeedbackRequest(BaseModel):
|
||
query: str
|
||
answer: str = None
|
||
rating: int = None # 1-5
|
||
feedback_type: str = None # thumbs_up/thumbs_down/correction
|
||
feedback_text: str = None
|
||
session_id: str = None
|
||
|
||
|
||
class EvaluationRecordRequest(BaseModel):
|
||
question_id: str
|
||
query: str
|
||
query_type: str
|
||
ndcg: float
|
||
answer_quality: float
|
||
|
||
|
||
@router.post("/rag/feedback")
|
||
async def submit_feedback(request: FeedbackRequest):
|
||
"""提交用户反馈"""
|
||
try:
|
||
init_feedback_db()
|
||
conn = sqlite3.connect(FEEDBACK_DB)
|
||
cursor = conn.cursor()
|
||
|
||
cursor.execute("""
|
||
INSERT INTO feedback (query, answer, rating, feedback_type, feedback_text, session_id, timestamp)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?)
|
||
""", (
|
||
request.query,
|
||
request.answer,
|
||
request.rating,
|
||
request.feedback_type,
|
||
request.feedback_text,
|
||
request.session_id,
|
||
datetime.now().isoformat()
|
||
))
|
||
|
||
conn.commit()
|
||
conn.close()
|
||
|
||
return {"status": "success", "message": "反馈已保存"}
|
||
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=f"保存反馈失败: {str(e)}")
|
||
|
||
|
||
@router.post("/rag/evaluation/record")
|
||
async def record_evaluation(request: EvaluationRecordRequest):
|
||
"""记录评估结果"""
|
||
try:
|
||
init_feedback_db()
|
||
conn = sqlite3.connect(FEEDBACK_DB)
|
||
cursor = conn.cursor()
|
||
|
||
cursor.execute("""
|
||
INSERT INTO evaluation_history (question_id, query, query_type, ndcg, answer_quality, timestamp)
|
||
VALUES (?, ?, ?, ?, ?, ?)
|
||
""", (
|
||
request.question_id,
|
||
request.query,
|
||
request.query_type,
|
||
request.ndcg,
|
||
request.answer_quality,
|
||
datetime.now().isoformat()
|
||
))
|
||
|
||
conn.commit()
|
||
conn.close()
|
||
|
||
return {"status": "success", "message": "评估记录已保存"}
|
||
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=f"保存评估记录失败: {str(e)}")
|
||
|
||
|
||
@router.get("/rag/feedback/stats")
|
||
async def get_feedback_stats():
|
||
"""获取反馈统计"""
|
||
try:
|
||
init_feedback_db()
|
||
conn = sqlite3.connect(FEEDBACK_DB)
|
||
cursor = conn.cursor()
|
||
|
||
# 总反馈数
|
||
cursor.execute("SELECT COUNT(*) FROM feedback")
|
||
total = cursor.fetchone()[0]
|
||
|
||
# 评分分布
|
||
cursor.execute("""
|
||
SELECT rating, COUNT(*)
|
||
FROM feedback
|
||
WHERE rating IS NOT NULL
|
||
GROUP BY rating
|
||
""")
|
||
rating_dist = {row[0]: row[1] for row in cursor.fetchall()}
|
||
|
||
# 反馈类型分布
|
||
cursor.execute("""
|
||
SELECT feedback_type, COUNT(*)
|
||
FROM feedback
|
||
WHERE feedback_type IS NOT NULL
|
||
GROUP BY feedback_type
|
||
""")
|
||
feedback_type_dist = {row[0]: row[1] for row in cursor.fetchall()}
|
||
|
||
# 平均评分
|
||
cursor.execute("SELECT AVG(rating) FROM feedback WHERE rating IS NOT NULL")
|
||
avg_rating = cursor.fetchone()[0] or 0
|
||
|
||
conn.close()
|
||
|
||
return {
|
||
"total_feedback": total,
|
||
"rating_distribution": rating_dist,
|
||
"feedback_type_distribution": feedback_type_dist,
|
||
"average_rating": round(avg_rating, 2)
|
||
}
|
||
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=f"获取统计失败: {str(e)}")
|
||
|
||
|
||
@router.get("/rag/evaluation/history")
|
||
async def get_evaluation_history(limit: int = 50):
|
||
"""获取评估历史"""
|
||
try:
|
||
init_feedback_db()
|
||
conn = sqlite3.connect(FEEDBACK_DB)
|
||
cursor = conn.cursor()
|
||
|
||
cursor.execute("""
|
||
SELECT question_id, query, query_type, ndcg, answer_quality, timestamp
|
||
FROM evaluation_history
|
||
ORDER BY timestamp DESC
|
||
LIMIT ?
|
||
""", (limit,))
|
||
|
||
results = []
|
||
for row in cursor.fetchall():
|
||
results.append({
|
||
"question_id": row[0],
|
||
"query": row[1],
|
||
"query_type": row[2],
|
||
"ndcg": row[3],
|
||
"answer_quality": row[4],
|
||
"timestamp": row[5]
|
||
})
|
||
|
||
conn.close()
|
||
|
||
return {"history": results, "count": len(results)}
|
||
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=f"获取历史失败: {str(e)}")
|
||
|
||
|
||
# ===== 高级RAG功能 =====
|
||
|
||
from src.rag.advanced import (
|
||
MultiGranularChunker,
|
||
HierarchicalRetriever,
|
||
AdaptiveRetriever,
|
||
RetrievalConfig,
|
||
QueryDecomposer,
|
||
ContextCompressor
|
||
)
|
||
|
||
# 高级组件实例
|
||
_advanced_chunker = None
|
||
_advanced_retriever = None
|
||
_query_decomposer = None
|
||
_context_compressor = None
|
||
|
||
|
||
def get_advanced_components():
|
||
"""获取高级组件"""
|
||
global _advanced_chunker, _advanced_retriever, _query_decomposer, _context_compressor
|
||
|
||
if _advanced_chunker is None:
|
||
_advanced_chunker = MultiGranularChunker()
|
||
_query_decomposer = QueryDecomposer()
|
||
_context_compressor = ContextCompressor()
|
||
|
||
# retriever需要vector store，暂不初始化
|
||
_advanced_retriever = None
|
||
|
||
return {
|
||
"chunker": _advanced_chunker,
|
||
"retriever": _advanced_retriever,
|
||
"decomposer": _query_decomposer,
|
||
"compressor": _context_compressor
|
||
}
|
||
|
||
|
||
class AdvancedSearchRequest(BaseModel):
|
||
query: str
|
||
top_k: int = 10
|
||
use_hierarchical: bool = False # 使用分层检索
|
||
use_query_decompose: bool = False # 使用查询分解
|
||
use_compression: bool = True # 使用上下文压缩
|
||
query_type: str = "general" # 查询类型
|
||
# Self-RAG + MultiHopRAG 参数
|
||
use_agentic: bool = False # 使用Agent编排
|
||
enable_self_rag: bool = False # 启用Self-RAG反思
|
||
enable_multi_hop: bool = False # 启用MultiHopRAG多跳检索
|
||
max_iterations: int = 3 # Self-RAG最大迭代次数
|
||
|
||
|
||
class QueryDecomposeRequest(BaseModel):
|
||
query: str
|
||
use_llm: bool = False
|
||
|
||
|
||
@router.post("/rag/advanced/search")
|
||
async def advanced_search(request: AdvancedSearchRequest):
|
||
"""高级检索接口 - 支持Self-RAG和MultiHopRAG"""
|
||
try:
|
||
comps = get_advanced_components()
|
||
|
||
results = {
|
||
"query": request.query,
|
||
"features": {
|
||
"hierarchical": request.use_hierarchical,
|
||
"query_decompose": request.use_query_decompose,
|
||
"compression": request.use_compression,
|
||
"agentic": request.use_agentic,
|
||
"self_rag": request.enable_self_rag,
|
||
"multi_hop": request.enable_multi_hop
|
||
}
|
||
}
|
||
|
||
# 1. 多跳推理模式 (直接使用MultiHopRAG)
|
||
if request.enable_multi_hop:
|
||
# 直接调用multi_hop_rag
|
||
try:
|
||
from src.rag.multi_hop_rag import MultiHopRAG
|
||
from src.rag.hybrid_retriever import HybridRetriever, SearchConfig
|
||
|
||
store = get_vector_store()
|
||
llm_mgr = get_llm_manager()
|
||
|
||
# 创建HybridRetriever包装store（MultiHopRAG需要正确接口）
|
||
config = SearchConfig(top_k=5, enable_query_rewrite=True)
|
||
hybrid_retriever = HybridRetriever(
|
||
vector_store=store,
|
||
config=config,
|
||
embedding_model=_embedding_model
|
||
)
|
||
|
||
# 创建MultiHopRAG实例
|
||
multi_hop = MultiHopRAG(
|
||
retriever=hybrid_retriever, # 使用HybridRetriever
|
||
llm_client=llm_mgr,
|
||
max_hops=request.max_iterations,
|
||
enable_web_search=False
|
||
)
|
||
|
||
# 执行多跳推理
|
||
hop_result = await multi_hop.query(request.query)
|
||
|
||
results["multi_hop_result"] = {
|
||
"final_answer": hop_result.final_answer,
|
||
"total_hops": hop_result.total_hops,
|
||
"sources": hop_result.sources[:5]
|
||
}
|
||
|
||
return results
|
||
|
||
except Exception as e:
|
||
results["error"] = f"MultiHopRAG failed: {str(e)}"
|
||
# 回退到普通检索
|
||
|
||
# 2. Self-RAG模式（暂不支持，agentic已移除）
|
||
if request.enable_self_rag or request.use_agentic:
|
||
results["message"] = "Self-RAG/Agentic模式已移除，请使用enable_multi_hop"
|
||
return results
|
||
|
||
# 2. 查询分解
|
||
if request.use_query_decompose:
|
||
sub_queries = comps["decomposer"].decompose(request.query)
|
||
results["sub_queries"] = [
|
||
{
|
||
"query": sq.query,
|
||
"reason": sq.reason,
|
||
"expected_docs": sq.expected_docs
|
||
}
|
||
for sq in sub_queries
|
||
]
|
||
|
||
# 3. 分层检索（需要实际索引）
|
||
if request.use_hierarchical:
|
||
results["message"] = "分层检索需要先构建多粒度索引"
|
||
|
||
# 4. 上下文压缩
|
||
if request.use_compression:
|
||
results["compression"] = "已启用上下文压缩"
|
||
|
||
# 5. 普通检索结果：保持与 gfa_core.mas.rag.RAGRetriever.search 的契约一致。
|
||
store = get_vector_store()
|
||
search_results = await _hybrid_search(store, request.query, max(1, min(request.top_k, 50)))
|
||
results["results"] = [
|
||
{
|
||
"doc_id": item.chunk.doc_id,
|
||
"content": item.chunk.content[:1000],
|
||
"score": item.score,
|
||
"metadata": item.chunk.metadata,
|
||
"pmid": item.chunk.metadata.get("pmid"),
|
||
}
|
||
for item in search_results
|
||
]
|
||
results["count"] = len(results["results"])
|
||
|
||
return results
|
||
|
||
except Exception as e:
|
||
import traceback
|
||
traceback.print_exc()
|
||
raise HTTPException(status_code=500, detail=f"高级检索失败: {str(e)}")
|
||
|
||
|
||
@router.post("/rag/advanced/decompose")
|
||
async def decompose_query(request: QueryDecomposeRequest):
|
||
"""查询分解接口"""
|
||
try:
|
||
comps = get_advanced_components()
|
||
|
||
if request.use_llm:
|
||
sub_queries = await comps["decomposer"].decompose_with_llm(request.query)
|
||
else:
|
||
sub_queries = comps["decomposer"].decompose(request.query)
|
||
|
||
return {
|
||
"original_query": request.query,
|
||
"sub_queries": [
|
||
{
|
||
"query": sq.query,
|
||
"reason": sq.reason,
|
||
"expected_docs": sq.expected_docs,
|
||
"priority": sq.priority
|
||
}
|
||
for sq in sub_queries
|
||
],
|
||
"count": len(sub_queries)
|
||
}
|
||
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=f"查询分解失败: {str(e)}")
|
||
|
||
|
||
def _extract_knowledge_graph(sources: List[RAGSource], query: str) -> KnowledgeGraphData:
|
||
"""
|
||
从RAG检索结果动态提取知识图谱
|
||
提取基因名、物种、表型、过程等实体和关系
|
||
"""
|
||
import re
|
||
|
||
nodes = []
|
||
edges = []
|
||
node_ids = set()
|
||
|
||
# 基因名模式 (支持任意物种)
|
||
gene_pattern = r'\b([A-Z][a-z]?[A-Z]{1,4}\d{1,4}[a-z]?)\b'
|
||
|
||
# 数据集 ID 模式 (V15.1 新增)
|
||
dataset_pattern = r'\b(PRJNA\d+|SRP\d+|GSE\d+|GSM\d+|SRX\d+|SRR\d+)\b'
|
||
|
||
# 物种模式
|
||
species_pattern = r'\b(Arabidopsis thaliana|Oryza sativa|Zea mays|Glycine max|Solanum lycopersicum|Malus domestica|Carya illinoinensis|pecan|Arabidopsis|rice|maize|soybean|tomato|apple)\b'
|
||
|
||
# 过程/功能关键词
|
||
process_keywords = ['development', 'growth', 'regulation', 'signaling', 'biosynthesis',
|
||
'response', 'metabolism', 'expression', 'flowering', 'fruit development']
|
||
|
||
# 组织/器官关键词
|
||
tissue_keywords = ['root', 'shoot', 'leaf', 'flower', 'fruit', 'seed', 'embryo',
|
||
'endosperm', 'pollen', 'ovule', 'meristem', 'bud']
|
||
|
||
# 处理每个source
|
||
for source in sources:
|
||
content = source.content
|
||
doc_id = source.doc_id
|
||
|
||
# 提取基因
|
||
genes = re.findall(gene_pattern, content)
|
||
for gene in genes:
|
||
if gene not in node_ids:
|
||
nodes.append({
|
||
"id": gene,
|
||
"type": "gene"
|
||
})
|
||
node_ids.add(gene)
|
||
|
||
# 提取数据集 (V15.1 新增)
|
||
datasets = re.findall(dataset_pattern, content)
|
||
for ds_id in datasets:
|
||
if ds_id not in node_ids:
|
||
# 判断数据集类型
|
||
if ds_id.startswith('GSE') or ds_id.startswith('GSM'):
|
||
ds_type = 'geo_dataset'
|
||
else:
|
||
ds_type = 'sra_dataset'
|
||
nodes.append({
|
||
"id": ds_id,
|
||
"type": ds_type
|
||
})
|
||
node_ids.add(ds_id)
|
||
|
||
# 提取物种
|
||
species_list = re.findall(species_pattern, content, re.IGNORECASE)
|
||
for sp in species_list:
|
||
sp_key = sp.lower()
|
||
if sp_key not in node_ids:
|
||
nodes.append({
|
||
"id": sp,
|
||
"type": "species"
|
||
})
|
||
node_ids.add(sp_key)
|
||
|
||
# 提取过程
|
||
for proc in process_keywords:
|
||
if proc.lower() in content.lower():
|
||
proc_key = f"{proc}_{doc_id[:8]}"
|
||
if proc_key not in node_ids:
|
||
nodes.append({
|
||
"id": proc,
|
||
"type": "process"
|
||
})
|
||
node_ids.add(proc_key)
|
||
|
||
# 提取组织
|
||
for tissue in tissue_keywords:
|
||
if tissue.lower() in content.lower():
|
||
tissue_key = f"{tissue}_{doc_id[:8]}"
|
||
if tissue_key not in node_ids:
|
||
nodes.append({
|
||
"id": tissue,
|
||
"type": "tissue"
|
||
})
|
||
node_ids.add(tissue_key)
|
||
|
||
# 构建关系：基因-基因 (共现)
|
||
if len(genes) >= 2:
|
||
for i, g1 in enumerate(genes[:5]): # 限制避免过多边
|
||
for g2 in genes[i+1:6]:
|
||
edges.append({
|
||
"source": g1,
|
||
"target": g2,
|
||
"relation": "co_occurrence"
|
||
})
|
||
|
||
# 构建关系：基因-过程
|
||
for gene in genes[:3]:
|
||
for proc in process_keywords:
|
||
if proc.lower() in content.lower():
|
||
edges.append({
|
||
"source": gene,
|
||
"target": proc,
|
||
"relation": "involved_in"
|
||
})
|
||
|
||
# 构建关系：基因-组织 (表达)
|
||
for gene in genes[:3]:
|
||
for tissue in tissue_keywords:
|
||
if tissue.lower() in content.lower():
|
||
edges.append({
|
||
"source": gene,
|
||
"target": tissue,
|
||
"relation": "expressed_in"
|
||
})
|
||
|
||
# 去重边
|
||
unique_edges = []
|
||
edge_keys = set()
|
||
for edge in edges:
|
||
key = f"{edge['source']}->{edge['target']}"
|
||
if key not in edge_keys:
|
||
unique_edges.append(edge)
|
||
edge_keys.add(key)
|
||
|
||
# V15.1: 为检测到的数据集节点添加文章连接
|
||
try:
|
||
kg_path = PROJECT_ROOT / "frontend" / "knowledge_graph_with_omics.json"
|
||
if kg_path.exists():
|
||
with open(kg_path, 'r', encoding='utf-8') as f:
|
||
full_kg = json.load(f)
|
||
|
||
# 构建数据集到文章连接的映射
|
||
dataset_connections = {}
|
||
for edge in full_kg.get('edges', []):
|
||
if edge.get('relation') in ['from_article', 'from_paper', 'has_data_type', 'from_species']:
|
||
source = edge.get('source', '')
|
||
if 'PRJNA' in source or 'SRP' in source or 'GSE' in source:
|
||
if source not in dataset_connections:
|
||
dataset_connections[source] = []
|
||
dataset_connections[source].append(edge)
|
||
|
||
# 为当前知识图谱中的数据集节点添加连接
|
||
dataset_ids = {n['id'] for n in nodes if 'dataset' in n.get('type', '')}
|
||
for ds_id in dataset_ids:
|
||
if ds_id in dataset_connections:
|
||
for conn in dataset_connections[ds_id]:
|
||
target = conn.get('target', '')
|
||
target_type = conn.get('target_type', 'unknown')
|
||
|
||
# 添加目标节点（如果不存在）
|
||
if target and target not in node_ids:
|
||
nodes.append({
|
||
"id": target,
|
||
"type": target_type
|
||
})
|
||
node_ids.add(target)
|
||
|
||
# 添加连接边
|
||
edge_key = f"{ds_id}->{target}"
|
||
if edge_key not in edge_keys:
|
||
unique_edges.append({
|
||
"source": ds_id,
|
||
"target": target,
|
||
"relation": conn.get('relation', 'connected_to')
|
||
})
|
||
edge_keys.add(edge_key)
|
||
except Exception as e:
|
||
logger.warning(f"添加数据集文章连接失败: {e}")
|
||
|
||
return KnowledgeGraphData(nodes=nodes, edges=unique_edges)
|
||
|
||
|
||
def _query_omics_datasets(query: str) -> List[OmicsDataset]:
|
||
"""
|
||
检测用户查询意图，从知识图谱中查找多组学数据集
|
||
根据查询中的物种关键词过滤相关数据集
|
||
"""
|
||
import re
|
||
import json
|
||
|
||
# 数据集查询关键词（中英文）
|
||
omics_keywords = [
|
||
'转录组', 'rna-seq', 'rnaseq', 'rna seq', 'transcriptome',
|
||
'sra', 'geo', '数据集', 'dataset', '编号', 'accession',
|
||
'prjna', 'gse', 'gsm', 'srp', 'srr', 'srx'
|
||
]
|
||
|
||
# 检查查询是否包含数据集相关关键词
|
||
query_lower = query.lower()
|
||
is_omics_query = any(kw in query_lower for kw in omics_keywords)
|
||
|
||
if not is_omics_query:
|
||
return []
|
||
|
||
# 从查询中提取物种关键词
|
||
species_keywords = {
|
||
'pecan': 'Carya illinoinensis',
|
||
'山核桃': 'Carya illinoinensis',
|
||
'carya': 'Carya illinoinensis',
|
||
'谷子': 'Setaria italica',
|
||
'foxtail millet': 'Setaria italica',
|
||
'setaria': 'Setaria italica',
|
||
'向日葵': 'Helianthus annuus',
|
||
'sunflower': 'Helianthus annuus',
|
||
'柑橘': 'Citrus',
|
||
'citrus': 'Citrus',
|
||
}
|
||
|
||
# 确定查询的物种
|
||
target_species = None
|
||
for keyword, species in species_keywords.items():
|
||
if keyword in query_lower:
|
||
target_species = species
|
||
break
|
||
|
||
# 从知识图谱文件中查找数据集
|
||
try:
|
||
kg_path = PROJECT_ROOT / "frontend" / "knowledge_graph_with_omics.json"
|
||
if not kg_path.exists():
|
||
kg_path = PROJECT_ROOT / "frontend" / "knowledge_graph.json"
|
||
|
||
with open(kg_path, 'r', encoding='utf-8') as f:
|
||
kg_data = json.load(f)
|
||
|
||
datasets = []
|
||
for node in kg_data.get('nodes', []):
|
||
node_type = node.get('type', '')
|
||
node_id = node.get('id', '')
|
||
|
||
# 查找 sra_dataset 类型的节点
|
||
if node_type == 'sra_dataset':
|
||
# 推断物种（从节点ID或名称）
|
||
species = 'Unknown'
|
||
if 'PRJNA' in node_id:
|
||
species = 'Carya illinoinensis' # PRJNA435846 是 pecan
|
||
elif any(x in node_id for x in ['SRP028', 'SRP051']):
|
||
species = 'Setaria italica' # 谷子
|
||
elif 'GSE221055' in node_id:
|
||
species = 'Helianthus annuus' # 向日葵
|
||
elif any(x in node_id for x in ['SRP162', 'SRP355', 'SRP357']):
|
||
species = 'Citrus' # 柑橘
|
||
|
||
# 如果指定了目标物种，只返回匹配的数据集
|
||
if target_species and species != target_species:
|
||
continue
|
||
|
||
datasets.append(OmicsDataset(
|
||
dataset_id=node_id,
|
||
dataset_type='transcriptome',
|
||
species=species,
|
||
source_paper='V15 Knowledge Graph'
|
||
))
|
||
|
||
return datasets
|
||
except Exception as e:
|
||
logger.warning(f"查询多组学数据集失败: {e}")
|
||
return []
|
||
|
||
|
||
def _add_omics_info_to_response(response: str, datasets: List[OmicsDataset]) -> str:
|
||
"""
|
||
在回答中添加多组学数据集信息
|
||
"""
|
||
if not datasets:
|
||
return response
|
||
|
||
# 构建数据集信息文本
|
||
omics_info = "\n\n📊 **多组学数据集信息**:\n\n"
|
||
omics_info += "| 数据集编号 | 类型 | 物种 |\n"
|
||
omics_info += "|-----------|------|------|\n"
|
||
|
||
for ds in datasets:
|
||
omics_info += f"| {ds.dataset_id} | {ds.dataset_type} | {ds.species} |\n"
|
||
|
||
omics_info += "\n💡 您可以在 NCBI SRA/GEO 数据库中使用这些编号检索原始数据。"
|
||
|
||
# 在回答末尾添加（在参考文献之前）
|
||
if "📚 参考文献" in response:
|
||
response = response.replace("📚 参考文献", omics_info + "\n📚 参考文献")
|
||
else:
|
||
response += omics_info
|
||
|
||
return response
|
||
|
||
|
||
def _extract_kg_entities(sources: List[RAGSource]) -> List[KGEntityInfo]:
|
||
"""
|
||
从检索结果中提取可用于Neo4j可视化的实体
|
||
优先从metadata中提取gene_family和species
|
||
"""
|
||
entities = []
|
||
seen = set()
|
||
|
||
for source in sources:
|
||
metadata = source.metadata
|
||
|
||
# 提取基因家族
|
||
if 'gene_family' in metadata:
|
||
family = metadata['gene_family']
|
||
if family and family not in seen:
|
||
entities.append(KGEntityInfo(
|
||
name=family,
|
||
type='family',
|
||
display_name=family
|
||
))
|
||
seen.add(family)
|
||
|
||
# 提取物种
|
||
if 'species' in metadata:
|
||
species = metadata['species']
|
||
if species and species not in seen:
|
||
entities.append(KGEntityInfo(
|
||
name=species,
|
||
type='species',
|
||
display_name=species
|
||
))
|
||
seen.add(species)
|
||
|
||
# 如果metadata中没有，尝试从内容中提取（简化版）
|
||
if not entities:
|
||
import re
|
||
for source in sources[:3]: # 只看前3个source
|
||
content = source.content
|
||
|
||
# 尝试匹配基因家族名（大写字母组合）
|
||
family_match = re.search(r'\b([A-Z]{2,10})\s+gene family', content, re.IGNORECASE)
|
||
if family_match:
|
||
family = family_match.group(1).upper()
|
||
if family not in seen:
|
||
entities.append(KGEntityInfo(
|
||
name=family,
|
||
type='family',
|
||
display_name=family
|
||
))
|
||
seen.add(family)
|
||
break
|
||
|
||
return entities
|
||
|
||
|
||
@router.get("/rag/advanced/status")
|
||
async def advanced_status():
|
||
"""高级功能状态"""
|
||
comps = get_advanced_components()
|
||
|
||
return {
|
||
"status": "ready",
|
||
"features": {
|
||
"multi_granular_chunking": True,
|
||
"hierarchical_retrieval": comps["retriever"] is not None,
|
||
"query_decomposition": True,
|
||
"context_compression": True,
|
||
"citation_extraction": True,
|
||
"phylogeny_extraction": True # v5新增
|
||
},
|
||
"chunk_levels": ["summary", "paragraph", "sentence", "keyword"]
|
||
}
|
||
|
||
|
||
# ========== v5 系统发育提取接口 ==========
|
||
|
||
@router.post("/rag/phylogeny/extract")
|
||
async def extract_phylogeny(request: PhylogenyExtractionRequest):
|
||
"""
|
||
系统发育信息提取接口 (v5.2 - TAIR关键词库增强版)
|
||
|
||
自动从PDF提取系统发育分组信息，集成TAIR权威基因家族数据，自动消歧
|
||
包含4轮审查流程和关键词匹配详情
|
||
"""
|
||
try:
|
||
# 导入v5.2提取器 (TAIR关键词库增强版)
|
||
sys.path.insert(0, str(PROJECT_ROOT))
|
||
from extract_phylogeny_v5_universal_v3 import extract_phylogeny_v52
|
||
|
||
print(f"[Phylogeny Extract v5.2] 开始提取: {request.pdf_path}")
|
||
print(f"[Phylogeny Extract v5.2] DOI: {request.doi}")
|
||
print(f"[Phylogeny Extract v5.2] PMCID: {request.pmcid}")
|
||
|
||
# 读取PDF文本
|
||
import fitz
|
||
doc = fitz.open(request.pdf_path)
|
||
text = ""
|
||
for page in doc:
|
||
text += page.get_text()
|
||
doc.close()
|
||
|
||
# 执行提取（v5.2包含TAIR关键词库）
|
||
result = extract_phylogeny_v52(text)
|
||
|
||
# 返回结果（包含审查详情和TAIR关键词匹配）
|
||
ext = result.get("extraction", {})
|
||
return {
|
||
"status": "success",
|
||
"pdf_name": Path(request.pdf_path).name,
|
||
"gene_family": ext.get("gene_family"),
|
||
"species": ext.get("species"),
|
||
"total_genes": ext.get("total_genes"),
|
||
"total_groups": ext.get("total_groups"),
|
||
"confidence_score": result.get("confidence"),
|
||
"confidence_level": result.get("assessment", {}).get("level"),
|
||
"source": "TAIR+v5.2",
|
||
# v5.2新增：TAIR关键词匹配详情
|
||
"matched_keywords": ext.get("matched_keywords", {}),
|
||
# v5.2新增：审查详情
|
||
"scrutiny": result.get("scrutiny", {}),
|
||
"validation": result.get("validation", {}),
|
||
"assessment": result.get("assessment", {}),
|
||
"version": result.get("version", "v5.2-TAIR")
|
||
}
|
||
|
||
except Exception as e:
|
||
import traceback
|
||
print(f"[Phylogeny Extract v5.1] 错误: {e}")
|
||
print(traceback.format_exc())
|
||
raise HTTPException(status_code=500, detail=f"提取失败: {str(e)}")
|
||
|
||
|
||
@router.get("/rag/phylogeny/status/{pdf_name}")
|
||
async def get_phylogeny_status(pdf_name: str):
|
||
"""
|
||
查询系统发育提取状态
|
||
"""
|
||
try:
|
||
result_path = PROJECT_ROOT / "evaluation" / "auto_results" / f"{pdf_name}_v5_report.json"
|
||
|
||
if not result_path.exists():
|
||
return {"status": "not_found", "message": "提取结果不存在"}
|
||
|
||
with open(result_path) as f:
|
||
result = json.load(f)
|
||
|
||
return {
|
||
"status": "success",
|
||
"pdf_name": pdf_name,
|
||
"confidence_score": result.get("confidence_score"),
|
||
"confidence_level": result.get("confidence_level"),
|
||
"total_groups": result.get("total_groups"),
|
||
"has_supplementary": result.get("has_supplementary")
|
||
}
|
||
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=f"查询失败: {str(e)}")
|