class EmbeddingFunc:
    """
    LightRAG 的核心 Embedding 抽象
    
    特点：
    1. 维度与 token 限制明确
    2. 支持自定义 embedding 函数
    3. 内置批处理优化
    4. 异步处理支持
    """
    def __init__(
        self,
        embedding_dim: int,
        max_token_size: int = 8192,
        func: callable = None,
    ):
        self.embedding_dim = embedding_dim
        self.max_token_size = max_token_size
        self.func = func
        self.embedding_batch_num = 64  # 批处理大小
    
    async def __call__(self, texts: list[str]) -> np.ndarray:
        """
        执行 embedding，自动批处理
        
        Returns:
            shape: (len(texts), embedding_dim)
        """
        if self.func is None:
            raise ValueError("Embedding function not configured")
        
        # 批处理优化（避免 API 限流）
        batches = [
            texts[i : i + self.embedding_batch_num]
            for i in range(0, len(texts), self.embedding_batch_num)
        ]
        
        # 异步并发处理
        results = await asyncio.gather(*[self.func(batch) for batch in batches])
        
        # 展平结果
        embeddings = [emb for batch_result in results for emb in batch_result]
        
        # 验证维度
        assert len(embeddings) == len(texts), "Embedding count mismatch"
        assert all(len(emb) == self.embedding_dim for emb in embeddings), \
            "Embedding dimension mismatch"
        
        return np.array(embeddings)

# 示例：OpenAI Embedding 函数
async def openai_embedding_func(texts: list[str]) -> list[list[float]]:
    """OpenAI API 调用"""
    from openai import AsyncOpenAI
    
    client = AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    
    response = await client.embeddings.create(
        model="text-embedding-3-large",
        input=texts,
        dimensions=3072  # 可配置维度（3072/1536/256）
    )
    
    return [item.embedding for item in response.data]

# 示例：本地 HuggingFace 模型
async def huggingface_embedding_func(texts: list[str]) -> list[list[float]]:
    """本地 sentence-transformers 模型"""
    from sentence_transformers import SentenceTransformer
    
    # 模型可缓存复用
    if not hasattr(huggingface_embedding_func, "model"):
        huggingface_embedding_func.model = SentenceTransformer(
            "BAAI/bge-base-en-v1.5",
            device="cuda" if torch.cuda.is_available() else "cpu"
        )
    
    embeddings = huggingface_embedding_func.model.encode(
        texts,
        batch_size=32,
        show_progress_bar=False,
        normalize_embeddings=True  # 归一化便于余弦相似度
    )
    
    return embeddings.tolist()

# 使用示例
embedding_func = EmbeddingFunc(
    embedding_dim=3072,
    max_token_size=8192,
    func=openai_embedding_func
)

# 执行 embedding
texts = ["What is RAG?", "Explain vector search"]
embeddings = await embedding_func(texts)

