Update chatbot.py

chatbot.py

@@ -1,408 +1,437 @@
-import re
-import json
-from pymilvus import MilvusClient, model
-from openai import OpenAI
-import time
-
-class LegalChatbot:
-    def __init__(self, milvus_db_path, collection_name, openai_api_key, openai_base_url=None, model_name="deepseek-reasoner"):
-        """
-        初始化法律RAG聊天机器人
-        
-        参数:
-            milvus_db_path: Milvus数据库路径
-            collection_name: 要搜索的集合名称
-            openai_api_key: OpenAI API密钥
-            openai_base_url: 可选的API基础URL（对于使用DeepSeek等服务）
-            model_name: 要使用的LLM模型名称
-        """
-        # 初始化Milvus客户端
-        self.milvus_client = MilvusClient(milvus_db_path)
-        self.collection_name = collection_name
-        
-        # 检查集合是否存在，如果不存在则创建它
-        if not self.milvus_client.has_collection(collection_name=collection_name):
-            print(f"Collection '{collection_name}' does not exist. Creating it...")
-            # 初始化嵌入模型
-            self.embedding_fn = model.DefaultEmbeddingFunction()
-            vector_dim = self.embedding_fn.dim
-            
-            # 创建新的集合
-            self.milvus_client.create_collection(
-                collection_name=collection_name,
-                dimension=vector_dim
-            )
-            print(f"Collection '{collection_name}' created successfully.")
-        
-        # 初始化嵌入模型
-        self.embedding_fn = model.DefaultEmbeddingFunction()
-        
-        # 初始化OpenAI客户端
-        if openai_base_url:
-            self.llm_client = OpenAI(api_key=openai_api_key, base_url=openai_base_url)
-        else:
-            self.llm_client = OpenAI(api_key=openai_api_key)
-        
-        self.model_name = model_name
-        self.conversation_history = [
-            {"role": "system", "content": """You are a helpful paralegal assistant with expertise in Canadian and U.S. law.
-            
-            You will help users with their legal questions. When answering, you should be helpful, accurate, and cite specific legal sources when possible.
-            
-            Users are members of the general public and may ask questions in Chinese or English. Please respond in the same language as the user's question.
-            """}
-        ]
-    
-    def search_legal_database(self, query, limit=5):
-        """
-        使用Milvus搜索法律数据库
-        
-        参数:
-            query: 搜索查询
-            limit: 返回结果的数量
-            
-        返回:
-            格式化的搜索结果字符串
-        """
-        if not query or query.strip() == "" or query.strip().lower() == "query":
-            return "无效的搜索查询。请提供具体的搜索内容。"
-            
-        # 检查数据库中是否有数据
-        collection_stats = self.milvus_client.get_collection_stats(self.collection_name)
-        row_count = collection_stats.get("row_count", 0)
-        
-        if row_count == 0:
-            # 如果集合为空，添加一些示例数据
-            print("集合为空，添加示例数据...")
-            self._add_sample_data()
-            
-        # 生成查询向量
-        query_vector = self.embedding_fn.encode_queries([query])
-        
-        # 执行搜索
-        search_results = self.milvus_client.search(
-            collection_name=self.collection_name,
-            data=query_vector,
-            limit=limit,
-            output_fields=["text", "page_num", "source"]
-        )
-        
-        # 检查是否有结果
-        if not search_results or len(search_results[0]) == 0:
-            return "没有找到与此查询相关的结果。"
-            
-        # 格式化搜索结果
-        formatted_results = []
-        for i, result in enumerate(search_results[0]):
-            similarity = 1 - result['distance']
-            source = result['entity'].get('source', 'Unknown source')
-            page_num = result['entity'].get('page_num', 'Unknown page')
-            text = result['entity'].get('text', '')
-            
-            formatted_result = f"[结果 {i+1}] 来源: {source}, 页码: {page_num}, 相关度: {similarity:.4f}\n"
-            formatted_result += f"内容: {text}\n\n"
-            formatted_results.append(formatted_result)
-        
-        return "\n".join(formatted_results)
-            
-    def _add_sample_data(self):
-        """添加示例法律文本数据到空集合中"""
-        # 简单的法律文本示例
-        docs = [
-            "Ontario Regulation 213/91 (Construction Projects) under the Occupational Health and Safety Act contains provisions for construction safety. Section 26 requires that every worker who may be exposed to the hazard of falling more than 3 metres shall use a fall protection system.",
-            "Under the Canada Labour Code, employers have a duty to ensure that the health and safety at work of every person employed by the employer is protected (Section 124). This includes providing proper training and supervision.",
-            "The Criminal Code of Canada Section 217.1 states that everyone who undertakes, or has the authority, to direct how another person does work or performs a task is under a legal duty to take reasonable steps to prevent bodily harm to that person, or any other person, arising from that work or task.",
-            "British Columbia's Workers Compensation Act requires employers to ensure the health and safety of all workers and comply with occupational health and safety regulations. This includes providing proper equipment, training, and supervision for construction activities.",
-            "Alberta's Occupational Health and Safety Code (Part 9) contains specific requirements for fall protection systems when workers are at heights of 3 metres or more, including the use of guardrails, safety nets, or personal fall arrest systems."
-        ]
-        
-        # 生成向量
-        vectors = self.embedding_fn.encode_documents(docs)
-        
-        # 准备数据
-        data = []
-        for i in range(len(docs)):
-            source_name = f"Sample Legal Text {i+1}"
-            data.append({
-                "id": i,
-                "vector": vectors[i],
-                "text": docs[i],
-                "page_num": 1,
-                "source": source_name
-            })
-            
-        # 插入数据
-        self.milvus_client.insert(collection_name=self.collection_name, data=data)
-        print(f"已添加 {len(data)} 条示例数据到集合")
-    
-    def _analyze_query_need(self, user_message):
-        """
-        分析用户消息，判断是否需要搜索法律数据库
-        
-        参数:
-            user_message: 用户的消息
-            
-        返回:
-            dict: {"needs_search": bool, "queries": list} 
-        """
-        # 预处理：检查用户是否明确要求搜索
-        search_keywords = [
-            "search in database", "search in the database", "search database", 
-            "look up", "find in database", "search for", "after searching",
-            "查询数据库", "搜索数据库", "数据库查找"
-        ]
-        
-        user_message_lower = user_message.lower()
-        explicit_search_request = any(keyword in user_message_lower for keyword in search_keywords)
-        
-        if explicit_search_request:
-            print("检测到用户明确要求搜索数据库")
-            # 清理查询内容，去掉搜索相关的指令（大小写不敏感）
-            clean_query = user_message
-            
-            # 所有需要去掉的短语
-            all_phrases_to_remove = search_keywords + [
-                "Answer me after searching in the database", "answer me after", 
-                "please search", "search and tell me", "look up and answer",
-                "tell me", "what is", "what are", "explain"
-            ]
-            
-            for phrase in all_phrases_to_remove:
-                # 大小写不敏感的替换
-                import re
-                pattern = re.compile(re.escape(phrase), re.IGNORECASE)
-                clean_query = pattern.sub("", clean_query)
-            
-            clean_query = clean_query.strip(".,?! ")
-            
-            if not clean_query or len(clean_query) < 3:
-                clean_query = "legal information"
-                
-            return {
-                "needs_search": True,
-                "reasoning": "用户明确要求搜索数据库",
-                "queries": [clean_query]
-            }
-        analysis_prompt = [
-            {"role": "system", "content": """You are an AI assistant that analyzes user questions to determine if they need legal database searches.
-            
-            Your task is to analyze the user's question and determine:
-            1. Whether this question requires searching a legal database
-            2. If yes, what specific search queries would be most helpful
-            
-            Respond in JSON format:
-            {
-                "needs_search": true/false,
-                "reasoning": "brief explanation of why search is or isn't needed",
-                "queries": ["query1", "query2"] // only if needs_search is true
-            }
-            
-            IMPORTANT RULES:
-            1. If the user explicitly requests database search (phrases like "search in database", "look up", "find in database"), always set needs_search to true
-            2. For ANY legal topic question, default to needs_search = true unless it's clearly a simple greeting or completely non-legal
-            3. Legal topics include: laws, regulations, legal procedures, legal documents, legal concepts, legal rights, etc.
-            
-            Search should be needed for:
-            - ANY legal question (wills, trusts, contracts, rights, procedures, etc.)
-            - Questions about specific laws, regulations, or legal codes  
-            - Requests for legal precedents or case law
-            - Questions about legal procedures or requirements
-            - Legal document comparisons (like will vs trust)
-            - When user explicitly asks to search database
-            
-            Search should NOT be needed ONLY for:
-            - Simple greetings ("hello", "how are you")
-            - Completely non-legal topics (weather, sports, etc.)
-            - Technical issues with the system itself
-            """}
-        ]
-        
-        # 添加当前对话历史的最后几条消息作为上下文
-        context_messages = self.conversation_history[-3:] if len(self.conversation_history) > 3 else self.conversation_history[1:]
-        for msg in context_messages:
-            analysis_prompt.append(msg)
-            
-        analysis_prompt.append({"role": "user", "content": f"Analyze this question: {user_message}"})
-        
-        response = self.llm_client.chat.completions.create(
-            model=self.model_name,
-            messages=analysis_prompt,
-            stream=False,
-            temperature=0.1
-        )
-        
-        response_content = response.choices[0].message.content.strip()
-        print(f"LLM原始响应: {response_content}")
-        
-        # 尝试提取JSON内容（如果响应包含其他文本）
-        import re
-        json_match = re.search(r'\{.*\}', response_content, re.DOTALL)
-        if json_match:
-            json_content = json_match.group(0)
-        else:
-            json_content = response_content
-        
-        analysis_result = json.loads(json_content)
-        print(f"查询分析结果: {analysis_result}")
-        return analysis_result
-        
-    def process_message(self, user_message):
-        """
-        处理用户消息并生成响应（两阶段模式）
-        
-        参数:
-            user_message: 用户的消息
-            
-        返回:
-            助手的响应
-        """
-        # 将用户消息添加到对话历史
-        self.conversation_history.append({"role": "user", "content": user_message})
-        
-        # 第一阶段：分析是否需要搜索
-        analysis = self._analyze_query_need(user_message)
-        
-        search_results = ""
-        if analysis.get("needs_search", False) and analysis.get("queries"):
-            # 第二阶段：执行搜索
-            all_results = []
-            for query in analysis["queries"][:2]:  # 最多执行2个查询
-                print(f"执行搜索查询: {query}")
-                result = self.search_legal_database(query)
-                if result and result.strip():
-                    all_results.append(f"查询: {query}\n{result}")
-            
-            if all_results:
-                search_results = "\n\n" + "="*50 + "\n".join(all_results)
-        
-        # 第三阶段：基于搜索结果生成回答
-        final_prompt = self.conversation_history.copy()
-        
-        if search_results:
-            final_prompt.append({
-                "role": "system", 
-                "content": f"以下是相关的法律搜索结果，请在回答中引用这些信息：\n{search_results}\n\n请基于这些搜索结果回答用户的问题，并引用具体的来源和页码。"
-            })
-        response = self.llm_client.chat.completions.create(
-            model=self.model_name,
-            messages=final_prompt,
-            stream=False
-        )
-        
-        assistant_response = response.choices[0].message.content
-        
-        # 将最终响应添加到对话历史
-        self.conversation_history.append({"role": "assistant", "content": assistant_response})
-        
-        return assistant_response
-    
-    def process_message_stream(self, user_message):
-        """
-        处理用户消息并以流式方式返回响应，支持智能RAG查询
-        
-        Args:
-            user_message (str): 用户输入的消息
-            
-        Yields:
-            str: 响应文本的片段
-        """
-        # 添加用户消息到对话历史
-        self.conversation_history.append({"role": "user", "content": user_message})
-    
-        # 第一阶段：分析是否需要搜索
-        analysis = self._analyze_query_need(user_message)
-        
-        search_results = ""
-        if analysis.get("needs_search", False) and analysis.get("queries"):
-            # 输出搜索提示
-            yield "[正在搜索相关法律信息...]\n\n"
-            
-            # 第二阶段：执行搜索
-            all_results = []
-            for query in analysis["queries"][:2]:  # 最多执行2个查询
-                print(f"执行流式搜索查询: {query}")
-                result = self.search_legal_database(query)
-                if result and result.strip():
-                    all_results.append(f"查询: {query}\n{result}")
-            
-            if all_results:
-                search_results = "\n\n" + "="*50 + "\n".join(all_results)
-                yield "[搜索完成，正在生成回答...]\n\n"
-        
-        # 第三阶段：基于搜索结果生成流式回答
-        final_prompt = self.conversation_history.copy()
-        
-        if search_results:
-            final_prompt.append({
-                "role": "system", 
-                "content": f"以下是相关的法律搜索结果，请在回答中引用这些信息：\n{search_results}\n\n请基于这些搜索结果回答用户的问题，并引用具体的来源和页码。"
-            })
-        
-        # 创建流式完成请求
-        response = self.llm_client.chat.completions.create(
-            model=self.model_name,
-            messages=final_prompt,
-            stream=True,
-            temperature=0.3,
-            max_tokens=2048
-        )
-        
-        full_response = ""  # 存储完整响应
-        
-        # 处理流式响应
-        for chunk in response:
-            if hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content:
-                content = chunk.choices[0].delta.content
-                full_response += content
-                yield content
-        
-        # 将最终响应添加到对话历史
-        self.conversation_history.append({"role": "assistant", "content": full_response})
-    
-    def reset_conversation(self):
-        """重置对话历史"""
-        self.conversation_history = [self.conversation_history[0]]  # 保留系统消息
-
-def main():
-    # 配置信息
-    MILVUS_DB_PATH = "./milvus_legal_codes.db"  # 使用您之前创建的数据库名，或创建新的
-    COLLECTION_NAME = "legal_codes_collection"  # 使用您之前创建的集合名或新名称
-    # OPENAI_API_KEY = "sk-dad31a53a4684587aed060afc0e4d75b"  # 请替换为实际的API密钥
-    # OPENAI_BASE_URL = "https://api.deepseek.com"  # 如果使用OpenAI API，请移除此行
-    OPENAI_API_KEY = "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"  # 请替换为实际的API密钥
-    OPENAI_BASE_URL = ""  # 如果使用OpenAI API，请移除此行
-    
-    # 初始化聊天机器人
-    chatbot = LegalChatbot(
-        milvus_db_path=MILVUS_DB_PATH,
-        collection_name=COLLECTION_NAME,
-        openai_api_key=OPENAI_API_KEY,
-        openai_base_url=OPENAI_BASE_URL,
-        # model_name="deepseek-chat"
-        model_name="gpt-4o"
-    )
-    
-    print("法律RAG聊天机器人已初始化。输入'exit'或'quit'结束会话。")
-    
-    while True:
-        user_input = input("\n您: ")
-        
-        if user_input.lower() in ['exit', 'quit']:
-            print("会话结束。")
-            break
-        
-        if user_input.lower() in ['reset', 'clear']:
-            chatbot.reset_conversation()
-            print("已重置对话历史。")
-            continue
-        
-        print("\n正在思考...")
-        start_time = time.time()
-        
-        response = chatbot.process_message(user_input)
-        
-        end_time = time.time()
-        print(f"助手 ({end_time - start_time:.2f}秒): {response}")
-
-if __name__ == "__main__":
+import re
+import json
+from pymilvus import MilvusClient, model
+from openai import OpenAI
+import time
+
+class LegalChatbot:
+    def __init__(self, milvus_db_path, collection_name, openai_api_key, openai_base_url=None, model_name="deepseek-reasoner"):
+        """
+        Initialize Legal RAG Chatbot
+        
+        Args:
+            milvus_db_path: Milvus database path
+            collection_name: Collection name to search
+            openai_api_key: OpenAI API key
+            openai_base_url: Optional API base URL (for DeepSeek etc.)
+            model_name: LLM model name to use
+        """
+        # Initialize Milvus client
+        self.milvus_client = MilvusClient(milvus_db_path)
+        self.collection_name = collection_name
+        
+        # Check if collection exists, create if not
+        if not self.milvus_client.has_collection(collection_name=collection_name):
+            print(f"Collection '{collection_name}' does not exist. Creating it...")
+            # Initialize embedding model
+            self.embedding_fn = model.DefaultEmbeddingFunction()
+            vector_dim = self.embedding_fn.dim
+            
+            # Create new collection
+            self.milvus_client.create_collection(
+                collection_name=collection_name,
+                dimension=vector_dim
+            )
+            print(f"Collection '{collection_name}' created successfully.")
+        
+        # Initialize embedding model
+        self.embedding_fn = model.DefaultEmbeddingFunction()
+        
+        # Initialize OpenAI client
+        if openai_base_url:
+            self.llm_client = OpenAI(api_key=openai_api_key, base_url=openai_base_url)
+        else:
+            self.llm_client = OpenAI(api_key=openai_api_key)
+        
+        self.model_name = model_name
+        self.conversation_history = [
+            {"role": "system", "content": """You are a helpful paralegal assistant with expertise in Canadian and U.S. law.
+            
+            You will help users with their legal questions. When answering, you should be helpful, accurate, and cite specific legal sources when possible.
+            
+            Users are members of the general public and may ask questions in Chinese or English. Please respond in the same language as the user's question.
+            """}
+        ]
+    
+    def search_legal_database(self, query, limit=5):
+        """
+        Search legal database using Milvus
+        
+        Args:
+            query: Search query
+            limit: Number of results to return
+            
+        Returns:
+            Formatted search results string
+        """
+        if not query or query.strip() == "" or query.strip().lower() == "query":
+            return "Invalid search query. Please provide specific search content."
+            
+        # Check if database has data
+        collection_stats = self.milvus_client.get_collection_stats(self.collection_name)
+        row_count = collection_stats.get("row_count", 0)
+        
+        if row_count == 0:
+            # If collection is empty, add sample data
+            print("Collection is empty, adding sample data...")
+            self._add_sample_data()
+            
+        # Generate query vector
+        query_vector = self.embedding_fn.encode_queries([query])
+        
+        # Execute search
+        search_results = self.milvus_client.search(
+            collection_name=self.collection_name,
+            data=query_vector,
+            limit=limit,
+            output_fields=["text", "page_num", "source"]
+        )
+        
+        # Check if there are results
+        if not search_results or len(search_results[0]) == 0:
+            return "No results found related to this query."
+            
+        # Format search results
+        formatted_results = []
+        for i, result in enumerate(search_results[0]):
+            similarity = 1 - result['distance']
+            source = result['entity'].get('source', 'Unknown source')
+            page_num = result['entity'].get('page_num', 'Unknown page')
+            text = result['entity'].get('text', '')
+            
+            formatted_result = f"[Result {i+1}] Source: {source}, Page: {page_num}, Relevance: {similarity:.4f}\n"
+            formatted_result += f"Content: {text}\n\n"
+            formatted_results.append(formatted_result)
+        
+        return "\n".join(formatted_results)
+            
+    def _add_sample_data(self):
+        """Add sample legal text data to empty collection"""
+        # Simple legal text examples
+        docs = [
+            "Ontario Regulation 213/91 (Construction Projects) under the Occupational Health and Safety Act contains provisions for construction safety. Section 26 requires that every worker who may be exposed to the hazard of falling more than 3 metres shall use a fall protection system.",
+            "Under the Canada Labour Code, employers have a duty to ensure that the health and safety at work of every person employed by the employer is protected (Section 124). This includes providing proper training and supervision.",
+            "The Criminal Code of Canada Section 217.1 states that everyone who undertakes, or has the authority, to direct how another person does work or performs a task is under a legal duty to take reasonable steps to prevent bodily harm to that person, or any other person, arising from that work or task.",
+            "British Columbia's Workers Compensation Act requires employers to ensure the health and safety of all workers and comply with occupational health and safety regulations. This includes providing proper equipment, training, and supervision for construction activities.",
+            "Alberta's Occupational Health and Safety Code (Part 9) contains specific requirements for fall protection systems when workers are at heights of 3 metres or more, including the use of guardrails, safety nets, or personal fall arrest systems."
+        ]
+        
+        # Generate vectors
+        vectors = self.embedding_fn.encode_documents(docs)
+        
+        # Prepare data
+        data = []
+        for i in range(len(docs)):
+            source_name = f"Sample Legal Text {i+1}"
+            data.append({
+                "id": i,
+                "vector": vectors[i],
+                "text": docs[i],
+                "page_num": 1,
+                "source": source_name
+            })
+            
+        # Insert data
+        self.milvus_client.insert(collection_name=self.collection_name, data=data)
+        print(f"Added {len(data)} sample data entries to collection")
+    
+    def _analyze_query_need(self, user_message):
+        """
+        Analyze user message to determine if legal database search is needed
+        
+        Args:
+            user_message: User's message
+            
+        Returns:
+            dict: {"needs_search": bool, "queries": list} 
+        """
+        # Preprocessing: Check if user explicitly requests search
+        search_keywords = [
+            "search in database", "search in the database", "search database", 
+            "look up", "find in database", "search for", "after searching",
+            "query database", "database search", "database lookup"
+        ]
+        
+        user_message_lower = user_message.lower()
+        explicit_search_request = any(keyword in user_message_lower for keyword in search_keywords)
+        
+        if explicit_search_request:
+            print("Detected explicit user request for database search")
+            # Clean query content, remove search-related instructions (case insensitive)
+            clean_query = user_message
+            
+            # All phrases to remove
+            all_phrases_to_remove = search_keywords + [
+                "Answer me after searching in the database", "answer me after", 
+                "please search", "search and tell me", "look up and answer",
+                "tell me", "what is", "what are", "explain"
+            ]
+            
+            for phrase in all_phrases_to_remove:
+                # Case insensitive replacement
+                import re
+                pattern = re.compile(re.escape(phrase), re.IGNORECASE)
+                clean_query = pattern.sub("", clean_query)
+            
+            clean_query = clean_query.strip(".,?! ")
+            
+            if not clean_query or len(clean_query) < 3:
+                clean_query = "legal information"
+                
+            return {
+                "needs_search": True,
+                "reasoning": "User explicitly requested database search",
+                "queries": [clean_query]
+            }
+        
+        analysis_prompt = [
+            {"role": "system", "content": """You are an AI assistant that analyzes user questions to determine if they need legal database searches.
+            
+            Your task is to analyze the user's question and determine:
+            1. Whether this question requires searching a legal database
+            2. If yes, what specific search queries would be most helpful
+            
+            Respond in JSON format:
+            {
+                "needs_search": true/false,
+                "reasoning": "brief explanation of why search is or isn't needed",
+                "queries": ["query1", "query2"] // only if needs_search is true
+            }
+            
+            IMPORTANT RULES:
+            1. If the user explicitly requests database search (phrases like "search in database", "look up", "find in database"), always set needs_search to true
+            2. For ANY legal topic question, default to needs_search = true unless it's clearly a simple greeting or completely non-legal
+            3. Legal topics include: laws, regulations, legal procedures, legal documents, legal concepts, legal rights, etc.
+            
+            Search should be needed for:
+            - ANY legal question (wills, trusts, contracts, rights, procedures, etc.)
+            - Questions about specific laws, regulations, or legal codes  
+            - Requests for legal precedents or case law
+            - Questions about legal procedures or requirements
+            - Legal document comparisons (like will vs trust)
+            - When user explicitly asks to search database
+            
+            Search should NOT be needed ONLY for:
+            - Simple greetings ("hello", "how are you")
+            - Completely non-legal topics (weather, sports, etc.)
+            - Technical issues with the system itself
+            """}
+        ]
+        
+        # Add recent conversation history as context
+        context_messages = self.conversation_history[-3:] if len(self.conversation_history) > 3 else self.conversation_history[1:]
+        for msg in context_messages:
+            analysis_prompt.append(msg)
+            
+        analysis_prompt.append({"role": "user", "content": f"Analyze this question: {user_message}"})
+        
+        # Display the analysis prompt
+        print("\n<prompt>")
+        print("Query Analysis Prompt:")
+        print(f"User Message: {user_message}")
+        print("System: Analyzing if legal database search is needed...")
+        print("</prompt>\n")
+        
+        response = self.llm_client.chat.completions.create(
+            model=self.model_name,
+            messages=analysis_prompt,
+            stream=False,
+            temperature=0.1
+        )
+        
+        response_content = response.choices[0].message.content.strip()
+        print(f"LLM Raw Response: {response_content}")
+        
+        # Try to extract JSON content (if response contains other text)
+        import re
+        json_match = re.search(r'\{.*\}', response_content, re.DOTALL)
+        if json_match:
+            json_content = json_match.group(0)
+        else:
+            json_content = response_content
+        
+        analysis_result = json.loads(json_content)
+        print(f"Query Analysis Result: {analysis_result}")
+        return analysis_result
+        
+    def process_message(self, user_message):
+        """
+        Process user message and generate response (two-stage mode)
+        
+        Args:
+            user_message: User's message
+            
+        Returns:
+            Assistant's response
+        """
+        # Add user message to conversation history
+        self.conversation_history.append({"role": "user", "content": user_message})
+        
+        # Stage 1: Analyze if search is needed
+        analysis = self._analyze_query_need(user_message)
+        
+        search_results = ""
+        if analysis.get("needs_search", False) and analysis.get("queries"):
+            # Stage 2: Execute search
+            all_results = []
+            for query in analysis["queries"][:2]:  # Execute max 2 queries
+                print(f"Executing search query: {query}")
+                result = self.search_legal_database(query)
+                if result and result.strip():
+                    all_results.append(f"Query: {query}\n{result}")
+            
+            if all_results:
+                search_results = "\n\n" + "="*50 + "\n".join(all_results)
+                
+                # Display RAG results with tags
+                print("\n<RAG_result>")
+                print("Search Results from Legal Database:")
+                print(search_results)
+                print("</RAG_result>\n")
+        
+        # Stage 3: Generate answer based on search results
+        final_prompt = self.conversation_history.copy()
+        
+        if search_results:
+            final_prompt.append({
+                "role": "system", 
+                "content": f"The following are relevant legal search results, please reference this information in your answer:\n{search_results}\n\nPlease answer the user's question based on these search results, and cite specific sources and page numbers."
+            })
+            
+        response = self.llm_client.chat.completions.create(
+            model=self.model_name,
+            messages=final_prompt,
+            stream=False
+        )
+        
+        assistant_response = response.choices[0].message.content
+        
+        # Add final response to conversation history
+        self.conversation_history.append({"role": "assistant", "content": assistant_response})
+        
+        return assistant_response
+    
+    def process_message_stream(self, user_message):
+        """
+        Process user message and return streaming response with intelligent RAG queries
+        
+        Args:
+            user_message (str): User input message
+            
+        Yields:
+            str: Response text fragments
+        """
+        # Add user message to conversation history
+        self.conversation_history.append({"role": "user", "content": user_message})
+    
+        # Stage 1: Analyze if search is needed
+        analysis = self._analyze_query_need(user_message)
+        
+        search_results = ""
+        if analysis.get("needs_search", False) and analysis.get("queries"):
+            # Output search prompt
+            yield "\n<prompt>\n"
+            yield "🔍 Analyzing query for legal database search...\n"
+            yield f"Query Analysis: {analysis.get('reasoning', 'Legal topic detected')}\n"
+            yield f"Search needed: {analysis.get('needs_search', False)}\n"
+            yield "</prompt>\n\n"
+            
+            yield "[🔍 Searching relevant legal information...]\n\n"
+            
+            # Stage 2: Execute search
+            all_results = []
+            for query in analysis["queries"][:2]:  # Execute max 2 queries
+                print(f"Executing streaming search query: {query}")
+                result = self.search_legal_database(query)
+                if result and result.strip():
+                    all_results.append(f"Query: {query}\n{result}")
+            
+            if all_results:
+                search_results = "\n\n" + "="*50 + "\n".join(all_results)
+                
+                # Output RAG results with tags
+                yield "\n<RAG_result>\n"
+                yield "📚 Search Results from Legal Database:\n\n"
+                for i, result in enumerate(all_results, 1):
+                    yield f"Search {i}:\n{result}\n\n"
+                yield "</RAG_result>\n\n"
+                
+                yield "[✅ Search completed, generating answer...]\n\n"
+        
+        # Stage 3: Generate streaming answer based on search results
+        final_prompt = self.conversation_history.copy()
+        
+        if search_results:
+            final_prompt.append({
+                "role": "system", 
+                "content": f"The following are relevant legal search results, please reference this information in your answer:\n{search_results}\n\nPlease answer the user's question based on these search results, and cite specific sources and page numbers."
+            })
+        
+        # Create streaming completion request
+        response = self.llm_client.chat.completions.create(
+            model=self.model_name,
+            messages=final_prompt,
+            stream=True,
+            temperature=0.3,
+            max_tokens=2048
+        )
+        
+        full_response = ""  # Store complete response
+        
+        # Process streaming response
+        for chunk in response:
+            if hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content:
+                content = chunk.choices[0].delta.content
+                full_response += content
+                yield content
+        
+        # Add final response to conversation history
+        self.conversation_history.append({"role": "assistant", "content": full_response})
+    
+    def reset_conversation(self):
+        """Reset conversation history"""
+        self.conversation_history = [self.conversation_history[0]]  # Keep system message
+
+def main():
+    # Configuration
+    MILVUS_DB_PATH = "./milvus_legal_codes.db"  # Use your existing database name or create new
+    COLLECTION_NAME = "legal_codes_collection"  # Use your existing collection name or new name
+    # OPENAI_API_KEY = "sk-dad31a53a4684587aed060afc0e4d75b"  # Replace with actual API key
+    # OPENAI_BASE_URL = "https://api.deepseek.com"  # Remove this line if using OpenAI API
+    OPENAI_API_KEY = "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"  # Replace with actual API key
+    OPENAI_BASE_URL = ""  # Remove this line if using OpenAI API
+    
+    # Initialize chatbot
+    chatbot = LegalChatbot(
+        milvus_db_path=MILVUS_DB_PATH,
+        collection_name=COLLECTION_NAME,
+        openai_api_key=OPENAI_API_KEY,
+        openai_base_url=OPENAI_BASE_URL,
+        # model_name="deepseek-chat"
+        model_name="gpt-4o"
+    )
+    
+    print("Legal RAG Chatbot initialized. Type 'exit' or 'quit' to end session.")
+    
+    while True:
+        user_input = input("\nYou: ")
+        
+        if user_input.lower() in ['exit', 'quit']:
+            print("Session ended.")
+            break
+        
+        if user_input.lower() in ['reset', 'clear']:
+            chatbot.reset_conversation()
+            print("Conversation history reset.")
+            continue
+        
+        print("\nThinking...")
+        start_time = time.time()
+        
+        response = chatbot.process_message(user_input)
+        
+        end_time = time.time()
+        print(f"Assistant ({end_time - start_time:.2f}s): {response}")
+
+if __name__ == "__main__":
     main()