Update app.py

app.py

@@ -1,212 +1,22 @@
 #!/usr/bin/env python3
 """
 Legal RAG Chatbot for Hugging Face Spaces
-基于原有chatbot.py的简化版本，使用Gradio ChatInterface
+直接使用原有的chatbot.py，只添加Gradio界面
 """
 
 import os
 import gradio as gr
-from pymilvus import MilvusClient, model
-from openai import OpenAI
-import json
+from chatbot import LegalChatbot
 
-class LegalChatbot:
-    def __init__(self, milvus_db_path, collection_name, openai_api_key, openai_base_url=None, model_name="gpt-4o"):
-        """
-        初始化法律RAG聊天机器人
-        """
-        print(f"🚀 Initializing Legal Chatbot...")
-        print(f"📁 Database path: {milvus_db_path}")
-        print(f"📚 Collection: {collection_name}")
-        print(f"🤖 Model: {model_name}")
-        
-        # 初始化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 with sample data...")
-            # 初始化嵌入模型
-            self.embedding_fn = model.DefaultEmbeddingFunction()
-            vector_dim = self.embedding_fn.dim
-            
-            # 创建新的集合
-            self.milvus_client.create_collection(
-                collection_name=collection_name,
-                dimension=vector_dim
-            )
-            self._add_sample_data()
-        else:
-            print(f"✅ Found existing collection '{collection_name}'")
-        
-        # 初始化嵌入模型
-        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
-        
-        print("✅ Legal Chatbot initialized successfully!")
-    
-    def search_legal_database(self, query, limit=5):
-        """使用Milvus搜索法律数据库"""
-        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("⚠️ Collection is empty, adding sample data...")
-            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.",
-            "Employment Standards Act provides minimum standards for wages, hours of work, overtime pay, vacation time, and termination notice. Employers must provide at least two weeks written notice or pay in lieu for termination.",
-            "Canadian Charter of Rights and Freedoms Section 15 guarantees equality rights. Every individual is equal before and under the law and has the right to equal protection and benefit of the law without discrimination.",
-            "Personal Information Protection and Electronic Documents Act (PIPEDA) governs how private sector organizations collect, use and disclose personal information in the course of commercial activities.",
-            "Competition Act prohibits anti-competitive practices including conspiracies, bid-rigging, price maintenance, and abuse of dominant position. The Act is enforced by the Competition Bureau.",
-            "Consumer Protection Acts across Canada provide consumers with rights including cooling-off periods for certain contracts, protection from unfair business practices, and warranties for goods and services."
-        ]
-        
-        # 生成向量
-        vectors = self.embedding_fn.encode_documents(docs)
-        
-        # 准备数据
-        data = []
-        for i in range(len(docs)):
-            source_name = f"Canadian Legal Code {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"✅ Added {len(data)} sample legal documents to collection")
-    
-    def _analyze_query_need(self, user_message):
-        """分析用户消息，判断是否需要搜索法律数据库"""
-        # 简化的搜索判断逻辑
-        legal_keywords = [
-            "law", "legal", "regulation", "act", "code", "section", "rights", 
-            "protection", "safety", "employment", "contract", "liability",
-            "charter", "constitution", "statute", "criminal", "civil",
-            "搜索", "查询", "数据库", "法律", "法规", "条例"
-        ]
-        
-        user_message_lower = user_message.lower()
-        
-        # 检查是否包含法律关键词
-        needs_search = any(keyword in user_message_lower for keyword in legal_keywords)
-        
-        if needs_search:
-            # 清理查询内容
-            clean_query = user_message.strip()
-            return {
-                "needs_search": True,
-                "reasoning": "检测到法律相关查询",
-                "queries": [clean_query]
-            }
-        else:
-            return {
-                "needs_search": False,
-                "reasoning": "未检测到法律关键词",
-                "queries": []
-            }
-        
-    def process_message_with_history(self, user_message, history, system_message):
-        """
-        处理用户消息，兼容Gradio ChatInterface的历史格式
-        """
-        # 构建完整的对话历史
-        messages = [{"role": "system", "content": system_message}]
-        
-        # 添加历史对话
-        for user_msg, assistant_msg in history:
-            if user_msg:
-                messages.append({"role": "user", "content": user_msg})
-            if assistant_msg:
-                messages.append({"role": "assistant", "content": assistant_msg})
-        
-        # 分析是否需要搜索
-        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"🔍 Searching for: {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)
-        
-        # 添加当前用户消息
-        if search_results:
-            enhanced_message = f"{user_message}\n\n以下是相关的法律搜索结果，请在回答中引用这些信息：\n{search_results}"
-            messages.append({"role": "user", "content": enhanced_message})
-        else:
-            messages.append({"role": "user", "content": user_message})
-        
-        return messages
+# 配置信息 - 从环境变量获取
+MILVUS_DB_PATH = os.environ.get("MILVUS_DB_PATH", "./milvus_legal_codes.db")
+COLLECTION_NAME = os.environ.get("COLLECTION_NAME", "legal_codes_collection")
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+OPENAI_BASE_URL = os.environ.get("OPENAI_BASE_URL", "")
+MODEL_NAME = os.environ.get("MODEL_NAME", "gpt-4o")
 
 # 初始化聊天机器人
 def initialize_chatbot():
