app.py

from flask import Flask, render_template, request, jsonify, Response
from flask_cors import CORS
import os
import time
from ss import SearchEngine
from bs4 import BeautifulSoup
from rag_qa_system import RAGQASystem
from llm_evaluator import LLMEvaluator
from recommendation_system import RecommendationSystem, SearchAnalytics
import json

app = Flask(__name__)
CORS
# 初始化搜索引擎
engine = SearchEngine()
engine_loaded = False

# RAG问答系统
rag_system = None

# LLM评价器
evaluator = None

# 推荐系统
recommendation_system = None

# 搜索统计
search_analytics = SearchAnalytics()

def init_search_engine():
    """初始化搜索引擎"""
    global engine_loaded, rag_system, evaluator, recommendation_system
    try:
        print("正在加载搜索引擎...")
        engine.load_index('search_index.pkl')
        engine.load_pagerank('pagerank_results.json')
        
        # 加载LTR模型（如果存在）
        if os.path.exists('ltr_model.pkl'):
            import pickle
            print("加载LTR模型...")
            with open('ltr_model.pkl', 'rb') as f:
                model_data = pickle.load(f)
                engine.ltr_model = model_data['model']
                engine.scaler = model_data['scaler']
            print("LTR模型加载完成")
        else:
            print("未找到LTR模型，将使用BM25排序")
        
        print("初始化RAG问答系统...")
        rag_system = RAGQASystem(engine)
        
        print("初始化LLM评价器...")
        evaluator = LLMEvaluator(engine)
        
        print("初始化推荐系统...")
        recommendation_system = RecommendationSystem(engine)
        
        print("预加载今日推荐和热门话题...")
        try:
            recommendation_system.get_daily_recommendations(top_k=10)
            recommendation_system.get_trending_topics(top_k=15)
            print("预加载完成！")
        except Exception as e:
            print(f"预加载失败（不影响使用）: {e}")
        
        engine_loaded = True
        print("搜索引擎初始化完成！")
        return True
    except Exception as e:
        print(f"搜索引擎初始化失败: {e}")
        return False

@app.route('/')
def index():
    """返回主页"""
    return render_template('index.html')

@app.route('/api/search', methods=['GET'])
def search():
    """搜索API接口（支持高级搜索选项）"""
    if not engine_loaded:
        return jsonify({
            'success': False,
            'error': '搜索引擎未初始化'
        }), 500
    
    query = request.args.get('q', '').strip()
    page = int(request.args.get('page', 1))
    page_size = int(request.args.get('size', 10))
    sort_by = request.args.get('sort', 'relevance')  # relevance, pagerank, length
    min_pagerank = float(request.args.get('min_pagerank', 0.0))
    diversity = float(request.args.get('diversity', 0.0))  # 0-1
    
    if not query:
        return jsonify({
            'success': False,
            'error': '查询内容不能为空'
        }), 400
    
    try:
        start_time = time.time()
        
        # 获取更多结果以支持分页
        total_results = page * page_size + 20
        
        # 使用高级搜索功能
        if sort_by != 'relevance' or min_pagerank > 0 or diversity > 0:
            results = engine.search_with_filters(
                query, 
                top_k=total_results, 
                use_ltr=True,
                sort_by=sort_by,
                min_pagerank=min_pagerank,
                diversity_factor=diversity
            )
        else:
            results = engine.search(query, top_k=total_results, use_ltr=True)
        
        search_time = time.time() - start_time
        
        # 记录搜索统计
        search_analytics.log_search(query, len(results), search_time)
        
        # 分页处理
        start_idx = (page - 1) * page_size
        end_idx = start_idx + page_size
        page_results = results[start_idx:end_idx]
        
        # 格式化结果
        formatted_results = []
        for idx, (url, score, title, html_content) in enumerate(page_results, start=start_idx + 1):
            # 从HTML中提取纯文本
            try:
                soup = BeautifulSoup(html_content, 'html.parser')
                # 移除script和style标签
                for script in soup(["script", "style"]):
                    script.decompose()
                # 获取纯文本
                text = soup.get_text()
                # 清理多余的空白字符
                text = ' '.join(text.split())
                preview = text[:200].strip()
                if len(text) > 200:
                    preview += '...'
            except Exception as e:
                # 如果解析失败，使用原始方法
                preview = html_content[:200].replace('\n', ' ').strip()
                if len(html_content) > 200:
                    preview += '...'
            
            # 获取PageRank分数
            pagerank_score = engine.pagerank_scores.get(url, 0.0)
            
            formatted_results.append({
                'rank': idx,
                'url': url,
                'title': title if title else url,
                'score': float(score),
                'pagerank': float(pagerank_score),
                'preview': preview
            })
        
        # 获取相关搜索建议
        related_searches = []
        if recommendation_system:
            related_searches = recommendation_system.get_related_searches(query, top_k=5)
        
        return jsonify({
            'success': True,
            'query': query,
            'total': len(results),
            'page': page,
            'page_size': page_size,
            'search_time': round(search_time, 3),
            'results': formatted_results,
            'related_searches': related_searches
        })
    
    except Exception as e:
        print(f"搜索出错: {e}")
        import traceback
        traceback.print_exc()
        return jsonify({
            'success': False,
            'error': f'搜索时发生错误: {str(e)}'
        }), 500

