TanujInsane/document-classification-env

1import pickle2import os3from datetime import datetime4 5CATEGORIES = {6    "easy":   {0:"General", 1:"Billing", 2:"Support", 3:"Technical", 4:"HR"},7    "medium": {0:"General", 1:"Billing", 2:"Support", 3:"Technical", 4:"HR",8               5:"Legal", 6:"Sales", 7:"Marketing", 8:"Operations", 9:"Complaints"},9}10 11CATEGORY_OVERRIDE = {12    "Technical":  ["crash", "crashes", "bug", "error", "not working", "broken", "upload fail",13                   "application", "software", "login fail", "404", "500", "exception",14                   "loading", "slow", "freeze", "hang", "install", "update failed", "glitch"],15    "Billing":    ["invoice", "charge", "charged", "payment", "refund", "subscription", "billing",16                   "amount", "price", "cost", "fee", "receipt", "transaction", "overcharged"],17    "HR":         ["leave", "salary", "policy", "maternity", "paternity", "holiday", "resignation",18                   "offer letter", "payslip", "appraisal", "human resources", "benefits"],19    "Complaints": ["complaint", "unacceptable", "disgusted", "terrible", "worst", "rude", "awful",20                   "disappointed", "frustrated", "escalate"],21    "Legal":      ["legal", "lawsuit", "compliance", "gdpr", "contract", "terms", "court",22                   "attorney", "lawyer", "liability", "copyright", "trademark"],23}24 25PRIORITY_RULES = {26    "urgent": ["urgent", "asap", "immediately", "critical", "emergency", "not working",27               "broken", "down", "error", "failed", "crashes", "crash"],28    "high":   ["incorrect", "wrong", "missing", "cannot", "can't", "unable",29               "problem", "issue", "complaint"],30    "medium": ["review", "check", "update", "change", "request", "need", "want", "please"],31    "low":    ["inquiry", "question", "info", "information", "general", "curious", "wondering"]32}33 34DEPARTMENT_MAP = {35    "Billing":    {"team": "Billing & Finance",   "email": "billing@company.com",    "sla": "24h"},36    "Technical":  {"team": "Tech Support",         "email": "techsupport@company.com","sla": "4h"},37    "HR":         {"team": "Human Resources",      "email": "hr@company.com",         "sla": "48h"},38    "Support":    {"team": "Customer Support",     "email": "support@company.com",    "sla": "12h"},39    "General":    {"team": "General Enquiries",    "email": "hello@company.com",      "sla": "48h"},40    "Legal":      {"team": "Legal & Compliance",   "email": "legal@company.com",      "sla": "72h"},41    "Sales":      {"team": "Sales",                "email": "sales@company.com",      "sla": "6h"},42    "Marketing":  {"team": "Marketing",            "email": "marketing@company.com",  "sla": "48h"},43    "Operations": {"team": "Operations",           "email": "ops@company.com",        "sla": "24h"},44    "Complaints": {"team": "Customer Relations",   "email": "complaints@company.com", "sla": "8h"},45}46 47REPLY_TEMPLATES = {48    "Billing": """Dear Customer,49 50Thank you for contacting us regarding your billing inquiry.51 52We have received your request and our Billing & Finance team will review your account details within 24 hours.53 54What we'll do:55  - Review your invoice/payment details56  - Verify the transaction in question57  - Send you a detailed resolution58 59Reference ID: {ref_id}60Priority: {priority}61Expected resolution: {sla}62 63If urgent, call: +1-800-BILLING64 65Best regards,66Billing & Finance Team""",67 68    "Technical": """Dear Customer,69 70Thank you for reporting this technical issue.71 72Our Tech Support team has been notified and will respond within 4 hours.73 74Immediate steps you can try:75  - Clear browser cache and retry76  - Check our status page: status.company.com77  - Restart the application78 79Reference ID: {ref_id}80Priority: {priority}81Expected resolution: {sla}82 83For critical issues: techsupport@company.com84 85Best regards,86Tech Support Team""",87 88    "HR": """Dear Team Member,89 90Thank you for reaching out to Human Resources.91 92Your query has been received and will be handled with full confidentiality.93 94Our HR team will respond within 48 hours regarding:95  - Policy clarifications96  - Benefits & compensation97  - Workplace concerns98 99Reference ID: {ref_id}100Priority: {priority}101Expected resolution: {sla}102 103Best regards,104Human Resources Team""",105 106    "Support": """Dear Customer,107 108Thank you for contacting Customer Support.109 110We have logged your request and a support agent will be in touch within 12 hours.111 112Reference ID: {ref_id}113Priority: {priority}114Expected resolution: {sla}115 116Track your ticket: support.company.com/track/{ref_id}117 118Best regards,119Customer Support Team""",120 121    "General": """Dear Customer,122 123Thank you for your inquiry.124 125We have received your message and will respond within 48 hours.126 127Reference ID: {ref_id}128Priority: {priority}129 130Best regards,131General Enquiries Team""",132 133    "Complaints": """Dear Customer,134 135We sincerely apologize for your experience.136 137Your complaint has been escalated to our Customer Relations team who will personally handle this within 8 hours.138 139Reference ID: {ref_id}140Priority: {priority}141Expected resolution: {sla}142 143We take all complaints seriously and are committed to resolving this promptly.144 145Best regards,146Customer Relations Team""",147 148    "Sales": """Dear Customer,149 150Thank you for your interest!151 152Our Sales team will reach out within 6 hours with a personalized proposal.153 154Reference ID: {ref_id}155Priority: {priority}156 157Best regards,158Sales Team""",159 160    "Legal": """Dear Customer,161 162Thank you for contacting our Legal & Compliance team.163 164Your matter will be reviewed by our legal team within 72 hours.165 166Reference ID: {ref_id}167Priority: {priority}168Expected resolution: {sla}169 170Best regards,171Legal & Compliance Team""",172 173    "Marketing": """Dear Customer,174 175Thank you for reaching out to Marketing.176 177We will respond within 48 hours.178 179Reference ID: {ref_id}180Priority: {priority}181 182Best regards,183Marketing Team""",184 185    "Operations": """Dear Customer,186 187Thank you for contacting Operations.188 189Your request has been logged and will be addressed within 24 hours.190 191Reference ID: {ref_id}192Priority: {priority}193Expected resolution: {sla}194 195Best regards,196Operations Team""",197}198 199 200class TicketAgent:201    def __init__(self, model_dir="."):202        self.models = {}203        self.model_dir = model_dir204        self._load()205 206    def _load(self):207        for d in ["easy", "medium", "hard"]:208            p = os.path.join(self.model_dir, f"model_{d}.pkl")209            if os.path.exists(p):210                with open(p, "rb") as f:211                    self.models[d] = pickle.load(f)212 213    def _name(self, c, diff):214        m = CATEGORIES.get(diff, {})215        try:216            return m.get(int(c), str(c))217        except:218            return str(c)219 220    def classify(self, text, difficulty="easy"):221        text = text.lower()222        223        for cat, kw_list in CATEGORY_OVERRIDE.items():224            if any(kw in text for kw in kw_list):225                if difficulty in self.models:226                    model = self.models[difficulty]227                    classes = model.classes_228                    proba = model.predict_proba([text])[0]229                    scores = {self._name(c, difficulty): round(float(p), 4) for c, p in zip(classes, proba)}230                else:231                    scores = {cat: 1.0}232                scores[cat] = max(scores.get(cat, 0), 0.85)233                return cat, scores234        235        if difficulty not in self.models:236            return "General", {"General": 1.0}237        model = self.models[difficulty]238        classes = model.classes_239        proba = model.predict_proba([text])[0]240        pred = classes[proba.argmax()]241        top = self._name(pred, difficulty)242        scores = {self._name(c, difficulty): round(float(p), 4) for c, p in zip(classes, proba)}243        return top, scores244 245    def get_priority(self, text):246        text = text.lower()247        for pri, kw_list in PRIORITY_RULES.items():248            if any(kw in text for kw in kw_list):249                return pri250        return "low"251 252    def route(self, category):253        return DEPARTMENT_MAP.get(category, DEPARTMENT_MAP["General"])254 255    def generate_reply(self, category, priority, ref_id):256        dept = self.route(category)257        tmpl = REPLY_TEMPLATES.get(category, REPLY_TEMPLATES["General"])258        return tmpl.format(ref_id=ref_id, priority=priority.upper(), sla=dept["sla"])259 260    def process_ticket(self, text, difficulty="easy", ticket_id=None):261        ref_id = ticket_id or f"TKT-{datetime.now().strftime('%Y%m%d%H%M%S')}"262        cat, scores = self.classify(text, difficulty)263        pri = self.get_priority(text)264        dept = self.route(cat)265        reply = self.generate_reply(cat, pri, ref_id)266        top3 = sorted(scores.items(), key=lambda x: -x[1])[:3]267        return {268            "ref_id": ref_id,269            "category": cat,270            "priority": pri,271            "department": dept["team"],272            "email": dept["email"],273            "sla": dept["sla"],274            "confidence": scores.get(cat, 0.0),275            "top3": top3,276            "reply": reply,277            "timestamp": datetime.now().isoformat()278        }279