class EmbeddingModel:
    """
    onyx 企业级 Embedding 模型
    
    企业特性：
    1. 模型服务器架构（独立部署）
    2. 精度控制（BFLOAT16/FLOAT32）
    3. 多租户隔离
    4. 性能监控与追踪
    5. 自动重试与降级
    """
    def __init__(
        self,
        model_name: str,
        provider_type: EmbeddingProvider,
        # 企业特性
        server_host: str = "localhost",
        server_port: int = 9000,
        normalize: bool = True,
        precision: str = "FLOAT32",  # BFLOAT16/FLOAT16/FLOAT32
        reduced_dimension: int | None = None,
        # 多租户
        tenant_id: str | None = None,
        # 监控
        enable_tracing: bool = True,
        callback: IndexingHeartbeatInterface | None = None,
    ):
        self.model_name = model_name
        self.provider_type = provider_type
        self.server_host = server_host
        self.server_port = server_port
        self.normalize = normalize
        self.precision = precision
        self.reduced_dimension = reduced_dimension
        self.tenant_id = tenant_id
        self.enable_tracing = enable_tracing
        self.callback = callback
        
        # 初始化模型服务器连接
        self._init_model_server()
    
    def _init_model_server(self):
        """连接到模型服务器（独立进程）"""
        self.client = ModelServerClient(
            host=self.server_host,
            port=self.server_port,
            timeout=30.0
        )
        
        # 健康检查
        if not self.client.health_check():
            raise ConnectionError(f"Model server unavailable: {self.server_host}:{self.server_port}")
    
    def encode(
        self,
        texts: list[str],
        text_type: EmbedTextType,  # PASSAGE / QUERY
        large_chunks_present: bool = False,
        request_id: str | None = None,
    ) -> list[Embedding]:
        """
        执行 embedding with 企业特性
        """
        # 1. 追踪开始
        trace_id = None
        if self.enable_tracing:
            trace_id = self._start_trace(request_id, len(texts))
        
        try:
            # 2. 预处理（token 限制检查）
            processed_texts = self._preprocess_texts(
                texts, 
                large_chunks_present=large_chunks_present
            )
            
            # 3. 调用模型服务器
            embeddings = self.client.embed(
                texts=processed_texts,
                model_name=self.model_name,
                text_type=text_type.value,
                tenant_id=self.tenant_id,
                normalize=self.normalize
            )
            
            # 4. 精度管理（降低存储成本）
            if self.precision != "FLOAT32":
                embeddings = self._convert_precision(embeddings)
            
            # 5. 维度缩减（可选，进一步降低成本）
            if self.reduced_dimension:
                embeddings = self._reduce_dimensions(embeddings)
            
            # 6. 记录指标
            self._record_metrics(trace_id, len(texts), embeddings[0].shape[0])
            
            return embeddings
            
        except Exception as e:
            # 错误追踪
            self._record_error(trace_id, e)
            
            # 重试逻辑
            if self._should_retry(e):
                return self._retry_with_backoff(texts, text_type, request_id)
            
            raise
    
    def _convert_precision(self, embeddings: list[np.ndarray]) -> list[np.ndarray]:
        """精度转换（降低内存/存储）"""
        if self.precision == "BFLOAT16":
            # bfloat16: 节省 50% 空间，精度损失小
            return [emb.astype(np.bfloat16) for emb in embeddings]
        elif self.precision == "FLOAT16":
            # float16: 节省 50% 空间，精度损失中等
            return [emb.astype(np.float16) for emb in embeddings]
        return embeddings
    
    def _reduce_dimensions(self, embeddings: list[np.ndarray]) -> list[np.ndarray]:
        """
        维度缩减（Matryoshka Representation Learning）
        
        某些模型（如 OpenAI embedding-3）支持动态维度：
        - 3072 维：最高精度
        - 1536 维：平衡
        - 256 维：高效存储
        """
        if self.reduced_dimension and embeddings[0].shape[0] > self.reduced_dimension:
            return [emb[:self.reduced_dimension] for emb in embeddings]
        return embeddings

# 模型服务器（独立部署）
class ModelServer:
    """
    独立的 Embedding 模型服务器
    
    优势：
    1. GPU 资源集中管理
    2. 多个 API 服务器共享模型
    3. 热加载/热更新模型
    4. 批处理优化
    """
    def __init__(self, host: str = "0.0.0.0", port: int = 9000):
        self.host = host
        self.port = port
        self.models = {}  # 模型缓存
    
    def load_model(self, model_name: str, device: str = "cuda"):
        """预加载模型到显存"""
        from sentence_transformers import SentenceTransformer
        
        if model_name not in self.models:
            self.models[model_name] = SentenceTransformer(
                model_name,
                device=device
            )
    
    async def embed_batch(
        self, 
        texts: list[str], 
        model_name: str,
        batch_size: int = 64
    ) -> np.ndarray:
        """批处理 embedding（优化 GPU 利用率）"""
        model = self.models.get(model_name)
        if model is None:
            raise ValueError(f"Model not loaded: {model_name}")
        
        # 动态批处理
        embeddings = model.encode(
            texts,
            batch_size=batch_size,
            show_progress_bar=False,
            normalize_embeddings=True,
            convert_to_numpy=True
        )
        
        return embeddings

from sentence_transformers import SentenceTransformer
import numpy as np

class LocalEmbedding:
    """
    极简本地 Embedding 方案
    
    优势：
    1. 完全本地（无 API 成本，隐私保护）
    2. 开箱即用（无复杂配置）
    3. 支持 GPU 加速
    4. 模型可缓存复用
    """
    def __init__(
        self,
        model_name: str = "BAAI/bge-m3",  # 多语言模型
        device: str = None,
        normalize: bool = True
    ):
        """
        推荐模型：
        - 英文：BAAI/bge-base-en-v1.5 (768维, 高质量)
        - 中文：BAAI/bge-large-zh-v1.5 (1024维)
        - 多语言：BAAI/bge-m3 (1024维, 100+语言)
        - 轻量：all-MiniLM-L6-v2 (384维, 快速)
        """
        self.model_name = model_name
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.normalize = normalize
        
        # 加载模型（首次会下载）
        print(f"Loading embedding model: {model_name} on {self.device}")
        self.model = SentenceTransformer(model_name, device=self.device)
        
        # 获取模型信息
        self.embedding_dim = self.model.get_sentence_embedding_dimension()
        self.max_seq_length = self.model.max_seq_length
        
        print(f"Model loaded: dim={self.embedding_dim}, max_seq={self.max_seq_length}")
    
    def embed(
        self, 
        texts: list[str] | str,
        batch_size: int = 32,
        show_progress: bool = False
    ) -> np.ndarray:
        """
        生成 embeddings
        
        Returns:
            shape: (len(texts), embedding_dim)
        """
        if isinstance(texts, str):
            texts = [texts]
        
        embeddings = self.model.encode(
            texts,
            batch_size=batch_size,
            show_progress_bar=show_progress,
            normalize_embeddings=self.normalize,
            convert_to_numpy=True
        )
        
        return embeddings
    
    def similarity(self, text1: str, text2: str) -> float:
        """计算两个文本的余弦相似度"""
        emb1, emb2 = self.embed([text1, text2])
        return float(np.dot(emb1, emb2))

