ai-LiaoWangweb-app/app/routes/deep_routes.py

from flask import Blueprint, render_template, request, jsonify
from flask_login import login_required
from app import db, create_app
from app.models import SpiderResult, DeepCollection, AIModel, TokenUsageLog
from sqlalchemy.exc import OperationalError
import time
import asyncio
import threading
import json
import requests
from crawl4ai import AsyncWebCrawler
import sys

if sys.platform == 'win32':
    asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())

bp = Blueprint('deep', __name__, url_prefix='/deep')

@bp.route('/dashboard')
@login_required
def dashboard():
    return render_template('deep_management.html')

@bp.route('/api/list')
@login_required
def list_data():
    page = request.args.get('page', 1, type=int)
    per_page = request.args.get('per_page', 10, type=int)
    query = request.args.get('query', '')
    
    q = DeepCollection.query
    
    if query:
        q = q.filter(DeepCollection.url.like(f'%{query}%') | DeepCollection.content.like(f'%{query}%'))
        
    pagination = q.order_by(DeepCollection.updated_at.desc()).paginate(page=page, per_page=per_page, error_out=False)
    
    items = []
    for item in pagination.items:
        items.append({
            'id': item.id,
            'title': item.title,
            'url': item.url,
            'summary': item.summary,
            'status': item.status,
            'created_at': item.created_at.strftime('%Y-%m-%d %H:%M:%S'),
            'updated_at': item.updated_at.strftime('%Y-%m-%d %H:%M:%S'),
            'content_length': len(item.content) if item.content else 0
        })
        
    return jsonify({
        'items': items,
        'total': pagination.total,
        'pages': pagination.pages,
        'current_page': page
    })

async def run_crawl(url):
    async with AsyncWebCrawler(verbose=True) as crawler:
        result = await crawler.arun(url=url)
        return result.markdown

def safe_commit(session, max_retries=5, delay=0.5):
    for i in range(max_retries):
        try:
            session.commit()
            return True
        except OperationalError as e:
            if "locked" in str(e):
                print(f"Database locked, retrying {i+1}/{max_retries}...")
                session.rollback()
                time.sleep(delay)
            else:
                raise e
    print("Database commit failed after max retries")
    return False

def generate_summary(content, model):
    if not model or not content:
        return None
        
    try:
        url = model.api_base
        if not url.endswith('/chat/completions'):
            if not url.endswith('/'):
                url += '/'
            url += 'chat/completions'

        headers = {
            "Authorization": f"Bearer {model.api_key}",
            "Content-Type": "application/json"
        }
        
        # Truncate content to avoid token limits
        truncated_content = content[:10000]
        
        payload = {
            "model": model.model_name,
            "messages": [
                {"role": "system", "content": "You are a helpful assistant. Please summarize the following web page content in Chinese within 200 words."},
                {"role": "user", "content": truncated_content}
            ],
            "stream": False,
            "max_tokens": 500,
            "temperature": 0.5
        }
        
        response = requests.post(url, json=payload, headers=headers, timeout=30)
        if response.status_code == 200:
            res_json = response.json()
            
            # Extract and log usage
            if 'usage' in res_json:
                usage = res_json['usage']
                print(f"Token Usage: {usage}")
                prompt_tokens = usage.get('prompt_tokens', 0)
                completion_tokens = usage.get('completion_tokens', 0)
                total_tokens = usage.get('total_tokens', 0)
                
                log = TokenUsageLog(
                    model_id=model.id,
                    prompt_tokens=prompt_tokens,
                    completion_tokens=completion_tokens,
                    total_tokens=total_tokens,
                    request_type='deep_collect'
                )
                model.total_tokens = (model.total_tokens or 0) + total_tokens
                db.session.add(log)
                safe_commit(db.session)
            else:
                print(f"Warning: No 'usage' field in AI response. Keys: {list(res_json.keys())}")
                
            return res_json['choices'][0]['message']['content']
        else:
            print(f"AI Summary Error: {response.text}")
            return None
    except Exception as e:
        print(f"AI Summary Exception: {e}")
        return None

