┌──────────────────────────────────────────────────────────────┐
│                    应用层 (Application Layer)                   │
│   AI Chatbot  ·  企业搜索  ·  代码助手  ·  文档分析  ·  客服系统  │
├──────────────────────────────────────────────────────────────┤
│                    编排层 (Orchestration Layer)                │
│   Query Processing → Retrieval → Augmentation → Generation    │
│   LangChain / LlamaIndex / Haystack / 自定义Pipeline          │
├──────────────────────────────────────────────────────────────┤
│                     数据层 (Data Layer)                        │
│   数据管道: 文档解析 → 分块 → 向量化 → 索引                    │
│   存储: 向量数据库 + 文档存储 + 缓存层                         │
└──────────────────────────────────────────────────────────────┘

import os
from typing import List, Optional
from dataclasses import dataclass

@dataclass
class Document:
    content: str
    metadata: dict
    source: str

class DocumentParser:
    """统一文档解析器，支持多种文档格式"""

    SUPPORTED_FORMATS = {'.pdf', '.docx', '.pptx', '.html', '.md', '.txt', '.csv', '.json'}

    def parse(self, file_path: str) -> Optional[Document]:
        ext = os.path.splitext(file_path)[1].lower()
        if ext not in self.SUPPORTED_FORMATS:
            print(f"⚠️ 不支持的格式: {ext}")
            return None

        metadata = {
            "source": file_path,
            "format": ext,
            "size": os.path.getsize(file_path)
        }

        try:
            if ext == '.pdf':
                return self._parse_pdf(file_path, metadata)
            elif ext == '.docx':
                return self._parse_docx(file_path, metadata)
            elif ext in {'.md', '.txt', '.csv', '.json'}:
                return self._parse_text(file_path, metadata)
            # ... 其他格式处理
        except Exception as e:
            print(f"❌ 解析失败 {file_path}: {e}")
            return None

    def _parse_pdf(self, path: str, metadata: dict) -> Document:
        # 实际项目中使用 PyMuPDF / pdfplumber / Unstructured
        # 这里展示接口设计模式
        import PyMuPDF  # fitz
        doc = fitz.open(path)
        content = "\n\n".join([page.get_text() for page in doc])
        metadata["pages"] = len(doc)
        return Document(content=content, metadata=metadata, source=path)

    def _parse_text(self, path: str, metadata: dict) -> Document:
        with open(path, 'r', encoding='utf-8', errors='ignore') as f:
            content = f.read()
        return Document(content=content, metadata=metadata, source=path)

from typing import List
import re

class ChunkingStrategy:
    """多种分块策略实现"""

    @staticmethod
    def recursive_character_split(
        text: str,
        chunk_size: int = 512,
        chunk_overlap: int = 128,
        separators: List[str] = ["\n\n", "\n", ".", "!", "?", ",", " ", ""]
    ) -> List[str]:
        """递归字符分割——LangChain/LLamaIndex经典策略"""
        chunks = []
        current = text

        for sep in separators:
            if sep == "":  # 最后回退到字符级别
                while len(current) > chunk_size:
                    chunks.append(current[:chunk_size])
                    current = current[chunk_size - chunk_overlap:]
                if current:
                    chunks.append(current)
                break

            parts = current.split(sep)
            merged = []
            current_part = ""

            for part in parts:
                if len(current_part) + len(sep) + len(part) <= chunk_size:
                    current_part = f"{current_part}{sep}{part}" if current_part else part
                else:
                    if current_part:
                        merged.append(current_part)
                    current_part = part

            if current_part:
                merged.append(current_part)

            # 检查是否所有块都在chunk_size内
            if all(len(c) <= chunk_size * 1.2 for c in merged) and len(merged) > 1:
                # 处理overlap
                result = []
                for i, m in enumerate(merged):
                    if i > 0 and chunk_overlap > 0:
                        prev_tail = merged[i-1][-chunk_overlap:]
                        m = prev_tail + m[:chunk_size - len(prev_tail)]
                    result.append(m)
                return result

            current = "\n\n".join(merged)

        return chunks if chunks else [text]

    @staticmethod
    def semantic_chunking(text: str, max_chunk_size: int = 512) -> List[str]:
        """语义分块——按主题段落分割"""
        # 检测主题变化（标题、空行、话题转换）
        paragraphs = re.split(r'\n\s*\n', text)
        chunks = []
        current_chunk = ""

        for para in paragraphs:
            if len(current_chunk) + len(para) > max_chunk_size and current_chunk:
                chunks.append(current_chunk.strip())
                current_chunk = para
            else:
                current_chunk = f"{current_chunk}\n\n{para}" if current_chunk else para

        if current_chunk:
            chunks.append(current_chunk.strip())

        return chunks

from typing import List
import numpy as np

class EmbeddingService:
    """嵌入服务抽象层——支持多种Embedding提供方"""

    def __init__(self, provider: str = "openai"):
        self.provider = provider
        self.client = self._init_client()

    def _init_client(self):
        if self.provider == "openai":
            from openai import OpenAI
            return OpenAI()
        elif self.provider == "text2vec":
            # 本地模型: text2vec-large-chinese / BAAI/bge-large-zh-v1.5
            from sentence_transformers import SentenceTransformer
            return SentenceTransformer('BAAI/bge-large-zh-v1.5')
        elif self.provider == "voyage":
            import voyageai
            return voyageai.Client()

    def embed_documents(self, texts: List[str]) -> np.ndarray:
        """批量嵌入文档"""
        if self.provider == "openai":
            resp = self.client.embeddings.create(
                model="text-embedding-3-large",
                input=texts,
                dimensions=1024  # 可降维
            )
            return np.array([d.embedding for d in resp.data])
        elif self.provider == "text2vec":
            return self.client.encode(texts, normalize_embeddings=True)
        # ... 更多提供方

    def embed_query(self, query: str) -> np.ndarray:
        """嵌入查询——推荐使用不同的prompt前缀"""
        if self.provider.startswith("BAAI"):
            return self.client.encode(
                [f"为这个句子生成表示以用于检索相关文章：{query}"],
                normalize_embeddings=True
            )[0]
        return self.embed_documents([query])[0]

class VectorStoreInterface:
    """向量存储抽象接口"""

    def create_collection(self, name: str, dimension: int, **kwargs): ...
    def insert(self, collection: str, vectors: np.ndarray, documents: List[str], metadatas: List[dict]): ...
    def search(self, collection: str, query_vector: np.ndarray, top_k: int = 10, **kwargs) -> List[dict]: ...
    def delete(self, collection: str, ids: List[str]): ...

class HNSWConfig:
    """HNSW索引参数配置"""

    def __init__(
        self,
        M: int = 16,           # 每个节点的最大连接数（默认16，范围4-64）
        ef_construction: int = 200,  # 构建时的动态列表大小
        ef_search: int = 50,   # 查询时的动态列表大小
        metric: str = "cosine"  # cosine / l2 / ip
    ):
        self.M = M
        self.ef_construction = ef_construction
        self.ef_search = ef_search
        self.metric = metric

    def to_params(self) -> dict:
        return {
            "M": self.M,
            "ef_construction": self.ef_construction,
            "ef_search": self.ef_search,
            "metric": self.metric
        }

    @staticmethod
    def suggest(dataset_size: int, dim: int) -> 'HNSWConfig':
        """根据数据集规模推荐参数"""
        if dataset_size < 100_000:
            return HNSWConfig(M=16, ef_construction=100, ef_search=50)
        elif dataset_size < 1_000_000:
            return HNSWConfig(M=24, ef_construction=200, ef_search=100)
        else:  # 百万级以上
            return HNSWConfig(M=32, ef_construction=400, ef_search=200)

from typing import List, Tuple
import numpy as np
from rank_bm25 import BM25Okapi

class HybridRetriever:
    """混合检索器——向量检索+BM25关键词检索"""

    def __init__(
        self,
        vector_weight: float = 0.7,
        top_k: int = 10
    ):
        self.vector_weight = vector_weight
        self.top_k = top_k
        self.documents: List[str] = []
        self.embeddings: np.ndarray = None
        self.bm25 = None
        self.embed_fn = None

    def build_index(self, documents: List[str], embeddings: np.ndarray, embed_fn):
        self.documents = documents
        self.embeddings = embeddings
        self.embed_fn = embed_fn

        tokenized = [doc.lower().split() for doc in documents]
        self.bm25 = BM25Okapi(tokenized)

    def search(self, query: str, **kwargs) -> List[Tuple[int, float, str]]:
        # 1. 向量检索
        query_vector = self.embed_fn(query)
        vector_scores = np.dot(self.embeddings, query_vector)
        vector_scores = np.argsort(-vector_scores)  # 降序排列

        # 2. BM25检索
        tokenized_query = query.lower().split()
        bm25_scores = self.bm25.get_scores(tokenized_query)
        bm25_ranking = np.argsort(-bm25_scores)

        # 3. RRF（Reciprocal Rank Fusion）融合排序
        from collections import defaultdict
        scores = defaultdict(float)

        for rank, idx in enumerate(vector_scores[:self.top_k * 2]):
            scores[idx] += self.vector_weight * (1.0 / (60 + rank + 1))

        for rank, idx in enumerate(bm25_ranking[:self.top_k * 2]):
            scores[idx] += (1 - self.vector_weight) * (1.0 / (60 + rank + 1))

        # 4. 排序取Top-K
        scored = sorted(scores.items(), key=lambda x: -x[1])[:self.top_k]

        return [(doc_id, score, self.documents[doc_id][:100] + "...") 
                for doc_id, score in scored]

class Reranker:
    """Cross-Encoder重排序"""

    def __init__(self, model_name: str = "BAAI/bge-reranker-v2-m3"):
        from transformers import AutoModelForSequenceClassification, AutoTokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)

    def rerank(self, query: str, candidates: List[Tuple[str, float]], top_k: int = 5) -> List[Tuple[str, float]]:
        pairs = [(query, doc_text) for doc_text, _ in candidates[:50]]  # 最多重排50个候选

        inputs = self.tokenizer(
            pairs, 
            padding=True, 
            truncation=True, 
            max_length=512, 
            return_tensors="pt"
        )

        import torch
        with torch.no_grad():
            outputs = self.model(**inputs)
            scores = outputs.logits.squeeze(-1).tolist()

        if isinstance(scores, float):
            scores = [scores]

        # 按Reranker分数重排
        ranked = sorted(zip(candidates, scores), key=lambda x: -x[1])
        return [(doc_text, score) for (doc_text, _), score in ranked[:top_k]]

class QueryTransformer:
    """智能查询转换"""

    def __init__(self, llm_client):
        self.llm = llm_client

    def expand_query(self, query: str) -> List[str]:
        """查询扩展——生成多个相关查询"""
        prompt = f"""请针对以下用户查询，生成3个不同角度的子查询，
        每个子查询应覆盖不同的语义维度，用于检索知识库：

        用户查询：{query}

        请直接输出3个子查询，每行一个："""

        response = self.llm.generate(prompt)
        sub_queries = [q.strip() for q in response.split('\n') if q.strip()]
        return [query] + sub_queries

    def rewrite_query_for_search(self, query: str) -> str:
        """将对话查询改写为独立搜索查询

        示例：
          「上一个问题的第二部分是什么」→ "RAG系统分块策略详解"
          「那个模型的具体参数」→ "BAAI/bge-large-zh-v1.5 嵌入模型 参数配置"
        """
        prompt = f"""你是一个搜索查询改写专家。将以下对话中的用户查询改写成
        适合向量检索的独立搜索查询（返回仅搜索查询本身）：

        用户查询：{query}
        改写："""

        return self.llm.generate(prompt).strip()

    def decompose_complex_query(self, query: str) -> List[str]:
        """复杂查询分解"""
        prompt = f"""将以下复杂查询分解为多个简单的原子查询：

        原始查询：{query}
        原子查询（每行一个）："""

        response = self.llm.generate(prompt)
        return [q.strip() for q in response.split('\n') if q.strip()]

class ContextAssembler:
    """将检索结果组装为LLM上下文"""

    def __init__(self, max_tokens: int = 4096):
        self.max_tokens = max_tokens
        from tiktoken import encoding_for_model
        self.tokenizer = encoding_for_model("gpt-4")

    def count_tokens(self, text: str) -> int:
        return len(self.tokenizer.encode(text))

    def assemble(
        self,
        query: str,
        retrieved_docs: List[dict],
        system_prompt: str = None
    ) -> dict:
        """智能组装上下文——在Token预算内最大化信息密度"""

        # 默认System Prompt
        if system_prompt is None:
            system_prompt = """你是一个基于知识库的AI助手。请基于以下提供的参考文档，
            准确、专业地回答用户问题。请注意：
            1. 如果参考文档中没有相关信息，请明确告知
            2. 引用信息来源时使用 [来源: doc_id] 格式
            3. 不要编造不在参考文档中的信息
            4. 如果多个文档的信息存在矛盾，请指出不同来源"""

        # 固定消耗
        system_tokens = self.count_tokens(system_prompt)
        query_tokens = self.count_tokens(f"## 用户问题\n{query}")
        format_tokens = 200  # 输出格式、结构字符

        max_context_tokens = self.max_tokens - system_tokens - query_tokens - format_tokens

        # 从最重要的文档开始填充
        context_parts = []
        used_tokens = 0

        # 第一遍：所有文档的摘要
        for i, doc in enumerate(retrieved_docs):
            doc_header = f"[文档 {i+1}] 来源: {doc.get('source', 'unknown')}"
            doc_content = doc['content']
            doc_block = f"{doc_header}\n{doc_content}"
            doc_tokens = self.count_tokens(doc_block)

            if used_tokens + doc_tokens <= max_context_tokens:
                context_parts.append(doc_block)
                used_tokens += doc_tokens
            else:
                # 截断剩余空间
                remaining = max_context_tokens - used_tokens
                if remaining > 100:  # 至少100 tokens才有意义
                    truncated = self.tokenizer.decode(
                        self.tokenizer.encode(doc_content)[:remaining - 50]
                    )
                    context_parts.append(f"{doc_header}\n{truncated}...")
                break

        context = "\n\n---\n\n".join(context_parts)

        full_prompt = f"{system_prompt}\n\n## 参考文档\n{context}\n\n## 用户问题\n{query}"

        return {
            "prompt": full_prompt,
            "total_tokens": self.count_tokens(full_prompt),
            "context_tokens": used_tokens,
            "doc_count": len(retrieved_docs),
            "actual_docs_used": len(context_parts)
        }

class CitationFormatter:
    """格式化的引用溯源"""

    @staticmethod
    def format_answer_with_citations(
        answer: str,
        citations: List[dict]
    ) -> str:
        """将LLM的回答附加上引用标记"""
        formatted = answer + "\n\n---\n## 参考文献\n"

        for i, cit in enumerate(citations, 1):
            formatted += f"\n[{i}] {cit['source']} (段落: {cit.get('chunk_id', 'N/A')})"

        return formatted

    @staticmethod
    def citation_aware_prompt() -> str:
        return """在回答中，当引用特定文档内容时，请在句子末尾添加 
        [来源: N] 标记（N为文档编号）。
        例如: "根据技术文档，RAG系统的检索延迟通常在100ms以内[来源: 3]"

        回答的最后添加"参考文献"部分，列出所有引用的来源。"""

import time
from dataclasses import dataclass, field
from typing import List, Optional

@dataclass
class RAGResult:
    query: str
    answer: str
    contexts: List[str]
    citations: List[dict]
    latency_ms: float
    retrieval_latency: float
    generation_latency: float
    tokens_used: int

class ProductionRAGPipeline:
    """生产级RAG管道"""

    def __init__(
        self,
        retriever: HybridRetriever,
        reranker: Reranker,
        assembler: ContextAssembler,
        llm_client,
        query_transformer: QueryTransformer = None,
        top_k_retrieval: int = 20,
        top_k_rerank: int = 5
    ):
        self.retriever = retriever
        self.reranker = reranker
        self.assembler = assembler
        self.llm = llm_client
        self.query_transformer = query_transformer
        self.top_k_retrieval = top_k_retrieval
        self.top_k_rerank = top_k_rerank

    def query(self, user_query: str, system_prompt: str = None) -> RAGResult:
        start = time.time()
        metrics = {}

        # Step 1: Query Transformation
        t0 = time.time()
        if self.query_transformer:
            queries = self.query_transformer.expand_query(user_query)
        else:
            queries = [user_query]
        metrics['query_transform'] = time.time() - t0

        # Step 2: Retrieval
        t0 = time.time()
        all_results = []
        for q in queries:
            results = self.retriever.search(q, top_k=self.top_k_retrieval)
            all_results.extend(results)
        # 去重
        seen = set()
        unique_results = []
        for r in all_results:
            if r[0] not in seen:
                seen.add(r[0])
                unique_results.append(r)
        metrics['retrieval'] = time.time() - t0

        # Step 3: Rerank
        t0 = time.time()
        reranked = self.reranker.rerank(
            user_query, 
            [(r[2], r[1]) for r in unique_results[:50]],
            top_k=self.top_k_rerank
        )
        metrics['rerank'] = time.time() - t0
        retrieval_total = metrics.get('retrieval', 0) + metrics.get('rerank', 0) + metrics.get('query_transform', 0)

        # Step 4: Context Assembly
        t0 = time.time()
        # 获取完整文档内容
        doc_texts = [doc for doc, score in reranked]
        assembled = self.assembler.assemble(user_query, [{"content": t, "source": f"doc_{i}"} for i, t in enumerate(doc_texts)], system_prompt)
        metrics['assembly'] = time.time() - t0

        # Step 5: Generation
        t0 = time.time()
        response = self.llm.generate(assembled['prompt'])
        metrics['generation'] = time.time() - t0

        total_latency = (time.time() - start) * 1000

        return RAGResult(
            query=user_query,
            answer=response,
            contexts=doc_texts,
            citations=[{"source": f"doc_{i}"} for i in range(len(doc_texts))],
            latency_ms=total_latency,
            retrieval_latency=retrieval_total * 1000,
            generation_latency=metrics.get('generation', 0) * 1000,
            tokens_used=assembled['total_tokens']
        )

