327 lines
11 KiB
Python
327 lines
11 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
将 vector_chunks.json 向量化并上传到 Milvus 和 PostgreSQL
|
||
使用中转站的 OpenAI 兼容 API
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import time
|
||
from pathlib import Path
|
||
from typing import List, Dict
|
||
|
||
import psycopg2
|
||
from psycopg2.extras import execute_values
|
||
from pymilvus import (
|
||
connections,
|
||
Collection,
|
||
FieldSchema,
|
||
CollectionSchema,
|
||
DataType,
|
||
utility,
|
||
)
|
||
from openai import OpenAI
|
||
|
||
# ===================== 配置 =====================
|
||
# 中转站配置
|
||
RELAY_BASE_URL = "http://6.86.80.4:30080/v1"
|
||
RELAY_API_KEY = "sk-5HeY7gfSIlyZMacfuXOf5cphpymsNqufEu1ou4U3avbULcyY"
|
||
EMBEDDING_MODEL = "text-embedding-v3" # 中转站支持的 embedding 模型
|
||
|
||
# Milvus 配置
|
||
MILVUS_HOST = "localhost"
|
||
MILVUS_PORT = "19530"
|
||
COLLECTION_NAME = "regulation_chunks"
|
||
|
||
# PostgreSQL 配置
|
||
PG_HOST = "6.86.80.10"
|
||
PG_PORT = 5432
|
||
PG_USER = "postgresql"
|
||
PG_PASSWORD = "postgresql123456"
|
||
PG_DATABASE = "postgres"
|
||
|
||
|
||
# ===================== Embedding =====================
|
||
def get_openai_client(api_key: str, base_url: str) -> OpenAI:
|
||
"""创建 OpenAI 客户端连接到中转站"""
|
||
return OpenAI(api_key=api_key, base_url=base_url)
|
||
|
||
|
||
def get_embeddings_batch(client: OpenAI, texts: List[str], batch_size: int = 10) -> List[List[float]]:
|
||
"""批量获取文本向量"""
|
||
all_embeddings = []
|
||
|
||
for i in range(0, len(texts), batch_size):
|
||
batch = texts[i:i + batch_size]
|
||
print(f"Embedding batch {i // batch_size + 1}/{(len(texts) - 1) // batch_size + 1}...")
|
||
|
||
response = client.embeddings.create(
|
||
model=EMBEDDING_MODEL,
|
||
input=batch,
|
||
)
|
||
|
||
embeddings = [item.embedding for item in response.data]
|
||
all_embeddings.extend(embeddings)
|
||
|
||
return all_embeddings
|
||
|
||
|
||
# ===================== Milvus =====================
|
||
def init_milvus(host: str, port: str):
|
||
connections.connect("default", host=host, port=port)
|
||
print(f"已连接 Milvus: {host}:{port}")
|
||
|
||
|
||
def create_collection(name: str, dim: int) -> Collection:
|
||
"""创建或获取 collection"""
|
||
if utility.has_collection(name):
|
||
print(f"Collection '{name}' 已存在，删除重建")
|
||
utility.drop_collection(name)
|
||
|
||
fields = [
|
||
FieldSchema(name="chunk_id", dtype=DataType.VARCHAR, max_length=64, is_primary=True),
|
||
FieldSchema(name="doc_id", dtype=DataType.VARCHAR, max_length=128),
|
||
FieldSchema(name="doc_title", dtype=DataType.VARCHAR, max_length=512),
|
||
FieldSchema(name="chunk_index", dtype=DataType.INT64),
|
||
FieldSchema(name="semantic_id", dtype=DataType.VARCHAR, max_length=64),
|
||
FieldSchema(name="chunk_type", dtype=DataType.VARCHAR, max_length=32),
|
||
FieldSchema(name="page_start", dtype=DataType.INT64),
|
||
FieldSchema(name="page_end", dtype=DataType.INT64),
|
||
FieldSchema(name="section_title", dtype=DataType.VARCHAR, max_length=512),
|
||
FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=2048),
|
||
FieldSchema(name="source_ids", dtype=DataType.VARCHAR, max_length=4096), # JSON 字符串
|
||
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
||
]
|
||
|
||
schema = CollectionSchema(fields, description="法规文档检索 chunks")
|
||
collection = Collection(name, schema)
|
||
|
||
# 创建向量索引（IVF_FLAT，适合中小规模）
|
||
index_params = {
|
||
"metric_type": "COSINE",
|
||
"index_type": "IVF_FLAT",
|
||
"params": {"nlist": 128},
|
||
}
|
||
collection.create_index("embedding", index_params)
|
||
print(f"Collection '{name}' 创建完成，索引已建立")
|
||
|
||
return collection
|
||
|
||
|
||
def insert_chunks(collection: Collection, chunks: List[Dict], embeddings: List[List[float]]):
|
||
"""插入 chunks 到 Milvus"""
|
||
data = [
|
||
[c["chunk_id"] for c in chunks],
|
||
[c["doc_id"] for c in chunks],
|
||
[c["doc_title"] for c in chunks],
|
||
[c["chunk_index"] for c in chunks],
|
||
[c["semantic_id"] for c in chunks],
|
||
[c["chunk_type"] for c in chunks],
|
||
[c["page_start"] for c in chunks],
|
||
[c["page_end"] for c in chunks],
|
||
[c["section_title"] for c in chunks],
|
||
[c["text"] for c in chunks],
|
||
[json.dumps(c.get("source_ids", [])) for c in chunks], # JSON 字符串
|
||
embeddings,
|
||
]
|
||
|
||
collection.insert(data)
|
||
collection.flush()
|
||
print(f"已插入 {len(chunks)} 个 chunks")
|
||
|
||
|
||
def load_collection(collection: Collection):
|
||
"""加载 collection 到内存（搜索前必须）"""
|
||
collection.load()
|
||
print(f"Collection 已加载到内存")
|
||
|
||
|
||
# ===================== PostgreSQL =====================
|
||
def get_pg_connection(host: str, port: int, user: str, password: str, database: str):
|
||
"""获取 PostgreSQL 连接"""
|
||
conn = psycopg2.connect(
|
||
host=host,
|
||
port=port,
|
||
user=user,
|
||
password=password,
|
||
database=database,
|
||
)
|
||
print(f"已连接 PostgreSQL: {host}:{port}/{database}")
|
||
return conn
|
||
|
||
|
||
def insert_chunks_to_pg(conn, chunks: List[Dict], doc_data: Dict):
|
||
"""插入 chunks 和相关数据到 PostgreSQL"""
|
||
cursor = conn.cursor()
|
||
|
||
try:
|
||
# 1. 插入文档
|
||
cursor.execute("""
|
||
INSERT INTO documents (doc_id, title, standard_number, upload_time)
|
||
VALUES (%s, %s, %s, NOW())
|
||
ON CONFLICT (doc_id) DO UPDATE SET title = EXCLUDED.title, updated_at = NOW()
|
||
""", (doc_data["doc_id"], doc_data["doc_title"], doc_data.get("standard_number")))
|
||
|
||
# 2. 插入语义块
|
||
semantic_blocks = doc_data.get("semantic_blocks", [])
|
||
if semantic_blocks:
|
||
block_rows = [
|
||
(
|
||
doc_data["doc_id"],
|
||
block["semantic_id"],
|
||
block["block_type"],
|
||
block["page_start"],
|
||
block["page_end"],
|
||
block.get("section_title"),
|
||
block.get("section_level"),
|
||
json.dumps(block.get("source_ids", [])),
|
||
block["text"],
|
||
)
|
||
for block in semantic_blocks
|
||
]
|
||
execute_values(
|
||
cursor,
|
||
"""
|
||
INSERT INTO semantic_blocks
|
||
(doc_id, semantic_id, block_type, page_start, page_end, section_title, section_level, source_ids, text)
|
||
VALUES %s
|
||
ON CONFLICT (doc_id, semantic_id) DO UPDATE SET text = EXCLUDED.text
|
||
""",
|
||
block_rows,
|
||
)
|
||
print(f"已插入 {len(semantic_blocks)} 个语义块")
|
||
|
||
# 3. 插入向量块元数据
|
||
chunk_rows = [
|
||
(
|
||
doc_data["doc_id"],
|
||
chunk["chunk_id"],
|
||
chunk["semantic_id"],
|
||
chunk["chunk_index"],
|
||
chunk.get("piece_index"),
|
||
chunk["page_start"],
|
||
chunk["page_end"],
|
||
chunk.get("section_title"),
|
||
chunk["text"],
|
||
json.dumps(chunk.get("source_ids", [])),
|
||
)
|
||
for chunk in chunks
|
||
]
|
||
execute_values(
|
||
cursor,
|
||
"""
|
||
INSERT INTO vector_chunks
|
||
(doc_id, chunk_id, semantic_id, chunk_index, piece_index, page_start, page_end, section_title, text, source_ids)
|
||
VALUES %s
|
||
ON CONFLICT (doc_id, chunk_id) DO UPDATE SET text = EXCLUDED.text
|
||
""",
|
||
chunk_rows,
|
||
)
|
||
print(f"已插入 {len(chunks)} 个向量块元数据")
|
||
|
||
conn.commit()
|
||
print("PostgreSQL 数据插入完成")
|
||
|
||
except Exception as e:
|
||
conn.rollback()
|
||
raise e
|
||
finally:
|
||
cursor.close()
|
||
|
||
|
||
# ===================== 主流程 =====================
|
||
def load_data(file_path: Path) -> Dict:
|
||
"""加载 vector_chunks.json，返回完整数据"""
|
||
data = json.loads(file_path.read_text(encoding="utf-8"))
|
||
return data
|
||
|
||
|
||
def upload_to_milvus_and_pg(
|
||
chunks_file: str,
|
||
api_key: str,
|
||
base_url: str,
|
||
milvus_host: str,
|
||
milvus_port: str,
|
||
collection_name: str,
|
||
batch_size: int,
|
||
pg_host: str,
|
||
pg_port: int,
|
||
pg_user: str,
|
||
pg_password: str,
|
||
pg_database: str,
|
||
):
|
||
# 1. 加载完整数据
|
||
chunks_path = Path(chunks_file).expanduser().resolve()
|
||
if not chunks_path.exists():
|
||
raise FileNotFoundError(f"文件不存在: {chunks_path}")
|
||
|
||
data = load_data(chunks_path)
|
||
chunks = data.get("vector_chunks", [])
|
||
if not chunks:
|
||
raise ValueError("vector_chunks 为空")
|
||
print(f"加载 {len(chunks)} 个 chunks")
|
||
|
||
# 2. 初始化连接
|
||
client = get_openai_client(api_key, base_url)
|
||
init_milvus(milvus_host, milvus_port)
|
||
pg_conn = get_pg_connection(pg_host, pg_port, pg_user, pg_password, pg_database)
|
||
|
||
# 3. 获取 embeddings
|
||
texts = [c["embedding_text"] for c in chunks]
|
||
embeddings = get_embeddings_batch(client, texts, batch_size)
|
||
print(f"生成 {len(embeddings)} 个向量")
|
||
|
||
# 4. 获取 embedding 维度
|
||
embedding_dim = len(embeddings[0])
|
||
print(f"Embedding 维度: {embedding_dim}")
|
||
|
||
# 5. 创建 collection 并插入 Milvus
|
||
collection = create_collection(collection_name, embedding_dim)
|
||
insert_chunks(collection, chunks, embeddings)
|
||
load_collection(collection)
|
||
|
||
# 6. 插入 PostgreSQL
|
||
insert_chunks_to_pg(pg_conn, chunks, data)
|
||
|
||
# 7. 关闭连接
|
||
pg_conn.close()
|
||
|
||
print("上传完成！")
|
||
|
||
|
||
# ===================== CLI =====================
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="将 vector_chunks 向量化并上传到 Milvus 和 PostgreSQL")
|
||
parser.add_argument("chunks_file", help="vector_chunks.json 文件路径")
|
||
parser.add_argument("--api-key", default=RELAY_API_KEY, help="中转站 API Key")
|
||
parser.add_argument("--base-url", default=RELAY_BASE_URL, help="中转站 Base URL")
|
||
parser.add_argument("--milvus-host", default=MILVUS_HOST, help="Milvus host")
|
||
parser.add_argument("--milvus-port", default=MILVUS_PORT, help="Milvus port")
|
||
parser.add_argument("--collection", default=COLLECTION_NAME, help="Milvus collection 名称")
|
||
parser.add_argument("--batch-size", type=int, default=10, help="Embedding 批量大小（中转站限制最大10）")
|
||
parser.add_argument("--pg-host", default=PG_HOST, help="PostgreSQL host")
|
||
parser.add_argument("--pg-port", type=int, default=PG_PORT, help="PostgreSQL port")
|
||
parser.add_argument("--pg-user", default=PG_USER, help="PostgreSQL user")
|
||
parser.add_argument("--pg-password", default=PG_PASSWORD, help="PostgreSQL password")
|
||
parser.add_argument("--pg-database", default=PG_DATABASE, help="PostgreSQL database")
|
||
args = parser.parse_args()
|
||
|
||
upload_to_milvus_and_pg(
|
||
chunks_file=args.chunks_file,
|
||
api_key=args.api_key,
|
||
base_url=args.base_url,
|
||
milvus_host=args.milvus_host,
|
||
milvus_port=args.milvus_port,
|
||
collection_name=args.collection,
|
||
batch_size=args.batch_size,
|
||
pg_host=args.pg_host,
|
||
pg_port=args.pg_port,
|
||
pg_user=args.pg_user,
|
||
pg_password=args.pg_password,
|
||
pg_database=args.pg_database,
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |