""" Legal RAG Agent - Core Orchestrator This module contains the main LegalRagAgent class that serves as the central coordinator for the Legal RAG Pipeline system. It integrates language models, vector databases, and legal research capabilities to provide intelligent legal assistance. Key Features: - Multi-model LLM support (Mistral, OpenAI, Ollama) - RAG pipeline for contextual legal advice - CanLII integration for Canadian legal research - Document upload and semantic search - Case and client management - Conversation history tracking Author: Legal RAG Pipeline Team Version: 1.0.0 """ from dotenv import load_dotenv import re import os import json from datetime import datetime from collections import deque from typing import List, Optional, Dict, Any import logging # Core RAG components from src.core.rag_pipeline import RAGPipeline, build_workflow from src.storage.database_manager import DatabaseManager # LLM providers from langchain_mistralai import ChatMistralAI from langchain_ollama import ChatOllama # Legal-specific prompt tools from src.generation.rag_prompt_tools import ( ChatSummarizer, PassiveLegalAdvice, CaseSimilaritySearch, LegislationSimilaritySearch, CanliiQueryWriter, ) # Utilities from src.utils.canlii_search import search_canlii from src.generation.model_provider import get_llm_model # Load environment variables load_dotenv() # Set up logging logging.basicConfig( level=logging.INFO) logger = logging.getLogger(__name__) class LegalRagAgent: """ Legal RAG Agent - Main orchestrator for the legal assistance system. This class coordinates interactions between the database manager, RAG pipeline, and language models to provide intelligent legal information and advice based on client cases, documents, and legal knowledge databases. The agent supports: - Case and client management - Document upload and semantic search - Legal research via CanLII integration - Contextual question answering - Conversation history tracking - Multi-model LLM support Attributes: database_manager (DatabaseManager): Handles data persistence and retrieval rag_pipeline (RAGPipeline): Manages document retrieval and answer generation llm: Current language model instance graph: Workflow graph for RAG operations """ def __init__(self, database_manager: DatabaseManager, model_name: str = "mistral-large-2411"): """ Initialize the Legal RAG Agent. Args: database_manager: Database manager instance for data operations model_name: Name of the language model to use (default: mistral-large-2411) Raises: ValueError: If required environment variables are missing """ self.database_manager = database_manager self.rag_pipeline = None self.graph = None # Validate required environment variables api_key = os.getenv('MISTRAL_API_KEY') if not api_key and 'mistral' in model_name.lower(): raise ValueError( "MISTRAL_API_KEY environment variable is required for Mistral models. " "Please set it in your .env file." ) # Configure language model self.llm_config = { "temperature": 0.0, # Deterministic responses for legal advice "num_ctx": 100000, # Large context window for legal documents "extract_reasoning": False # Focus on final answers } self.llm = get_llm_model(model_name, True, **self.llm_config) # Initialize RAG pipeline and workflow self.rag_pipeline = RAGPipeline(llm=self.llm) self.graph = build_workflow(self.rag_pipeline) # Initialize chat summarizer for conversation management self.chat_summarizer = ChatSummarizer(self.llm) logger.info(f"Legal RAG Agent initialized with {model_name}") def set_llm(self, model_name: str) -> bool: """ Update the language model used by the agent. Args: model_name: Name of the new language model Returns: bool: True if model was successfully updated, False otherwise """ try: self.llm = get_llm_model(model_name, True, **self.llm_config) self.rag_pipeline.llm = self.llm self.graph = build_workflow(self.rag_pipeline) logger.info(f"Language model updated to {model_name}") return True except Exception as e: logger.info(f"Failed to update language model: {str(e)}") return False def _get_object_handler(self, search_type, query): """ Search for an existing client or case by name or ID. If not found, returns None. """ if not query: query = input(f"Enter name to search for an existing {search_type}, or press Enter to create a new {search_type}: ").strip() if search_type == "client": matches = self.database_manager.get_client_by_id(query) elif search_type == "case": matches = self.database_manager.get_cases_by_name(query) else: logger.info(f"Unknown search_type: {search_type}") return None # # WEB-SAFE SEARCH METHODS # def get_case_by_name_web_safe(self, case_name: str) -> Optional[Any]: """ Web-safe version of case search that doesn't use input() calls. This method is designed for use by web APIs where interactive input is not possible. Args: case_name: Name or ID of the case to find Returns: Case object if found, None otherwise """ if not case_name: return None # Try to get case by ID first (if it's a number) if case_name.isdigit(): case = self.database_manager.get_case_by_id(case_name) if case: return case # Search by name matches = self.database_manager.get_cases_by_name(case_name) if matches: return matches[0] # Return first match for web safety return None def get_client_by_name_web_safe(self, client_name: str) -> Optional[Any]: """ Web-safe version of client search that doesn't use input() calls. This method is designed for use by web APIs where interactive input is not possible. Args: client_name: Name or ID of the client to find Returns: Client object if found, None otherwise """ if not client_name: return None # Try to get client by ID first (if it's a number) if client_name.isdigit(): client = self.database_manager.get_client_by_id(client_name) if client: return client # Search by name matches = self.database_manager.get_clients_by_name(client_name) if matches: return matches[0] # Return first match for web safety return None def _new_object_from_json(self, search_type: str, json_path: Optional[str] = None) -> Optional[Any]: """ Create a new client or case object from JSON file data. Args: search_type: Type of object to create ("client" or "case") json_path: Path to JSON file containing object data Returns: Newly created client or case object, or None if creation fails """ if not json_path or not os.path.exists(json_path): logger.info(f"JSON file not found: {json_path}") return None try: with open(json_path, 'r') as f: data = json.load(f) logger.info(f"Loaded {search_type} data from {json_path}") if search_type == "client": return self.database_manager.add_client(**data) elif search_type == "case": if self.database_manager.client: return self.database_manager.add_case( client_id=self.database_manager.client.id, **data ) else: return self.database_manager.add_case(**data) else: logger.info(f"Unknown search_type: {search_type}") return None except Exception as e: logger.info(f"Error creating {search_type} from JSON: {str(e)}") return None # # DOCUMENT MANAGEMENT METHODS # def _handle_documents(self, message: str) -> str: """ Handle document processing commands embedded in messages. Args: message: Input message that may contain document commands Returns: Response message about document processing """ if isinstance(message, str) and "UPLOAD_DOCUMENTS" in message: file_path = os.path.join(os.getcwd(), 'srcdoc.pdf') logger.info(f"Uploading document from {file_path}") self.upload_documents(doc_files=[file_path]) return "Document uploaded and processed successfully." return message def validate_urls(self, urls: List[str]) -> List[str]: """ Validate and filter URL list for document upload. Args: urls: List of URLs to validate Returns: List of valid URLs """ valid_urls = [] url_pattern = re.compile( r'^https?://' # http:// or https:// r'(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+[A-Z]{2,6}\.?|' # domain... r'localhost|' # localhost... r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})' # ...or ip r'(?::\d+)?' # optional port r'(?:/?|[/?]\S+)$', re.IGNORECASE) for url in urls: if url_pattern.match(url): valid_urls.append(url) else: logger.info(f"Invalid URL skipped: {url}") return valid_urls def validate_file_paths(self, file_paths: List[str]) -> List[str]: """ Validate and filter file path list for document upload. Args: file_paths: List of file paths to validate Returns: List of valid file paths that exist """ valid_paths = [] for path in file_paths: if os.path.exists(path): valid_paths.append(path) else: logger.info(f"File not found, skipped: {path}") return valid_paths def upload_documents(self, web_urls: List[str] = None, doc_files: List[str] = None) -> int: """ Upload and index documents from web URLs and local files. This method processes documents through the vector store for semantic search and retrieval during legal consultations. Args: web_urls: List of web URLs to scrape and index doc_files: List of local file paths to index Returns: Number of documents successfully uploaded """ web_urls = web_urls or [] doc_files = doc_files or [] # Validate inputs web_urls = self.validate_urls(web_urls) doc_files = self.validate_file_paths(doc_files) documents_added = 0 try: # Index web pages if any if web_urls: logger.info(f"Indexing {len(web_urls)} web pages...") web_results = self.database_manager.add_web_documents(web_urls) documents_added += len(web_results) logger.info(f"Indexed {len(web_results)} web documents") # Index local files if any if doc_files: logger.info(f"Indexing {len(doc_files)} local files...") file_results = self.database_manager.add_file_documents(doc_files) documents_added += len(file_results) logger.info(f"Indexed {len(file_results)} file documents") # Save the updated vector store if documents_added > 0: self.database_manager.case_vector_store_manager.save_vector_store() logger.info(f"Successfully uploaded and indexed {documents_added} documents") else: logger.info("No valid documents found to upload") except Exception as e: logger.info(f"Error uploading documents: {str(e)}") return documents_added # # RAG QUESTION ANSWERING METHODS # def ask_rag_question(self, question: str, uuids_for_retrieval: List[str] = None, included_web_pages: List[str] = None) -> Optional[Any]: """ Process a legal question through the RAG pipeline. This method orchestrates the entire question-answering process, including document retrieval, context preparation, and answer generation. Args: question: The user's legal question uuids_for_retrieval: Specific document UUIDs to include in retrieval included_web_pages: Web pages to include in context Returns: Chat interaction object with question and answer, or None if failed """ uuids_for_retrieval = uuids_for_retrieval or [] included_web_pages = included_web_pages or [] try: # Prepare conversation history context if not self.database_manager.conversation_history: history_text = "This is the first question in this conversation." else: # Use last 3 interactions for context recent_history = list(self.database_manager.conversation_history)[-3:] history_text = "\n".join([f"Q: {q}\nA: {a}" for q, a in recent_history]) # Prepare client context client_notes = "" if self.database_manager.client and hasattr(self.database_manager.client, 'notes'): client_notes = f"Client: {self.database_manager.client.notes or ''}" # Prepare case context case_details = "" if self.database_manager.case: case_details = ( f"Case: {self.database_manager.case.name}, " f"Type: {getattr(self.database_manager.case, 'case_type', 'Unknown')}, " f"Jurisdiction: {self.database_manager.case.jurisdiction_code}" ) # Prepare state for RAG pipeline state = { "question": question, "case_details": case_details, "client_notes": client_notes, "conversation_history": history_text, "documents_to_retrieve": uuids_for_retrieval, "web_pages": included_web_pages, } # Run the RAG pipeline logger.info(f"Processing question: {question[:50]}...") result = self.graph.invoke(state) answer = result.get("generation", "Sorry, I couldn't generate an answer.") # Save the interaction to database if case exists if self.database_manager.case: chat_interaction = self.database_manager.add_case_chat_history(question, answer) logger.info("Question processed and saved to chat history") return chat_interaction else: logger.info("No case loaded - conversation not saved") return None except Exception as e: logger.info(f"Error processing question: {str(e)}") return None def execute_actions(self, question: str) -> Optional[Any]: """ Execute actions based on the user's question. This method handles special commands and routes questions to appropriate processing methods. Args: question: User's input question or command Returns: Chat interaction object or None if processing failed """ # Handle document upload commands processed_message = self._handle_documents(question) if processed_message != question: # Document command was processed return None # Check for special commands commands = self.match_commands(question) if commands: logger.info(f"Executing command: {commands}") # Handle commands here if needed # Process as regular RAG question return self.ask_rag_question(question) def match_commands(self, question: str) -> Optional[str]: """ Match question against known command patterns. Args: question: Input question to check for commands Returns: Matched command string or None """ # Define command patterns patterns = { r'#summary': 'GENERATE_SUMMARY', r'#doc': 'DOCUMENT_SEARCH', r'#web': 'WEB_SEARCH', r'#save': 'SAVE_SESSION', } for pattern, command in patterns.items(): if re.search(pattern, question, re.IGNORECASE): return command return None def clear_conversation_history(self) -> None: """ Clear the current conversation history. This removes all cached conversation history but does not affect the persistent database records. """ self.database_manager.conversation_history.clear() logger.info("Conversation history cleared") def format_conversation_history(self) -> str: """ Format conversation history for display or context. Returns: Formatted string of recent conversation history """ if not self.database_manager.conversation_history: return "No conversation history available." formatted_history = [] for i, (question, answer) in enumerate(self.database_manager.conversation_history, 1): formatted_history.append(f"Q{i}: {question}") formatted_history.append(f"A{i}: {answer}") formatted_history.append("") # Empty line for readability return "\n".join(formatted_history) # # LEGAL RESEARCH AND ADVICE METHODS # def generate_passive_legal_information(self) -> Dict[str, Any]: """ Generate contextual legal information and resources. This method provides general legal guidance that may be relevant to the current case and conversation context. It combines AI-generated advice with relevant web resources. Returns: Dictionary containing advice, resources, and search results """ try: # Generate AI advice based on conversation history passive_legal_generator = PassiveLegalAdvice(self.llm) formatted_conversation_history = self.format_conversation_history() advice_result = passive_legal_generator.invoke({ 'conversation_history': formatted_conversation_history }) advice = advice_result.content if hasattr(advice_result, 'content') else str(advice_result) # Search for relevant web resources query = "legal information general guidance" if self.database_manager.case and hasattr(self.database_manager.case, 'case_type'): case_type = getattr(self.database_manager.case, 'case_type', '') if case_type: query = f"legal information {case_type}" # Perform web search for relevant resources links = [] titles = [] search_results = self.rag_pipeline.web_search({"question": query}) for doc in search_results.get("documents", []): if hasattr(doc, 'metadata'): url = doc.metadata.get("url", "") title = doc.metadata.get("title", "Legal Resource") if url: links.append(url) titles.append(title) return { "advice": advice, "links": links, "link_titles": titles } except Exception as e: logger.info(f"Error generating passive legal information: {str(e)}") return { "advice": "Please consult with a qualified legal professional for specific legal advice.", "links": ["https://www.canlii.org/", "https://laws-lois.justice.gc.ca/", "https://www.justice.gc.ca/"], "link_titles": [ "CanLII - Canadian Legal Information Institute", "Justice Laws Website - Government of Canada", "Department of Justice Canada"] } def set_legal_references(self, top_k: Optional[int] = 10) -> None: """ Set up legal references for the current case by searching CanLII. This method automatically searches Canadian legal databases for relevant case law and legislation based on the current case details. The AI evaluates and filters the most relevant legal precedents. Args: top_k: Maximum number of references to retrieve (default: use system default) """ if not self.database_manager.case: logger.info("No case is currently selected. Cannot set legal references.") return try: # Prepare case details for search case_details = self.database_manager.case.name if hasattr(self.database_manager.case, 'notes') and self.database_manager.case.notes: case_details += "\n" + self.database_manager.case.notes # Generate search query using AI query_writer = CanliiQueryWriter(self.llm) query_result = query_writer.invoke({"case_details": case_details}) search_query = query_result.query if hasattr(query_result, 'query') else str(query_result) logger.info(f"Searching CanLII with query: {search_query}") # Search for both legislation and case law for search_type in ['LEGISLATION', 'CASE']: try: logger.info(f"Searching for relevant {search_type.lower()}...") results = search_canlii( jurisdiction_code=self.database_manager.case.jurisdiction_code, search_term=search_query, search_type=search_type, top_k=top_k ) if results: logger.info(f"Found {len(results)} {search_type.lower()} results") # Use AI to evaluate relevance prompt = { "case_details": case_details, "reference_list": [r.get('description', r.get('title', '')) for r in results] } if search_type == 'CASE': logger.info("Evaluating and filtering relevant cases...") relevance_evaluator = CaseSimilaritySearch(llm=self.llm) else: logger.info("Evaluating and filtering relevant legislation...") relevance_evaluator = LegislationSimilaritySearch(llm=self.llm) relevant_items = relevance_evaluator.invoke(prompt=prompt).items # Index relevant references for retrieval if relevant_items: reference_urls = [] for i, item in enumerate(relevant_items): if item.binary_score == '1': url = results[i].get('path', '') if url: reference_urls.append(url) reference_urls = [results[0].get('path', '')] if reference_urls: logger.info(f"Indexing {len(reference_urls)} reference documents...") self.database_manager.add_file_documents( reference_urls, source_type="legal_reference" ) else: logger.info(f"No {search_type.lower()} results found") except Exception as search_error: logger.info(f"Error searching {search_type}: {search_error}") logger.info("Legal references setup completed") except Exception as e: logger.info(f"Error setting legal references: {str(e)}") def get_chat_summary(self) -> str: """ Generate a summary of the current chat conversation. Returns: Formatted summary of the conversation history """ if not self.database_manager.conversation_history: return "No conversation history available to summarize." try: # Format conversation for summarization chat_text = "\n".join([ f"Q: {q}\nA: {a}" for q, a in self.database_manager.conversation_history ]) # Generate summary using AI summary_result = self.chat_summarizer.invoke(chat_text) summary = getattr(summary_result, 'summary', str(summary_result)) return f"## Chat Summary\n\n{summary}" except Exception as e: logger.info(f"Error generating chat summary: {str(e)}") return f"Error generating summary: {str(e)}"