# 使用示例（极简）
embedding_model = LocalEmbedding(model_name="BAAI/bge-base-en-v1.5")

# 生成 embeddings
texts = ["What is machine learning?", "Explain deep learning"]
embeddings = embedding_model.embed(texts)

print(f"Generated {len(embeddings)} embeddings of dimension {embeddings.shape[1]}")

from sentence_transformers import SentenceTransformer, InputExample, losses
from torch.utils.data import DataLoader

def finetune_embedding_model(
    base_model: str = "BAAI/bge-base-en-v1.5",
    train_data: list[tuple[str, str, float]],  # (query, doc, similarity)
    output_path: str = "./finetuned_model",
    epochs: int = 3
):
    """
    微调 embedding 模型
    
    训练数据格式：
    [
        ("query 1", "relevant doc", 1.0),
        ("query 1", "irrelevant doc", 0.0),
        ...
    ]
    """
    # 1. 加载基础模型
    model = SentenceTransformer(base_model)
    
    # 2. 准备训练样本
    train_examples = [
        InputExample(texts=[query, doc], label=score)
        for query, doc, score in train_data
    ]
    
    # 3. 创建 DataLoader
    train_dataloader = DataLoader(
        train_examples, 
        shuffle=True, 
        batch_size=16
    )
    
    # 4. 定义损失函数（Cosine Similarity Loss）
    train_loss = losses.CosineSimilarityLoss(model)
    
    # 5. 训练
    model.fit(
        train_objectives=[(train_dataloader, train_loss)],
        epochs=epochs,
        warmup_steps=100,
        output_path=output_path,
        show_progress_bar=True
    )
    
    print(f"Model fine-tuned and saved to: {output_path}")
    return model

# 使用微调后的模型
finetuned_model = SentenceTransformer("./finetuned_model")
embeddings = finetuned_model.encode(["domain-specific query"])

# OpenAI Matryoshka 支持动态维度
from openai import OpenAI

client = OpenAI()

# 场景1：高精度检索（昂贵但准确）
high_precision = client.embeddings.create(
    model="text-embedding-3-large",
    input=texts,
    dimensions=3072  # 全维度
)

# 场景2：平衡（推荐）
balanced = client.embeddings.create(
    model="text-embedding-3-large",
    input=texts,
    dimensions=1536  # 一半维度，存储减半
)

# 场景3：高效存储（粗排）
efficient = client.embeddings.create(
    model="text-embedding-3-large",
    input=texts,
    dimensions=256   # 仅 8% 存储，适合粗排
)

# 混合策略：粗排 + 精排
# 1. 用 256 维粗排召回 top-100
# 2. 用 3072 维精排选 top-10

import hashlib
from functools import lru_cache

class CachedEmbedding:
    """带缓存的 Embedding（避免重复计算）"""
    
    def __init__(self, embedding_model, cache_dir: str = "./emb_cache"):
        self.model = embedding_model
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(exist_ok=True)
    
    def _cache_key(self, text: str) -> str:
        """生成缓存键"""
        return hashlib.md5(text.encode()).hexdigest()
    
    def embed_with_cache(self, texts: list[str]) -> np.ndarray:
        """检索缓存 + 批量计算未缓存"""
        embeddings = []
        uncached_texts = []
        uncached_indices = []
        
        # 1. 检查缓存
        for i, text in enumerate(texts):
            cache_file = self.cache_dir / f"{self._cache_key(text)}.npy"
            if cache_file.exists():
                embeddings.append(np.load(cache_file))
            else:
                uncached_texts.append(text)
                uncached_indices.append(i)
        
        # 2. 批量计算未缓存
        if uncached_texts:
            new_embeddings = self.model.embed(uncached_texts)
            
            # 保存到缓存
            for text, emb in zip(uncached_texts, new_embeddings):
                cache_file = self.cache_dir / f"{self._cache_key(text)}.npy"
                np.save(cache_file, emb)
            
            # 插入到结果中
            for idx, emb in zip(uncached_indices, new_embeddings):
                embeddings.insert(idx, emb)
        
        return np.array(embeddings)