From b0b6ef014f3b06621c05f1f45cf46b1999c0ff38 Mon Sep 17 00:00:00 2001 From: kx_guo Date: Sat, 5 Sep 2026 19:00:25 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=9D=E7=95=99RAG=E6=A3=80=E7=B4=A2?= =?UTF-8?q?=E5=B9=B6=E4=BF=AE=E5=A4=8D=E7=9B=B8=E4=BC=BC=E6=A1=88=E4=BE=8B?= =?UTF-8?q?=E8=B0=83=E7=94=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 保留 pgvector 与混合检索能力，返回 TOP 5 相似案例\n- 将相似案例上下文传入 AI 预防建议生成\n- 修复风险评估调用缺少 similar_cases 参数导致的运行报错\n- 保留 Embedding 配置、数据库迁移、回填脚本及 RAG 测试\n- 恢复基础字符串数组格式的 AI 预防建议协议 --- .env.example | 10 + app/core/config.py | 36 ++ app/schemas/risk_prediction.py | 8 + app/services/conversations.py | 76 ++++- app/services/rag.py | 311 ++++++++++++++++++ .../20260905_0014_add_optional_pgvector.py | 51 +++ scripts/backfill_embeddings.py | 175 ++++++++++ scripts/enable_pgvector.py | 38 +++ tests/test_rag.py | 87 +++++ 9 files changed, 784 insertions(+), 8 deletions(-) create mode 100644 app/services/rag.py create mode 100644 migrations/versions/20260905_0014_add_optional_pgvector.py create mode 100644 scripts/backfill_embeddings.py create mode 100644 scripts/enable_pgvector.py create mode 100644 tests/test_rag.py diff --git a/.env.example b/.env.example index 89e8730..3e19b2c 100644 --- a/.env.example +++ b/.env.example @@ -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 diff --git a/app/core/config.py b/app/core/config.py index 2dbf7f8..24d59c1 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -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: diff --git a/app/schemas/risk_prediction.py b/app/schemas/risk_prediction.py index a02dad4..6a5a364 100644 --- a/app/schemas/risk_prediction.py +++ b/app/schemas/risk_prediction.py @@ -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): diff --git a/app/services/conversations.py b/app/services/conversations.py index 03b57db..1ddd7e1 100644 --- a/app/services/conversations.py +++ b/app/services/conversations.py @@ -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. diff --git a/app/services/rag.py b/app/services/rag.py new file mode 100644 index 0000000..5e50a1b --- /dev/null +++ b/app/services/rag.py @@ -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 "请查看原始案例记录"), + ) diff --git a/migrations/versions/20260905_0014_add_optional_pgvector.py b/migrations/versions/20260905_0014_add_optional_pgvector.py new file mode 100644 index 0000000..1598ce6 --- /dev/null +++ b/migrations/versions/20260905_0014_add_optional_pgvector.py @@ -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") + ) diff --git a/scripts/backfill_embeddings.py b/scripts/backfill_embeddings.py new file mode 100644 index 0000000..2da7356 --- /dev/null +++ b/scripts/backfill_embeddings.py @@ -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)) diff --git a/scripts/enable_pgvector.py b/scripts/enable_pgvector.py new file mode 100644 index 0000000..b452ef7 --- /dev/null +++ b/scripts/enable_pgvector.py @@ -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()) diff --git a/tests/test_rag.py b/tests/test_rag.py new file mode 100644 index 0000000..d026682 --- /dev/null +++ b/tests/test_rag.py @@ -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