84 lines
2.8 KiB
Python
84 lines
2.8 KiB
Python
|
"""构造并 upsert NL2SQL 元数据行。"""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import time
|
||
|
|
from collections.abc import Awaitable, Callable
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from nl2sql.embedding import embed_texts
|
||
|
|
from nl2sql.metadata import build_metadata_chunks
|
||
|
|
from nl2sql.milvus_collections import NL2SQL_COLLECTION
|
||
|
|
|
||
|
|
|
||
|
|
def _escape_filter_value(value: str) -> str:
|
||
|
|
return value.replace("\\", "\\\\").replace('"', '\\"')
|
||
|
|
|
||
|
|
|
||
|
|
async def prepare_metadata_rows(
|
||
|
|
chunks: list[dict[str, Any]],
|
||
|
|
*,
|
||
|
|
embedder: Callable[[list[str]], Awaitable[list[list[float]]]] = embed_texts,
|
||
|
|
timestamp: int | None = None,
|
||
|
|
) -> list[dict[str, Any]]:
|
||
|
|
if not chunks:
|
||
|
|
return []
|
||
|
|
vectors = await embedder([chunk["text"] for chunk in chunks])
|
||
|
|
if len(vectors) != len(chunks):
|
||
|
|
raise ValueError("Embedding result count must match metadata chunk count")
|
||
|
|
created_at = int(time.time()) if timestamp is None else timestamp
|
||
|
|
return [
|
||
|
|
{**chunk, "vector": vector, "created_at": created_at}
|
||
|
|
for chunk, vector in zip(chunks, vectors)
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
async def sync_metadata(
|
||
|
|
milvus_client,
|
||
|
|
table_rows: list[dict[str, Any]],
|
||
|
|
column_rows: list[dict[str, Any]],
|
||
|
|
*,
|
||
|
|
embedder: Callable[[list[str]], Awaitable[list[list[float]]]] = embed_texts,
|
||
|
|
timestamp: int | None = None,
|
||
|
|
) -> int:
|
||
|
|
chunks = build_metadata_chunks(table_rows, column_rows)
|
||
|
|
chunks_by_table: dict[str, list[dict[str, Any]]] = {}
|
||
|
|
for chunk in chunks:
|
||
|
|
chunks_by_table.setdefault(chunk["table_name"], []).append(chunk)
|
||
|
|
|
||
|
|
updated_count = 0
|
||
|
|
for table_name, table_chunks in chunks_by_table.items():
|
||
|
|
escaped_name = _escape_filter_value(table_name)
|
||
|
|
existing_rows = await milvus_client.query(
|
||
|
|
collection_name=NL2SQL_COLLECTION,
|
||
|
|
filter=f'table_name == "{escaped_name}"',
|
||
|
|
output_fields=["id", "content_hash", "is_valid", "is_deprecated"],
|
||
|
|
)
|
||
|
|
desired_signature = {
|
||
|
|
(chunk["id"], chunk["content_hash"], chunk["is_valid"], chunk["is_deprecated"])
|
||
|
|
for chunk in table_chunks
|
||
|
|
}
|
||
|
|
existing_signature = {
|
||
|
|
(
|
||
|
|
row.get("id"),
|
||
|
|
row.get("content_hash"),
|
||
|
|
row.get("is_valid", True),
|
||
|
|
row.get("is_deprecated", False),
|
||
|
|
)
|
||
|
|
for row in existing_rows or []
|
||
|
|
}
|
||
|
|
if desired_signature == existing_signature:
|
||
|
|
continue
|
||
|
|
|
||
|
|
if existing_rows:
|
||
|
|
await milvus_client.delete(
|
||
|
|
collection_name=NL2SQL_COLLECTION,
|
||
|
|
filter=f'table_name == "{escaped_name}"',
|
||
|
|
)
|
||
|
|
rows = await prepare_metadata_rows(
|
||
|
|
table_chunks, embedder=embedder, timestamp=timestamp
|
||
|
|
)
|
||
|
|
if rows:
|
||
|
|
await milvus_client.upsert(collection_name=NL2SQL_COLLECTION, data=rows)
|
||
|
|
updated_count += len(rows)
|
||
|
|
return updated_count
|