@app.route('/api/qa', methods=['POST'])
def question_answer():
    """RAG问答API接口（非流式）"""
    if not engine_loaded or not rag_system:
        return jsonify({
            'success': False,
            'error': 'RAG问答系统未初始化'
        }), 500
    
    data = request.get_json()
    question = data.get('question', '').strip()
    top_k = data.get('top_k', 5)
    search_mode = data.get('search_mode', 'smart')  # 搜索模式
    use_context_extract = data.get('use_context_extract', True)  # 是否使用上下文提取（聚焦模式）
    api_token = data.get('api_token', '').strip()
    model_name = data.get('model_name', '').strip()
    api_url = data.get('api_url', '').strip()
    
    if not question:
        return jsonify({
            'success': False,
            'error': '问题不能为空'
        }), 400
    
    try:
        print(f"\n收到问答请求: {question}")
        print(f"  搜索模式: {search_mode}")
        print(f"  上下文提取: {'启用' if use_context_extract else '禁用'}")
        if api_token:
            print("  使用自定义API Token")
            
        result = rag_system.answer_question(
            question, 
            top_k=top_k, 
            search_mode=search_mode, 
            use_context_extract=use_context_extract,
            api_token=api_token if api_token else None,
            model_name=model_name if model_name else None,
            api_url=api_url if api_url else None
        )
        return jsonify(result)
    
    except Exception as e:
        print(f"问答出错: {e}")
        return jsonify({
            'success': False,
            'error': f'问答时发生错误: {str(e)}'
        }), 500

@app.route('/api/qa/stream', methods=['POST'])
def question_answer_stream():
    """RAG问答API接口（流式输出）"""
    if not engine_loaded or not rag_system:
        return jsonify({
            'success': False,
            'error': 'RAG问答系统未初始化'
        }), 500
    
    data = request.get_json()
    question = data.get('question', '').strip()
    top_k = data.get('top_k', 5)
    search_mode = data.get('search_mode', 'smart')
    use_context_extract = data.get('use_context_extract', True)
    api_token = data.get('api_token', '').strip()
    model_name = data.get('model_name', '').strip()
    api_url = data.get('api_url', '').strip()
    
    if not question:
        return jsonify({
            'success': False,
            'error': '问题不能为空'
        }), 400
    
    def generate():
        try:
            print(f"\n收到流式问答请求: {question}")
            print(f"  搜索模式: {search_mode}")
            print(f"  上下文提取: {'启用' if use_context_extract else '禁用'}")
            if api_token:
                print("  使用自定义API Token")
            
            # 1. 检索相关文档
            retrieved_docs = rag_system.retrieve_context(
                question, 
                top_k=top_k, 
                search_mode=search_mode, 
                use_context_extract=use_context_extract
            )
            
            if not retrieved_docs:
                # 发送错误事件
                yield f"data: {json.dumps({'type': 'error', 'message': '没有找到相关文档'}, ensure_ascii=False)}\n\n"
                return
            
            # 2. 发送检索到的文档信息
            sources = [
                {
                    'rank': doc['rank'],
                    'title': doc['title'],
                    'url': doc['url'],
                    'score': doc['score']
                }
                for doc in retrieved_docs
            ]
            yield f"data: {json.dumps({'type': 'sources', 'sources': sources, 'model': model_name or rag_system.model_name}, ensure_ascii=False)}\n\n"
            
            # 3. 构建提示词
            prompt = rag_system.build_prompt(question, retrieved_docs)
            
            # 4. 流式生成答案
            for chunk in rag_system.generate_answer_stream(
                prompt,
                api_token=api_token if api_token else None,
                model_name=model_name if model_name else None,
                api_url=api_url if api_url else None
            ):
                # 发送文本块
                yield f"data: {json.dumps({'type': 'content', 'content': chunk}, ensure_ascii=False)}\n\n"
            
            # 5. 发送完成信号
            yield f"data: {json.dumps({'type': 'done'}, ensure_ascii=False)}\n\n"
            
        except Exception as e:
            print(f"流式问答出错: {e}")
            import traceback
            traceback.print_exc()
            yield f"data: {json.dumps({'type': 'error', 'message': str(e)}, ensure_ascii=False)}\n\n"
    
    return Response(generate(), mimetype='text/event-stream', headers={
        'Cache-Control': 'no-cache',
        'X-Accel-Buffering': 'no'
    })

