feat:客服agent接入nl2sql

nl2sql/metadata_sync.py

@@ -2,18 +2,127 @@
 from __future__ import annotations
 
 import time
-from collections.abc import Awaitable, Callable
+from collections.abc import Awaitable, Callable, Iterable
 from typing import Any
 
 from nl2sql.embedding import embed_texts
 from nl2sql.metadata import build_metadata_chunks
 from nl2sql.milvus_collections import NL2SQL_COLLECTION
 
+# ORM 已定义但数据库尚未建表的表，在表说明后追加此标注，
+# 让 NL2SQL 生成侧知道该表暂不可查询（权威 Schema 校验也会兜底拒绝）。
+MISSING_TABLE_MARKER = "（注意：当前数据库中尚未建表）"
+
 
 def _escape_filter_value(value: str) -> str:
     return value.replace("\\", "\\\\").replace('"', '\\"')
 
 
+def _row_table_name(row: dict[str, Any]) -> str:
+    """兼容 information_schema 大小写键名，取行内表名。"""
+    value = (
+        row.get("table_name")
+        or row.get("TABLE_NAME")
+        or row.get("Table_name")
+        or ""
+    )
+    return str(value).strip()
+
+
+def _filter_rows_by_tables(
+    rows: list[dict[str, Any]], allowed_tables: set[str]
+) -> list[dict[str, Any]]:
+    return [row for row in rows if _row_table_name(row) in allowed_tables]
+
+
+async def _delete_tables_outside_allowlist(
+    milvus_client, allowed_tables: set[str]
+) -> list[str]:
+    """删除集合中不在允许名单内的表元数据 chunk，返回被清理的表名。"""
+    existing_rows = await milvus_client.query(
+        collection_name=NL2SQL_COLLECTION,
+        filter="is_valid == true",
+        output_fields=["table_name"],
+        limit=16384,
+    )
+    existing_tables = {
+        str(row.get("table_name") or "").strip()
+        for row in existing_rows or []
+        if str(row.get("table_name") or "").strip()
+    }
+    removed: list[str] = []
+    for table_name in sorted(existing_tables - allowed_tables):
+        await milvus_client.delete(
+            collection_name=NL2SQL_COLLECTION,
+            filter=f'table_name == "{_escape_filter_value(table_name)}"',
+        )
+        removed.append(table_name)
+    return removed
+
+
+def build_orm_metadata_rows(
+    tables: Iterable[Any],
+) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
+    """把 SQLAlchemy Table 定义合成 information_schema 风格的元数据行。
+
+    用于 model/ 已定义、但数据库尚未建表的表：
+    - 表注释取 ORM ``__table_args__`` 的 comment（ORM 无则空串）；
+    - 字段注释取列 comment（ORM 通常未标注，则为空串）；
+    - 字段类型/可空性从 ORM 列定义推导。
+    """
+    table_rows: list[dict[str, Any]] = []
+    column_rows: list[dict[str, Any]] = []
+    for table in tables:
+        table_rows.append(
+            {
+                "TABLE_NAME": str(table.name).strip(),
+                "TABLE_COMMENT": str(table.comment or "").strip(),
+                "TABLE_TYPE": "BASE TABLE",
+            }
+        )
+        for position, column in enumerate(table.columns, start=1):
+            column_rows.append(
+                {
+                    "TABLE_NAME": str(table.name).strip(),
+                    "COLUMN_NAME": str(column.name).strip(),
+                    "COLUMN_COMMENT": str(column.comment or "").strip(),
+                    "DATA_TYPE": str(column.type).strip().lower(),
+                    "IS_NULLABLE": "YES" if column.nullable else "NO",
+                    "ORDINAL_POSITION": position,
+                }
+            )
+    return table_rows, column_rows
+
+
+def merge_orm_metadata_rows(
+    table_rows: list[dict[str, Any]],
+    column_rows: list[dict[str, Any]],
+    *,
+    allowed_tables: set[str],
+    missing_table_objects: Iterable[Any] = (),
+) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
+    """以 model/ ORM 全量为名单主体合并元数据行。
+
+    - 库里真实存在的 ORM 表：沿用 information_schema 行（列注释更全）；
+    - 库里没有的 ORM 表（missing_table_objects）：用 ORM 定义合成行，
+      并在表说明追加 MISSING_TABLE_MARKER 标注"尚未建表"；
+    - 名单之外（DB-only）的行在此丢弃，交由 sync_metadata 从 Milvus 清理。
+    """
+    allowed = {str(name).strip() for name in allowed_tables if str(name).strip()}
+    merged_tables = [row for row in table_rows if _row_table_name(row) in allowed]
+    merged_columns = [row for row in column_rows if _row_table_name(row) in allowed]
+
+    missing_table_rows, missing_column_rows = build_orm_metadata_rows(
+        missing_table_objects
+    )
+    for row in missing_table_rows:
+        comment = str(row.get("TABLE_COMMENT") or "").strip()
+        row["TABLE_COMMENT"] = f"{comment}{MISSING_TABLE_MARKER}".strip()
+    merged_tables.extend(missing_table_rows)
+    merged_columns.extend(missing_column_rows)
+    return merged_tables, merged_columns
+
+
 async def prepare_metadata_rows(
     chunks: list[dict[str, Any]],
     *,
@@ -39,7 +148,18 @@ async def sync_metadata(
     *,
     embedder: Callable[[list[str]], Awaitable[list[list[float]]]] = embed_texts,
     timestamp: int | None = None,
+    allowed_tables: set[str] | None = None,
 ) -> int:
+    """构造并 upsert 元数据 chunk。
+
+    allowed_tables 提供时：仅同步名单内的表，并删除集合中名单之外的
+    表元数据（ Milvus 里只保留主数据源认可的表，如 model/ ORM 覆盖的表）。
+    """
+    if allowed_tables is not None:
+        allowed = {str(name).strip() for name in allowed_tables if str(name).strip()}
+        table_rows = _filter_rows_by_tables(table_rows, allowed)
+        column_rows = _filter_rows_by_tables(column_rows, allowed)
+
     chunks = build_metadata_chunks(table_rows, column_rows)
     chunks_by_table: dict[str, list[dict[str, Any]]] = {}
     for chunk in chunks:
@@ -80,4 +200,7 @@ async def sync_metadata(
         if rows:
             await milvus_client.upsert(collection_name=NL2SQL_COLLECTION, data=rows)
             updated_count += len(rows)
+
+    if allowed_tables is not None:
+        await _delete_tables_outside_allowlist(milvus_client, allowed)
     return updated_count