317 lines
13 KiB
Python
317 lines
13 KiB
Python
|
"""组合召回：MySQL 权威召回 + 可选向量召回，Redis 只做可重建的加速层。
|
||
|
|
|
||
|
|
设计取舍（对应 P3 缺口 1）：
|
||
|
|
1. `VectorMemoryAdapter` 只有"向量 → 命中"的能力，既不做文本向量化，也不回表。
|
||
|
|
而本项目没有 text-embedding 端点配置（`Settings` 里只有 Milvus 连接信息），
|
||
|
|
因此语义召回缺少"文本 → 向量"的第一步。这里不猜造 embedding：调用方须注入
|
||
|
|
`embed`（生产接模型网关），未注入时向量通道显式降级并记录原因，
|
||
|
|
MySQL 结构化召回照常返回。适配器侧补 `parse_hits`，把 pymilvus 的嵌套命中
|
||
|
|
折叠成可回表的 `VectorHit`，避免服务层去猜第三方的返回结构。
|
||
|
|
2. `MemoryCacheAdapter` 的 get/set 语义可用，但缺"按前缀删除"，而客户级联失效
|
||
|
|
需要清缓存键，因此在适配器上新增 `delete(*keys)`；写入失败仍只返回 False。
|
||
|
|
"""
|
||
|
|
|
||
|
|
import hashlib
|
||
|
|
import json
|
||
|
|
import logging
|
||
|
|
from collections.abc import Awaitable, Callable
|
||
|
|
from dataclasses import asdict, dataclass, field
|
||
|
|
from typing import Any, Protocol
|
||
|
|
|
||
|
|
from sqlalchemy import select
|
||
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
|
|
||
|
|
from app.infrastructure.memory_cache import MemoryCacheAdapter
|
||
|
|
from app.infrastructure.vector_memory import VectorMemoryAdapter
|
||
|
|
from app.model.memory import MemoryUnit
|
||
|
|
from app.service.memory_service import MemoryService
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
CACHE_KEY_PREFIX = "mem:recall"
|
||
|
|
CACHE_TTL_SECONDS = 300
|
||
|
|
# 客户级缓存失效需要枚举出真实存在的键：limit 与查询摘要都由调用方决定，
|
||
|
|
# 因此把常用组合固化为常量，供 `cache_keys` 与生命周期服务共用。
|
||
|
|
COMMON_LIMITS = (5, 10, 20)
|
||
|
|
# 空查询与 None 归一化后同键，只保留一个代表值。
|
||
|
|
COMMON_QUERIES: tuple[str | None, ...] = (None,)
|
||
|
|
|
||
|
|
# 向量相似度在综合置信度中的权重上限：语义召回的置信度不得与结构化证据等价。
|
||
|
|
VECTOR_WEIGHT = 0.7
|
||
|
|
STRUCTURED_SOURCE = "mysql"
|
||
|
|
VECTOR_SOURCE = "milvus"
|
||
|
|
|
||
|
|
EmbeddingProvider = Callable[[str], Awaitable[list[float]]]
|
||
|
|
|
||
|
|
|
||
|
|
class CacheAdapter(Protocol):
|
||
|
|
async def get(self, key: str) -> Any: ...
|
||
|
|
|
||
|
|
async def set(self, key: str, value: str, ttl_seconds: int = 300) -> bool: ...
|
||
|
|
|
||
|
|
async def delete(self, *keys: str) -> int: ...
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class RecallItem:
|
||
|
|
"""一条召回结果；`sources` 记录它来自哪些通道，`confidence` 是综合置信度。"""
|
||
|
|
|
||
|
|
memory_uuid: str
|
||
|
|
memory_key: str
|
||
|
|
content: str
|
||
|
|
memory_type: str
|
||
|
|
confidence: float
|
||
|
|
sources: tuple[str, ...]
|
||
|
|
evidence: dict[str, Any] = field(default_factory=dict)
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class RecallResult:
|
||
|
|
customer_id: int
|
||
|
|
items: tuple[RecallItem, ...]
|
||
|
|
degraded: bool = False
|
||
|
|
degraded_reasons: tuple[str, ...] = ()
|
||
|
|
from_cache: bool = False
|
||
|
|
|
||
|
|
|
||
|
|
class MemoryRecallService:
|
||
|
|
"""组合召回入口；任何外部依赖（Milvus/Redis）故障都只降级，不阻塞主流程。"""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
session: AsyncSession,
|
||
|
|
*,
|
||
|
|
vector: VectorMemoryAdapter | None = None,
|
||
|
|
cache: MemoryCacheAdapter | CacheAdapter | None = None,
|
||
|
|
embed: EmbeddingProvider | None = None,
|
||
|
|
cache_ttl_seconds: int = CACHE_TTL_SECONDS,
|
||
|
|
) -> None:
|
||
|
|
self.session = session
|
||
|
|
self.memory = MemoryService(session)
|
||
|
|
self.vector = vector
|
||
|
|
self.cache = cache
|
||
|
|
self.embed = embed
|
||
|
|
self.cache_ttl_seconds = cache_ttl_seconds
|
||
|
|
|
||
|
|
async def recall(
|
||
|
|
self,
|
||
|
|
customer_id: int,
|
||
|
|
query: str | None = None,
|
||
|
|
*,
|
||
|
|
limit: int = 10,
|
||
|
|
use_cache: bool = True,
|
||
|
|
) -> RecallResult:
|
||
|
|
bounded = max(1, min(limit, 100))
|
||
|
|
reasons: list[str] = []
|
||
|
|
cache_key = self.cache_key(customer_id, query, bounded)
|
||
|
|
if use_cache:
|
||
|
|
cached, cache_degraded = await self._cache_read(cache_key)
|
||
|
|
if cached is not None:
|
||
|
|
return cached
|
||
|
|
if cache_degraded:
|
||
|
|
reasons.append("redis_unavailable")
|
||
|
|
structured = await self._structured(customer_id, query, bounded)
|
||
|
|
vector_items, vector_reasons = await self._vector(customer_id, query, bounded)
|
||
|
|
reasons.extend(vector_reasons)
|
||
|
|
merged = self._merge(structured, vector_items, bounded)
|
||
|
|
degraded = bool(reasons)
|
||
|
|
result = RecallResult(
|
||
|
|
customer_id=customer_id,
|
||
|
|
items=tuple(merged),
|
||
|
|
degraded=degraded,
|
||
|
|
degraded_reasons=tuple(dict.fromkeys(reasons)),
|
||
|
|
)
|
||
|
|
if use_cache and self.cache is not None:
|
||
|
|
await self._cache_write(cache_key, result)
|
||
|
|
return result
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def cache_key(customer_id: int, query: str | None, limit: int) -> str:
|
||
|
|
digest = hashlib.sha256((query or "").strip().lower().encode("utf-8")).hexdigest()[:16]
|
||
|
|
return f"{CACHE_KEY_PREFIX}:{customer_id}:{limit}:{digest}"
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def cache_keys(cls, customer_id: int) -> list[str]:
|
||
|
|
"""某客户常见参数组合下的热缓存键；供客户级失效枚举使用。"""
|
||
|
|
return [
|
||
|
|
cls.cache_key(customer_id, query, limit)
|
||
|
|
for limit in COMMON_LIMITS
|
||
|
|
for query in COMMON_QUERIES
|
||
|
|
]
|
||
|
|
|
||
|
|
async def invalidate_cache(self, customer_id: int) -> int:
|
||
|
|
if self.cache is None:
|
||
|
|
return 0
|
||
|
|
return await self.cache.delete(*self.cache_keys(customer_id))
|
||
|
|
|
||
|
|
async def _structured(
|
||
|
|
self, customer_id: int, query: str | None, limit: int
|
||
|
|
) -> list[MemoryUnit]:
|
||
|
|
"""复用 `MemoryService.recall_with_decay`（含关键词过滤与时间衰减）。"""
|
||
|
|
return await self.memory.recall_with_decay(customer_id, query, limit=limit)
|
||
|
|
|
||
|
|
async def _vector(
|
||
|
|
self, customer_id: int, query: str | None, limit: int
|
||
|
|
) -> tuple[list[RecallItem], list[str]]:
|
||
|
|
if self.vector is None:
|
||
|
|
return [], []
|
||
|
|
text = (query or "").strip()
|
||
|
|
if not text:
|
||
|
|
return [], []
|
||
|
|
if self.embed is None:
|
||
|
|
# 缺少文本向量化能力时，语义通道不可用；这是配置缺口，不是运行时故障。
|
||
|
|
logger.warning("vector recall skipped: embedding provider not configured")
|
||
|
|
return [], ["embedding_unavailable"]
|
||
|
|
try:
|
||
|
|
embedding = await self.embed(text)
|
||
|
|
except Exception:
|
||
|
|
logger.warning("vector recall degraded: embedding failed")
|
||
|
|
return [], ["embedding_failed"]
|
||
|
|
if not embedding:
|
||
|
|
return [], ["embedding_failed"]
|
||
|
|
# 纵深防御：适配器契约上不抛异常，但召回服务自己对外承诺"外部依赖故障只降级、
|
||
|
|
# 不阻塞主流程"，这条承诺不能寄托在适配器实现上。
|
||
|
|
try:
|
||
|
|
found = self.vector.search(embedding, limit=limit)
|
||
|
|
if found.degraded:
|
||
|
|
logger.warning("vector recall degraded: milvus unavailable collection=%s",
|
||
|
|
self.vector.collection)
|
||
|
|
return [], ["milvus_unavailable"]
|
||
|
|
hits = self.vector.parse_hits(found.hits)
|
||
|
|
except Exception:
|
||
|
|
logger.warning("vector recall degraded: vector backend raised collection=%s",
|
||
|
|
self.vector.collection)
|
||
|
|
return [], ["milvus_unavailable"]
|
||
|
|
if not hits:
|
||
|
|
return [], []
|
||
|
|
scores = {hit.memory_uuid: hit.score for hit in hits}
|
||
|
|
rows: list[RecallItem] = []
|
||
|
|
memories = await self._load_by_uuids(customer_id, list(scores))
|
||
|
|
for memory in memories:
|
||
|
|
score = scores.get(memory.memory_uuid, 0.0)
|
||
|
|
rows.append(self._item(
|
||
|
|
memory,
|
||
|
|
sources=(VECTOR_SOURCE,),
|
||
|
|
confidence=min(1.0, abs(score) * VECTOR_WEIGHT),
|
||
|
|
evidence={"vector_score": round(score, 6)},
|
||
|
|
))
|
||
|
|
return rows, []
|
||
|
|
|
||
|
|
async def _load_by_uuids(self, customer_id: int, uuids: list[str]) -> list[MemoryUnit]:
|
||
|
|
if not uuids:
|
||
|
|
return []
|
||
|
|
found = await self.session.scalars(
|
||
|
|
select(MemoryUnit).where(
|
||
|
|
MemoryUnit.customer_id == customer_id,
|
||
|
|
MemoryUnit.memory_uuid.in_(uuids),
|
||
|
|
MemoryUnit.status == "active",
|
||
|
|
)
|
||
|
|
)
|
||
|
|
return list(found)
|
||
|
|
|
||
|
|
def _merge(
|
||
|
|
self, structured: list[MemoryUnit], vector_items: list[RecallItem], limit: int
|
||
|
|
) -> list[RecallItem]:
|
||
|
|
"""以 memory_uuid 为身份合并：结构化结果保留，向量命中追加并取较高置信度。"""
|
||
|
|
merged: dict[str, RecallItem] = {}
|
||
|
|
for memory in structured:
|
||
|
|
item = self._item(memory, sources=(STRUCTURED_SOURCE,), confidence=float(
|
||
|
|
memory.confidence))
|
||
|
|
merged[item.memory_uuid] = item
|
||
|
|
for item in vector_items:
|
||
|
|
existing = merged.get(item.memory_uuid)
|
||
|
|
if existing is None:
|
||
|
|
merged[item.memory_uuid] = item
|
||
|
|
continue
|
||
|
|
merged[item.memory_uuid] = RecallItem(
|
||
|
|
memory_uuid=existing.memory_uuid,
|
||
|
|
memory_key=existing.memory_key,
|
||
|
|
content=existing.content,
|
||
|
|
memory_type=existing.memory_type,
|
||
|
|
confidence=max(existing.confidence, item.confidence),
|
||
|
|
sources=(STRUCTURED_SOURCE, VECTOR_SOURCE),
|
||
|
|
evidence={**existing.evidence, **item.evidence},
|
||
|
|
)
|
||
|
|
ordered = sorted(merged.values(), key=lambda item: item.confidence, reverse=True)
|
||
|
|
return ordered[:limit]
|
||
|
|
|
||
|
|
def _item(
|
||
|
|
self,
|
||
|
|
memory: MemoryUnit,
|
||
|
|
*,
|
||
|
|
sources: tuple[str, ...],
|
||
|
|
confidence: float,
|
||
|
|
evidence: dict[str, Any] | None = None,
|
||
|
|
) -> RecallItem:
|
||
|
|
return RecallItem(
|
||
|
|
memory_uuid=memory.memory_uuid,
|
||
|
|
memory_key=memory.memory_key,
|
||
|
|
content=memory.content,
|
||
|
|
memory_type=memory.memory_type,
|
||
|
|
confidence=round(min(1.0, max(0.0, confidence)), 6),
|
||
|
|
sources=sources,
|
||
|
|
evidence=evidence or {"mysql_confidence": float(memory.confidence)},
|
||
|
|
)
|
||
|
|
|
||
|
|
async def _cache_read(self, key: str) -> tuple[RecallResult | None, bool]:
|
||
|
|
"""返回 (命中结果, 是否降级)；`degraded` 只在适配器区分出故障时为真。"""
|
||
|
|
if self.cache is None:
|
||
|
|
return None, False
|
||
|
|
try:
|
||
|
|
cached = await self.cache.get(key)
|
||
|
|
except Exception:
|
||
|
|
# 自定义适配器未按契约吞异常时，缓存故障仍不得阻塞主流程。
|
||
|
|
logger.warning("recall cache read degraded key=%s", key)
|
||
|
|
return None, True
|
||
|
|
if getattr(cached, "degraded", False):
|
||
|
|
logger.warning("recall cache read degraded key=%s", key)
|
||
|
|
return None, True
|
||
|
|
raw = getattr(cached, "value", None)
|
||
|
|
if raw is None:
|
||
|
|
return None, False
|
||
|
|
try:
|
||
|
|
payload = json.loads(raw if isinstance(raw, str) else str(raw))
|
||
|
|
items = tuple(
|
||
|
|
RecallItem(
|
||
|
|
memory_uuid=str(entry["memory_uuid"]),
|
||
|
|
memory_key=str(entry["memory_key"]),
|
||
|
|
content=str(entry["content"]),
|
||
|
|
memory_type=str(entry["memory_type"]),
|
||
|
|
confidence=float(entry["confidence"]),
|
||
|
|
sources=tuple(str(source) for source in entry["sources"]),
|
||
|
|
evidence=dict(entry.get("evidence") or {}),
|
||
|
|
)
|
||
|
|
for entry in payload["items"]
|
||
|
|
)
|
||
|
|
except (KeyError, TypeError, ValueError):
|
||
|
|
# 缓存格式不可信：按未命中处理并覆盖回填，绝不返回半解析结果。
|
||
|
|
logger.warning("recall cache payload ignored key=%s", key)
|
||
|
|
return None, False
|
||
|
|
# 降级状态必须与结果一起从缓存恢复，否则缓存命中会把"降级数据"伪装成正常结果。
|
||
|
|
return RecallResult(
|
||
|
|
customer_id=int(payload["customer_id"]),
|
||
|
|
items=items,
|
||
|
|
degraded=bool(payload.get("degraded")),
|
||
|
|
degraded_reasons=tuple(
|
||
|
|
str(reason) for reason in (payload.get("degraded_reasons") or ())
|
||
|
|
),
|
||
|
|
from_cache=True,
|
||
|
|
), False
|
||
|
|
|
||
|
|
async def _cache_write(self, key: str, result: RecallResult) -> bool:
|
||
|
|
if self.cache is None:
|
||
|
|
return False
|
||
|
|
payload = {
|
||
|
|
"customer_id": result.customer_id,
|
||
|
|
"degraded": result.degraded,
|
||
|
|
"degraded_reasons": list(result.degraded_reasons),
|
||
|
|
"items": [asdict(item) for item in result.items],
|
||
|
|
}
|
||
|
|
try:
|
||
|
|
return await self.cache.set(
|
||
|
|
key, json.dumps(payload, ensure_ascii=False), ttl_seconds=self.cache_ttl_seconds
|
||
|
|
)
|
||
|
|
except Exception:
|
||
|
|
logger.warning("recall cache write degraded key=%s", key)
|
||
|
|
return False
|