322 lines
12 KiB
Python
322 lines
12 KiB
Python
"""Hybrid retrieval for historical anomaly cases."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import re
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import httpx
|
|
from sqlalchemy import select
|
|
from sqlalchemy import text as sql_text
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.core.config import get_settings
|
|
from app.models.entities import AnomalyCase
|
|
from app.schemas.risk_prediction import SimilarCase
|
|
from app.services.supplier_identity import resolve_supplier
|
|
|
|
logger = logging.getLogger("dieya_risk_prediction.rag")
|
|
_embedding_semaphore = asyncio.Semaphore(1)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class RetrievedCase:
|
|
case: SimilarCase
|
|
score: float
|
|
case_id: int
|
|
|
|
|
|
class AnomalyCaseRAG:
|
|
"""Retrieve historical cases with business, text, and optional vector signals."""
|
|
|
|
def __init__(
|
|
self,
|
|
session: AsyncSession,
|
|
*,
|
|
top_k: int | None = None,
|
|
transport: httpx.AsyncBaseTransport | None = None,
|
|
) -> None:
|
|
self.session = session
|
|
self.top_k = top_k or get_settings().rag_top_k
|
|
self.transport = transport
|
|
|
|
async def search(
|
|
self,
|
|
*,
|
|
query: str,
|
|
category: str | None = None,
|
|
material: str | None = None,
|
|
supplier: str | None = None,
|
|
factory: str | None = None,
|
|
stage: str | None = None,
|
|
anomaly_type: str | None = None,
|
|
rows: list[dict[str, Any]] | None = None,
|
|
) -> list[RetrievedCase]:
|
|
identity = await resolve_supplier(self.session, supplier)
|
|
supplier_match = (
|
|
identity.short_name or identity.aliases[0]
|
|
if identity
|
|
else supplier
|
|
)
|
|
terms = self._terms(query, category, material, supplier, factory, stage, anomaly_type)
|
|
rows = rows if rows is not None else await self._load_rows(
|
|
query,
|
|
category=category,
|
|
material=material,
|
|
supplier=supplier,
|
|
factory=factory,
|
|
stage=stage,
|
|
anomaly_type=anomaly_type,
|
|
)
|
|
ranked = [
|
|
RetrievedCase(
|
|
case=self._to_case(row),
|
|
score=self._score(
|
|
row,
|
|
query=query,
|
|
category=category,
|
|
material=material,
|
|
supplier=supplier_match,
|
|
factory=factory,
|
|
stage=stage,
|
|
anomaly_type=anomaly_type,
|
|
terms=terms,
|
|
),
|
|
case_id=int(row["id"]),
|
|
)
|
|
for row in rows
|
|
]
|
|
ranked = [item for item in ranked if item.score > 0]
|
|
ranked.sort(key=lambda item: (item.score, item.case_id), reverse=True)
|
|
for item in ranked[: self.top_k]:
|
|
vector_similarity = next(
|
|
(
|
|
float(row["vector_score"])
|
|
for row in rows
|
|
if int(row["id"]) == item.case_id and row.get("vector_score") is not None
|
|
),
|
|
None,
|
|
)
|
|
item.case.similarity = round(
|
|
max(0.0, min(1.0, vector_similarity if vector_similarity is not None else item.score / 100)),
|
|
3,
|
|
)
|
|
return ranked[: self.top_k]
|
|
|
|
async def _load_rows(
|
|
self,
|
|
query: str,
|
|
*,
|
|
category: str | None,
|
|
material: str | None,
|
|
supplier: str | None,
|
|
factory: str | None,
|
|
stage: str | None,
|
|
anomaly_type: str | None,
|
|
) -> list[dict[str, Any]]:
|
|
fields = (
|
|
AnomalyCase.id,
|
|
AnomalyCase.date,
|
|
AnomalyCase.product_name,
|
|
AnomalyCase.product_category,
|
|
AnomalyCase.material_type,
|
|
AnomalyCase.supplier,
|
|
AnomalyCase.factory,
|
|
AnomalyCase.stage_new,
|
|
AnomalyCase.anomaly_type,
|
|
AnomalyCase.description,
|
|
AnomalyCase.solution,
|
|
AnomalyCase.root_cause,
|
|
AnomalyCase.defect_level,
|
|
AnomalyCase.result,
|
|
)
|
|
vector_ready = await self._vector_column_ready()
|
|
has_vectors = await self._has_stored_vectors() if vector_ready else False
|
|
vector = await self._query_embedding(query) if has_vectors else None
|
|
if vector is None:
|
|
return [dict(row) for row in (await self.session.execute(select(*fields).limit(5000))).mappings().all()]
|
|
|
|
vector_literal = "[" + ",".join(str(float(value)) for value in vector) + "]"
|
|
statement = sql_text(
|
|
"""
|
|
SELECT id, date, product_name, product_category, material_type, supplier,
|
|
factory, stage_new, anomaly_type, description, solution, root_cause,
|
|
defect_level, result, 1 - (embedding_vector <=> CAST(:embedding AS halfvec)) AS vector_score
|
|
FROM anomaly_cases
|
|
WHERE embedding_vector IS NOT NULL
|
|
ORDER BY embedding_vector <=> CAST(:embedding AS halfvec)
|
|
LIMIT 5000
|
|
"""
|
|
)
|
|
result = await self.session.execute(statement, {"embedding": vector_literal})
|
|
rows = [dict(row) for row in result.mappings().all()]
|
|
if rows:
|
|
return rows
|
|
return [dict(row) for row in (await self.session.execute(select(*fields).limit(5000))).mappings().all()]
|
|
|
|
async def _vector_column_ready(self) -> bool:
|
|
try:
|
|
result = await self.session.execute(
|
|
sql_text(
|
|
"""
|
|
SELECT EXISTS (
|
|
SELECT 1
|
|
FROM information_schema.columns
|
|
WHERE table_name = 'anomaly_cases'
|
|
AND column_name = 'embedding_vector'
|
|
AND udt_name IN ('vector', 'halfvec')
|
|
)
|
|
"""
|
|
)
|
|
)
|
|
return bool(result.scalar())
|
|
except Exception: # noqa: BLE001 - a missing optional extension uses lexical retrieval
|
|
await self.session.rollback()
|
|
return False
|
|
|
|
async def _query_embedding(self, query: str) -> list[float] | None:
|
|
settings = get_settings()
|
|
api_key = (
|
|
settings.rag_embedding_api_key.get_secret_value()
|
|
if settings.rag_embedding_api_key
|
|
else settings.llm_api_key.get_secret_value() if settings.llm_api_key else ""
|
|
)
|
|
base_url = settings.rag_embedding_base_url or settings.llm_base_url
|
|
if not api_key or not base_url:
|
|
return None
|
|
endpoint = f"{base_url.rstrip('/')}/embeddings"
|
|
if not base_url.rstrip("/").endswith("/v1"):
|
|
endpoint = f"{base_url.rstrip('/')}/v1/embeddings"
|
|
timeout = httpx.Timeout(
|
|
settings.rag_embedding_timeout_seconds,
|
|
connect=min(5, settings.rag_embedding_timeout_seconds),
|
|
)
|
|
try:
|
|
async with _embedding_semaphore:
|
|
async with httpx.AsyncClient(timeout=timeout, transport=self.transport) as client:
|
|
response = await client.post(
|
|
endpoint,
|
|
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
|
|
json={"model": settings.rag_embedding_model, "input": query[:2000]},
|
|
)
|
|
response.raise_for_status()
|
|
body = response.json()
|
|
embedding = (body.get("data") or [{}])[0].get("embedding")
|
|
if not isinstance(embedding, list) or len(embedding) != settings.rag_embedding_dimensions:
|
|
return None
|
|
return [float(value) for value in embedding]
|
|
except (httpx.HTTPError, ValueError, TypeError, KeyError) as exc:
|
|
logger.warning(
|
|
"RAG embedding request failed error=%s; using lexical retrieval",
|
|
exc.__class__.__name__,
|
|
)
|
|
return None
|
|
|
|
async def _has_stored_vectors(self) -> bool:
|
|
"""Avoid a paid/network embedding call until case vectors exist."""
|
|
|
|
try:
|
|
result = await self.session.execute(
|
|
sql_text(
|
|
"SELECT EXISTS ("
|
|
"SELECT 1 FROM anomaly_cases WHERE embedding_vector IS NOT NULL LIMIT 1"
|
|
")"
|
|
)
|
|
)
|
|
return bool(result.scalar())
|
|
except Exception: # noqa: BLE001 - lexical retrieval is the safe fallback
|
|
await self.session.rollback()
|
|
return False
|
|
|
|
@staticmethod
|
|
def _terms(*values: str | None) -> list[str]:
|
|
terms: list[str] = []
|
|
for value in values:
|
|
if not value:
|
|
continue
|
|
cleaned = " ".join(str(value).strip().lower().split())
|
|
runs = re.findall(r"[a-z0-9]+|[\u4e00-\u9fff]{2,}", cleaned)
|
|
terms.extend(runs)
|
|
for run in runs:
|
|
if re.fullmatch(r"[\u4e00-\u9fff]+", run) and len(run) > 3:
|
|
terms.extend(
|
|
run[index : index + size]
|
|
for size in (2, 3, 4)
|
|
for index in range(len(run) - size + 1)
|
|
)
|
|
return list(dict.fromkeys(term for term in terms if len(term) > 1))
|
|
|
|
@classmethod
|
|
def _score(
|
|
cls,
|
|
row: dict[str, Any],
|
|
*,
|
|
query: str,
|
|
category: str | None,
|
|
material: str | None,
|
|
supplier: str | None,
|
|
factory: str | None,
|
|
stage: str | None,
|
|
anomaly_type: str | None,
|
|
terms: list[str],
|
|
) -> float:
|
|
score = 0.0
|
|
score += cls._field_score(row.get("product_category"), category, 7)
|
|
score += cls._field_score(row.get("product_name"), category, 5)
|
|
score += cls._field_score(row.get("material_type"), material, 7)
|
|
score += cls._field_score(row.get("supplier"), supplier, 6)
|
|
score += cls._field_score(row.get("factory"), factory, 6)
|
|
score += cls._field_score(row.get("stage_new"), stage, 3)
|
|
score += cls._field_score(row.get("anomaly_type"), anomaly_type, 7)
|
|
vector_score = row.get("vector_score")
|
|
if vector_score is not None:
|
|
score += max(0.0, float(vector_score)) * 40
|
|
searchable = " ".join(str(row.get(key) or "") for key in (
|
|
"product_name",
|
|
"product_category",
|
|
"material_type",
|
|
"anomaly_type",
|
|
"description",
|
|
"root_cause",
|
|
"solution",
|
|
)).lower()
|
|
query_lower = query.lower().strip()
|
|
if query_lower and query_lower in searchable:
|
|
score += 8
|
|
score += min(12.0, sum(1 for term in terms if term in searchable) * 1.5)
|
|
if str(row.get("defect_level") or "").upper() == "A":
|
|
score += 1.0
|
|
return round(score, 2)
|
|
|
|
@staticmethod
|
|
def _field_score(actual: Any, expected: str | None, weight: float) -> float:
|
|
if not actual or not expected:
|
|
return 0.0
|
|
actual_text = str(actual).lower()
|
|
expected_text = str(expected).lower()
|
|
if actual_text == expected_text:
|
|
return weight
|
|
if expected_text in actual_text or actual_text in expected_text:
|
|
return weight * 0.7
|
|
return 0.0
|
|
|
|
@staticmethod
|
|
def _to_case(row: dict[str, Any]) -> SimilarCase:
|
|
return SimilarCase(
|
|
case_id=int(row["id"]) if row.get("id") is not None else None,
|
|
date=row.get("date").isoformat() if row.get("date") else None,
|
|
product=row.get("product_name") or row.get("product_category") or "未命名产品",
|
|
category=row.get("product_category"),
|
|
material=row.get("material_type"),
|
|
anomaly_type=row.get("anomaly_type"),
|
|
description=row.get("description"),
|
|
root_cause=row.get("root_cause"),
|
|
defect_level=row.get("defect_level"),
|
|
result=row.get("result"),
|
|
problem=(row.get("anomaly_type") or row.get("description") or "历史异常未填写"),
|
|
solution=(row.get("solution") or row.get("root_cause") or "请查看原始案例记录"),
|
|
)
|