dev #1

Merged
kx_guo merged 24 commits from dev into main 2026-09-21 10:51:18 +08:00
9 changed files with 784 additions and 8 deletions
Showing only changes of commit b0b6ef014f - Show all commits

View File

@ -47,6 +47,16 @@ RISK_LLM_FALLBACK_MAX_RETRIES=1
RISK_LLM_FALLBACK_REASONING_EFFORT=high
RISK_LLM_FALLBACK_MAX_OUTPUT_TOKENS=16384
# ==================== RAG 向量检索配置 ====================
# 可选的 Embeddings 服务;留空 API Key 时复用主 LLM Key 和地址。
# 当前数据库未安装 pgvector 时,系统自动使用业务字段 + 文本相关性混合检索。
RISK_RAG_EMBEDDING_API_KEY=
RISK_RAG_EMBEDDING_BASE_URL=https://api.openlux.ai/v1
RISK_RAG_EMBEDDING_MODEL=gemini-embedding-2-preview
RISK_RAG_EMBEDDING_DIMENSIONS=3072
RISK_RAG_EMBEDDING_TIMEOUT_SECONDS=10
RISK_RAG_TOP_K=5
# ==================== 跨域配置 ====================
# 允许访问后端接口的前端地址,多个地址用英文逗号分隔。
RISK_CORS_ALLOW_ORIGINS=http://localhost:39158,http://127.0.0.1:39158

View File

