MetaSearch

"""
查询规划智能体 (QueryPlannerAgent)

负责分析用户查询，制定检索策略和执行计划
"""

import json
import re
from typing import Dict, List, Any
from .base_agent import BaseAgent

class QueryPlannerAgent(BaseAgent):
    """查询规划智能体 - 分析查询复杂度并制定检索策略"""
    
    def __init__(self, name: str = "QueryPlannerAgent", **kwargs):
        super().__init__(name, **kwargs)
        
        # 查询类型模式匹配
        self.query_patterns = {
            'factual': [
                r'什么是', r'谁是', r'哪里', r'何时', r'定义', r'概念',
                r'what is', r'who is', r'where', r'when', r'definition'
            ],
            'analytical': [
                r'为什么', r'如何', r'怎样', r'原因', r'分析', r'解释',
                r'why', r'how', r'analyze', r'explain', r'reason'
            ],
            'comparative': [
                r'比较', r'对比', r'区别', r'不同', r'相同', r'优缺点',
                r'compare', r'difference', r'similarity', r'vs', r'versus'
            ]
        }
    
    def execute(self, task: Dict) -> Dict:
        """执行查询规划任务"""
        query = task.get('query', '')
        
        self.log_step("开始分析用户查询")
        
        # 基础查询分析
        basic_analysis = self._analyze_query_basic(query)
        self.log_step("基础分析完成", 
                     f"长度={basic_analysis['length']}, 词数={basic_analysis['word_count']}")
        
        # 查询类型识别
        query_type = self._identify_query_type(query)
        self.log_step("查询类型识别", f"类型={query_type}")
        
        # 制定检索策略
        strategy = self._plan_retrieval_strategy(query, query_type, basic_analysis)
        self.log_step("检索策略制定", 
                     f"策略={strategy['method']}, 预计轮次={strategy['max_iterations']}")
        
        return {
            'query_analysis': basic_analysis,
            'query_type': query_type,
            'retrieval_strategy': strategy,
            'agent_name': self.name
        }

"""
检索执行智能体 (RetrievalAgent)

根据规划执行具体的检索任务，支持多种检索策略
"""

from typing import Dict, List, Any
from .base_agent import BaseAgent

class RetrievalAgent(BaseAgent):
    """检索执行智能体 - 执行具体的检索操作"""
    
    def __init__(self, name: str = "RetrievalAgent", searcher=None, **kwargs):
        super().__init__(name, **kwargs)
        self.searcher = searcher
        
        # 检索策略配置
        self.retrieval_methods = {
            'vector': self._vector_retrieval,
            'keyword': self._keyword_retrieval,
            'hybrid': self._hybrid_retrieval,
            'adaptive': self._adaptive_retrieval
        }
    
    def execute(self, task: Dict) -> Dict:
        """执行检索任务"""
        query = task.get('query', '')
        strategy = task.get('strategy', {})
        iteration = task.get('iteration', 1)
        
        self.log_step(f"开始第{iteration}轮检索", f"查询: {query[:50]}...")
        
        # 根据策略选择检索方法
        method = strategy.get('method', 'hybrid')
        retrieval_func = self.retrieval_methods.get(method, self._hybrid_retrieval)
        
        # 执行检索
        results = retrieval_func(query, strategy)
        
        # 结果后处理
        processed_results = self._process_results(results, strategy)
        
        self.log_step("检索完成", 
                     f"找到{len(processed_results)}个结果, "
                     f"平均相关度{self._calculate_avg_score(processed_results):.3f}")
        
        return {
            'results': processed_results,
            'iteration': iteration,
            'method': method,
            'total_found': len(processed_results),
            'agent_name': self.name
        }
    
    def _hybrid_retrieval(self, query: str, strategy: Dict) -> List[Dict]:
        """混合检索：结合向量检索和关键词检索"""
        # 向量检索
        vector_results = self.searcher.vector_search(
            query, 
            top_k=strategy.get('vector_candidates', 50)
        )
        
        # 关键词检索
        keyword_results = self.searcher.keyword_search(
            query,
            top_k=strategy.get('keyword_candidates', 50)
        )
        
        # 结果融合
        merged_results = self._merge_results(vector_results, keyword_results)
        
        return merged_results[:strategy.get('max_results', 20)]

