tools/user_management_tool.py

from typing import Any, Dict, Optional
from smolagents.tools import Tool
import json
import os

class UserManagementTool(Tool):
    name = "user_management"
    description = "Manages user information for the car sharing system."
    inputs = {
        'action': {'type': 'string', 'description': 'The action to perform (get, set, list)'},
        'user_name': {'type': 'string', 'description': 'Name of the user', 'nullable': True},
        'user_info': {'type': 'any', 'description': 'User information to store (for set action)', 'nullable': True}
    }
    output_type = "any"

    def __init__(self, users_file="car_sharing_users.json"):
        self.is_initialized = True
        self.users_file = users_file
        self.users = self._load_users()
        self.current_user = None
    
    def _load_users(self) -> Dict[str, Dict[str, Any]]:
        """Load users from file or create an empty dict if file doesn't exist."""
        if os.path.exists(self.users_file):
            try:
                with open(self.users_file, 'r') as f:
                    return json.load(f)
            except Exception:
                return {}
        return {}
    
    def _save_users(self) -> bool:
        """Save users to file."""
        try:
            with open(self.users_file, 'w') as f:
                json.dump(self.users, f, indent=2)
            return True
        except Exception:
            return False
    
    def forward(self, action: str, user_name: Optional[str] = None, user_info: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """
        Manage user information.
        
        Args:
            action: The action to perform (get, set, list, set_current)
            user_name: Name of the user
            user_info: User information to store (for set action)
            
        Returns:
            A dictionary with the result of the operation
        """
        try:
            action = action.lower()
            
            if action == "list":
                # List all users
                return {
                    "success": True,
                    "action": "list",
                    "users": list(self.users.keys()),
                    "current_user": self.current_user
                }
            
            elif action == "get":
                # Get user info
                if not user_name:
                    return {
                        "success": False,
                        "action": "get",
                        "error": "User name is required for 'get' action"
                    }
                
                if user_name not in self.users:
                    return {
                        "success": False,
                        "action": "get",
                        "error": f"User '{user_name}' not found"
                    }
                
                return {
                    "success": True,
                    "action": "get",
                    "user_name": user_name,
                    "user_info": self.users[user_name]
                }
            
            elif action == "set":
                # Set user info
                if not user_name:
                    return {
                        "success": False,
                        "action": "set",
                        "error": "User name is required for 'set' action"
                    }
                
                if user_info is None:
                    user_info = {}
                
                # Create or update user
                if user_name in self.users:
                    self.users[user_name].update(user_info)
                else:
                    self.users[user_name] = user_info
                
                self._save_users()
                
                return {
                    "success": True,
                    "action": "set",
                    "user_name": user_name,
                    "user_info": self.users[user_name]
                }
            
            elif action == "set_current":
                # Set current user
                if not user_name:
                    return {
                        "success": False,
                        "action": "set_current",
                        "error": "User name is required for 'set_current' action"
                    }
                
                # Create user if it doesn't exist
                if user_name not in self.users:
                    self.users[user_name] = {}
                    self._save_users()
                
                self.current_user = user_name
                
                return {
                    "success": True,
                    "action": "set_current",
                    "user_name": user_name,
                    "current_user": self.current_user
                }
            
            else:
                return {
                    "success": False,
                    "error": f"Unknown action: {action}"
                }
        
        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }