# semantic_chunker.py
import re
from typing import List, Dict
from sentence_transformers import SentenceTransformer
import numpy as np

class SemanticRecursiveChunker:
    def __init__(self, 
                 model_name: str = "BAAI/bge-m3",
                 min_chunk_size: int = 200,
                 max_chunk_size: int = 1500,
                 split_threshold: float = 0.35):
        self.encoder = SentenceTransformer(model_name)
        self.min_chunk_size = min_chunk_size
        self.max_chunk_size = max_chunk_size
        self.split_threshold = split_threshold
    
    def _split_sentences(self, text: str) -> List[str]:
        """增强句子分割，保留缩写和数字"""
        # 保护常见缩写
        text = re.sub(r'\b(\w\.){2,}', lambda m: m.group().replace('.', '@@@'), text)
        sentences = re.split(r'(?<=[。！？!?])\s*|\n{2,}', text)
        return [s.replace('@@@', '.') for s in sentences if s.strip()]
    
    def _calc_cosine_sim(self, vec1: np.ndarray, vec2: np.ndarray) -> float:
        return float(np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2)))
    
    def chunk(self, document: str) -> List[Dict]:
        """递归语义分块主逻辑"""
        sentences = self._split_sentences(document)
        if not sentences:
            return []
        
        # 预计算所有句子向量
        sent_vecs = self.encoder.encode(sentences, normalize_embeddings=True)
        
        # 递归切分
        chunks = []
        start_idx = 0
        current_chunk = []
        current_len = 0
        
        for i in range(len(sentences)):
            current_chunk.append(sentences[i])
            current_len += len(sentences[i])
            
            # 检查是否应该切分
            if current_len >= self.min_chunk_size:
                if i == len(sentences) - 1:
                    # 最后一块直接收尾
                    chunks.append({
                        "text": "".join(current_chunk),
                        "start_idx": start_idx,
                        "end_idx": i
                    })
                    break
                
                # 计算与下一句的语义相似度
                sim = self._calc_cosine_sim(sent_vecs[i], sent_vecs[i+1])
                
                if current_len >= self.max_chunk_size or sim < self.split_threshold:
                    # 切分当前块
                    chunks.append({
                        "text": "".join(current_chunk),
                        "start_idx": start_idx,
                        "end_idx": i
                    })
                    # 重置状态
                    current_chunk = []
                    current_len = 0
                    start_idx = i + 1
        
        # 后处理：合并过小的块
        merged_chunks = []
        for chunk in chunks:
            if merged_chunks and len(chunk["text"]) < self.min_chunk_size * 0.5:
                merged_chunks[-1]["text"] += chunk["text"]
                merged_chunks[-1]["end_idx"] = chunk["end_idx"]
            else:
                merged_chunks.append(chunk)
        
        return merged_chunks

# 使用示例
if __name__ == "__main__":
    chunker = SemanticRecursiveChunker()
    with open("sample_doc.txt", "r", encoding="utf-8") as f:
        content = f.read()
    
    chunks = chunker.chunk(content)
    print(f"共生成 {len(chunks)} 个语义块")
    for i, c in enumerate(chunks[:3]):
        print(f"\n--- 块 {i+1} (长度 {len(c['text'])}) ---")
        print(c["text"][:200] + "...")

# hybrid_retriever.py
from pymilvus import (
    connections, Collection, CollectionSchema, FieldSchema, DataType,
    utility, AnnSearchRequest, RRFRanker, WeightedRanker
)
import numpy as np
from typing import List, Dict, Optional
from sentence_transformers import SentenceTransformer
import jieba  # 中文分词

