feat:新增投顾agent和nl2sqlagent

nl2sql/metadata_sync.py

@@ -0,0 +1,83 @@
+"""构造并 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