from WxMini.Milvus.Config.MulvusConfig import *
from WxMini.Milvus.Utils.MilvusCollectionManager import MilvusCollectionManager
from WxMini.Milvus.Utils.MilvusConnectionPool import *
from gensim.models import KeyedVectors
import jieba
import time

# 1. 加载预训练的 Word2Vec 模型
model_path = MS_MODEL_PATH
model = KeyedVectors.load_word2vec_format(model_path, binary=False, limit=MS_MODEL_LIMIT)
print(f"模型加载成功，词向量维度: {model.vector_size}")

# 功能：将文本转换为嵌入向量
def text_to_embedding(text):
    words = jieba.lcut(text)  # 使用 jieba 分词
    print(f"文本: {text}, 分词结果: {words}")
    embeddings = [model[word] for word in words if word in model]
    print(f"有效词向量数量: {len(embeddings)}")
    if embeddings:
        avg_embedding = sum(embeddings) / len(embeddings)
        print(f"生成的平均向量: {avg_embedding[:5]}...")  # 打印前 5 维
        return avg_embedding
    else:
        print("未找到有效词，返回零向量")
        return [0.0] * model.vector_size

# 2. 使用连接池管理 Milvus 连接
milvus_pool = MilvusConnectionPool(host=MS_HOST, port=MS_PORT, max_connections=MS_MAX_CONNECTIONS)

# 3. 从连接池中获取一个连接
connection = milvus_pool.get_connection()

# 4. 初始化集合管理器
collection_name = MS_COLLECTION_NAME
collection_manager = MilvusCollectionManager(collection_name)

# 5. 输入一个用户问题
user_input = input("请输入一句话：")  # 例如：“我今天心情不太好”
model_response = "我没听懂，能再说一遍吗？"  # 大模型的固定回复
timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())  # 当前时间
session_id = "session_001"  # 会话 ID（可以根据需要动态生成）

# 6. 将用户问题转换为嵌入向量
user_embedding = text_to_embedding(user_input)

# 7. 插入数据，确保字段顺序与集合定义一致
entities = [
    [session_id],  # session_id
    [user_input],  # user_input
    [model_response],  # model_response
    [timestamp],  # timestamp
    [user_embedding]  # embedding
]
collection_manager.insert_data(entities)
print("数据插入成功。")

# 8. 释放连接
milvus_pool.release_connection(connection)

# 9. 关闭连接池
milvus_pool.close()