Spaces:
Running
Running
File size: 3,694 Bytes
8e453ef 1a1eddd 8e453ef |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
# 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)}")
|