@ -92,6 +92,37 @@ class Settings(BaseSettings):
le=393_216,
validation_alias=AliasChoices("RISK_LLM_FALLBACK_MAX_OUTPUT_TOKENS", "DEEPSEEK_MAX_OUTPUT_TOKENS"),
)
rag_embedding_api_key: SecretStr | None = Field(
default=None,
repr=False,
validation_alias=AliasChoices("RISK_RAG_EMBEDDING_API_KEY", "EMBEDDING_API_KEY"),
)
rag_embedding_base_url: str = Field(
default="",
validation_alias=AliasChoices("RISK_RAG_EMBEDDING_BASE_URL", "EMBEDDING_BASE_URL"),
)
rag_embedding_model: str = Field(
default="bge-m3",
validation_alias=AliasChoices("RISK_RAG_EMBEDDING_MODEL", "EMBEDDING_MODEL"),
)
rag_embedding_dimensions: int = Field(
default=3072,
ge=1,
le=4096,
validation_alias=AliasChoices("RISK_RAG_EMBEDDING_DIMENSIONS", "EMBEDDING_DIMENSIONS"),
)
rag_embedding_timeout_seconds: int = Field(
default=10,
gt=0,
le=60,
validation_alias=AliasChoices("RISK_RAG_EMBEDDING_TIMEOUT_SECONDS", "EMBEDDING_TIMEOUT"),
)
rag_top_k: int = Field(
default=5,
ge=1,
le=20,
validation_alias=AliasChoices("RISK_RAG_TOP_K", "RAG_TOP_K"),
)
cors_allow_origins: Annotated[tuple[str, ...], NoDecode] = (
"http://localhost:39158",
"http://127.0.0.1:39158",
@ -121,6 +152,11 @@ class Settings(BaseSettings):
def normalize_llm_fallback_base_url(cls, value: str) -> str:
return value.rstrip("/")
@field_validator("rag_embedding_base_url")
@classmethod
def normalize_rag_embedding_base_url(cls, value: str) -> str:
return value.rstrip("/")
@field_validator("cors_allow_origins", mode="before")
@classmethod
def split_cors_origins(cls, value: object) -> object:

View File

@ -67,10 +67,18 @@ class RiskDimensionScore(ApiSchema):
class SimilarCase(ApiSchema):
case_id: int | None = None
date: str | None = None
product: str | None = None
category: str | None = None
material: str | None = None
anomaly_type: str | None = None
description: str | None = None
root_cause: str | None = None
defect_level: str | None = None
problem: str | None = None
solution: str | None = None
similarity: float | None = None
class RiskAssessmentResult(ApiSchema):

View File

@ -28,6 +28,7 @@ from app.schemas.risk_prediction import (
)
from app.services.database_risk_scoring import DatabaseRiskScoringService
from app.services.llm_client import LLMClient
from app.services.rag import AnomalyCaseRAG
logger = logging.getLogger("dieya_risk_prediction.conversations")
@ -233,11 +234,7 @@ class ConversationService:
show_risk = self._is_risk_request(text, semantics)
response_type = self._response_type(text, semantics)
response_type = await self._use_solution_card_without_history(
response_type,
text,
semantics,
)
response_type = await self._use_solution_card_without_history(response_type, text, semantics)
risk_result = (
await self._enrich_risk_result(text, semantics, self._predict(text))
if show_risk and response_type == "risk_assessment"
@ -657,6 +654,18 @@ class ConversationService:
) or "当前新品"
direct_count = await self._count_category_anomalies(category)
standards = await self._standards_for(category, text)
similar_cases = await self._similar_cases(
text,
category=category,
material=self._first_match(
text,
("软管", "PET瓶", "PET", "铝瓶", "泵头", "彩盒", "瓶子", "喷头", "盖子"),
),
supplier=self._semantic_text(semantics, "supplier"),
factory=self._semantic_text(semantics, "factory"),
stage=self._semantic_text(semantics, "stage"),
anomaly_type=self._semantic_text(semantics, "anomaly_type"),
)
standard_fallback_enabled = direct_count == 0
if standard_fallback_enabled and not standards:
standards = await self._fallback_standards_for(text)
@ -700,7 +709,10 @@ class ConversationService:
"standard_fallback_enabled": standard_fallback_enabled,
"matched_standards": [self._standard_dict(item) for item in standards[:5]],
"identified_risks": risks,
"rag_similar_cases": [],
"rag_similar_cases": [
{"similarity": item.case.similarity, **item.case.model_dump(mode="json", by_alias=False)}
for item in similar_cases
],
}
)
except Exception as exc: # noqa: BLE001 - the solution card must remain available
@ -722,6 +734,7 @@ class ConversationService:
),
"standardFallbackEnabled": standard_fallback_enabled,
"standards": [self._standard_dict(item) for item in standards[:4]],
"similarCases": [item.case.model_dump(mode="json", by_alias=False) for item in similar_cases],
"risks": risks,
"suggestions": suggestions,
"warning": "以上方案由 AI 推理生成,实际执行前请咨询品控部门确认。",
@ -908,6 +921,18 @@ class ConversationService:
factory_count = await self._count_anomalies(AnomalyCase.factory, factory)
anomaly_count = supplier_count if supplier else factory_count
standard_rows = await self._standards_for(category, text)
similar_cases = await self._similar_cases(
text,
category=category,
material=self._first_match(
text,
("软管", "PET瓶", "PET", "铝瓶", "泵头", "彩盒", "瓶子", "喷头", "盖子"),
),
supplier=supplier,
factory=factory,
stage=self._semantic_text(semantics, "stage"),
anomaly_type=self._semantic_text(semantics, "anomaly_type"),
)
alpha_count = supplier_count if supplier else 0
alpha = round(min(1.5, 1 + (alpha_count / 100) * 0.3), 2) if alpha_count else 1.0
@ -952,6 +977,7 @@ class ConversationService:
database_score=database_score,
dimensions=dimensions,
standard_rows=standard_rows,
similar_cases=similar_cases,
)
return result.model_copy(
@ -986,7 +1012,7 @@ class ConversationService:
),
"standard_judgment": standard_rows,
"dimension_scores": dimensions,
"similar_cases": [],
"similar_cases": [item.case.model_dump(mode="json", by_alias=False) for item in similar_cases],
"ai_suggestions": suggestions,
}
)
@ -1012,6 +1038,7 @@ class ConversationService:
database_score,
dimensions: Sequence[RiskDimensionScore],
standard_rows: Sequence[StandardJudgment],
similar_cases: Sequence[Any],
) -> Sequence[str]: # noqa: ANN001
"""Generate suggestions from live scoring data, with a deterministic fallback."""
@ -1042,7 +1069,13 @@ class ConversationService:
for item in dimensions
],
"matched_standards": [self._standard_dict(item) for item in standard_rows[:5]],
"rag_similar_cases": [],
"rag_similar_cases": [
{
"similarity": item.case.similarity,
**item.case.model_dump(mode="json", by_alias=False),
}
for item in similar_cases
],
}
try:
suggestions = await self.llm.generate_prevention_suggestions(context)
@ -1055,6 +1088,33 @@ class ConversationService:
)
return self._suggestions(text, supplier_or_factory)
async def _similar_cases(
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,
) -> Sequence[Any]:
"""Return historical cases for both the card and the AI context."""
try:
return await AnomalyCaseRAG(self.session).search(
query=query,
category=category,
material=material,
supplier=supplier,
factory=factory,
stage=stage,
anomaly_type=anomaly_type,
)
except Exception as exc: # noqa: BLE001 - RAG must not block risk scoring
logger.warning("RAG historical case search failed error=%s", exc.__class__.__name__)
return []
async def _count_category_anomalies(self, category: str) -> int:
"""The imported workbook has sparse product_category values.

311
app/services/rag.py Normal file
View File

@ -0,0 +1,311 @@
"""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
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,
) -> list[RetrievedCase]:
terms = self._terms(query, category, material, supplier, factory, stage, anomaly_type)
rows = 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,
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,
)
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, 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"),
problem=(row.get("anomaly_type") or row.get("description") or "历史异常未填写"),
solution=(row.get("solution") or row.get("root_cause") or "请查看原始案例记录"),
)