-    # 从环境变量获取配置
-    MILVUS_DB_PATH = os.environ.get("MILVUS_DB_PATH", "./milvus_legal_codes.db")
-    COLLECTION_NAME = os.environ.get("COLLECTION_NAME", "legal_codes_collection")
-    OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
-    OPENAI_BASE_URL = os.environ.get("OPENAI_BASE_URL", "")
-    MODEL_NAME = os.environ.get("MODEL_NAME", "gpt-4o")
-    
     if not OPENAI_API_KEY:
         raise ValueError("❌ OPENAI_API_KEY environment variable is required!")
     
@@ -225,9 +35,11 @@ def initialize_chatbot():
 try:
     chatbot = initialize_chatbot()
     print("✅ Chatbot initialized successfully!")
+    chatbot_status = f"✅ **Status**: Connected to database with {chatbot.collection_name} collection"
 except Exception as e:
     print(f"❌ Failed to initialize chatbot: {e}")
     chatbot = None
+    chatbot_status = f"❌ **Status**: Configuration error - {str(e)}"
 
 def respond(
     message,
@@ -237,28 +49,34 @@ def respond(
     temperature,
     top_p,
 ):
+    """
+    Gradio ChatInterface响应函数
+    使用原有chatbot的流式处理功能
+    """
     if chatbot is None:
         yield "❌ Chatbot not initialized. Please check the configuration."
         return
     
     try:
-        # 使用chatbot处理消息和历史
-        messages = chatbot.process_message_with_history(message, history, system_message)
+        # 重置聊天机器人的对话历史
+        chatbot.reset_conversation()
         
-        response = ""
+        # 设置系统消息
+        if system_message.strip():
+            chatbot.conversation_history[0]["content"] = system_message
+        
+        # 添加历史对话到聊天机器人
+        for user_msg, assistant_msg in history:
+            if user_msg:
+                chatbot.conversation_history.append({"role": "user", "content": user_msg})
+            if assistant_msg:
+                chatbot.conversation_history.append({"role": "assistant", "content": assistant_msg})
         
-        # 调用LLM API
-        for chunk in chatbot.llm_client.chat.completions.create(
-            model=chatbot.model_name,
-            messages=messages,
-            max_tokens=max_tokens,
-            stream=True,
-            temperature=temperature,
-            top_p=top_p,
-        ):
-            if hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content:
-                token = chunk.choices[0].delta.content
-                response += token
+        # 使用原有的流式处理功能
+        response = ""
+        for chunk in chatbot.process_message_stream(message):
+            if chunk:
+                response += chunk
                 yield response
                 
     except Exception as e:
@@ -266,14 +84,22 @@ def respond(
         yield f"抱歉，处理您的消息时出现错误：{str(e)}"
 
 # 准备描述信息
-base_description = "🤖 AI法律助手，结合向量数据库搜索和大语言模型，为您提供准确的法律信息咨询。"
+base_description = """🤖 **AI法律助手** - 结合向量数据库搜索和大语言模型的智能法律咨询系统
+
+🔍 **核心功能:**
+- 智能查询分析 - 自动判断是否需要搜索法律数据库
+- 向量相似度搜索 - 基于Milvus的高效法律文档检索  
+- RAG增强生成 - 结合搜索结果提供准确回答
+- 实时流式回复 - 支持打字机效果的实时响应
+
+💡 **试试这些问题:**
+• "What are the fall protection requirements in Ontario construction?"
+• "Tell me about employer duties under Canada Labour Code"  
+• "Search for information about workplace safety regulations"
+• "What are my rights under the Charter of Rights and Freedoms?"
+"""
 
-if chatbot:
-    status_info = f"\n\n✅ **Status**: Connected to database with {chatbot.collection_name} collection"
-    examples_info = "\n\n💡 **Try asking:**\n• What are the fall protection requirements in Ontario construction?\n• Tell me about employer duties under Canada Labour Code\n• What are my rights under the Charter of Rights and Freedoms?\n• Search for information about workplace safety regulations"
-    full_description = base_description + status_info + examples_info
-else:
-    full_description = base_description + "\n\n❌ **Status**: Configuration error - please check environment variables"
+full_description = base_description + f"\n\n{chatbot_status}"
 
 # 创建Gradio ChatInterface
 demo = gr.ChatInterface(
@@ -284,23 +110,39 @@ demo = gr.ChatInterface(
         gr.Textbox(
             value="You are a helpful legal assistant with expertise in Canadian law. You have access to a legal database and should provide accurate, well-sourced legal information. Always cite specific legal sources when possible. Remember to include appropriate disclaimers that this is for informational purposes only and not legal advice.",
             label="System Message",
-            lines=3
+            lines=3,
+            max_lines=5
+        ),
+        gr.Slider(
+            minimum=1, 
+            maximum=2048, 
+            value=1024, 
+            step=1, 
+            label="Max new tokens"
+        ),
+        gr.Slider(
+            minimum=0.1, 
+            maximum=2.0, 
+            value=0.7, 
+            step=0.1, 
+            label="Temperature"
         ),
-        gr.Slider(minimum=1, maximum=2048, value=1024, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(
             minimum=0.1,
             maximum=1.0,
             value=0.95,
             step=0.05,
-            label="Top-p (nucleus sampling)",
+            label="Top-p (nucleus sampling)"
         ),
-    ]
+    ],
+    theme=gr.themes.Soft(),
+    analytics_enabled=False
 )
 
 if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        share=False
+        share=False,
+        show_error=True
     )