# query_understanding.py
import jieba
import jieba.analyse
from typing import Dict, List, Tuple
import re

class QueryUnderstandingEngine:
    """查询理解引擎：意图识别 + 查询重写"""
    
    def __init__(self):
        # 意图分类模型
        self.intent_keywords = {
            'conceptual': ['是什么', '什么是', '如何理解', '概念', '原理'],
            'code_example': ['代码', '示例', 'demo', '怎么写', '实现'],
            'api_reference': ['API', '参数', '接口', '调用', '文档'],
            'error_solution': ['错误', '报错', '失败', '怎么办', '解决'],
            'comparison': ['对比', '区别', '哪个好', 'vs', '比较']
        }
        
    def analyze_intent(self, query: str) -> Dict:
        """分析查询意图"""
        intent_scores = {intent: 0.0 for intent in self.intent_keywords.keys()}
        
        # 基于关键词的意图识别
        for intent, keywords in self.intent_keywords.items():
            for keyword in keywords:
                if keyword in query:
                    intent_scores[intent] += 1.0
        
        # 基于长度的启发式规则：长查询更可能是概念性查询
        query_length = len(query)
        if query_length > 20:
            intent_scores['conceptual'] += 0.5
        
        # 归一化分数
        total = sum(intent_scores.values())
        if total > 0:
            intent_scores = {k: v/total for k, v in intent_scores.items()}
        
        primary_intent = max(intent_scores.items(), key=lambda x: x[1])
        
        return {
            'primary_intent': primary_intent[0],
            'confidence': primary_intent[1],
            'intent_breakdown': intent_scores
        }
    
    def query_rewrite(self, query: str, intent: str) -> List[str]:
        """查询重写：生成多个查询变体"""
        base_queries = [query]
        
        # 根据意图进行特定的查询扩展
        if intent == 'code_example':
            # 代码查询添加语言标签
            languages = ['Python', 'Java', 'JavaScript', 'Go', 'C++']
            for lang in languages:
                if any(word in query for word in [lang.lower(), lang]):
                    base_queries.append(f"{query} {lang}代码示例")
        
        elif intent == 'api_reference':
            # API查询添加参数说明
            if 'API' in query:
                base_queries.extend([
                    f"{query} 参数说明",
                    f"{query} 使用示例",
                    f"{query} 返回值"
                ])
        
        return base_queries

# intelligent_router.py
from typing import Dict, List
import numpy as np

class IntelligentRouter:
    """智能路由：根据意图选择最优检索策略"""
    
    def __init__(self):
        # 意图到检索策略的映射
        self.intent_strategy_map = {
            'conceptual': {
                'strategies': ['vector', 'hybrid'],
                'weights': [0.7, 0.3],
                'description': '概念查询以向量检索为主'
            },
            'code_example': {
                'strategies': ['code_search', 'vector', 'keyword'],
                'weights': [0.6, 0.25, 0.15],
                'description': '代码查询需要专用代码检索'
            },
            'api_reference': {
                'strategies': ['keyword', 'vector', 'api_search'],
                'weights': [0.5, 0.3, 0.2],
                'description': 'API文档需要精确匹配'
            },
            'error_solution': {
                'strategies': ['hybrid', 'vector', 'keyword'],
                'weights': [0.4, 0.4, 0.2],
                'description': '错误解决方案需要综合检索'
            }
        }
        
    def route_strategy(self, intent_info: Dict, query: str) -> Dict:
        """路由决策"""
        intent = intent_info['primary_intent']
        confidence = intent_info['confidence']
        
        # 获取基础策略配置
        strategy_config = self.intent_strategy_map.get(
            intent, 
            self.intent_strategy_map['conceptual']  # 默认策略
        )
        
        strategies = strategy_config['strategies']
        base_weights = strategy_config['weights']
        
        # 根据置信度调整权重
        adjusted_weights = self._adjust_weights_by_confidence(
            base_weights, confidence
        )
        
        # 根据查询特征微调权重
        final_weights = self._adjust_weights_by_query_features(
            adjusted_weights, strategies, query
        )
        
        return {
            'strategies': strategies,
            'weights': final_weights,
            'reasoning': f"意图:{intent}, 置信度:{confidence:.2f}"
        }
    
    def _adjust_weights_by_confidence(self, weights: List[float], 
                                    confidence: float) -> List[float]:
        """根据意图置信度调整权重"""
        # 高置信度时强化主要策略，低置信度时更均衡
        if confidence > 0.7:
            # 强化主要策略
            main_weight = weights[0] * 1.2
            other_weights = [w * 0.8 for w in weights[1:]]
            # 重新归一化
            total = main_weight + sum(other_weights)
            return [main_weight/total] + [w/total for w in other_weights]
        else:
            return weights
    
    def _adjust_weights_by_query_features(self, weights: List[float],
                                        strategies: List[str], query: str) -> List[float]:
        """根据查询特征微调权重"""
        # 查询长度特征
        query_len = len(query)
        if query_len < 10:
            # 短查询更适合关键词检索
            if 'keyword' in strategies:
                idx = strategies.index('keyword')
                weights[idx] *= 1.3
        elif query_len > 50:
            # 长查询更适合向量检索
            if 'vector' in strategies:
                idx = strategies.index('vector')
                weights[idx] *= 1.2
        
        # 归一化
        total = sum(weights)
        return [w/total for w in weights]