@app.route('/api/evaluate', methods=['POST'])
def evaluate_search():
    """搜索结果评价API接口"""
    if not engine_loaded or not evaluator:
        return jsonify({
            'success': False,
            'error': 'LLM评价器未初始化'
        }), 500
    
    data = request.get_json()
    query = data.get('query', '').strip()
    top_k = data.get('top_k', 5)
    api_token = data.get('api_token', '').strip()
    model_name = data.get('model_name', '').strip()
    api_url = data.get('api_url', '').strip()
    
    if not query:
        return jsonify({
            'success': False,
            'error': '查询内容不能为空'
        }), 400
    
    try:
        print(f"\n收到评价请求: {query}")
        if api_token:
            print("  使用自定义API Token")
            
        result = evaluator.evaluate_search_results(
            query, 
            top_k=top_k,
            api_token=api_token if api_token else None,
            model_name=model_name if model_name else None,
            api_url=api_url if api_url else None
        )
        return jsonify(result)
    
    except Exception as e:
        print(f"评价出错: {e}")
        return jsonify({
            'success': False,
            'error': f'评价时发生错误: {str(e)}'
        }), 500

@app.route('/api/evaluate/stream', methods=['POST'])
def evaluate_search_stream():
    """搜索结果评价API接口（流式输出）"""
    if not engine_loaded or not evaluator:
        return jsonify({
            'success': False,
            'error': 'LLM评价器未初始化'
        }), 500
    
    data = request.get_json()
    query = data.get('query', '').strip()
    top_k = data.get('top_k', 5)
    api_token = data.get('api_token', '').strip()
    model_name = data.get('model_name', '').strip()
    api_url = data.get('api_url', '').strip()
    
    if not query:
        return jsonify({
            'success': False,
            'error': '查询内容不能为空'
        }), 400
    
    def generate():
        try:
            print(f"\n收到流式评价请求: {query}")
            if api_token:
                print("  使用自定义API Token")
            
            # 流式评价搜索结果
            for event in evaluator.evaluate_search_results_stream(
                query,
                top_k=top_k,
                api_token=api_token if api_token else None,
                model_name=model_name if model_name else None,
                api_url=api_url if api_url else None
            ):
                yield f"data: {json.dumps(event, ensure_ascii=False)}\n\n"
            
        except Exception as e:
            print(f"流式评价出错: {e}")
            import traceback
            traceback.print_exc()
            yield f"data: {json.dumps({'type': 'error', 'message': str(e)}, ensure_ascii=False)}\n\n"
    
    return Response(generate(), mimetype='text/event-stream', headers={
        'Cache-Control': 'no-cache',
        'X-Accel-Buffering': 'no'
    })

@app.route('/api/suggestions', methods=['GET'])
def get_suggestions():
    """获取搜索建议"""
    if not engine_loaded:
        return jsonify({
            'success': False,
            'suggestions': []
        })
    
    query = request.args.get('q', '').strip()
    if not query or len(query) < 2:  # 至少2个字符才显示建议
        return jsonify({
            'success': True,
            'suggestions': []
        })
    
    try:
        # 从索引中获取相关的词条作为建议
        suggestions = []
        query_lower = query.lower()
        
        # 从倒排索引中查找包含查询词的词条
        seen = set()
        for term in engine.inverted_index.keys():
            if query_lower in term.lower() and term not in seen:
                suggestions.append(term)
                seen.add(term)
                if len(suggestions) >= 10:  # 最多返回10条建议
                    break
        
        # 如果建议不足，从文档标题中查找
        if len(suggestions) < 10:
            for url, data in engine.url_to_id.items():
                doc_id = data['id']
                title = engine.doc_titles.get(doc_id, '')
                if query_lower in title.lower():
                    # 提取标题中包含查询词的部分
                    words = title.split()
                    for word in words:
                        if query_lower in word.lower() and word not in seen:
                            suggestions.append(word)
                            seen.add(word)
                            if len(suggestions) >= 10:
                                break
                if len(suggestions) >= 10:
                    break
        
        return jsonify({
            'success': True,
            'suggestions': suggestions[:10]
        })
    
    except Exception as e:
        print(f"获取建议出错: {e}")
        return jsonify({
            'success': False,
            'suggestions': []
        })