def execute_deep_task(deep_id, app_context=None):
    """Background task for deep collection"""
    app = create_app()
    with app.app_context():
        try:
            deep_item = DeepCollection.query.get(deep_id)
            if not deep_item:
                return

            deep_item.status = 'running'
            deep_item.progress = 10
            deep_item.progress_msg = 'Initializing task...'
            safe_commit(db.session)
            
            url = deep_item.url
            
            # 1. Crawl
            deep_item.progress = 20
            deep_item.progress_msg = f'Connecting to {url}...'
            safe_commit(db.session)
            
            print(f"Starting crawl for {url}")
            try:
                markdown = asyncio.run(run_crawl(url))
            except Exception as e:
                deep_item.status = 'failed'
                deep_item.error_msg = f"Crawl failed: {str(e)}"
                deep_item.progress_msg = 'Crawl failed'
                safe_commit(db.session)
                return

            if not markdown:
                deep_item.status = 'failed'
                deep_item.error_msg = 'Failed to crawl content (empty result)'
                deep_item.progress_msg = 'Crawl returned empty'
                safe_commit(db.session)
                return

            deep_item.content = markdown
            deep_item.progress = 50
            deep_item.progress_msg = 'Content crawled successfully. Analyzing...'
            safe_commit(db.session)

            # 2. AI Summary
            model = AIModel.query.filter_by(is_active=True).first()
            summary = None
            if model:
                deep_item.progress = 60
                deep_item.progress_msg = f'Generating summary with {model.name}...'
                safe_commit(db.session)
                
                print(f"Generating summary with model {model.name}")
                summary = generate_summary(markdown, model)
                
                if summary:
                    deep_item.progress = 90
                    deep_item.progress_msg = 'Summary generated.'
                else:
                    deep_item.progress_msg = 'Summary generation skipped or failed.'
            else:
                deep_item.progress_msg = 'No active AI model, skipping summary.'
            
            # 3. Save final result
            deep_item.summary = summary
            deep_item.status = 'completed'
            deep_item.progress = 100
            deep_item.progress_msg = 'Deep collection completed successfully.'
            deep_item.error_msg = None
            safe_commit(db.session)
            
        except Exception as e:
            print(f"Deep task failed: {e}")
            # Re-query in case of session issues
            with app.app_context():
                 deep_item = DeepCollection.query.get(deep_id)
                 if deep_item:
                     deep_item.status = 'failed'
                     deep_item.error_msg = str(e)
                     deep_item.progress_msg = 'Internal error occurred'
                     safe_commit(db.session)

@bp.route('/api/collect', methods=['POST'])
@login_required
def deep_collect():
    data = request.json
    data_id = data.get('id') # SpiderResult ID
    
    source_data = SpiderResult.query.get(data_id)
    if not source_data:
        return jsonify({'error': 'Source data not found'}), 404
        
    url = source_data.link
    if not url:
        return jsonify({'error': 'No URL in source data'}), 400
        
    # Check or create DeepCollection record
    deep_item = DeepCollection.query.filter_by(url=url).first()
    if not deep_item:
        deep_item = DeepCollection(
            url=url,
            title=source_data.title,
            status='pending',
            progress=0,
            progress_msg='Queued...'
        )
        db.session.add(deep_item)
    else:
        # Reset for re-run
        deep_item.status = 'pending'
        deep_item.title = source_data.title # Update title in case it changed
        deep_item.progress = 0
        deep_item.progress_msg = 'Queued...'
        deep_item.error_msg = None
        
    # Update SpiderResult to show it has deep collection
    source_data.has_deep_collection = True
        
    db.session.commit()
    
    # Start background thread
    thread = threading.Thread(target=execute_deep_task, args=(deep_item.id,))
    thread.start()
    
    return jsonify({
        'message': 'Deep collection started',
        'task_id': deep_item.id,
        'status': 'started'
    })

@bp.route('/api/status/<int:id>', methods=['GET'])
@login_required
def check_status(id):
    item = DeepCollection.query.get(id)
    if not item:
        return jsonify({'error': 'Task not found'}), 404
        
    return jsonify({
        'id': item.id,
        'status': item.status,
        'progress': item.progress,
        'progress_msg': item.progress_msg,
        'error': item.error_msg,
        'summary': item.summary if item.status == 'completed' else None,
        'url': item.url
    })

@bp.route('/api/get/<int:id>')
@login_required
def get_deep_data(id):
    item = DeepCollection.query.get_or_404(id)
    return jsonify({
        'id': item.id,
        'url': item.url,
        'content': item.content,
        'summary': item.summary,
        'status': item.status,
        'updated_at': item.updated_at.strftime('%Y-%m-%d %H:%M:%S')
    })

@bp.route('/api/delete', methods=['POST'])
@login_required
def delete_deep_data():
    data = request.json
    ids = data.get('ids', [])
    if not ids:
         return jsonify({'error': 'No IDs provided'}), 400
         
    try:
        DeepCollection.query.filter(DeepCollection.id.in_(ids)).delete(synchronize_session=False)
        db.session.commit()
        return jsonify({'message': 'Deleted successfully'})
    except Exception as e:
        db.session.rollback()
        return jsonify({'error': str(e)}), 500