"""
质量评估智能体 (EvaluatorAgent)

评估检索结果质量，决定是否继续检索或生成答案
"""

import json
from typing import Dict, List, Any
from .base_agent import BaseAgent

class EvaluatorAgent(BaseAgent):
    """质量评估智能体 - 评估信息质量并决策下一步"""
    
    def __init__(self, name: str = "EvaluatorAgent", ranker=None, **kwargs):
        super().__init__(name, **kwargs)
        self.ranker = ranker
        
        # 评估维度权重
        self.evaluation_weights = {
            'coverage': 0.25,      # 覆盖度
            'depth': 0.25,         # 深度
            'consistency': 0.25,   # 一致性
            'novelty': 0.25        # 新颖性
        }
    
    def execute(self, task: Dict) -> Dict:
        """执行质量评估任务"""
        original_query = task.get('original_query', '')
        current_results = task.get('current_results', [])
        accumulated_chunks = task.get('accumulated_chunks', set())
        iteration = task.get('iteration', 1)
        max_iterations = task.get('max_iterations', 3)
        
        self.log_step(f"开始第{iteration}轮质量评估")
        
        # 重排序优化结果
        if self.ranker and current_results:
            reranked_results = self._rerank_results(original_query, current_results)
        else:
            reranked_results = current_results
        
        # 多维度质量评估
        quality_scores = self._evaluate_quality(
            original_query, reranked_results, accumulated_chunks
        )
        
        # 决策下一步行动
        decision = self._make_decision(
            quality_scores, iteration, max_iterations, accumulated_chunks
        )
        
        self.log_step(f"评估完成", 
                     f"整体质量={quality_scores['overall']:.3f}, "
                     f"决策={decision['action']}")
        
        return {
            'reranked_results': reranked_results,
            'quality_scores': quality_scores,
            'decision': decision,
            'iteration': iteration,
            'agent_name': self.name
        }

"""
答案生成智能体 (GeneratorAgent)

整合多轮检索结果，生成高质量的综合答案
"""

from typing import Dict, List, Any
from .base_agent import BaseAgent

class GeneratorAgent(BaseAgent):
    """答案生成智能体 - 整合信息并生成最终答案"""
    
    def __init__(self, name: str = "GeneratorAgent", **kwargs):
        super().__init__(name, **kwargs)
        
        # 生成模板
        self.generation_templates = {
            'factual': "基于检索到的信息，直接回答用户的事实性问题",
            'analytical': "深入分析检索结果，提供全面的解释和分析",
            'comparative': "对比不同观点和信息，提供平衡的比较分析",
            'procedural': "整理步骤信息，提供清晰的操作指南",
            'comprehensive': "综合所有信息，生成结构化的全面回答"
        }
    
    def execute(self, task: Dict) -> Dict:
        """执行答案生成任务"""
        original_query = task.get('original_query', '')
        final_results = task.get('final_results', [])
        query_type = task.get('query_type', 'comprehensive')
        total_iterations = task.get('total_iterations', 1)
        
        self.log_step("开始生成最终答案", 
                     f"基于{len(final_results)}个信息块, {total_iterations}轮检索")
        
        # 信息预处理和去重
        processed_info = self._preprocess_information(final_results)
        
        # 结构化组织信息
        organized_info = self._organize_information(processed_info, query_type)
        
        # 生成答案
        final_answer = self._generate_answer(
            original_query, organized_info, query_type
        )
        
        # 答案后处理
        polished_answer = self._polish_answer(final_answer)
        
        self.log_step("答案生成完成", f"长度={len(polished_answer)}字符")
        
        return {
            'final_answer': polished_answer,
            'information_sources': len(final_results),
            'query_type': query_type,
            'total_iterations': total_iterations,
            'agent_name': self.name
        }