View File

@ -0,0 +1,51 @@
"""Add the optional pgvector column for anomaly-case retrieval."""
from collections.abc import Sequence
import sqlalchemy as sa
from alembic import op
revision: str = "20260905_0014"
down_revision: str | None = "20260902_0013"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def upgrade() -> None:
connection = op.get_bind()
available = connection.execute(
sa.text("SELECT EXISTS (SELECT 1 FROM pg_available_extensions WHERE name = 'vector')")
).scalar()
if not available:
# The application has a structured/text retrieval fallback. Keep the
# migration non-blocking on managed PostgreSQL instances without the
# optional extension; install pgvector and rerun this revision later.
return
is_superuser = connection.execute(
sa.text("SELECT rolsuper FROM pg_roles WHERE rolname = current_user")
).scalar()
if not is_superuser:
# Managed PostgreSQL may expose the package but reserve CREATE
# EXTENSION for an administrator. Lexical RAG remains usable.
return
connection.execute(sa.text("CREATE EXTENSION IF NOT EXISTS vector"))
connection.execute(
sa.text(
"ALTER TABLE anomaly_cases "
"ADD COLUMN IF NOT EXISTS embedding_vector halfvec(3072)"
)
)
connection.execute(
sa.text(
"CREATE INDEX IF NOT EXISTS ix_anomaly_cases_embedding_vector "
"ON anomaly_cases USING hnsw (embedding_vector halfvec_cosine_ops)"
)
)
def downgrade() -> None:
connection = op.get_bind()
connection.execute(sa.text("DROP INDEX IF EXISTS ix_anomaly_cases_embedding_vector"))
connection.execute(
sa.text("ALTER TABLE anomaly_cases DROP COLUMN IF EXISTS embedding_vector")
)

View File

