175 lines
4.9 KiB
Python
Executable File
175 lines
4.9 KiB
Python
Executable File
"""
|
|
RAG v10专用路由 - 使用苦荞SPL附表增强版索引
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import json
|
|
import pickle
|
|
import numpy as np
|
|
from pathlib import Path
|
|
from fastapi import APIRouter, HTTPException
|
|
from pydantic import BaseModel
|
|
from typing import List, Optional
|
|
import requests
|
|
from src.rag.project_config import PROJECT_ROOT
|
|
|
|
sys.path.insert(0, str(PROJECT_ROOT))
|
|
|
|
router = APIRouter(prefix="/v10")
|
|
|
|
# ========== 配置 ==========
|
|
OLLAMA_URL = "http://localhost:11434"
|
|
EMBEDDING_MODEL = "bge-m3:567m"
|
|
LLM_MODEL = "qwen3:0.6b"
|
|
|
|
# v10索引路径
|
|
V10_INDEX_PATH = PROJECT_ROOT / "paper" / "v10" / "rag_v10_faiss_final.index"
|
|
V10_CHUNKS_PATH = PROJECT_ROOT / "paper" / "v10" / "rag_v10_chunks_with_summary.json"
|
|
V10_ENHANCED_PATH = PROJECT_ROOT / "paper" / "v10" / "rag_v10_chunks_enhanced.json"
|
|
|
|
# 全局变量
|
|
_v10_index = None
|
|
_v10_chunks = None
|
|
|
|
def get_embedding(text):
|
|
"""获取文本向量"""
|
|
response = requests.post(
|
|
f"{OLLAMA_URL}/api/embeddings",
|
|
json={"model": EMBEDDING_MODEL, "prompt": text}
|
|
)
|
|
vec = np.array(response.json()['embedding'], dtype=np.float32)
|
|
return vec / np.linalg.norm(vec)
|
|
|
|
def load_v10_index():
|
|
"""加载v10索引"""
|
|
global _v10_index, _v10_chunks
|
|
|
|
if _v10_index is not None:
|
|
return _v10_index, _v10_chunks
|
|
|
|
print(f"[V10 RAG] 加载v10索引...")
|
|
|
|
# 加载FAISS索引
|
|
import faiss
|
|
_v10_index = faiss.read_index(str(V10_INDEX_PATH))
|
|
print(f"[V10 RAG] 索引加载完成: {_v10_index.ntotal} 个向量")
|
|
|
|
# 加载文本块
|
|
with open(V10_CHUNKS_PATH, 'r') as f:
|
|
chunks1 = json.load(f)
|
|
with open(V10_ENHANCED_PATH, 'r') as f:
|
|
chunks2 = json.load(f)
|
|
|
|
_v10_chunks = chunks1 + chunks2
|
|
print(f"[V10 RAG] 文本块加载完成: {len(_v10_chunks)} 个")
|
|
|
|
return _v10_index, _v10_chunks
|
|
|
|
# ========== 请求/响应模型 ==========
|
|
class V10RAGQueryRequest(BaseModel):
|
|
query: str
|
|
top_k: int = 3
|
|
use_llm: bool = True
|
|
|
|
class V10RAGSource(BaseModel):
|
|
chunk_id: str
|
|
chunk_type: str
|
|
content: str
|
|
score: float
|
|
|
|
class V10RAGResponse(BaseModel):
|
|
query: str
|
|
answer: str
|
|
sources: List[V10RAGSource]
|
|
total_chunks: int
|
|
|
|
# ========== API端点 ==========
|
|
@router.post("/chat", response_model=V10RAGResponse)
|
|
async def v10_rag_chat(request: V10RAGQueryRequest):
|
|
"""V10专用RAG查询"""
|
|
try:
|
|
# 加载索引
|
|
index, chunks = load_v10_index()
|
|
|
|
# 1. 查询向量化
|
|
query_vec = get_embedding(request.query).reshape(1, -1)
|
|
|
|
# 2. FAISS检索
|
|
distances, indices = index.search(query_vec, request.top_k)
|
|
|
|
# 3. 构建sources
|
|
sources = []
|
|
context_parts = []
|
|
for i, (dist, idx) in enumerate(zip(distances[0], indices[0])):
|
|
chunk = chunks[idx]
|
|
sources.append(V10RAGSource(
|
|
chunk_id=chunk['id'],
|
|
chunk_type=chunk['metadata']['type'],
|
|
content=chunk['text'],
|
|
score=float(dist)
|
|
))
|
|
context_parts.append(f"[{i+1}] {chunk['text']}")
|
|
|
|
# 4. LLM生成（可选）
|
|
if request.use_llm:
|
|
context = "\n\n".join(context_parts)
|
|
prompt = f"""你是基因家族分析助手，基于以下资料回答问题。
|
|
请准确回答，并标注引用来源（如[1]、[2]等）。
|
|
|
|
参考资料：
|
|
{context}
|
|
|
|
用户问题：{request.query}
|
|
|
|
请基于以上资料回答："""
|
|
|
|
response = requests.post(
|
|
f"{OLLAMA_URL}/api/generate",
|
|
json={
|
|
"model": LLM_MODEL,
|
|
"prompt": prompt,
|
|
"stream": False,
|
|
"options": {"temperature": 0.3}
|
|
}
|
|
)
|
|
answer = response.json()['response']
|
|
else:
|
|
# 直接返回检索结果
|
|
answer = "\n\n".join([f"[{i+1}] {s.content}" for i, s in enumerate(sources)])
|
|
|
|
return V10RAGResponse(
|
|
query=request.query,
|
|
answer=answer,
|
|
sources=sources,
|
|
total_chunks=len(chunks)
|
|
)
|
|
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"V10 RAG查询失败: {str(e)}")
|
|
|
|
@router.get("/info")
|
|
async def v10_info():
|
|
"""获取v10索引信息"""
|
|
try:
|
|
index, chunks = load_v10_index()
|
|
|
|
# 统计
|
|
type_counts = {}
|
|
for chunk in chunks:
|
|
t = chunk['metadata']['type']
|
|
type_counts[t] = type_counts.get(t, 0) + 1
|
|
|
|
return {
|
|
"index_type": "FAISS v10专用",
|
|
"total_chunks": len(chunks),
|
|
"embedding_model": EMBEDDING_MODEL,
|
|
"llm_model": LLM_MODEL,
|
|
"chunk_types": type_counts,
|
|
"data_source": "苦荞SPL附表增强版"
|
|
}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
print("[V10 RAG] 路由加载完成")
|