# multi_route_retriever.py
import asyncio
from concurrent.futures import ThreadPoolExecutor
from typing import Dict, List, Any
import numpy as np
from sentence_transformers import SentenceTransformer

class MultiRouteRetriever:
    """多路检索引擎：并行执行多种检索策略"""
    
    def __init__(self, config: Dict):
        self.config = config
        self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
        
        # 初始化各检索器
        self.vector_retriever = VectorRetriever(config)
        self.keyword_retriever = KeywordRetriever(config)
        self.code_retriever = CodeRetriever(config)
        self.hybrid_retriever = HybridRetriever(config)
        
        # 线程池用于并行检索
        self.executor = ThreadPoolExecutor(max_workers=4)
    
    async def retrieve(self, query: str, route_plan: Dict) -> List[Dict]:
        """并行执行多路检索"""
        strategies = route_plan['strategies']
        weights = route_plan['weights']
        
        # 创建检索任务
        tasks = []
        for strategy in strategies:
            if strategy == 'vector':
                task = self.executor.submit(self.vector_retriever.search, query, 10)
            elif strategy == 'keyword':
                task = self.executor.submit(self.keyword_retriever.search, query, 10)
            elif strategy == 'code_search':
                task = self.executor.submit(self.code_retriever.search, query, 10)
            elif strategy == 'hybrid':
                task = self.executor.submit(self.hybrid_retriever.search, query, 10)
            else:
                continue
            tasks.append((strategy, task))
        
        # 等待所有检索完成
        results = {}
        for strategy, task in tasks:
            try:
                results[strategy] = await asyncio.get_event_loop().run_in_executor(
                    None, lambda: task.result(timeout=5.0)
                )
            except Exception as e:
                print(f"Strategy {strategy} failed: {e}")
                results[strategy] = []
        
        # 结果融合
        fused_results = self._fuse_results(results, weights, strategies)
        return fused_results
    
    def _fuse_results(self, results: Dict[str, List], weights: List[float],
                     strategies: List[str]) -> List[Dict]:
        """多路结果融合算法"""
        all_docs = {}
        doc_scores = {}
        
        # 收集所有文档并计算加权分数
        for strategy, weight in zip(strategies, weights):
            if strategy not in results:
                continue
                
            for i, doc in enumerate(results[strategy]):
                doc_id = doc['id']
                
                if doc_id not in all_docs:
                    all_docs[doc_id] = doc
                    doc_scores[doc_id] = 0.0
                
                # 排名衰减权重：排名越靠前，权重越高
                rank_weight = 1.0 / (i + 1)  
                doc_scores[doc_id] += weight * rank_weight * doc.get('score', 0.5)
        
        # 按加权分数排序
        sorted_docs = sorted(all_docs.values(), 
                           key=lambda x: doc_scores[x['id']], reverse=True)
        
        return sorted_docs[:20]  # 返回Top 20

