menu_app.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from ragflow_client import RAGFlowClient
import os
import time
import sys

class RAGFlowMenuApp:
    def __init__(self):
        """初始化菜单应用"""
        print("正在连接RAGFlow服务...")
        self.client = RAGFlowClient(
            api_key="ragflow-Q4ZWNiMjE4MWE4NTExZjBiODlmMzZiNj",
            base_url="http://localhost:8080"
        )
        self.current_dataset = None
        self.current_chat = None
        self.chat_session = None
        print("连接成功！")

    def show_main_menu(self):
        """显示主菜单"""
        print("\n" + "="*50)
        print("RAGFlow 客户端菜单")
        print("="*50)
        print("1. 知识库管理")
        print("2. 聊天助手管理")
        print("3. 检索与对话")
        print("4. 系统设置")
        print("0. 退出程序")
        print("="*50)
        choice = input("请选择功能 (0-4): ")
        return choice

    def show_dataset_menu(self):
        """显示知识库管理菜单"""
        print("\n" + "="*50)
        print("知识库管理")
        print("="*50)
        print("1. 创建新知识库")
        print("2. 列出所有知识库")
        print("3. 上传文档到知识库")
        print("4. 解析知识库文档")
        print("5. 删除知识库")
        print("0. 返回主菜单")
        print("="*50)
        if self.current_dataset:
            print(f"当前选择的知识库: {self.current_dataset.name}")
        print("="*50)
        choice = input("请选择功能 (0-5): ")
        return choice

    def show_chat_menu(self):
        """显示聊天助手管理菜单"""
        print("\n" + "="*50)
        print("聊天助手管理")
        print("="*50)
        print("1. 创建新聊天助手")
        print("2. 列出所有聊天助手")
        print("3. 开始对话")
        print("4. 删除聊天助手")
        print("0. 返回主菜单")
        print("="*50)
        if self.current_chat:
            print(f"当前选择的助手: {self.current_chat.name}")
        print("="*50)
        choice = input("请选择功能 (0-4): ")
        return choice

    def show_retrieve_menu(self):
        """显示检索与对话菜单"""
        print("\n" + "="*50)
        print("检索与对话")
        print("="*50)
        print("1. 从知识库检索内容")
        print("2. 开始聊天会话")
        print("0. 返回主菜单")
        print("="*50)
        if self.current_dataset:
            print(f"当前知识库: {self.current_dataset.name}")
        if self.current_chat:
            print(f"当前聊天助手: {self.current_chat.name}")
        print("="*50)
        choice = input("请选择功能 (0-2): ")
        return choice

    def show_settings_menu(self):
        """显示系统设置菜单"""
        print("\n" + "="*50)
        print("系统设置")
        print("="*50)
        print("1. 修改API密钥")
        print("2. 修改服务器地址")
        print("0. 返回主菜单")
        print("="*50)
        choice = input("请选择功能 (0-2): ")
        return choice

    def create_dataset(self):
        """创建新知识库"""
        name = input("请输入知识库名称: ")
        description = input("请输入知识库描述 (可选): ")
        
        print("可用的嵌入模型:")
        print("1. BAAI/bge-large-zh-v1.5 (默认)")
        print("2. BAAI/bge-base-zh-v1.5")
        print("3. 其他...")
        model_choice = input("请选择嵌入模型 (默认使用1): ")
        
        embedding_model = "BAAI/bge-large-zh-v1.5"
        if model_choice == "2":
            embedding_model = "BAAI/bge-base-zh-v1.5"
        elif model_choice == "3":
            embedding_model = input("请输入嵌入模型名称: ")
        
        print("创建知识库中...")
        dataset = self.client.create_dataset(
            name=name,
            description=description,
            embedding_model=embedding_model
        )
        
        if dataset:
            self.current_dataset = dataset
            input("\n知识库创建成功！按Enter键继续...")
        else:
            input("\n知识库创建失败！按Enter键继续...")
    
    def list_datasets(self):
        """列出所有知识库"""
        print("\n获取知识库列表中...")
        datasets = self.client.list_datasets()
        
        if not datasets:
            input("没有找到知识库！按Enter键继续...")
            return
        
        print("\n知识库列表:")
        print("-"*50)
        print(f"{'序号':<6}{'名称':<20}{'ID':<25}{'创建时间':<20}")
        print("-"*50)
        
        for i, ds in enumerate(datasets):
            print(f"{i+1:<6}{ds.name:<20}{ds.id:<25}{getattr(ds, 'create_time', 'N/A'):<20}")
        
        print("-"*50)
        choice = input("选择一个知识库的序号 (或直接按Enter返回): ")
        
        if choice.isdigit() and 1 <= int(choice) <= len(datasets):
            self.current_dataset = datasets[int(choice)-1]
            print(f"\n已选择知识库: {self.current_dataset.name}")
        
        input("按Enter键继续...")
    
    def upload_document(self):
        """上传文档到知识库"""
        if not self.current_dataset:
            input("\n请先选择一个知识库！按Enter键继续...")
            return
        
        file_path = input("\n请输入文件路径: ")
        
        if not os.path.exists(file_path):
            input(f"文件不存在: {file_path}！按Enter键继续...")
            return
        
        print(f"\n上传文件 {file_path} 到知识库 {self.current_dataset.name} 中...")
        success = self.client.upload_document(self.current_dataset, file_path)
        
        if success:
            input("\n文件上传成功！按Enter键继续...")
        else:
            input("\n文件上传失败！按Enter键继续...")
    
    def parse_documents(self):
        """解析知识库文档"""
        if not self.current_dataset:
            input("\n请先选择一个知识库！按Enter键继续...")
            return
        
        print("\n获取文档列表中...")
        documents = self.current_dataset.list_documents()
        
        if not documents:
            input("知识库中没有找到文档！按Enter键继续...")
            return
        
        print("\n文档列表:")
        print("-"*70)
        print(f"{'序号':<6}{'名称':<30}{'ID':<25}{'状态':<10}")
        print("-"*70)
        
        unparsed_docs = []
        for i, doc in enumerate(documents):
            print(f"{i+1:<6}{doc.name:<30}{doc.id:<25}{doc.run:<10}")
            if doc.run not in ["DONE", "FAIL"]:
                unparsed_docs.append(doc)
        
        print("-"*70)
        
        if not unparsed_docs:
            input("所有文档已解析完成！按Enter键继续...")
            return
        
        choice = input("是否解析所有未完成的文档？(y/n): ")
        
        if choice.lower() == 'y':
            document_ids = [doc.id for doc in unparsed_docs]
            print(f"\n开始解析 {len(document_ids)} 个文档...")
            self.client.parse_documents(self.current_dataset, document_ids)
            
            # 等待解析
            wait = input("是否等待解析完成？(y/n): ")
            if wait.lower() == 'y':
                print("\n等待文档解析完成...")
                max_wait_time = 300  # 最多等待5分钟
                start_time = time.time()
                
                while True:
                    time.sleep(5)
                    # 重新获取文档状态
                    docs = self.current_dataset.list_documents()
                    all_done = all(doc.run == "DONE" or doc.run == "FAIL" for doc in docs)
                    
                    if all_done:
                        input("\n所有文档解析完成！按Enter键继续...")
                        break
                        
                    if time.time() - start_time > max_wait_time:
                        input("\n等待超时！按Enter键继续...")
                        break
                        
                    print("文档仍在解析中，继续等待...")
            else:
                input("\n解析任务已启动，将在后台运行。按Enter键继续...")
        else:
            input("\n操作已取消。按Enter键继续...")
    
    def delete_dataset(self):
        """删除知识库"""
        if not self.current_dataset:
            print("\n获取知识库列表中...")
            datasets = self.client.list_datasets()
            
            if not datasets:
                input("没有找到知识库！按Enter键继续...")
                return
            
            print("\n知识库列表:")
            print("-"*50)
            print(f"{'序号':<6}{'名称':<20}{'ID':<25}")
            print("-"*50)
            
            for i, ds in enumerate(datasets):
                print(f"{i+1:<6}{ds.name:<20}{ds.id:<25}")
            
            print("-"*50)
            choice = input("选择要删除的知识库序号 (或直接按Enter取消): ")
            
            if not choice.isdigit() or int(choice) < 1 or int(choice) > len(datasets):
                input("\n操作已取消。按Enter键继续...")
                return
                
            dataset_to_delete = datasets[int(choice)-1]
        else:
            dataset_to_delete = self.current_dataset
        
        confirm = input(f"\n确定要删除知识库 '{dataset_to_delete.name}' 吗？此操作不可恢复！(y/n): ")
        
        if confirm.lower() != 'y':
            input("\n操作已取消。按Enter键继续...")
            return
        
        print(f"\n正在删除知识库 {dataset_to_delete.name}...")
        self.client.delete_datasets(ids=[dataset_to_delete.id])
        
        if self.current_dataset and self.current_dataset.id == dataset_to_delete.id:
            self.current_dataset = None
        
        input("\n知识库已删除！按Enter键继续...")
    
    def create_chat(self):
        """创建新聊天助手"""
        if not self.current_dataset:
            print("\n获取知识库列表中...")
            datasets = self.client.list_datasets()
            
            if not datasets:
                input("没有找到知识库！请先创建知识库。按Enter键继续...")
                return
            
            print("\n知识库列表:")
            print("-"*50)
            print(f"{'序号':<6}{'名称':<20}{'ID':<25}")
            print("-"*50)
            
            for i, ds in enumerate(datasets):
                print(f"{i+1:<6}{ds.name:<20}{ds.id:<25}")
            
            print("-"*50)
            choice = input("选择要关联的知识库序号 (或直接按Enter取消): ")
            
            if not choice.isdigit() or int(choice) < 1 or int(choice) > len(datasets):
                input("\n操作已取消。按Enter键继续...")
                return
                
            self.current_dataset = datasets[int(choice)-1]
        
        name = input("\n请输入聊天助手名称: ")
        
        print("\n创建聊天助手中...")
        chat = self.client.create_chat(
            name=name,
            dataset_ids=[self.current_dataset.id]
        )
        
        if chat:
            self.current_chat = chat
            input("\n聊天助手创建成功！按Enter键继续...")
        else:
            input("\n聊天助手创建失败！按Enter键继续...")
    
    def list_chats(self):
        """列出所有聊天助手"""
        print("\n获取聊天助手列表中...")
        chats = self.client.rag.list_chats()
        
        if not chats:
            input("没有找到聊天助手！按Enter键继续...")
            return
        
        print("\n聊天助手列表:")
        print("-"*60)
        print(f"{'序号':<6}{'名称':<20}{'ID':<25}{'创建时间':<20}")
        print("-"*60)
        
        for i, chat in enumerate(chats):
            print(f"{i+1:<6}{chat.name:<20}{chat.id:<25}{getattr(chat, 'create_time', 'N/A'):<20}")
        
        print("-"*60)
        choice = input("选择一个聊天助手的序号 (或直接按Enter返回): ")
        
        if choice.isdigit() and 1 <= int(choice) <= len(chats):
            self.current_chat = chats[int(choice)-1]
            print(f"\n已选择聊天助手: {self.current_chat.name}")
        
        input("按Enter键继续...")
    
    def start_chat(self):
        """开始聊天会话"""
        if not self.current_chat:
            self.list_chats()
            if not self.current_chat:
                return
        
        print(f"\n开始与 {self.current_chat.name} 聊天，输入'退出'结束对话")
        session = self.current_chat.create_session("新会话")
        self.chat_session = session
        
        while True:
            try:
                question = input("\n用户: ")
                if question.lower() in ['退出', 'quit', 'exit']:
                    break

                print("\n助手: ", end='')
                response = ""
                for chunk in session.ask(question, stream=True):
                    new_content = chunk.content[len(response):]
                    print(new_content, end='', flush=True)
                    response = chunk.content
                print()
            except KeyboardInterrupt:
                print("\n用户中断对话")
                break
            except Exception as e:
                print(f"\n聊天过程中出错: {e}")
                continue
        
        input("\n聊天会话已结束。按Enter键继续...")
    
    def delete_chat(self):
        """删除聊天助手"""
        if not self.current_chat:
            print("\n获取聊天助手列表中...")
            chats = self.client.rag.list_chats()
            
            if not chats:
                input("没有找到聊天助手！按Enter键继续...")
                return
            
            print("\n聊天助手列表:")
            print("-"*50)
            print(f"{'序号':<6}{'名称':<20}{'ID':<25}")
            print("-"*50)
            
            for i, chat in enumerate(chats):
                print(f"{i+1:<6}{chat.name:<20}{chat.id:<25}")
            
            print("-"*50)
            choice = input("选择要删除的聊天助手序号 (或直接按Enter取消): ")
            
            if not choice.isdigit() or int(choice) < 1 or int(choice) > len(chats):
                input("\n操作已取消。按Enter键继续...")
                return
                
            chat_to_delete = chats[int(choice)-1]
        else:
            chat_to_delete = self.current_chat
        
        confirm = input(f"\n确定要删除聊天助手 '{chat_to_delete.name}' 吗？此操作不可恢复！(y/n): ")
        
        if confirm.lower() != 'y':
            input("\n操作已取消。按Enter键继续...")
            return
        
        print(f"\n正在删除聊天助手 {chat_to_delete.name}...")
        self.client.rag.delete_chats(ids=[chat_to_delete.id])
        
        if self.current_chat and self.current_chat.id == chat_to_delete.id:
            self.current_chat = None
        
        input("\n聊天助手已删除！按Enter键继续...")
    
    def retrieve_from_dataset(self):
        """从知识库检索内容"""
        if not self.current_dataset:
            input("\n请先选择一个知识库！按Enter键继续...")
            return
        
        question = input("\n请输入查询问题: ")
        similarity_threshold = input("请输入相似度阈值 (0-1，默认0.2): ")
        
        try:
            similarity_threshold = float(similarity_threshold) if similarity_threshold else 0.2
        except ValueError:
            similarity_threshold = 0.2
        
        print("\n正在检索内容...")
        chunks = self.client.retrieve(
            question=question,
            dataset_ids=[self.current_dataset.id],
            similarity_threshold=similarity_threshold
        )
        
        if not chunks:
            input("\n没有找到相关内容！按Enter键继续...")
            return
        
        print(f"\n检索到{len(chunks)}个相关片段：")
        for i, chunk in enumerate(chunks):
            print(f"\n片段 {i+1}:")
            print(f"内容: {chunk.content}")
            print(f"文档: {chunk.document_name}")
            if hasattr(chunk, 'similarity'):
                print(f"相似度: {chunk.similarity}")
        
        input("\n检索完成。按Enter键继续...")
    
    def change_api_key(self):
        """修改API密钥"""
        new_key = input("\n请输入新的API密钥: ")
        if not new_key:
            input("\n操作已取消。按Enter键继续...")
            return
        
        confirm = input("\n确定要修改API密钥吗？(y/n): ")
        if confirm.lower() != 'y':
            input("\n操作已取消。按Enter键继续...")
            return
        
        self.client = RAGFlowClient(
            api_key=new_key,
            base_url=self.client.rag.base_url
        )
        
        # 重置状态
        self.current_dataset = None
        self.current_chat = None
        self.chat_session = None
        
        input("\nAPI密钥已更新！按Enter键继续...")
    
    def change_server_url(self):
        """修改服务器地址"""
        new_url = input("\n请输入新的服务器地址 (例如: http://localhost:8080): ")
        if not new_url:
            input("\n操作已取消。按Enter键继续...")
            return
        
        confirm = input("\n确定要修改服务器地址吗？(y/n): ")
        if confirm.lower() != 'y':
            input("\n操作已取消。按Enter键继续...")
            return
        
        self.client = RAGFlowClient(
            api_key=self.client.rag.api_key,
            base_url=new_url
        )
        
        # 重置状态
        self.current_dataset = None
        self.current_chat = None
        self.chat_session = None
        
        input("\n服务器地址已更新！按Enter键继续...")
    
    def run(self):
        """运行菜单应用"""
        while True:
            choice = self.show_main_menu()
            
            if choice == '0':
                print("\n程序已退出，感谢使用！")
                break
            
            elif choice == '1':
                # 知识库管理
                while True:
                    dataset_choice = self.show_dataset_menu()
                    
                    if dataset_choice == '0':
                        break
                    elif dataset_choice == '1':
                        self.create_dataset()
                    elif dataset_choice == '2':
                        self.list_datasets()
                    elif dataset_choice == '3':
                        self.upload_document()
                    elif dataset_choice == '4':
                        self.parse_documents()
                    elif dataset_choice == '5':
                        self.delete_dataset()
            
            elif choice == '2':
                # 聊天助手管理
                while True:
                    chat_choice = self.show_chat_menu()
                    
                    if chat_choice == '0':
                        break
                    elif chat_choice == '1':
                        self.create_chat()
                    elif chat_choice == '2':
                        self.list_chats()
                    elif chat_choice == '3':
                        self.start_chat()
                    elif chat_choice == '4':
                        self.delete_chat()
            
            elif choice == '3':
                # 检索与对话
                while True:
                    retrieve_choice = self.show_retrieve_menu()
                    
                    if retrieve_choice == '0':
                        break
                    elif retrieve_choice == '1':
                        self.retrieve_from_dataset()
                    elif retrieve_choice == '2':
                        self.start_chat()
            
            elif choice == '4':
                # 系统设置
                while True:
                    settings_choice = self.show_settings_menu()
                    
                    if settings_choice == '0':
                        break
                    elif settings_choice == '1':
                        self.change_api_key()
                    elif settings_choice == '2':
                        self.change_server_url()
            
            else:
                print("\n无效选择，请重试！")

if __name__ == "__main__":
    try:
        app = RAGFlowMenuApp()
        app.run()
    except KeyboardInterrupt:
        print("\n\n程序已被用户中断，感谢使用！")
    except Exception as e:
        print(f"\n程序出现错误: {e}")
    sys.exit(0)