@app.route('/api/status', methods=['GET'])
def status():
    """返回搜索引擎状态"""
    return jsonify({
        'loaded': engine_loaded,
        'total_docs': engine.total_docs if engine_loaded else 0,
        'has_ltr': engine.ltr_model is not None if engine_loaded else False,
        'has_rag': rag_system is not None,
        'has_evaluator': evaluator is not None
    })

@app.route('/api/recommendations', methods=['GET'])
def get_recommendations():
    """获取今日推荐"""
    if not engine_loaded or not recommendation_system:
        return jsonify({
            'success': False,
            'error': '推荐系统未初始化'
        }), 500
    
    try:
        top_k = int(request.args.get('top_k', 10))
        recommendations = recommendation_system.get_daily_recommendations(top_k=top_k)
        
        return jsonify({
            'success': True,
            'recommendations': recommendations
        })
    
    except Exception as e:
        print(f"获取推荐失败: {e}")
        return jsonify({
            'success': False,
            'error': f'获取推荐时发生错误: {str(e)}'
        }), 500

@app.route('/api/trending', methods=['GET'])
def get_trending():
    """获取热门话题"""
    if not engine_loaded or not recommendation_system:
        return jsonify({
            'success': False,
            'error': '推荐系统未初始化'
        }), 500
    
    try:
        top_k = int(request.args.get('top_k', 10))
        topics = recommendation_system.get_trending_topics(top_k=top_k)
        
        return jsonify({
            'success': True,
            'topics': topics
        })
    
    except Exception as e:
        print(f"获取热门话题失败: {e}")
        return jsonify({
            'success': False,
            'error': f'获取热门话题时发生错误: {str(e)}'
        }), 500

@app.route('/api/analytics/popular', methods=['GET'])
def get_popular_searches():
    """获取热门搜索统计"""
    try:
        days = int(request.args.get('days', 7))
        top_k = int(request.args.get('top_k', 10))
        
        popular = search_analytics.get_popular_queries(days=days, top_k=top_k)
        
        return jsonify({
            'success': True,
            'popular_queries': popular
        })
    
    except Exception as e:
        print(f"获取热门搜索失败: {e}")
        return jsonify({
            'success': False,
            'error': f'获取热门搜索时发生错误: {str(e)}'
        }), 500

@app.route('/api/analytics/trends', methods=['GET'])
def get_search_trends():
    """获取搜索趋势统计"""
    try:
        days = int(request.args.get('days', 7))
        trends = search_analytics.get_search_trends(days=days)
        
        return jsonify({
            'success': True,
            'trends': trends
        })
    
    except Exception as e:
        print(f"获取搜索趋势失败: {e}")
        return jsonify({
            'success': False,
            'error': f'获取搜索趋势时发生错误: {str(e)}'
        }), 500

@app.route('/api/analytics/stats', methods=['GET'])
def get_analytics_stats():
    """获取搜索统计信息"""
    try:
        stats = search_analytics.get_statistics()
        
        return jsonify({
            'success': True,
            'statistics': stats
        })
    
    except Exception as e:
        print(f"获取统计信息失败: {e}")
        return jsonify({
            'success': False,
            'error': f'获取统计信息时发生错误: {str(e)}'
        }), 500

if __name__ == '__main__':
    # 启动前初始化搜索引擎
    init_success = init_search_engine()
    
    if not init_success:
        print("\n警告: 搜索引擎初始化失败，请确保以下文件存在：")
        print("  - search_index.pkl")
        print("  - pagerank_results.json")
        print("  - ltr_model.pkl (可选)")
        print("\n服务器仍会启动，但搜索功能将不可用。\n")
    
    print("\n" + "="*60)
    print("RUC智能搜索引擎Web服务启动中...")
    print("功能: 搜索 + RAG问答 + LLM评价")
    print("访问地址: http://localhost:5000")
    print("="*60 + "\n")
    
    app.run(debug=False, host='0.0.0.0', port=5000)