Spaces:
Running
Running
| # query_classifier.py | |
| import google.generativeai as genai | |
| import json | |
| import os | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| genai.configure(api_key=os.getenv('GOOGLE_API_KEY')) | |
| class QueryClassifier: | |
| def __init__(self): | |
| self.model = genai.GenerativeModel('gemini-2.5-flash-lite') | |
| self.process_mapping = { | |
| 'legal_aid': 'Legal Aid Application', | |
| 'consumer_complaint': 'Consumer Complaint Filing', | |
| 'court_status': 'Court Case Status Check', | |
| 'general': 'General Legal Query' | |
| } | |
| def classify_query(self, user_query: str) -> dict: | |
| """Classify user query to determine legal process and current step""" | |
| classification_prompt = f""" | |
| Analyze this user query and determine: | |
| 1. Which legal process they're asking about | |
| 2. What step they might be at in the process | |
| 3. Whether they need step-by-step guidance or just information | |
| Query: "{user_query}" | |
| Available processes: | |
| - legal_aid: Questions about getting free legal help | |
| - consumer_complaint: Filing complaints about products/services | |
| - court_status: Checking case status or court procedures | |
| - general: General legal questions not fitting above categories | |
| Possible steps for legal_aid: | |
| - check_eligibility: Asking about who qualifies | |
| - gather_documents: Asking what documents needed | |
| - choose_application_method: How to apply | |
| - online_application: Help with online form | |
| - wait_processing: Application submitted, waiting for response | |
| - contact_lawyer: Have lawyer assigned, need next steps | |
| Return ONLY a JSON object with this format: | |
| {{ | |
| "process": "process_name", | |
| "step": "step_id_or_null", | |
| "intent": "process_guidance|information_only", | |
| "confidence": 0.8, | |
| "reasoning": "Brief explanation" | |
| }} | |
| """ | |
| try: | |
| response = self.model.generate_content(classification_prompt) | |
| # Extract JSON from response | |
| response_text = response.text.strip() | |
| if response_text.startswith('```'): | |
| response_text = response_text[7:-3].strip() | |
| elif response_text.startswith('```'): | |
| response_text = response_text[3:-3].strip() | |
| result = json.loads(response_text) | |
| # Validate result | |
| if result.get('process') not in self.process_mapping: | |
| result['process'] = 'general' | |
| result['process_name'] = self.process_mapping.get(result['process']) | |
| return result | |
| except Exception as e: | |
| print(f"Classification error: {e}") | |
| return { | |
| "process": "general", | |
| "step": None, | |
| "intent": "information_only", | |
| "confidence": 0.3, | |
| "reasoning": "Classification failed, defaulting to general", | |
| "process_name": "General Legal Query" | |
| } | |
| # Test the classifier | |
| if __name__ == "__main__": | |
| classifier = QueryClassifier() | |
| test_queries = [ | |
| "I need free legal help but don't know if I qualify", | |
| "My landlord is trying to evict me illegally", | |
| "How do I check my court case status online?", | |
| "I bought a defective phone and the company won't refund" | |
| ] | |
| for query in test_queries: | |
| result = classifier.classify_query(query) | |
| print(f"\nQuery: {query}") | |
| print(f"Classification: {json.dumps(result, indent=2)}") | |