100 lines
3.0 KiB
Python
100 lines
3.0 KiB
Python
"""NL2SQL 元数据召回与授权过滤。"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Awaitable, Callable, Iterable
|
|
from typing import Any
|
|
|
|
from nl2sql.embedding import embed_texts
|
|
from nl2sql.milvus_collections import NL2SQL_COLLECTION
|
|
|
|
|
|
def _flatten_hits(results: Any) -> Iterable[dict[str, Any]]:
|
|
"""兼容 Milvus 返回的批次列表和单条字典结构。"""
|
|
for batch in results or []:
|
|
if isinstance(batch, dict):
|
|
yield batch
|
|
else:
|
|
yield from batch or []
|
|
|
|
|
|
def filter_authorized_metadata(
|
|
tables: list[dict[str, Any]],
|
|
columns: list[dict[str, Any]],
|
|
permission: dict[str, Any],
|
|
) -> dict[str, list[dict[str, Any]]]:
|
|
"""过滤失效实体、无表权限实体和无字段权限实体。"""
|
|
authorized_tables = set(permission.get("tables", set()))
|
|
authorized_columns = permission.get("columns", {})
|
|
valid_table_names = {
|
|
table.get("table_name")
|
|
for table in tables
|
|
if table.get("is_valid", True) and table.get("table_name") in authorized_tables
|
|
}
|
|
filtered_tables = [
|
|
table
|
|
for table in tables
|
|
if table.get("is_valid", True)
|
|
and table.get("table_name") in valid_table_names
|
|
]
|
|
filtered_columns = [
|
|
column
|
|
for column in columns
|
|
if column.get("is_valid", True)
|
|
and column.get("table_name") in valid_table_names
|
|
and column.get("field_name")
|
|
in set(authorized_columns.get(column.get("table_name"), set()))
|
|
]
|
|
return {"tables": filtered_tables, "columns": filtered_columns}
|
|
|
|
|
|
async def retrieve_metadata(
|
|
query: str,
|
|
milvus_client,
|
|
*,
|
|
embedder: Callable[[list[str]], Awaitable[list[list[float]]]] = embed_texts,
|
|
top_k: int = 5,
|
|
) -> list[dict[str, Any]]:
|
|
"""分别召回表级和字段级元数据，并过滤无效向量。"""
|
|
if not query or not query.strip():
|
|
return []
|
|
if top_k <= 0:
|
|
return []
|
|
|
|
vector = (await embedder([query]))[0]
|
|
output_fields = [
|
|
"id",
|
|
"chunk_type",
|
|
"table_name",
|
|
"field_name",
|
|
"text",
|
|
"is_valid",
|
|
"is_deprecated",
|
|
]
|
|
async def search_chunk(chunk_type: str):
|
|
return await milvus_client.search(
|
|
collection_name=NL2SQL_COLLECTION,
|
|
data=[vector],
|
|
limit=top_k,
|
|
filter=f'chunk_type == "{chunk_type}" and is_valid == true',
|
|
output_fields=output_fields,
|
|
)
|
|
|
|
table_hits, field_hits = await asyncio.gather(
|
|
search_chunk("table_meta"),
|
|
search_chunk("field_meta"),
|
|
)
|
|
results: list[dict[str, Any]] = []
|
|
for hits in (table_hits, field_hits):
|
|
for hit in _flatten_hits(hits):
|
|
entity = hit.get("entity") or hit
|
|
if not entity.get("is_valid", True) or entity.get("is_deprecated", False):
|
|
continue
|
|
results.append(
|
|
{
|
|
**entity,
|
|
"distance": hit.get("distance", hit.get("score")),
|
|
}
|
|
)
|
|
return results
|