import hashlib
import json
from functools import lru_cache
import diskcache as dc

class RAGCache:
    """RAG多级缓存——降低延迟和API成本"""

    def __init__(self, cache_dir: str = "./rag_cache"):
        self.disk_cache = dc.Cache(cache_dir)
        self.hit_count = 0
        self.miss_count = 0

    def _hash_query(self, query: str) -> str:
        return hashlib.sha256(query.encode()).hexdigest()[:16]

    def get(self, query: str, retrieval_threshold: float = 0.95) -> Optional[str]:
        """缓存命中策略"""
        cache_key = self._hash_query(query)

        if cache_key in self.disk_cache:
            self.hit_count += 1
            return self.disk_cache[cache_key]

        self.miss_count += 1
        return None

    def set(self, query: str, answer: str, ttl: int = 86400):
        cache_key = self._hash_query(query)
        self.disk_cache.set(cache_key, answer, expire=ttl)

    def get_stats(self) -> dict:
        total = self.hit_count + self.miss_count
        return {
            "total_requests": total,
            "hit_count": self.hit_count,
            "miss_count": self.miss_count,
            "hit_rate": self.hit_count / total if total > 0 else 0,
            "cache_size": len(self.disk_cache)
        }

# production-rag-config.yaml
api:
  rate_limit: 1000 req/min
  timeout: 30s

retrieval:
  chunk_size: 512
  chunk_overlap: 128
  embedding_model: text-embedding-3-large
  embedding_dimensions: 1024
  vector_store: milvus
  index_type: HNSW
  index_params:
    M: 24
    ef_construction: 200
    ef_search: 100

search:
  hybrid_weight: 0.7  # 向量 vs BM25
  top_k_retrieval: 20
  top_k_rerank: 5

generation:
  model: gpt-4o
  max_tokens: 2048
  temperature: 0.3  # 知识问答使用低温度
  context_limit: 8192  # 上下文Token上限

cache:
  type: diskcache
  ttl_hours: 24
  max_size_gb: 10

monitoring:
  metrics_port: 9090
  log_level: INFO
  trace_all_queries: true

from dataclasses import dataclass
from typing import List

@dataclass
class RAGEvaluationMetrics:
    """RAG系统评估指标"""

    # 检索质量
    recall: float          # Top-K召回率
    precision: float       # 精确率
    mrr: float             # Mean Reciprocal Rank
    ndcg: float            # NDCG

    # 生成质量
    faithfulness: float    # 忠实度（不编造）
    relevance: float       # 相关性
    completeness: float    # 完整性

    # 性能
    p50_latency_ms: float
    p99_latency_ms: float
    throughput_qps: float

class RAGEvaluator:
    """RAG系统评估器"""

    def __init__(self, test_dataset: List[dict]):
        """
        test_dataset: [{"query": "...", "golden_docs": [...], "golden_answer": "..."}]
        """
        self.dataset = test_dataset

    def evaluate_retrieval(self, rag_pipeline) -> dict:
        recalls = []
        mrrs = []

        for item in self.dataset:
            results = rag_pipeline.retrieve(item['query'])

            retrieved_ids = [r['doc_id'] for r in results]
            golden_ids = item['golden_docs']

            # Recall@K
            hits = sum(1 for d in golden_ids if d in retrieved_ids[:10])
            recalls.append(hits / len(golden_ids))

            # MRR
            for rank, r in enumerate(results, 1):
                if r['doc_id'] in golden_ids:
                    mrrs.append(1.0 / rank)
                    break
            else:
                mrrs.append(0.0)

        return {
            "recall@10": sum(recalls) / len(recalls),
            "mrr": sum(mrrs) / len(mrrs)
        }

    def evaluate_generation(self, rag_pipeline) -> dict:
        """评估生成质量——使用LLM-as-Judge"""
        from judge_llm import FaithfulnessJudge

        judge = FaithfulnessJudge()
        faithful_scores = []

        for item in self.dataset[:50]:  # 子集评估
            result = rag_pipeline.query(item['query'])
            score = judge.evaluate_faithfulness(
                answer=result.answer,
                contexts=result.contexts
            )
            faithful_scores.append(score)

        return {
            "faithfulness": sum(faithful_scores) / len(faithful_scores),
            "sample_size": len(faithful_scores)
        }