app.py

import gradio as gr
import llama_cpp
import base64
from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad


def decrypt_file(input_path, key):
    # 读取加密文件
    with open(input_path, "rb") as f:
        encrypted_data = base64.b64decode(f.read())

    key = key.ljust(32, "0")[:32].encode("utf-8")
    iv = encrypted_data[:16]
    ciphertext = encrypted_data[16:]

    cipher = AES.new(key, AES.MODE_CBC, iv)
    plaintext = unpad(cipher.decrypt(ciphertext), AES.block_size)

    return plaintext.decode("utf-8")


llm = llama_cpp.Llama.from_pretrained(
    repo_id="mradermacher/bge-large-zh-v1.5-GGUF",
    filename="bge-large-zh-v1.5.Q4_K_M.gguf",
    embedding=True,
)

# embedding_1 = llm.create_embedding("Hello, world!")
# embedding_2 = llm.create_embedding("你好, 世界!") # type(embedding_1['data'][0]['embedding']) list

from pymilvus import MilvusClient

client = MilvusClient("./books.db")

client.create_collection(collection_name="collection_1", dimension=1024)

import os, json

aeskey = os.getenv("aeskey")
decrypted_content = decrypt_file("encrypted.txt", aeskey)
raw_jsons = json.loads(decrypted_content)


with open("embeddings.json", mode="r") as embedding_file:
    all_embs = json.load(embedding_file)


for vhjx_index, vhjx_item in enumerate(raw_jsons):
    chapter = vhjx_item[0]

    docs = []
    metas = []
    for jvvi_item in vhjx_item[1:]:
        content = jvvi_item["原文"]
        docs.append(content)
        metas.append(
            {
                "index": jvvi_item["index"],
                "text": content,
                "annotation": jvvi_item.get("注释", ""),
                "critique": jvvi_item.get("批判", ""),
                "chapter": chapter,
            }
        )

    # 一个章节一次
    # 批量生成 embeddings（每个为 list[float]）
    # emb_result = llm.create_embedding(docs)
    embeddings = all_embs[vhjx_index]  # List[List[float]]
    print(len(embeddings))
    # 准备数据
    milvus_data = []
    for i, emb in enumerate(embeddings):
        item = metas[i]
        milvus_data.append(
            {
                "id": vhjx_index * 100 + i,
                "index": item["index"],
                "vector": emb,
                "text": item["text"],
                "annotation": item["annotation"],
                "critique": item["critique"],
                "chapter": item["chapter"],
            }
        )
    print(f"✅ 共 {len(milvus_data)} 条数据")

    # 插入数据
    client.insert(collection_name="collection_1", data=milvus_data)
    print(f"✅ 插入完成：共 {len(milvus_data)} 条数据")


def greet(name):
    embeddings = llm.create_embedding(name)
    res = client.search(
        collection_name="collection_1",
        data=[embeddings["data"][0]["embedding"]],
        limit=5,
        output_fields=["index", "text", "annotation", "critique"],
    )
    return res


demo = gr.Interface(
    fn=greet,
    inputs=gr.Textbox(label="输入部分原文句子"),
    outputs=gr.JSON(label="查询结果"),
    title="论语批判MCP (Embedding版本)",
    description="输入模糊的论语原文，可以向量检索到对应的批判内容。",
)
demo.launch(mcp_server=True)