Spaces:
Sleeping
Sleeping
| 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() |