@ -0,0 +1,175 @@
"""Generate and store anomaly-case embeddings after pgvector is available.
Usage:
.venv/Scripts/python -m scripts.backfill_embeddings
The script is intentionally separate from application startup. Embedding 2,448
rows can take time and should never make the API unavailable.
"""
from __future__ import annotations
import argparse
import asyncio
import logging
import random
from typing import Any
import httpx
from app.core.config import get_settings
from app.database.session import SessionFactory
from sqlalchemy import text
logger = logging.getLogger("dieya_risk_prediction.embedding_backfill")
def _case_text(row: dict[str, Any]) -> str:
labels = (
("产品", row.get("product_name")),
("品类", row.get("product_category")),
("包材", row.get("material_type")),
("供应商", row.get("supplier")),
("加工厂", row.get("factory")),
("阶段", row.get("stage_new")),
("异常类型", row.get("anomaly_type")),
("异常描述", row.get("description")),
("原因", row.get("root_cause")),
("解决方案", row.get("solution")),
)
return ";".join(f"{label}:{value}" for label, value in labels if value)
async def _embedding(
client: httpx.AsyncClient,
endpoint: str,
api_key: str,
texts: list[str],
model: str,
) -> list[list[float]]:
for attempt in range(4):
response = await client.post(
endpoint,
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
json={"model": model, "input": texts},
)
if response.status_code not in {429, 500, 502, 503, 504}:
response.raise_for_status()
data = response.json().get("data") or []
ordered = sorted(data, key=lambda item: item.get("index", 0))
return [list(map(float, item["embedding"])) for item in ordered]
if attempt == 3:
response.raise_for_status()
retry_after = response.headers.get("retry-after")
try:
delay = min(60.0, max(3.0, float(retry_after))) if retry_after else 3.0 * (2**attempt)
except ValueError:
delay = 3.0 * (2**attempt)
delay += random.uniform(0, 1)
logger.warning(
"Embedding API 限流或暂时不可用 status=%s attempt=%s delay=%.1fs",
response.status_code,
attempt + 1,
delay,
)
await asyncio.sleep(delay)
raise RuntimeError("Embedding 请求重试次数已用尽")
async def main(batch_size: int) -> 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:
raise RuntimeError("未配置 Embeddings API Key 或 Base URL")
endpoint = f"{base_url.rstrip('/')}/embeddings"
if not base_url.rstrip("/").endswith("/v1"):
endpoint = f"{base_url.rstrip('/')}/v1/embeddings"
async with SessionFactory() as session:
column_exists = await session.scalar(
text(
"""
SELECT EXISTS (
SELECT 1 FROM information_schema.columns
WHERE table_name = 'anomaly_cases'
AND column_name = 'embedding_vector'
AND udt_name = 'halfvec'
)
"""
)
)
if not column_exists:
raise RuntimeError("当前数据库没有 embedding_vector,请先安装 pgvector 并执行 alembic upgrade head")
rows = (
await session.execute(
text(
"""
SELECT id, product_name, product_category, material_type, supplier,
factory, stage_new, anomaly_type, description, root_cause, solution
FROM anomaly_cases
WHERE embedding_vector IS NULL
ORDER BY id
"""
)
)
).mappings().all()
logger.info("待生成案例向量数量=%s", len(rows))
print(f"待生成案例向量数量:{len(rows)},批次大小:{batch_size}", flush=True)
timeout = httpx.Timeout(
settings.rag_embedding_timeout_seconds,
connect=min(5, settings.rag_embedding_timeout_seconds),
)
async with httpx.AsyncClient(timeout=timeout) as client:
success_count = 0
failure_count = 0
for start in range(0, len(rows), batch_size):
batch = rows[start : start + batch_size]
try:
vectors = await _embedding(
client,
endpoint,
api_key,
[_case_text(dict(row)) for row in batch],
settings.rag_embedding_model,
)
except (httpx.HTTPError, ValueError, TypeError, KeyError) as exc:
failure_count += len(batch)
logger.exception("案例向量生成失败 case_id=%s", batch[0]["id"])
print(f"案例 {batch[0]['id']} 失败:{exc.__class__.__name__},已跳过", flush=True)
continue
if len(vectors) != len(batch):
raise RuntimeError("Embeddings 返回数量与请求数量不一致")
async with SessionFactory() as session:
for row, vector in zip(batch, vectors, strict=True):
literal = "[" + ",".join(str(value) for value in vector) + "]"
await session.execute(
text(
"UPDATE anomaly_cases "
"SET embedding_vector = CAST(:embedding AS halfvec) "
"WHERE id = :id"
),
{"embedding": literal, "id": row["id"]},
)
await session.commit()
success_count += len(batch)
logger.info("已完成案例向量回填 %s/%s", min(start + batch_size, len(rows)), len(rows))
print(
f"已完成:{min(start + batch_size, len(rows))}/{len(rows)},"
f"成功:{success_count},失败:{failure_count},当前案例:{batch[-1]['id']}",
flush=True,
)
await asyncio.sleep(1.5)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--batch-size", type=int, default=1)
args = parser.parse_args()
if not 1 <= args.batch_size <= 8:
raise SystemExit("--batch-size 必须在 1 到 8 之间")
asyncio.run(main(args.batch_size))