class MilvusHybridRetriever:
    def __init__(self, 
                 host: str = "localhost",
                 port: str = "19530",
                 collection_name: str = "enterprise_kb",
                 dense_model: str = "BAAI/bge-m3",
                 dim: int = 1024):
        connections.connect(host=host, port=port)
        self.collection_name = collection_name
        self.dense_encoder = SentenceTransformer(dense_model)
        self.dim = dim
        self._create_collection_if_not_exists()
    
    def _create_collection_if_not_exists(self):
        """创建支持混合检索的 Collection"""
        if utility.has_collection(self.collection_name):
            self.collection = Collection(self.collection_name)
            return
        
        fields = [
            FieldSchema(name="doc_id", dtype=DataType.VARCHAR, is_primary=True, max_length=64),
            FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=65535),
            FieldSchema(name="dense_vector", dtype=DataType.FLOAT_VECTOR, dim=self.dim),
            FieldSchema(name="sparse_vector", dtype=DataType.SPARSE_FLOAT_VECTOR),  # Milvus 2.4+ 支持
            FieldSchema(name="metadata", dtype=DataType.JSON)
        ]
        schema = CollectionSchema(fields, description="企业级混合检索知识库")
        self.collection = Collection(self.collection_name, schema)
        
        # 创建索引：IVF_FLAT + BM25
        self.collection.create_index(
            "dense_vector",
            {"index_type": "IVF_FLAT", "metric_type": "IP", "params": {"nlist": 1024}}
        )
        self.collection.create_index(
            "sparse_vector",
            {"index_type": "SPARSE_INVERTED_INDEX", "metric_type": "IP"}
        )
    
    def _bm25_tokenize(self, text: str) -> List[str]:
        """中文+英文混合分词"""
        tokens = jieba.lcut(text)
        # 过滤停用词和单字
        stopwords = set(["的", "了", "和", "是", "在", "有", "我", "你"])
        return [t for t in tokens if t.strip() and t not in stopwords and len(t) > 1]
    
    def _build_sparse_vector(self, text: str) -> Dict[int, float]:
        """构建 BM25 稀疏向量（简化版）"""
        tokens = self._bm25_tokenize(text)
        token_freq = {}
        for token in tokens:
            token_freq[token] = token_freq.get(token, 0) + 1
        
        # 使用哈希将 token 映射到固定维度
        sparse_vec = {}
        for token, freq in token_freq.items():
            idx = hash(token) % 1000000  # 1M 维度
            sparse_vec[idx] = freq * 1.0
        return sparse_vec
    
    def insert_documents(self, documents: List[Dict]):
        """批量插入文档"""
        rows = []
        for doc in documents:
            dense_vec = self.dense_encoder.encode(doc["content"]).tolist()
            sparse_vec = self._build_sparse_vector(doc["content"])
            rows.append({
                "doc_id": doc["doc_id"],
                "content": doc["content"],
                "dense_vector": dense_vec,
                "sparse_vector": sparse_vec,
                "metadata": doc.get("metadata", {})
            })
        
        self.collection.insert(rows)
        self.collection.flush()
    
    def hybrid_search(self, 
                      query: str,
                      top_k: int = 10,
                      alpha: float = 0.5,
                      ranker_type: str = "rrf") -> List[Dict]:
        """
        混合检索：Dense + Sparse
        
        Args:
            query: 查询文本
            top_k: 返回数量
            alpha: Dense 权重 (0-1)，仅 WeightedRanker 使用
            ranker_type: "rrf" (推荐) 或 "weighted"
        """
        # 生成查询向量
        dense_query = self.dense_encoder.encode(query).tolist()
        sparse_query = self._build_sparse_vector(query)
        
        # 构建检索请求
        dense_req = AnnSearchRequest(
            data=[dense_query],
            anns_field="dense_vector",
            param={"metric_type": "IP", "params": {"nprobe": 64}},
            limit=top_k * 2  # 扩大召回池
        )
        sparse_req = AnnSearchRequest(
            data=[sparse_query],
            anns_field="sparse_vector",
            param={"metric_type": "IP"},
            limit=top_k * 2
        )
        
        # 选择排序策略
        if ranker_type == "rrf":
            ranker = RRFRanker(k=60)  # RRF 的 k 参数
        else:
            ranker = WeightedRanker(alpha, 1 - alpha)
        
        # 执行混合检索
        results = self.collection.hybrid_search(
            reqs=[dense_req, sparse_req],
            ranker=ranker,
            limit=top_k,
            output_fields=["doc_id", "content", "metadata"]
        )
        
        # 格式化结果
        formatted = []
        for hits in results:
            for hit in hits:
                formatted.append({
                    "doc_id": hit.entity.get("doc_id"),
                    "content": hit.entity.get("content"),
                    "score": hit.score,
                    "metadata": hit.entity.get("metadata")
                })
        return formatted

# 使用示例
if __name__ == "__main__":
    retriever = MilvusHybridRetriever()
    
    # 插入示例文档
    docs = [
        {"doc_id": "001", "content": "Milvus 是一个开源的向量数据库，支持万亿级向量检索。"},
        {"doc_id": "002", "content": "BGE-Reranker 是北京智源研究院发布的重排序模型。"},
        {"doc_id": "003", "content": "GraphRAG 通过知识图谱增强大模型的推理能力。"}
    ]
    retriever.insert_documents(docs)
    
    # 混合检索
    results = retriever.hybrid_search("向量数据库如何实现混合检索？", top_k=3)
    for r in results:
        print(f"doc_id: {r['doc_id']}, score: {r['score']:.4f}")
        print(f"content: {r['content'][:100]}...")

# reranker_service.py
from FlagEmbedding import FlagReranker
from typing import List, Tuple
import numpy as np
import time

