diff --git a/src/add_chunk.py b/src/add_chunk.py index c167d5a..d5ea5bc 100644 --- a/src/add_chunk.py +++ b/src/add_chunk.py @@ -1,20 +1,50 @@ from tkinter import Tk, StringVar, Label, OptionMenu, Button, filedialog from ragflow_sdk import RAGFlow -api_key = "ragflow-I5ZDNjMWNhNTdlMjExZjBiOTEwMzI0ZT" -base_url = "http://192.168.107.165:8099" +base_url = "http://127.0.0.1" +api_key = "ragflow-MyMjM2ODE2NThlMTExZjBiMzJlNzY5Mj" rag_object = RAGFlow(api_key=api_key, base_url=base_url) + def add_chunk_to_document(): - dataset_id = dataset_var.get() - document_id = document_var.get() - file_path = filedialog.askopenfilename(filetypes=[("Text files", "*.txt")]) - - if file_path: - with open(file_path, 'r') as file: + try: + dataset_name = dataset_var.get() + document_name = document_var.get() + if not dataset_name or not document_name: + print("请选择数据集和文档!") + return + + file_path = filedialog.askopenfilename(filetypes=[("Text files", "*.txt")]) + if not file_path: + print("未选择文件") + return + + with open(file_path, 'r', encoding='utf-8') as file: content = file.read() - rag_object.add_chunk(dataset_id, document_id, content) + + datasets = rag_object.list_datasets(name=dataset_name) + if not datasets: + print(f"数据集 '{dataset_name}' 不存在!") + return + dataset = datasets[0] + + documents = dataset.list_documents(name=document_name) + if not documents: + print(f"文档 '{document_name}' 不存在!") + return + document = documents[0] + + content="test chunk" + chunk = document.add_chunk(content=content) + print(f"Chunk添加成功! ID: {chunk.id}") + + except Exception as e: + print(f"添加chunk时发生错误: {e}") + import traceback + traceback.print_exc() + + def update_documents(*args): dataset_name = dataset_var.get() @@ -28,11 +58,15 @@ def update_documents(*args): root = Tk() root.title("Add Chunk to Document") -dataset_var = StringVar(root) -document_var = StringVar(root) - +# 初始化变量时设置默认值 datasets = rag_object.list_datasets() -dataset_menu = OptionMenu(root, dataset_var, *[ds.name for ds in datasets], command=update_documents) +dataset_names = [ds.name for ds in datasets] +dataset_var = StringVar(root) +dataset_var.set(dataset_names[0] if dataset_names else "") +document_var = StringVar(root) +document_var.set("") + +dataset_menu = OptionMenu(root, dataset_var, *dataset_names, command=update_documents) dataset_menu.pack()