feat(flywheel): few-shot 在线检索注入打通风险规则编译链路
- 新增 FewShotStore：独立 Qdrant collection few_shot_samples，向量 upsert/search/delete， 全程失败降级不阻塞主链路 - 新增 FewShotIngestionService：RiskObservation confirmed/false_positive → FewShotSample + 向量，带 sample_key 幂等去重 - 新增 FewShotRetriever：按 case 特征检索相似历史样本，去重 + token 预算 + 单条字符上限裁剪 - risk_observations.create_feedback commit 后挂 hook 自动入库，带 feature flag 和 try/except 兜底 - risk_rule_generation_prompt 新增 few_shot_samples 可选参数，合并进 examples 并标 source=historical_confirmed；risk_rule_generation 构造 prompt 前调 retriever，失败降级为空
This commit is contained in:
caoxiaozhu
214
server/src/app/services/few_shot_store.py
Normal file
214
server/src/app/services/few_shot_store.py
Normal file
@@ -0,0 +1,214 @@
|
||||
"""Few-shot 样本的 Qdrant 向量存储。
|
||||
|
||||
独立于 LightRAG 的 Qdrant 客户端，使用专用 collection ``few_shot_samples``，
|
||||
与知识库 RAG 的 collection 隔离。所有操作失败都不抛异常（记日志返回空），
|
||||
保证主链路不阻塞。
|
||||
|
||||
向量来自 :class:`EmbeddingProvider`，payload 带业务过滤字段（scene/label/domain/risk_type），
|
||||
检索时按这些字段过滤 + 向量相似度排序。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from app.core.logging import get_logger
|
||||
from app.services.knowledge_rag import _resolve_default_qdrant_url
|
||||
|
||||
logger = get_logger("app.services.few_shot_store")
|
||||
|
||||
FEW_SHOT_COLLECTION = "few_shot_samples"
|
||||
|
||||
|
||||
def _resolve_qdrant_config() -> tuple[str, str]:
|
||||
"""复用 knowledge_rag 的 Qdrant URL/key 解析逻辑。"""
|
||||
|
||||
url = os.environ.get("QDRANT_URL", "").strip() or _resolve_default_qdrant_url()
|
||||
api_key = os.environ.get("QDRANT_API_KEY", "").strip()
|
||||
return url, api_key
|
||||
|
||||
|
||||
class FewShotStore:
|
||||
"""对 Qdrant 的轻量封装，专供 few-shot 样本检索使用。
|
||||
|
||||
设计要点：
|
||||
- 惰性创建 client 和 collection，首次操作时初始化。
|
||||
- 所有公共方法吞异常（返回空/False），主链路永远不被拖崩。
|
||||
- 向量写入和检索都依赖外部传入的 :class:`EmbeddingProvider`，
|
||||
由调用方保证与配置一致。
|
||||
"""
|
||||
|
||||
def __init__(self, embedding_provider: Any) -> None:
|
||||
self._embedding_provider = embedding_provider
|
||||
self._client: Any = None
|
||||
self._ensured = False
|
||||
|
||||
def _client_or_none(self) -> Any:
|
||||
"""惰性初始化 QdrantClient，失败返回 None。"""
|
||||
|
||||
if self._client is not None:
|
||||
return self._client
|
||||
try:
|
||||
from qdrant_client import QdrantClient
|
||||
|
||||
url, api_key = _resolve_qdrant_config()
|
||||
self._client = QdrantClient(url=url, api_key=api_key or None)
|
||||
except Exception:
|
||||
logger.warning("few-shot QdrantClient 初始化失败，本轮操作跳过", exc_info=True)
|
||||
self._client = None
|
||||
return self._client
|
||||
|
||||
def _ensure_collection(self) -> bool:
|
||||
"""确保 collection 存在，成功返回 True。"""
|
||||
|
||||
if self._ensured:
|
||||
return True
|
||||
client = self._client_or_none()
|
||||
if client is None:
|
||||
return False
|
||||
try:
|
||||
from qdrant_client.http.exceptions import UnexpectedResponse
|
||||
|
||||
try:
|
||||
client.get_collection(FEW_SHOT_COLLECTION)
|
||||
self._ensured = True
|
||||
return True
|
||||
except UnexpectedResponse as exc:
|
||||
if exc.status_code != 404:
|
||||
raise
|
||||
# collection 不存在则创建
|
||||
dim = self._embedding_provider.dimension()
|
||||
from qdrant_client.http.models import (
|
||||
Distance,
|
||||
VectorParams,
|
||||
PayloadSchemaType,
|
||||
)
|
||||
|
||||
client.create_collection(
|
||||
collection_name=FEW_SHOT_COLLECTION,
|
||||
vectors_config=VectorParams(size=dim, distance=Distance.COSINE),
|
||||
)
|
||||
for field, field_type in [
|
||||
("sample_id", PayloadSchemaType.KEYWORD),
|
||||
("scene", PayloadSchemaType.KEYWORD),
|
||||
("label", PayloadSchemaType.KEYWORD),
|
||||
("domain", PayloadSchemaType.KEYWORD),
|
||||
("risk_type", PayloadSchemaType.KEYWORD),
|
||||
("status", PayloadSchemaType.KEYWORD),
|
||||
]:
|
||||
try:
|
||||
client.create_payload_index(
|
||||
collection_name=FEW_SHOT_COLLECTION,
|
||||
field_name=field,
|
||||
field_schema=field_type,
|
||||
)
|
||||
except Exception:
|
||||
logger.debug("payload index 创建跳过 field=%s", field, exc_info=True)
|
||||
self._ensured = True
|
||||
logger.info("few-shot collection 创建成功 dim=%s", dim)
|
||||
return True
|
||||
except Exception:
|
||||
logger.warning("few-shot collection 初始化失败，本轮操作跳过", exc_info=True)
|
||||
return False
|
||||
|
||||
def upsert(self, sample: Any) -> str | None:
|
||||
"""把一条样本向量化并写入 Qdrant，返回 vector_id，失败返回 None。"""
|
||||
|
||||
if not self._ensure_collection():
|
||||
return None
|
||||
client = self._client
|
||||
try:
|
||||
vector = self._embedding_provider.embed([sample.case_text])[0]
|
||||
except Exception:
|
||||
logger.warning("few-shot embedding 失败 sample_key=%s", getattr(sample, "sample_key", ""), exc_info=True)
|
||||
return None
|
||||
vector_id = uuid.uuid4().hex
|
||||
payload = {
|
||||
"sample_id": sample.id,
|
||||
"scene": sample.scene,
|
||||
"label": sample.label,
|
||||
"domain": sample.domain,
|
||||
"risk_type": sample.risk_type,
|
||||
"risk_level": sample.risk_level,
|
||||
"status": getattr(sample, "status", "active"),
|
||||
"conclusion_text": sample.conclusion_text,
|
||||
"payload_json": sample.payload_json,
|
||||
}
|
||||
try:
|
||||
client.upsert(
|
||||
collection_name=FEW_SHOT_COLLECTION,
|
||||
points=[{"id": vector_id, "vector": vector, "payload": payload}],
|
||||
)
|
||||
return vector_id
|
||||
except Exception:
|
||||
logger.warning("few-shot upsert 失败 sample_key=%s", getattr(sample, "sample_key", ""), exc_info=True)
|
||||
return None
|
||||
|
||||
def search(
|
||||
self,
|
||||
case_text: str,
|
||||
*,
|
||||
scene: str | None = None,
|
||||
labels: list[str] | None = None,
|
||||
top_k: int = 3,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""按 case_text 检索相似样本，可按 scene/label 过滤。失败返回空列表。"""
|
||||
|
||||
if not case_text or not self._ensure_collection():
|
||||
return []
|
||||
client = self._client
|
||||
try:
|
||||
vector = self._embedding_provider.embed([case_text])[0]
|
||||
except Exception:
|
||||
logger.warning("few-shot 检索 embedding 失败", exc_info=True)
|
||||
return []
|
||||
must: list[dict[str, Any]] = [{"key": "status", "match": {"value": "active"}}]
|
||||
if scene:
|
||||
must.append({"key": "scene", "match": {"value": scene}})
|
||||
if labels:
|
||||
must.append({"key": "label", "match": {"any": labels}})
|
||||
try:
|
||||
from qdrant_client.http.models import Filter
|
||||
|
||||
results = client.query_points(
|
||||
collection_name=FEW_SHOT_COLLECTION,
|
||||
query=vector,
|
||||
query_filter=Filter(must=must),
|
||||
limit=top_k,
|
||||
with_payload=True,
|
||||
).points
|
||||
except Exception:
|
||||
logger.warning("few-shot 检索失败", exc_info=True)
|
||||
return []
|
||||
hits: list[dict[str, Any]] = []
|
||||
for point in results:
|
||||
payload = getattr(point, "payload", None) or {}
|
||||
hits.append(
|
||||
{
|
||||
"sample_id": payload.get("sample_id"),
|
||||
"score": float(getattr(point, "score", 0.0)),
|
||||
"label": payload.get("label"),
|
||||
"domain": payload.get("domain"),
|
||||
"risk_type": payload.get("risk_type"),
|
||||
"conclusion_text": payload.get("conclusion_text") or "",
|
||||
"payload_json": payload.get("payload_json") or {},
|
||||
}
|
||||
)
|
||||
return hits
|
||||
|
||||
def delete_by_vector_id(self, vector_id: str) -> bool:
|
||||
"""按 vector_id 删除向量，失败返回 False。"""
|
||||
|
||||
if not vector_id or not self._ensure_collection():
|
||||
return False
|
||||
try:
|
||||
self._client.delete(
|
||||
collection_name=FEW_SHOT_COLLECTION,
|
||||
points_selector=[vector_id],
|
||||
)
|
||||
return True
|
||||
except Exception:
|
||||
logger.warning("few-shot 删除失败 vector_id=%s", vector_id, exc_info=True)
|
||||
return False
|
||||
Reference in New Issue
Block a user