File size: 14,443 Bytes
d66ab65
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
import os
from flask import Blueprint, request, render_template, jsonify, current_app

# Import services
from .services.tokenizer_service import tokenizer_service
from .services.file_service import file_service
from .utils.validators import validators, ValidationError

# Create Blueprint
main_bp = Blueprint('main', __name__)


@main_bp.route('/tokenizer-info', methods=['GET'])
def tokenizer_info():
    """Endpoint to get tokenizer information without processing text."""
    model_id = request.args.get('model_id', '')
    is_custom = request.args.get('is_custom', 'false').lower() == 'true'
    
    if not model_id:
        return jsonify({"error": "No model ID provided"}), 400
    
    try:
        # Validate custom model path if it's a custom model
        if is_custom:
            try:
                validators.validate_model_path(model_id)
            except ValidationError as e:
                return jsonify({"error": str(e)}), 400
            
        # For predefined models, use the model name from the dictionary
        if not is_custom and tokenizer_service.is_predefined_model(model_id):
            model_id_or_name = model_id
        else:
            # For custom models, use the model ID directly
            model_id_or_name = model_id
            
        # Load the tokenizer and get info
        tokenizer, info, error = tokenizer_service.load_tokenizer(model_id_or_name)
        
        if error:
            return jsonify({"error": error}), 400
            
        return jsonify(info)
    except Exception as e:
        return jsonify({"error": f"Failed to get tokenizer info: {str(e)}"}), 500

@main_bp.route('/', methods=['GET', 'POST'])
def index():
    text = ""
    token_data = None
    error_message = ""
    selected_model = request.args.get('model', request.form.get('model', 'qwen3'))
    custom_model = request.args.get('custom_model', request.form.get('custom_model', ''))
    model_type = request.args.get('model_type', request.form.get('model_type', 'predefined'))
    
    # Determine which model to use based on model_type
    model_to_use = selected_model if model_type == 'predefined' else custom_model
    
    if request.method == 'POST':
        # Check if file upload
        if 'file' in request.files and request.files['file'].filename:
            uploaded_file = request.files['file']
            
            try:
                # Validate file
                validators.validate_filename(uploaded_file.filename)
                validators.validate_file_extension(uploaded_file.filename, file_service.ALLOWED_EXTENSIONS)
                
                # Validate custom model if needed
                if model_type == 'custom' and custom_model:
                    validators.validate_model_path(custom_model)
                
                # Save file securely
                file_path = file_service.save_uploaded_file(uploaded_file, current_app.config['UPLOAD_FOLDER'])
                
                # Read a small preview of the file
                preview_char_limit = current_app.config.get('PREVIEW_CHAR_LIMIT', 8096)
                with open(file_path, 'r', errors='replace') as f:
                    text = f.read(preview_char_limit)
                
                try:
                    # Process the file using file service
                    token_data = file_service.process_file_for_tokenization(
                        file_path=file_path,
                        model_id_or_name=model_to_use,
                        preview_char_limit=preview_char_limit,
                        max_display_tokens=current_app.config.get('MAX_DISPLAY_TOKENS', 50000),
                        chunk_size=current_app.config.get('CHUNK_SIZE', 1024 * 1024)
                    )
                    
                    # Clean up the file after processing
                    file_service.cleanup_file(file_path)
                    
                    # If request is AJAX, return JSON
                    if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                        return jsonify(token_data)
                        
                except Exception as e:
                    error_message = str(e)
                    file_service.cleanup_file(file_path)
                    
                    if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                        return jsonify({"error": error_message}), 400
                    return render_template(
                        'index.html',
                        text=text,
                        token_data=None,
                        models=tokenizer_service.TOKENIZER_MODELS,
                        selected_model=selected_model,
                        custom_model=custom_model,
                        model_type=model_type,
                        error=error_message
                    )
                    
            except ValidationError as e:
                error_message = str(e)
                if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                    return jsonify({"error": error_message}), 400
                return render_template(
                    'index.html',
                    text="",
                    token_data=None,
                    models=tokenizer_service.TOKENIZER_MODELS,
                    selected_model=selected_model,
                    custom_model=custom_model,
                    model_type=model_type,
                    error=error_message
                )

        # Regular text processing
        else:
            text = request.form.get('text', '')
            if text:
                try:
                    # Validate text input
                    validators.validate_text_input(text)
                    
                    # Validate custom model if needed
                    if model_type == 'custom' and custom_model:
                        validators.validate_model_path(custom_model)
                    
                    # Process text using file service
                    token_data = file_service.process_text_for_tokenization(
                        text=text,
                        model_id_or_name=model_to_use,
                        preview_char_limit=current_app.config.get('PREVIEW_CHAR_LIMIT', 8096),
                        max_display_tokens=current_app.config.get('MAX_DISPLAY_TOKENS', 50000)
                    )
                    
                    # If request is AJAX, return JSON
                    if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                        return jsonify(token_data)
                        
                except ValidationError as e:
                    error_message = str(e)
                    if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                        return jsonify({"error": error_message}), 400
                    return render_template(
                        'index.html',
                        text=text,
                        token_data=None,
                        models=tokenizer_service.TOKENIZER_MODELS,
                        selected_model=selected_model,
                        custom_model=custom_model,
                        model_type=model_type,
                        error=error_message
                    )
                except Exception as e:
                    error_message = str(e)
                    if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                        return jsonify({"error": error_message}), 400
                    return render_template(
                        'index.html',
                        text=text,
                        token_data=None,
                        models=tokenizer_service.TOKENIZER_MODELS,
                        selected_model=selected_model,
                        custom_model=custom_model,
                        model_type=model_type,
                        error=error_message
                    )
    
    return render_template(
        'index.html',
        text=text,
        token_data=token_data,
        models=tokenizer_service.TOKENIZER_MODELS,
        selected_model=selected_model,
        custom_model=custom_model,
        model_type=model_type,
        error=error_message
    )


