feat: QYLAW 法律法规知识库

- 语义检索(FAISS + embedding)+ 精确查找(法规名+条号)
- RAG 问答(SSE 流式,支持 thinking 折叠显示)
- 法规浏览(原文阅读)
- 历史记录(检索+对话持久化到 SQLite)
- 设置页(系统提示词/模板/LLM 参数可配置)
- 检索质量评估脚本

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
2026-08-07 14:55:25 +08:00
commit 641e33b834
39 changed files with 5254 additions and 0 deletions
+1
View File
@@ -0,0 +1 @@
"""服务层模块"""
+79
View File
@@ -0,0 +1,79 @@
"""embedding 服务客户端 — 调用 114:8003 Qwen3-Embedding 服务"""
import asyncio
import logging
from typing import List
import httpx
from app.config import EMBEDDING_URL, EMBEDDING_MODEL, EMBEDDING_BATCH_SIZE
logger = logging.getLogger(__name__)
_client: httpx.AsyncClient | None = None
def _get_client() -> httpx.AsyncClient:
"""复用 httpx 客户端"""
global _client
if _client is None or _client.is_closed:
_client = httpx.AsyncClient(timeout=60)
return _client
async def embed_batch(texts: List[str], batch_size: int = EMBEDDING_BATCH_SIZE) -> List[List[float]]:
"""批量向量化
Args:
texts: 文本列表
batch_size: 每批大小,默认 32
Returns:
向量列表,顺序与输入一致
Raises:
RuntimeError: embedding 服务不可用或返回错误
"""
if not texts:
return []
client = _get_client()
all_vecs: List[List[float]] = []
for i in range(0, len(texts), batch_size):
batch = texts[i : i + batch_size]
payload = {
"model": EMBEDDING_MODEL,
"input": batch,
}
# 重试 3 次,指数退避
last_err = None
for attempt in range(3):
try:
resp = await client.post(
f"{EMBEDDING_URL}/embeddings",
json=payload,
)
resp.raise_for_status()
data = resp.json()
# vLLM embeddings 响应格式: {data: [{embedding: [...]}], ...}
vecs = [item["embedding"] for item in data["data"]]
all_vecs.extend(vecs)
last_err = None
break
except Exception as e:
last_err = e
wait = 2 ** attempt
logger.warning(f"embedding 批次 {i//batch_size}{attempt+1} 次失败: {e}, {wait}s 后重试")
await asyncio.sleep(wait)
if last_err is not None:
raise RuntimeError(f"embedding 服务调用失败(重试 3 次): {last_err}")
return all_vecs
async def embed_text(text: str) -> List[float]:
"""单条文本向量化"""
vecs = await embed_batch([text])
return vecs[0]
+300
View File
@@ -0,0 +1,300 @@
"""历史记录服务 — 检索历史和对话历史持久化到 SQLite
表结构:
- search_history: 检索历史(查询词、过滤条件、结果数)
- chat_history: 对话历史(问题、模板、thinking、答案、引用条文)
"""
import sqlite3
import time
import json
import logging
from typing import Optional
from app.config import SQLITE_PATH
logger = logging.getLogger(__name__)
def _get_conn() -> sqlite3.Connection:
"""获取 SQLite 连接(自动建表)"""
SQLITE_PATH.parent.mkdir(parents=True, exist_ok=True)
conn = sqlite3.connect(str(SQLITE_PATH), check_same_thread=False)
conn.executescript("""
CREATE TABLE IF NOT EXISTS search_history (
id INTEGER PRIMARY KEY AUTOINCREMENT,
query TEXT NOT NULL,
category TEXT,
province TEXT,
city TEXT,
result_count INTEGER DEFAULT 0,
top_k INTEGER DEFAULT 20,
created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS chat_history (
id INTEGER PRIMARY KEY AUTOINCREMENT,
question TEXT NOT NULL,
template TEXT DEFAULT 'simple',
category TEXT,
province TEXT,
thinking TEXT,
answer TEXT,
citations TEXT,
thinking_enabled INTEGER DEFAULT 1,
status TEXT DEFAULT 'done',
created_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_search_history_created
ON search_history(created_at DESC);
CREATE INDEX IF NOT EXISTS idx_chat_history_created
ON chat_history(created_at DESC);
""")
return conn
def add_search_history(
query: str,
category: Optional[str] = None,
province: Optional[str] = None,
city: Optional[str] = None,
result_count: int = 0,
top_k: int = 20,
) -> int:
"""记录检索历史
Returns:
记录 ID,失败返回 0
"""
try:
conn = _get_conn()
ts = time.strftime("%Y-%m-%d %H:%M:%S")
cursor = conn.execute(
"""INSERT INTO search_history
(query, category, province, city, result_count, top_k, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?)""",
(query[:200], category, province, city, result_count, top_k, ts),
)
conn.commit()
rid = cursor.lastrowid
conn.close()
return rid
except Exception as e:
logger.error(f"记录检索历史失败: {e}")
return 0
def add_chat_history(
question: str,
template: str = "simple",
category: Optional[str] = None,
province: Optional[str] = None,
thinking: str = "",
answer: str = "",
citations: str = "[]",
thinking_enabled: bool = True,
status: str = "done",
) -> int:
"""记录对话历史
Args:
citations: JSON 字符串,引用条文列表
Returns:
记录 ID,失败返回 0
"""
try:
conn = _get_conn()
ts = time.strftime("%Y-%m-%d %H:%M:%S")
cursor = conn.execute(
"""INSERT INTO chat_history
(question, template, category, province, thinking, answer,
citations, thinking_enabled, status, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
(question[:1000], template, category, province,
thinking, answer, citations, 1 if thinking_enabled else 0,
status, ts),
)
conn.commit()
rid = cursor.lastrowid
conn.close()
return rid
except Exception as e:
logger.error(f"记录对话历史失败: {e}")
return 0
def update_chat_history(
chat_id: int,
thinking: str = "",
answer: str = "",
citations: str = "[]",
status: str = "done",
) -> bool:
"""更新对话历史(流式生成完成后更新完整内容)
Returns:
True 成功
"""
if not chat_id:
return False
try:
conn = _get_conn()
conn.execute(
"""UPDATE chat_history
SET thinking = ?, answer = ?, citations = ?, status = ?
WHERE id = ?""",
(thinking, answer, citations, status, chat_id),
)
conn.commit()
conn.close()
return True
except Exception as e:
logger.error(f"更新对话历史失败: {e}")
return False
def get_search_history(
page: int = 1,
page_size: int = 20,
) -> dict:
"""获取检索历史列表(分页,最新在前)
Returns:
{"items": [...], "total": int, "page": int, "page_size": int}
"""
try:
conn = _get_conn()
total = conn.execute("SELECT COUNT(*) FROM search_history").fetchone()[0]
offset = (page - 1) * page_size
rows = conn.execute(
"""SELECT id, query, category, province, city, result_count, top_k, created_at
FROM search_history
ORDER BY created_at DESC
LIMIT ? OFFSET ?""",
(page_size, offset),
).fetchall()
conn.close()
items = [
{
"id": r[0], "query": r[1], "category": r[2], "province": r[3],
"city": r[4], "result_count": r[5], "top_k": r[6], "created_at": r[7],
}
for r in rows
]
return {"items": items, "total": total, "page": page, "page_size": page_size}
except Exception as e:
logger.error(f"获取检索历史失败: {e}")
return {"items": [], "total": 0, "page": page, "page_size": page_size}
def get_chat_history(
page: int = 1,
page_size: int = 20,
) -> dict:
"""获取对话历史列表(分页,最新在前)
Returns:
{"items": [...], "total": int, "page": int, "page_size": int}
"""
try:
conn = _get_conn()
total = conn.execute("SELECT COUNT(*) FROM chat_history").fetchone()[0]
offset = (page - 1) * page_size
rows = conn.execute(
"""SELECT id, question, template, category, province,
thinking, answer, citations, thinking_enabled, status, created_at
FROM chat_history
ORDER BY created_at DESC
LIMIT ? OFFSET ?""",
(page_size, offset),
).fetchall()
conn.close()
items = []
for r in rows:
items.append({
"id": r[0], "question": r[1], "template": r[2], "category": r[3],
"province": r[4], "thinking": r[5], "answer": r[6],
"citations": json.loads(r[7]) if r[7] else [],
"thinking_enabled": bool(r[8]), "status": r[9], "created_at": r[10],
})
return {"items": items, "total": total, "page": page, "page_size": page_size}
except Exception as e:
logger.error(f"获取对话历史失败: {e}")
return {"items": [], "total": 0, "page": page, "page_size": page_size}
def get_chat_detail(chat_id: int) -> Optional[dict]:
"""获取单条对话历史详情"""
try:
conn = _get_conn()
row = conn.execute(
"""SELECT id, question, template, category, province,
thinking, answer, citations, thinking_enabled, status, created_at
FROM chat_history WHERE id = ?""",
(chat_id,),
).fetchone()
conn.close()
if not row:
return None
return {
"id": row[0], "question": row[1], "template": row[2], "category": row[3],
"province": row[4], "thinking": row[5], "answer": row[6],
"citations": json.loads(row[7]) if row[7] else [],
"thinking_enabled": bool(row[8]), "status": row[9], "created_at": row[10],
}
except Exception as e:
logger.error(f"获取对话详情失败: {e}")
return None
def delete_search_history(history_id: int) -> bool:
"""删除单条检索历史"""
try:
conn = _get_conn()
conn.execute("DELETE FROM search_history WHERE id = ?", (history_id,))
conn.commit()
conn.close()
return True
except Exception as e:
logger.error(f"删除检索历史失败: {e}")
return False
def delete_chat_history(history_id: int) -> bool:
"""删除单条对话历史"""
try:
conn = _get_conn()
conn.execute("DELETE FROM chat_history WHERE id = ?", (history_id,))
conn.commit()
conn.close()
return True
except Exception as e:
logger.error(f"删除对话历史失败: {e}")
return False
def clear_search_history() -> bool:
"""清空全部检索历史"""
try:
conn = _get_conn()
conn.execute("DELETE FROM search_history")
conn.commit()
conn.close()
return True
except Exception as e:
logger.error(f"清空检索历史失败: {e}")
return False
def clear_chat_history() -> bool:
"""清空全部对话历史"""
try:
conn = _get_conn()
conn.execute("DELETE FROM chat_history")
conn.commit()
conn.close()
return True
except Exception as e:
logger.error(f"清空对话历史失败: {e}")
return False
+147
View File
@@ -0,0 +1,147 @@
"""法规原文读取服务 — 从 law-pack 目录读取 markdown,从 SQLite 取列表"""
import sqlite3
import logging
from typing import Optional, Dict, Any, List
from pathlib import Path
from app.config import SQLITE_PATH, LAW_PACK_DIR
logger = logging.getLogger(__name__)
class LawReader:
"""法规原文读取单例"""
_instance: Optional["LawReader"] = None
@classmethod
def get_instance(cls) -> "LawReader":
if cls._instance is None:
cls._instance = cls()
return cls._instance
def __init__(self):
self._conn: Optional[sqlite3.Connection] = None
def _ensure_conn(self):
if self._conn is None and SQLITE_PATH.exists():
self._conn = sqlite3.connect(str(SQLITE_PATH), check_same_thread=False)
self._conn.row_factory = sqlite3.Row
def list_laws(
self,
category: Optional[str] = None,
province: Optional[str] = None,
keyword: Optional[str] = None,
page: int = 1,
page_size: int = 50,
) -> Dict[str, Any]:
"""法规列表(分页 + 过滤)"""
self._ensure_conn()
if self._conn is None:
return {"results": [], "total": 0, "page": page, "page_size": page_size}
where = []
params: List[Any] = []
if category:
where.append("category = ?")
params.append(category)
if province:
where.append("province = ?")
params.append(province)
if keyword:
where.append("name LIKE ?")
params.append(f"%{keyword}%")
where_clause = f"WHERE {' AND '.join(where)}" if where else ""
# 总数
total = self._conn.execute(
f"SELECT COUNT(*) FROM laws {where_clause}", params
).fetchone()[0]
# 分页
offset = (page - 1) * page_size
rows = self._conn.execute(
f"""
SELECT id, name, category, publish_date, province, city, clause_count, file_path
FROM laws {where_clause}
ORDER BY category, name
LIMIT ? OFFSET ?
""",
params + [page_size, offset],
).fetchall()
results = [
{
"law_id": row["id"],
"name": row["name"],
"category": row["category"],
"publish_date": row["publish_date"],
"province": row["province"],
"city": row["city"],
"clause_count": row["clause_count"],
"file_path": row["file_path"],
}
for row in rows
]
return {
"results": results,
"total": total,
"page": page,
"page_size": page_size,
}
def get_law(self, law_id: int) -> Optional[Dict[str, Any]]:
"""法规详情(含原文 + 条文列表)"""
self._ensure_conn()
if self._conn is None:
return None
row = self._conn.execute(
"SELECT * FROM laws WHERE id = ?", (law_id,)
).fetchone()
if row is None:
return None
# 读取原文(路径校验防穿越)
file_path = Path(row["file_path"])
try:
# 确保在 law-pack 目录内
abs_path = file_path.resolve()
pack_root = LAW_PACK_DIR.resolve()
if not str(abs_path).startswith(str(pack_root)):
logger.error(f"路径穿越尝试: {file_path}")
return None
content = abs_path.read_text(encoding="utf-8")
except Exception as e:
logger.error(f"读取原文失败 {file_path}: {e}")
content = f"[读取失败: {e}]"
# 条文列表
clauses = self._conn.execute(
"SELECT id, chapter, clause_no, content FROM clauses WHERE law_id = ? ORDER BY id",
(law_id,),
).fetchall()
clause_list = [
{
"clause_id": c["id"],
"chapter": c["chapter"],
"clause_no": c["clause_no"],
"content": c["content"],
}
for c in clauses
]
return {
"law_id": row["id"],
"name": row["name"],
"category": row["category"],
"publish_date": row["publish_date"],
"province": row["province"],
"city": row["city"],
"content": content,
"file_path": row["file_path"],
"clauses": clause_list,
}
+140
View File
@@ -0,0 +1,140 @@
"""LLM 客户端 — 调用 114:7000 Qwen3.5-35B 服务(流式)"""
import json
import logging
from typing import AsyncGenerator
import httpx
from app.config import LLM_URL, LLM_MODEL
logger = logging.getLogger(__name__)
_client: httpx.AsyncClient | None = None
def _get_client() -> httpx.AsyncClient:
global _client
if _client is None or _client.is_closed:
_client = httpx.AsyncClient(timeout=120)
return _client
DEFAULT_SYSTEM_PROMPT = "你是法律助手,根据提供的法规条文回答问题,必须引用法规名和条号。"
async def stream_chat(
prompt: str,
system: str | None = None,
temperature: float = 0.3,
max_tokens: int = 2048,
thinking_enabled: bool = True,
thinking_budget: int = 2048,
) -> AsyncGenerator[tuple[str, str], None]:
"""流式对话,区分 thinking 和 answer
模型输出格式:thinking 内容以 "Here's a thinking process:" 开头,
以 </thinking> 结尾,之后是正式回答。
Args:
prompt: 用户 prompt(已含 context)
system: system prompt
temperature: 温度
max_tokens: 最大生成 token(含 thinking)
thinking_enabled: 是否启用 thinking
thinking_budget: thinking 预算(thinking 最大 token 数)
Yields:
(type, content) 元组,type 为 "thinking""answer"
Raises:
RuntimeError: LLM 服务不可用
"""
client = _get_client()
sys_content = system if system else DEFAULT_SYSTEM_PROMPT
# thinking 关闭时,通过 system prompt 补充指示
if not thinking_enabled:
sys_content += "\n\n注意:不要输出思考过程,直接给出答案。"
# payload:使用 chat_template_kwargs.enable_thinking 控制 thinking 开关
# 这是 Qwen3.5 vLLM 的原生参数,比 system prompt 更可靠
payload = {
"model": LLM_MODEL,
"messages": [
{"role": "system", "content": sys_content},
{"role": "user", "content": prompt},
],
"temperature": temperature,
"max_tokens": max_tokens,
"stream": True,
"chat_template_kwargs": {"enable_thinking": thinking_enabled},
}
# 状态机:跟踪当前是否在 thinking 区域
in_thinking = thinking_enabled # thinking 开启时,输出通常以 thinking 开头
# 小缓冲:仅用于检测跨块的 thinking 结束标记
tail_buffer = ""
# Qwen3.5 的 thinking 结束标记: </think>
THINKING_END_TAG = "</think>"
last_err = None
for attempt in range(2): # 重试 1 次
try:
async with client.stream(
"POST",
f"{LLM_URL}/chat/completions",
json=payload,
timeout=120,
) as resp:
resp.raise_for_status()
async for line in resp.aiter_lines():
if not line or not line.startswith("data: "):
continue
data_str = line[6:]
if data_str.strip() == "[DONE]":
return
try:
chunk = json.loads(data_str)
delta = chunk["choices"][0].get("delta", {})
content = delta.get("content", "")
if not content:
continue
if not thinking_enabled:
# thinking 关闭,全部作为 answer
yield ("answer", content)
else:
# thinking 开启,需要解析 THINKING_END_TAG 标记
if in_thinking:
if THINKING_END_TAG in content:
# 分割:前面是 thinking,后面是 answer
parts = content.split(THINKING_END_TAG, 1)
thinking_part = parts[0]
answer_part = parts[1] if len(parts) > 1 else ""
if thinking_part:
yield ("thinking", thinking_part)
if answer_part:
yield ("answer", answer_part.lstrip())
in_thinking = False
tail_buffer = ""
else:
# 检查是否是 THINKING_END_TAG 的前缀(跨块情况)
tail_buffer = (tail_buffer + content)[-12:]
# 检查 tail_buffer 是否是 THINKING_END_TAG 的前缀
is_prefix = any(
tail_buffer.endswith(THINKING_END_TAG[:k])
for k in range(1, len(THINKING_END_TAG))
)
if is_prefix:
continue
yield ("thinking", content)
else:
yield ("answer", content)
except (json.JSONDecodeError, KeyError, IndexError):
continue
return
except Exception as e:
last_err = e
logger.warning(f"LLM 流式调用第 {attempt+1} 次失败: {e}")
raise RuntimeError(f"LLM 服务调用失败: {last_err}")
+190
View File
@@ -0,0 +1,190 @@
"""提示词配置服务 — 从 SQLite 读写系统提示词和模板正文
配置项(config_key):
- system_prompt: 系统提示词(覆盖默认)
- template_simple: 通俗解释模板正文
- template_professional: 专业分析模板正文
- template_compare: 对比条文模板正文
"""
import sqlite3
import time
import logging
from typing import Optional, Dict
from pathlib import Path
from app.config import SQLITE_PATH
logger = logging.getLogger(__name__)
# 默认值(与 rag.py 中 PROMPT_TEMPLATES 保持一致)
DEFAULTS = {
"system_prompt": "你是法律助手,根据提供的法规条文回答问题,必须引用法规名和条号。",
"template_simple": (
"你是一位耐心的法律科普助手。请根据以下法规条文,用通俗易懂的语言回答用户问题。\n"
"要求:\n"
"1. 避免法律术语,用日常语言解释\n"
'2. 必须引用依据的法规名和条号,格式如"根据《XX法》第X条"\n'
'3. 如果条文不足以回答,明确说明"根据现有法规无法完全回答"\n'
"4. 不要编造法规\n\n"
"【相关法规条文】\n{context}\n\n【用户问题】\n{question}"
),
"template_professional": (
"你是一位专业的法律分析助手。请根据以下法规条文,对用户问题进行专业分析。\n"
"要求:\n"
"1. 使用规范法律术语\n"
'2. 必须引用依据的法规名和条号,格式如"依据《XX法》第X条"\n'
"3. 分析条文的适用条件、法律后果\n"
"4. 如有多条相关,对比分析\n"
"5. 指出条文的适用边界和可能的争议点\n"
"6. 不要编造法规\n\n"
"【相关法规条文】\n{context}\n\n【用户问题】\n{question}"
),
"template_compare": (
"你是一位法规研究助手。请根据以下法规条文,对比分析不同法规对同一问题的规定。\n"
"要求:\n"
"1. 列出每条相关法规的具体规定\n"
"2. 对比规定的异同\n"
"3. 标注法规名、条号、发布日期\n"
"4. 指出适用范围差异(如全国性 vs 地方性)\n"
"5. 不要编造法规\n\n"
"【相关法规条文】\n{context}\n\n【用户问题】\n{question}"
),
# LLM 参数默认值
"llm_thinking_enabled": "true", # thinking 开关
"llm_max_tokens": "4096", # 最大生成 token 数
"llm_temperature": "0.3", # 温度
"llm_thinking_budget": "2048", # thinking 预算(thinking 最大 token 数)
}
# 模板 key 与 rag.py PROMPT_TEMPLATES 的映射
TEMPLATE_KEY_MAP = {
"simple": "template_simple",
"professional": "template_professional",
"compare": "template_compare",
}
def _get_conn() -> sqlite3.Connection:
"""获取 SQLite 连接(自动建表)"""
SQLITE_PATH.parent.mkdir(parents=True, exist_ok=True)
conn = sqlite3.connect(str(SQLITE_PATH), check_same_thread=False)
# 自动建表(若数据库不存在或表不存在)
conn.executescript("""
CREATE TABLE IF NOT EXISTS prompt_configs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
config_key TEXT NOT NULL UNIQUE,
config_value TEXT,
updated_at TEXT NOT NULL
);
""")
return conn
def get_all_configs() -> Dict[str, str]:
"""读取所有提示词配置,未配置的用默认值填充
Returns:
{config_key: config_value} 所有配置项(含默认值)
"""
configs = dict(DEFAULTS) # 先用默认值
try:
conn = _get_conn()
rows = conn.execute(
"SELECT config_key, config_value FROM prompt_configs"
).fetchall()
for row in rows:
key, val = row[0], row[1]
if val: # 空值视为未配置,用默认
configs[key] = val
conn.close()
except Exception as e:
logger.warning(f"读取提示词配置失败,使用默认值: {e}")
return configs
def get_config(key: str) -> str:
"""读取单个配置,未配置返回默认值"""
try:
conn = _get_conn()
row = conn.execute(
"SELECT config_value FROM prompt_configs WHERE config_key = ?", (key,)
).fetchone()
conn.close()
if row and row[0]:
return row[0]
except Exception as e:
logger.warning(f"读取配置 {key} 失败: {e}")
return DEFAULTS.get(key, "")
def set_config(key: str, value: str) -> bool:
"""写入单个配置(空值删除记录,回退到默认)
Returns:
True 成功
"""
try:
conn = _get_conn()
ts = time.strftime("%Y-%m-%d %H:%M:%S")
if value: # 有值则 upsert
conn.execute(
"""
INSERT INTO prompt_configs (config_key, config_value, updated_at)
VALUES (?, ?, ?)
ON CONFLICT(config_key) DO UPDATE SET
config_value = excluded.config_value,
updated_at = excluded.updated_at
""",
(key, value, ts),
)
else: # 空值则删除(回退默认)
conn.execute(
"DELETE FROM prompt_configs WHERE config_key = ?", (key,)
)
conn.commit()
conn.close()
return True
except Exception as e:
logger.error(f"写入配置 {key} 失败: {e}")
return False
def get_template_text(template: str) -> str:
"""获取指定模板的正文
Args:
template: 模板名(simple/professional/compare)
Returns:
模板正文(含 {context}{question} 占位符)
"""
key = TEMPLATE_KEY_MAP.get(template, "template_simple")
return get_config(key)
def get_llm_params() -> dict:
"""获取 LLM 参数配置
Returns:
{"thinking_enabled": bool, "max_tokens": int, "temperature": float, "thinking_budget": int}
"""
configs = get_all_configs()
try:
return {
"thinking_enabled": configs.get("llm_thinking_enabled", "true").lower() == "true",
"max_tokens": int(configs.get("llm_max_tokens", "4096")),
"temperature": float(configs.get("llm_temperature", "0.3")),
"thinking_budget": int(configs.get("llm_thinking_budget", "2048")),
}
except (ValueError, TypeError):
return {
"thinking_enabled": True,
"max_tokens": 4096,
"temperature": 0.3,
"thinking_budget": 2048,
}
def get_system_prompt() -> str:
"""获取系统提示词"""
return get_config("system_prompt")
+264
View File
@@ -0,0 +1,264 @@
"""FAISS 检索器 — 加载索引 + SQLite metadata,实现语义检索 + 过滤"""
import sqlite3
import logging
from typing import List, Optional, Dict, Any
import numpy as np
from app.config import (
FAISS_INDEX_PATH,
SQLITE_PATH,
HNSW_EF_SEARCH,
)
logger = logging.getLogger(__name__)
class Retriever:
"""检索器单例"""
_instance: Optional["Retriever"] = None
@classmethod
def get_instance(cls) -> "Retriever":
if cls._instance is None:
cls._instance = cls()
return cls._instance
def __init__(self):
self._faiss = None
self._conn: Optional[sqlite3.Connection] = None
self._dim: Optional[int] = None
self._ready = False
def is_ready(self) -> bool:
return self._ready
def load(self):
"""加载 FAISS 索引和 SQLite"""
if not FAISS_INDEX_PATH.exists():
logger.warning(f"FAISS 索引不存在: {FAISS_INDEX_PATH}")
return
if not SQLITE_PATH.exists():
logger.warning(f"SQLite 不存在: {SQLITE_PATH}")
return
import faiss
logger.info(f"加载 FAISS 索引: {FAISS_INDEX_PATH}")
self._faiss = faiss.read_index(str(FAISS_INDEX_PATH))
self._dim = self._faiss.d
# 设置 HNSW 搜索参数
if hasattr(self._faiss, "hnsw"):
self._faiss.hnsw.efSearch = HNSW_EF_SEARCH
logger.info(f"加载 SQLite: {SQLITE_PATH}")
self._conn = sqlite3.connect(str(SQLITE_PATH), check_same_thread=False)
self._conn.row_factory = sqlite3.Row
self._ready = True
count = self._faiss.ntotal
logger.info(f"索引加载完成: {count} 向量, dim={self._dim}")
def search(
self,
query_vec: List[float],
top_k: int = 20,
category: Optional[str] = None,
province: Optional[str] = None,
city: Optional[str] = None,
) -> List[Dict[str, Any]]:
"""语义检索 + metadata 过滤
Args:
query_vec: 查询向量
top_k: 返回数量
category: 法规类别过滤
province: 省份过滤
city: 市级过滤
Returns:
结果列表,每项含 law_id/law_name/category/chapter/clause_no/content/score/file_path/province
"""
if not self._ready:
return []
# FAISS 检索(取 top_k * 3 用于过滤后仍有足够结果)
vec = np.array([query_vec], dtype=np.float32)
fetch_k = min(top_k * 3 if (category or province or city) else top_k, self._faiss.ntotal)
scores, indices = self._faiss.search(vec, fetch_k)
results = []
for score, faiss_idx in zip(scores[0], indices[0]):
if faiss_idx < 0:
continue
# 从 SQLite 取 metadata
row = self._conn.execute(
"""
SELECT c.id, c.law_id, c.chapter, c.clause_no, c.content, c.faiss_idx,
l.name, l.category, l.publish_date, l.file_path, l.province, l.city
FROM clauses c JOIN laws l ON c.law_id = l.id
WHERE c.faiss_idx = ?
""",
(int(faiss_idx),),
).fetchone()
if row is None:
continue
# metadata 过滤
if category and row["category"] != category:
continue
if province and row["province"] != province:
continue
if city and row["city"] != city:
continue
results.append({
"clause_id": row["id"],
"law_id": row["law_id"],
"law_name": row["name"],
"category": row["category"],
"chapter": row["chapter"],
"clause_no": row["clause_no"],
"content": row["content"],
"score": float(score),
"file_path": row["file_path"],
"province": row["province"],
"city": row["city"],
"publish_date": row["publish_date"],
})
if len(results) >= top_k:
break
return results
def keyword_search(
self,
query: str,
top_k: int = 20,
category: Optional[str] = None,
province: Optional[str] = None,
city: Optional[str] = None,
) -> List[Dict[str, Any]]:
"""关键词精确搜索(按法规名 + 条号 + 条文内容匹配)
支持"法规名 条号"格式(如"郑州市劳动用工条例 第三十二条"),
也支持纯关键词搜索条文内容。
Returns:
结果列表,格式与 search() 一致,score 为匹配相关度(非 FAISS 距离)
"""
if not self._ready or not query.strip():
return []
query = query.strip()
# 尝试解析"法规名 条号"格式
# 中文条号模式:第X条
import re
clause_match = re.search(r'(第[一二三四五六七八九十百千零\d]+条)', query)
clause_no = clause_match.group(1) if clause_match else None
# 法规名 = 去掉条号后的部分
law_name_query = re.sub(r'\s*第[一二三四五六七八九十百千零\d]+条\s*', '', query).strip()
results = []
try:
if clause_no and law_name_query:
# 精确匹配:法规名 LIKE + 条号 =
rows = self._conn.execute(
"""
SELECT c.id, c.law_id, c.chapter, c.clause_no, c.content, c.faiss_idx,
l.name, l.category, l.publish_date, l.file_path, l.province, l.city
FROM clauses c JOIN laws l ON c.law_id = l.id
WHERE l.name LIKE ? AND c.clause_no = ?
""",
(f"%{law_name_query}%", clause_no),
).fetchall()
elif clause_no:
# 只按条号搜索
rows = self._conn.execute(
"""
SELECT c.id, c.law_id, c.chapter, c.clause_no, c.content, c.faiss_idx,
l.name, l.category, l.publish_date, l.file_path, l.province, l.city
FROM clauses c JOIN laws l ON c.law_id = l.id
WHERE c.clause_no = ?
""",
(clause_no,),
).fetchall()
else:
# 纯关键词:搜索法规名或条文内容
rows = self._conn.execute(
"""
SELECT c.id, c.law_id, c.chapter, c.clause_no, c.content, c.faiss_idx,
l.name, l.category, l.publish_date, l.file_path, l.province, l.city
FROM clauses c JOIN laws l ON c.law_id = l.id
WHERE l.name LIKE ? OR c.content LIKE ?
""",
(f"%{query}%", f"%{query}%"),
).fetchall()
for row in rows:
# metadata 过滤
if category and row["category"] != category:
continue
if province and row["province"] != province:
continue
if city and row["city"] != city:
continue
results.append({
"clause_id": row["id"],
"law_id": row["law_id"],
"law_name": row["name"],
"category": row["category"],
"chapter": row["chapter"],
"clause_no": row["clause_no"],
"content": row["content"],
"score": 0.0, # 关键词搜索无相似度分数
"file_path": row["file_path"],
"province": row["province"],
"city": row["city"],
"publish_date": row["publish_date"],
})
if len(results) >= top_k:
break
except Exception as e:
logger.error(f"关键词搜索失败: {e}")
return results
def get_stats(self) -> Dict[str, Any]:
"""返回索引统计"""
if not self._ready:
return {
"total_laws": 0,
"total_clauses": 0,
"category_stats": {},
"index_built_at": None,
"faiss_dim": None,
"ready": False,
}
total_laws = self._conn.execute("SELECT COUNT(*) FROM laws").fetchone()[0]
total_clauses = self._conn.execute("SELECT COUNT(*) FROM clauses").fetchone()[0]
cat_rows = self._conn.execute(
"SELECT category, COUNT(*) as cnt FROM laws GROUP BY category"
).fetchall()
category_stats = {row["category"]: row["cnt"] for row in cat_rows}
# 索引构建时间(取最新一条 law 的 indexed_at)
built_at_row = self._conn.execute(
"SELECT indexed_at FROM laws ORDER BY indexed_at DESC LIMIT 1"
).fetchone()
index_built_at = built_at_row["indexed_at"] if built_at_row else None
return {
"total_laws": total_laws,
"total_clauses": total_clauses,
"category_stats": category_stats,
"index_built_at": index_built_at,
"faiss_dim": self._dim,
"ready": True,
}