View File

@ -0,0 +1,38 @@
"""Enable pgvector schema objects with a PostgreSQL administrator account.
Run once after the pgvector package is installed on the server:
.venv/Scripts/python -m scripts.enable_pgvector
Then the normal application account can run ``scripts.backfill_embeddings``.
"""
from __future__ import annotations
import asyncio
from app.database.session import SessionFactory
from sqlalchemy import text
async def main() -> None:
async with SessionFactory() as session:
await session.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
await session.execute(
text(
"ALTER TABLE anomaly_cases "
"ADD COLUMN IF NOT EXISTS embedding_vector halfvec(3072)"
)
)
await session.execute(
text(
"CREATE INDEX IF NOT EXISTS ix_anomaly_cases_embedding_vector "
"ON anomaly_cases USING hnsw (embedding_vector halfvec_cosine_ops)"
)
)
await session.commit()
print("pgvector 已启用,embedding_vector 和 HNSW 索引已准备完成。")
if __name__ == "__main__":
asyncio.run(main())

87
tests/test_rag.py Normal file
View File

@ -0,0 +1,87 @@
from datetime import date
import pytest
from app.services.rag import AnomalyCaseRAG
def _row(**values):
return {
"id": 1,
"date": date(2026, 1, 10),
"product_name": "洗发水 500ml",
"product_category": "洗护类",
"material_type": "软管",
"supplier": "祥印",
"factory": "晶品",
"stage_new": "产品开发",
"anomaly_type": "密封性异常",
"description": "软管漏液",
"solution": "增加密封性测试",
"root_cause": "封口参数不稳定",
"defect_level": "B",
**values,
}
def test_rag_terms_and_structured_text_score() -> None:
terms = AnomalyCaseRAG._terms("软管密封性异常", "洗发水", "软管")
score = AnomalyCaseRAG._score(
_row(),
query="洗发水软管密封性异常",
category="洗发水",
material="软管",
supplier="祥印",
factory=None,
stage="产品开发",
anomaly_type="密封性异常",
terms=terms,
)
assert "密封性" in terms
assert score > 20
case = AnomalyCaseRAG._to_case(_row())
assert case.product == "洗发水 500ml"
assert case.problem == "密封性异常"
assert case.solution == "增加密封性测试"
@pytest.mark.asyncio
async def test_rag_does_not_call_embedding_until_vectors_are_stored() -> None:
class FakeMappings:
def all(self):
return [_row()]
class FakeResult:
def mappings(self):
return FakeMappings()
class FakeSession:
async def execute(self, _statement):
return FakeResult()
rag = AnomalyCaseRAG(FakeSession(), top_k=1)
async def vector_column_ready():
return True
async def has_stored_vectors():
return False
async def embedding_must_not_run(_query):
raise AssertionError("没有已回填向量时不应调用 Embedding")
rag._vector_column_ready = vector_column_ready
rag._has_stored_vectors = has_stored_vectors
rag._query_embedding = embedding_must_not_run
rows = await rag._load_rows(
"洗发水软管密封性异常",
category="洗发水",
material="软管",
supplier=None,
factory=None,
stage=None,
anomaly_type=None,
)
assert len(rows) == 1