sail_scripts/model/category_manager.py

"""
Category Manager for domain-specific training and inference.
Tracks trained categories and enables category-aware responses.
"""
import json
import os
from typing import List, Dict, Optional
from datetime import datetime

class CategoryManager:
    """Manages training categories and metadata."""
    
    PREDEFINED_CATEGORIES = [
        "text",      # General text
        "code",      # Programming
        "video",     # Video descriptions/scripts
        "image",     # Image captions/analysis
        "math",      # Mathematics
        "history",   # Historical content
        "science",   # Scientific content
        "grammar",   # Language/grammar
        "custom"     # User-defined
    ]
    
    def __init__(self, metadata_path="category_metadata.json"):
        self.metadata_path = metadata_path
        self.categories = {}
        self.load()
    
    def load(self):
        """Load category metadata from disk."""
        if os.path.exists(self.metadata_path):
            with open(self.metadata_path, 'r') as f:
                self.categories = json.load(f)
        else:
            self.categories = {}
    
    def save(self):
        """Save category metadata to disk."""
        with open(self.metadata_path, 'w') as f:
            json.dump(self.categories, f, indent=2)
    
    def add_training(self, category: str, details: str, vocab_size: int):
        """Record a training session for a category."""
        category = category.lower()
        
        if category not in self.categories:
            self.categories[category] = {
                "training_count": 0,
                "sessions": [],
                "vocab_size": 0
            }
        
        self.categories[category]["training_count"] += 1
        self.categories[category]["vocab_size"] = vocab_size
        self.categories[category]["sessions"].append({
            "timestamp": datetime.now().isoformat(),
            "details": details
        })
        
        self.save()
    
    def get_categories(self) -> List[str]:
        """Get list of all trained categories."""
        return list(self.categories.keys())
    
    def get_category_info(self, category: str) -> Optional[Dict]:
        """Get information about a specific category."""
        return self.categories.get(category.lower())
    
    def detect_category(self, text: str) -> str:
        """
        Detect the most likely category from user input.
        Uses keyword matching for now, can be enhanced with ML.
        """
        text_lower = text.lower()
        
        # Keyword-based detection
        category_keywords = {
            "math": ["calculate", "equation", "derivative", "integral", "algebra", "geometry", "math"],
            "code": ["function", "class", "variable", "python", "javascript", "code", "programming"],
            "history": ["war", "century", "historical", "ancient", "empire", "revolution"],
            "science": ["atom", "molecule", "physics", "chemistry", "biology", "experiment"],
            "grammar": ["noun", "verb", "sentence", "grammar", "syntax", "adjective"],
            "video": ["video", "scene", "frame", "footage", "clip"],
            "image": ["image", "picture", "photo", "visual", "pixel"]
        }
        
        # Count keyword matches
        scores = {}
        for category, keywords in category_keywords.items():
            if category in self.categories:  # Only consider trained categories
                score = sum(1 for keyword in keywords if keyword in text_lower)
                if score > 0:
                    scores[category] = score
        
        # Return category with highest score, or "text" as default
        if scores:
            return max(scores, key=scores.get)
        return "text"
    
    def get_summary(self) -> str:
        """Get a formatted summary of all trained categories."""
        if not self.categories:
            return "No categories trained yet."
        
        summary = "Trained Categories:\n"
        for cat, info in self.categories.items():
            summary += f"  - {cat.capitalize()}: {info['training_count']} session(s), {info['vocab_size']} vocab\n"
        return summary

# Singleton instance
_category_manager = None

def get_category_manager():
    """Get or create the singleton CategoryManager instance."""
    global _category_manager
    if _category_manager is None:
        _category_manager = CategoryManager()
    return _category_manager