# vector_retriever.py
import hnswlib
import numpy as np
from typing import List, Dict

class VectorRetriever:
    """基于HNSW的向量检索器"""
    
    def __init__(self, config: Dict):
        self.config = config
        self.index = hnswlib.Index(space='cosine', dim=384)
        try:
            self.index.load_index(config['vector_index_path'])
        except:
            # 初始化空索引
            self.index.init_index(max_elements=100000, ef_construction=200, M=16)
        
        self.doc_store = {}  # 文档存储
    
    def search(self, query: str, top_k: int = 10) -> List[Dict]:
        """向量检索"""
        # 生成查询向量
        query_embedding = self.embedding_model.encode([query])[0]
        
        # HNSW检索
        labels, distances = self.index.knn_query([query_embedding], k=top_k)
        
        results = []
        for label, distance in zip(labels[0], distances[0]):
            if label in self.doc_store:
                doc = self.doc_store[label]
                results.append({
                    'id': f"vec_{label}",
                    'content': doc['content'],
                    'type': doc['type'],
                    'score': 1 - distance,  # 余弦相似度转换
                    'source': 'vector_search'
                })
        
        return results

# code_retriever.py
import ast
from tree_sitter import Parser, Language
from typing import List, Dict
import hashlib

class CodeRetriever:
    """专用代码检索器：支持AST解析和语义检索"""
    
    def __init__(self, config: Dict):
        self.config = config
        # 初始化Tree-sitter解析器
        self.parser = Parser()
        try:
            LANGUAGE = Language('build/languages.so', 'python')
            self.parser.set_language(LANGUAGE)
        except:
            print("Tree-sitter not available, using fallback mode")
        
    def extract_code_features(self, code_text: str) -> Dict:
        """提取代码特征：函数、类、API调用等"""
        features = {
            'functions': [],
            'classes': [],
            'imports': [],
            'api_calls': [],
            'code_structure': {}
        }
        
        try:
            # 使用AST解析Python代码
            tree = ast.parse(code_text)
            
            for node in ast.walk(tree):
                if isinstance(node, ast.FunctionDef):
                    features['functions'].append(node.name)
                elif isinstance(node, ast.ClassDef):
                    features['classes'].append(node.name)
                elif isinstance(node, ast.Import):
                    for alias in node.names:
                        features['imports'].append(alias.name)
                elif isinstance(node, ast.ImportFrom):
                    module = node.module
                    for alias in node.names:
                        features['imports'].append(f"{module}.{alias.name}")
        
        except SyntaxError:
            # 非Python代码或语法错误，使用简单启发式规则
            lines = code_text.split('\n')
            for line in lines:
                if 'def ' in line and '(' in line:
                    # 简单函数提取
                    func_name = line.split('def ')[1].split('(')[0].strip()
                    features['functions'].append(func_name)
        
        return features
    
    def search(self, query: str, top_k: int = 10) -> List[Dict]:
        """代码专用检索"""
        # 基于代码特征的混合检索
        keyword_results = self._keyword_search(query, top_k*2)
        semantic_results = self._semantic_search(query, top_k*2)
        
        # 代码特定的重排逻辑
        reranked_results = self._rerank_for_code(
            keyword_results + semantic_results, query
        )
        
        return reranked_results[:top_k]
    
    def _rerank_for_code(self, results: List[Dict], query: str) -> List[Dict]:
        """代码特定的重排算法"""
        for doc in results:
            code_features = self.extract_code_features(doc.get('content', ''))
            
            # 计算代码质量分数
            quality_score = self._calculate_code_quality(doc['content'])
            
            # 计算查询匹配度
            query_match_score = self._calculate_query_match(doc['content'], query)
            
            # 综合评分
            doc['score'] = (doc.get('score', 0.5) * 0.4 + 
                          quality_score * 0.3 + 
                          query_match_score * 0.3)
        
        return sorted(results, key=lambda x: x['score'], reverse=True)

​

# multi_stage_reranker.py
from typing import List, Dict
import numpy as np