class RerankService:
    def __init__(self, 
                 model_path: str = "BAAI/bge-reranker-v2-m3",
                 use_fp16: bool = True,
                 device: str = "cuda:0"):
        self.reranker = FlagReranker(
            model_path,
            use_fp16=use_fp16,
            device=device
        )
        self.batch_size = 16  # 批处理大小，根据显存调整
    
    def rerank(self, 
               query: str, 
               candidates: List[dict],
               top_k: int = 5,
               threshold: float = 0.0) -> List[dict]:
        """
        对候选文档进行重排序
        
        Args:
            query: 查询文本
            candidates: 候选文档列表 [{"doc_id": str, "content": str, "score": float}]
            top_k: 返回数量
            threshold: 相关性阈值，低于此分数的结果将被过滤
        
        Returns:
            重排序后的文档列表
        """
        if not candidates:
            return []
        
        # 构建 query-passage 对
        pairs = [(query, doc["content"]) for doc in candidates]
        
        # 分批计算相关性分数
        scores = []
        for i in range(0, len(pairs), self.batch_size):
            batch_pairs = pairs[i:i+self.batch_size]
            batch_scores = self.reranker.compute_score(batch_pairs, normalize=True)
            if isinstance(batch_scores, float):
                batch_scores = [batch_scores]
            scores.extend(batch_scores)
        
        # 添加分数并排序
        for doc, score in zip(candidates, scores):
            doc["rerank_score"] = float(score)
        
        # 过滤低分结果
        filtered = [doc for doc in candidates if doc["rerank_score"] >= threshold]
        
        # 按重排序分数降序排列
        filtered.sort(key=lambda x: x["rerank_score"], reverse=True)
        
        return filtered[:top_k]
    
    def batch_rerank(self, 
                     queries: List[str], 
                     candidates: List[List[dict]],
                     top_k: int = 5) -> List[List[dict]]:
        """批量重排序，用于并发场景"""
        results = []
        for query, cands in zip(queries, candidates):
            results.append(self.rerank(query, cands, top_k))
        return results

# 集成到检索流程的示例
class RAGPipeline:
    def __init__(self, retriever, reranker):
        self.retriever = retriever
        self.reranker = reranker
    
    def retrieve_and_rerank(self, query: str, top_k: int = 5):
        # 第一阶段：混合检索召回 20 个候选
        candidates = self.retriever.hybrid_search(query, top_k=20)
        
        # 第二阶段：重排序取 top_k
        reranked = self.reranker.rerank(
            query, 
            candidates,
            top_k=top_k,
            threshold=0.3  # 过滤掉相关性低于 0.3 的结果
        )
        return reranked

# 性能测试
if __name__ == "__main__":
    reranker_service = RerankService()
    
    # 模拟候选集
    query = "如何优化 Milvus 的混合检索性能？"
    candidates = [
        {"doc_id": "001", "content": "Milvus 支持多种索引类型，包括 IVF_FLAT 和 HNSW。"},
        {"doc_id": "002", "content": "混合检索需要同时使用向量索引和 BM25 索引。"},
        {"doc_id": "003", "content": "优化 nprobe 参数可以提高召回率。"},
        {"doc_id": "004", "content": "Python 是一种流行的编程语言。"},
        {"doc_id": "005", "content": "RRF 算法用于融合多个检索结果。"}
    ]
    
    # 测试单次重排序延迟
    start = time.time()
    results = reranker_service.rerank(query, candidates, top_k=3)
    latency = (time.time() - start) * 1000
    
    print(f"重排序延迟: {latency:.1f} ms")
    for r in results:
        print(f"doc_id: {r['doc_id']}, rerank_score: {r['rerank_score']:.4f}")

# adaptive_context_router.py
from typing import Dict, List
import re