@main_bp.route('/health', methods=['GET'])
def health_check():
    """Basic health check endpoint."""
    import time
    import psutil
    from flask import __version__ as flask_version
    
    try:
        # Basic application status
        status = {
            "status": "healthy",
            "timestamp": int(time.time()),
            "version": "1.0.0",
            "flask_version": flask_version,
            "uptime": int(time.time()),  # Simple uptime since this request
        }
        
        return jsonify(status), 200
    except Exception as e:
        return jsonify({
            "status": "unhealthy",
            "error": str(e),
            "timestamp": int(time.time())
        }), 500


@main_bp.route('/health/detailed', methods=['GET'])
def detailed_health_check():
    """Detailed health check with system and service status."""
    import time
    import psutil
    import os
    from flask import __version__ as flask_version
    
    try:
        # System information
        cpu_percent = psutil.cpu_percent(interval=1)
        memory = psutil.virtual_memory()
        disk = psutil.disk_usage('/')
        
        # Check tokenizer service
        tokenizer_status = "healthy"
        tokenizer_cache_size = len(tokenizer_service.tokenizers) + len(tokenizer_service.custom_tokenizers)
        
        # Test basic tokenizer loading
        try:
            test_tokenizer, _, error = tokenizer_service.load_tokenizer('gpt2')
            if error:
                tokenizer_status = f"warning: {error}"
        except Exception as e:
            tokenizer_status = f"error: {str(e)}"
        
        # Check upload directory
        upload_folder = current_app.config.get('UPLOAD_FOLDER', '/tmp')
        upload_dir_exists = os.path.exists(upload_folder)
        upload_dir_writable = os.access(upload_folder, os.W_OK) if upload_dir_exists else False
        
        status = {
            "status": "healthy",
            "timestamp": int(time.time()),
            "version": "1.0.0",
            "flask_version": flask_version,
            "system": {
                "cpu_percent": round(cpu_percent, 1),
                "memory": {
                    "total": memory.total,
                    "available": memory.available,
                    "percent": memory.percent,
                    "used": memory.used
                },
                "disk": {
                    "total": disk.total,
                    "used": disk.used,
                    "free": disk.free,
                    "percent": round((disk.used / disk.total) * 100, 1)
                }
            },
            "services": {
                "tokenizer_service": {
                    "status": tokenizer_status,
                    "cached_tokenizers": tokenizer_cache_size,
                    "available_models": len(tokenizer_service.TOKENIZER_MODELS)
                },
                "file_service": {
                    "upload_directory": upload_folder,
                    "directory_exists": upload_dir_exists,
                    "directory_writable": upload_dir_writable,
                    "allowed_extensions": list(file_service.ALLOWED_EXTENSIONS)
                }
            },
            "configuration": {
                "max_content_length": current_app.config.get('MAX_CONTENT_LENGTH'),
                "cache_expiration": current_app.config.get('CACHE_EXPIRATION', 3600),
                "max_display_tokens": current_app.config.get('MAX_DISPLAY_TOKENS', 50000),
                "preview_char_limit": current_app.config.get('PREVIEW_CHAR_LIMIT', 8096)
            }
        }
        
        # Determine overall status
        overall_status = "healthy"
        if tokenizer_status.startswith("error"):
            overall_status = "unhealthy"
        elif tokenizer_status.startswith("warning") or not upload_dir_writable:
            overall_status = "degraded"
        
        status["status"] = overall_status
        
        return jsonify(status), 200 if overall_status == "healthy" else 503
        
    except Exception as e:
        return jsonify({
            "status": "unhealthy",
            "error": str(e),
            "timestamp": int(time.time())
        }), 500


@main_bp.route('/health/ready', methods=['GET'])
def readiness_check():
    """Readiness check - determines if the application is ready to serve requests."""
    try:
        # Check if core services are ready
        checks = {
            "tokenizer_service": False,
            "file_service": False,
            "configuration": False
        }
        
        # Test tokenizer service
        try:
            test_tokenizer, _, error = tokenizer_service.load_tokenizer('gpt2')
            checks["tokenizer_service"] = error is None
        except:
            checks["tokenizer_service"] = False
        
        # Test file service
        try:
            upload_folder = current_app.config.get('UPLOAD_FOLDER', '/tmp')
            checks["file_service"] = os.path.exists(upload_folder) and os.access(upload_folder, os.W_OK)
        except:
            checks["file_service"] = False
        
        # Check configuration
        required_configs = ['MAX_CONTENT_LENGTH', 'UPLOAD_FOLDER']
        checks["configuration"] = all(current_app.config.get(config) is not None for config in required_configs)
        
        all_ready = all(checks.values())
        
        return jsonify({
            "ready": all_ready,
            "checks": checks,
            "timestamp": int(time.time())
        }), 200 if all_ready else 503
        
    except Exception as e:
        return jsonify({
            "ready": False,
            "error": str(e),
            "timestamp": int(time.time())
        }), 500