class MultiStageReranker:
    """多轮重排引擎：精细化结果排序"""
    
    def __init__(self):
        self.stages = [
            'freshness_rerank',      # 新鲜度重排
            'authority_rerank',     # 权威性重排  
            'diversity_rerank',      # 多样性重排
            'personalization_rerank' # 个性化重排
        ]
    
    def rerank(self, documents: List[Dict], query: str, 
               user_context: Dict = None) -> List[Dict]:
        """多轮重排流水线"""
        reranked_docs = documents.copy()
        
        for stage in self.stages:
            if stage == 'freshness_rerank':
                reranked_docs = self._freshness_rerank(reranked_docs)
            elif stage == 'authority_rerank':
                reranked_docs = self._authority_rerank(reranked_docs)
            elif stage == 'diversity_rerank':
                reranked_docs = self._diversity_rerank(reranked_docs)
            elif stage == 'personalization_rerank' and user_context:
                reranked_docs = self._personalization_rerank(reranked_docs, user_context)
        
        return reranked_docs
    
    def _freshness_rerank(self, documents: List[Dict]) -> List[Dict]:
        """新鲜度重排：优先显示更新内容"""
        for doc in documents:
            freshness = self._calculate_freshness_score(doc.get('timestamp'))
            doc['score'] = doc.get('score', 0.5) * 0.7 + freshness * 0.3
        
        return sorted(documents, key=lambda x: x['score'], reverse=True)
    
    def _diversity_rerank(self, documents: List[Dict]) -> List[Dict]:
        """多样性重排：避免相同内容重复出现"""
        seen_content = set()
        diversified_docs = []
        
        for doc in documents:
            content_hash = hashlib.md5(doc['content'][:100].encode()).hexdigest()
            if content_hash not in seen_content:
                diversified_docs.append(doc)
                seen_content.add(content_hash)
            else:
                # 重复内容降权
                doc['score'] *= 0.5
        
        return diversified_docs

# data_consistency_manager.py
import time
from typing import Dict, List
from threading import Lock

class DataConsistencyManager:
    """数据一致性管理：保证多源数据实时性"""
    
    def __init__(self):
        self.last_update_time = {}
        self.update_locks = {}
        self.version_control = {}
    
    def check_consistency(self, data_sources: List[str]) -> Dict:
        """检查数据源一致性"""
        consistency_report = {}
        
        for source in data_sources:
            current_version = self._get_current_version(source)
            last_known_version = self.version_control.get(source)
            
            if last_known_version != current_version:
                consistency_report[source] = {
                    'status': 'stale',
                    'current_version': current_version,
                    'last_known_version': last_known_version,
                    'action': 'needs_refresh'
                }
            else:
                consistency_report[source] = {
                    'status': 'fresh',
                    'version': current_version
                }
        
        return consistency_report
    
    def synchronized_update(self, data_source: str, update_func) -> bool:
        """同步更新数据源，避免并发冲突"""
        if data_source not in self.update_locks:
            self.update_locks[data_source] = Lock()
        
        with self.update_locks[data_source]:
            # 双重检查，避免重复更新
            if not self._needs_update(data_source):
                return False
            
            try:
                update_func()
                self.version_control[data_source] = self._get_current_version(data_source)
                self.last_update_time[data_source] = time.time()
                return True
            except Exception as e:
                print(f"Update failed for {data_source}: {e}")
                return False

# performance_optimizer.py
from functools import lru_cache
import time
from statistics import mean

class PerformanceOptimizer:
    """性能优化器：实时调整系统参数"""
    
    def __init__(self):
        self.performance_metrics = {
            'response_times': [],
            'cache_hit_rates': [],
            'accuracy_scores': []
        }
    
    def adaptive_timeout_control(self, current_load: int) -> float:
        """自适应超时控制"""
        base_timeout = 2.0  # 基础超时2秒
        
        if current_load > 1000:  # 高负载
            return base_timeout * 0.7  # 缩短超时
        elif current_load < 100:  # 低负载
            return base_timeout * 1.5  # 延长超时获取更好结果
        else:
            return base_timeout
    
    @lru_cache(maxsize=1000)
    def cached_embedding(self, text: str) -> List[float]:
        """带缓存的嵌入计算"""
        return self.embedding_model.encode([text])[0].tolist()