class AdaptiveContextRouter:
    def __init__(self, 
                 short_context_limit: int = 30000,
                 long_context_limit: int = 150000):
        self.short_limit = short_context_limit
        self.long_limit = long_context_limit
    
    def analyze_query(self, query: str) -> Dict:
        """分析查询复杂度"""
        features = {
            "length": len(query),
            "has_code": bool(re.search(r'[{}();\[\]]', query)),
            "has_multiple_entities": len(re.findall(r'[A-Z][a-z]+', query)) > 2,
            "has_negation": bool(re.search(r'不是|没有|except|not', query, re.I)),
            "num_question_words": len(re.findall(r'为什么|如何|怎样|what|how|why', query, re.I))
        }
        
        complexity_score = sum([
            features["has_code"] * 2,
            features["has_multiple_entities"] * 1.5,
            features["has_negation"] * 1,
            features["num_question_words"] * 0.5
        ])
        
        return {
            "complexity_score": complexity_score,
            "needs_long_context": complexity_score > 3.0,
            "features": features
        }
    
    def assemble_context(self, 
                         retrieved_chunks: List[Dict],
                         graph_evidence: List[Dict] = None,
                         max_tokens: int = 30000) -> str:
        """智能组装上下文，优先保留高相关性内容"""
        # 按相关性分数排序
        sorted_chunks = sorted(retrieved_chunks, key=lambda x: x["score"], reverse=True)
        
        context_parts = []
        total_tokens = 0
        
        for chunk in sorted_chunks:
            chunk_tokens = len(chunk["content"]) // 2  # 粗略估计 Token 数
            if total_tokens + chunk_tokens > max_tokens:
                # 截断最后一块
                remaining = max_tokens - total_tokens
                context_parts.append(chunk["content"][:remaining * 2])
                break
            
            context_parts.append(chunk["content"])
            total_tokens += chunk_tokens
        
        # 添加知识图谱证据（如果有）
        if graph_evidence:
            evidence_text = "\n\n[知识图谱证据]\n" + "\n".join(
                f"- {ev['subject']} {ev['relation']} {ev['object']}" 
                for ev in graph_evidence[:5]
            )
            context_parts.append(evidence_text)
        
        return "\n\n".join(context_parts)

# graphrag_builder.py
from langchain_community.graphs import Neo4jGraph
from langchain_experimental.graph_transformers import LLMGraphTransformer
from langchain_openai import ChatOpenAI
from typing import List, Dict
import json

class GraphRAGBuilder:
    def __init__(self, 
                 neo4j_uri: str = "neo4j://localhost:7687",
                 neo4j_user: str = "neo4j",
                 neo4j_password: str = "password",
                 llm_model: str = "gpt-4o"):
        self.graph = Neo4jGraph(
            url=neo4j_uri,
            username=neo4j_user,
            password=neo4j_password
        )
        self.llm = ChatOpenAI(model=llm_model, temperature=0)
        self.transformer = LLMGraphTransformer(
            llm=self.llm,
            allowed_nodes=["Entity", "Person", "Organization", "Technology", "Product"],
            allowed_relationships=["RELATED_TO", "WORKS_AT", "DEVELOPS", "PARTNERS_WITH"]
        )
    
    def extract_and_store(self, documents: List[str]):
        """从文档中抽取实体关系并存入 Neo4j"""
        for i, doc in enumerate(documents):
            print(f"处理文档 {i+1}/{len(documents)}...")
            
            # 转换为 LangChain Document 格式
            from langchain_core.documents import Document
            lc_doc = Document(page_content=doc, metadata={"source": f"doc_{i}"})
            
            # 抽取知识图谱
            graph_docs = self.transformer.convert_to_graph_documents([lc_doc])
            
            # 存入 Neo4j
            self.graph.add_graph_documents(
                graph_docs,
                include_source=True,
                baseEntityLabel=True
            )
            print(f"  抽取到 {len(graph_docs[0].nodes)} 个实体，{len(graph_docs[0].relationships)} 条关系")
    
    def query_graph(self, cypher: str) -> List[Dict]:
        """执行 Cypher 查询"""
        result = self.graph.query(cypher)
        return result
    
    def get_entity_neighbors(self, entity_name: str, max_depth: int = 2) -> List[Dict]:
        """获取实体的邻居节点（用于多跳推理）"""
        cypher = f"""
        MATCH (n {{name: '{entity_name}'}})-[r*1..{max_depth}]-(m)
        RETURN n, r, m
        LIMIT 50
        """
        return self.query_graph(cypher)
    
    def find_shortest_path(self, entity1: str, entity2: str) -> List[Dict]:
        """查找两个实体间的最短路径"""
        cypher = f"""
        MATCH p = shortestPath(
            (a {{name: '{entity1}'}})-[*]-(b {{name: '{entity2}'}})
        )
        RETURN p
        """
        return self.query_graph(cypher)

# 实体关系抽取示例
if __name__ == "__main__":
    builder = GraphRAGBuilder()
    
    # 示例文档
    docs = [
        "微软和 OpenAI 是重要的合作伙伴关系。",
        "OpenAI 开发了 GPT-4 大语言模型。",
        "微软的 Azure 云平台提供了 GPT-4 的部署服务。",
        "Anthropic 是 OpenAI 的竞争对手，专注于 AI 安全。"
    ]
    
    # 构建知识图谱
    builder.extract_and_store(docs)
    
    # 查询实体关系
    print("\n=== 查询 '微软' 的邻居 ===")
    neighbors = builder.get_entity_neighbors("微软")
    for n in neighbors:
        print(f"  关系: {n}")
    
    print("\n=== 查询 'OpenAI' 到 'Azure' 的路径 ===")
    path = builder.find_shortest_path("OpenAI", "Azure")
    print(f"  最短路径: {path}")