Few-Shot Prompt 优化：示例选择的算法化方法

┌──────────────────────────────────┐
│        示例选择三要素             │
├──────────────────────────────────┤
│  相关性：示例与当前输入的关联度    │
│  多样性：示例集覆盖不同情况        │
│  一致性：示例间的标注风格统一      │
└──────────────────────────────────┘

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

class SimilarityBasedSelector:
    """基于语义相似度的示例选择"""
    
    def __init__(self, embedding_model="text-embedding-3-large"):
        self.embedding_model = embedding_model
        self.example_pool = []
        self.example_embeddings = []
    
    def add_examples(self, examples: list):
        """添加示例到池"""
        self.example_pool.extend(examples)
        embeddings = self._batch_embed([e['input'] for e in examples])
        self.example_embeddings.extend(embeddings)
    
    def select(self, query: str, k: int = 3) -> list:
        """选择最相似的 k 个示例"""
        query_emb = self._embed(query)
        
        # 计算与所有示例的相似度
        similarities = cosine_similarity(
            [query_emb], self.example_embeddings
        )[0]
        
        # 取 top-k
        top_indices = np.argsort(similarities)[-k:][::-1]
        
        return [self.example_pool[i] for i in top_indices]
    
    def _embed(self, text: str) -> np.ndarray:
        # 使用 embedding API
        pass
    
    def _batch_embed(self, texts: list) -> list:
        pass

class DiversityBasedSelector:
    """基于多样性的示例选择——确保示例覆盖不同情况"""
    
    def __init__(self, embedding_model="text-embedding-3-large"):
        self.embedder = embedding_model
    
    def select(self, query: str, examples: list, k: int = 3) -> list:
        """选择覆盖面最广的 k 个示例"""
        # 1. 计算所有示例与 query 的相似度
        query_emb = self._embed(query)
        example_embs = [self._embed(e['input']) for e in examples]
        
        similarities = [
            cosine_similarity([query_emb], [emb])[0][0]
            for emb in example_embs
        ]
        
        # 2. 使用 MMR (Maximal Marginal Relevance) 算法
        selected = []
        selected_indices = []
        
        # 第一个选最相似的
        first = np.argmax(similarities)
        selected.append(examples[first])
        selected_indices.append(first)
        
        # 后续选择：平衡相关性和多样性
        while len(selected) < k:
            best_score = -float('inf')
            best_idx = -1
            
            for i, example in enumerate(examples):
                if i in selected_indices:
                    continue
                
                # 相关性分数
                relevance = similarities[i]
                
                # 多样性分数（与已选示例的最大相似度的负值）
                diversity = min(
                    cosine_similarity(
                        [example_embs[i]], 
                        [example_embs[j]]
                    )[0][0]
                    for j in selected_indices
                )
                
                # MMR 分数
                mmr_score = 0.7 * relevance - 0.3 * (1 - diversity)
                
                if mmr_score > best_score:
                    best_score = mmr_score
                    best_idx = i
            
            selected.append(examples[best_idx])
            selected_indices.append(best_idx)
        
        return selected

class VotingBasedSelector:
    """基于投票的示例选择——多策略集成"""
    
    def __init__(self):
        self.selectors = [
            SimilarityBasedSelector(),
            DiversityBasedSelector(),
            ComplexityBasedSelector(),
            LabelBalancedSelector(),
        ]
    
    def select(self, query: str, k: int = 3) -> list:
        """多策略投票选择"""
        votes = {}
        
        for selector in self.selectors:
            selected = selector.select(query, k=k)
            for example in selected:
                ex_id = example['id']
                votes[ex_id] = votes.get(ex_id, 0) + 1
        
        # 按票数排序
        ranked = sorted(votes.items(), key=lambda x: -x[1])
        selected_ids = [ex_id for ex_id, _ in ranked[:k]]
        
        return [self.get_example(eid) for eid in selected_ids]

class ComplexityBasedSelector:
    """基于复杂度匹配的示例选择"""
    
    def __init__(self, llm_client):
        self.llm = llm_client
    
    def assess_complexity(self, text: str) -> dict:
        """评估输入的复杂度"""
        prompt = f"""
        评估以下输入的复杂度，返回JSON：
        {{
            "reasoning_depth": 1-5,     // 推理深度
            "knowledge_required": 1-5,  // 所需知识
            "ambiguity": 1-5,           // 歧义程度
            "length": 1-5               // 输入长度
        }}
        
        输入：{text}
        """
        return json.loads(self.llm.generate(prompt))
    
    def select(self, query: str, examples: list, k: int = 3) -> list:
        query_complexity = self.assess_complexity(query)
        
        # 为每个示例评估复杂度
        scored_examples = []
        for example in examples:
            ex_complexity = self.assess_complexity(example['input'])
            
            # 计算复杂度匹配度（欧氏距离的倒数）
            distance = sum(
                (query_complexity[dim] - ex_complexity[dim]) ** 2
                for dim in query_complexity
            ) ** 0.5
            
            score = 1 / (1 + distance)
            scored_examples.append((example, score))
        
        # 按匹配度排序
        scored_examples.sort(key=lambda x: -x[1])
        return [ex for ex, _ in scored_examples[:k]]

class LabelBalancedSelector:
    """标签平衡选择——确保示例标签分布合理"""
    
    def select(self, query: str, examples: list, k: int = 3) -> list:
        # 按 label 分组
        label_groups = {}
        for ex in examples:
            label = ex.get('label', 'unknown')
            if label not in label_groups:
                label_groups[label] = []
            label_groups[label].append(ex)
        
        # 计算每个组应选的数量
        n_labels = len(label_groups)
        per_label = max(1, k // n_labels)
        
        # 从每个组中选最相似的
        selected = []
        for label, group in label_groups.items():
            # 在组内按相似度排序
            scored = [(ex, self._similarity(query, ex['input'])) 
                      for ex in group]
            scored.sort(key=lambda x: -x[1])
            selected.extend([ex for ex, _ in scored[:per_label]])
        
        # 如果不够 k 个，从剩余中补充
        while len(selected) < k:
            remaining = [ex for ex in examples if ex not in selected]
            if not remaining:
                break
            selected.append(remaining[0])
        
        return selected[:k]

class ExampleOrderOptimizer:
    """示例顺序优化"""
    
    def __init__(self, llm_client):
        self.llm = llm_client
    
    def optimize_order(self, examples: list, query: str) -> list:
        """优化示例排列顺序"""
        # 策略1：复杂度递增（简单→复杂）
        ordered = sorted(examples, key=lambda e: e.get('complexity', 3))
        
        # 策略2：与query最相似的放最后（近因效应）
        similarities = [self._similarity(query, e['input']) for e in ordered]
        # 将最相似的移到最后
        max_sim_idx = similarities.index(max(similarities))
        ordered.append(ordered.pop(max_sim_idx))
        
        return ordered

EXAMPLE_FORMATS = {
    'minimal': '{input}\n→ {output}',
    
    'explained': '输入：{input}\n分析：{reasoning}\n输出：{output}',
    
    'structured': """
<example>
  <input>{input}</input>
  <reasoning>{reasoning}</reasoning>
  <output>{output}</output>
</example>""",
    
    'conversational': '用户：{input}\n助手：{output}',
    
    'annotated': '{input}\n[正确答案：{output}]\n[原因：{reasoning}]',
}

def select_format(task_type: str) -> str:
    """根据任务类型选择示例格式"""
    mapping = {
        'classification': 'minimal',      # 分类任务用最简格式
        'generation': 'explained',        # 生成任务需要推理过程
        'extraction': 'structured',       # 信息提取用结构化
        'conversation': 'conversational', # 对话任务用对话格式
        'reasoning': 'annotated',         # 推理任务需要标注
    }
    return mapping.get(task_type, 'explained')

class AdaptiveExampleSelector:
    """自适应示例选择系统——根据输入特征动态选择策略"""
    
    def __init__(self, llm_client, embedding_model):
        self.llm = llm_client
        self.similarity_selector = SimilarityBasedSelector(embedding_model)
        self.diversity_selector = DiversityBasedSelector(embedding_model)
        self.complexity_selector = ComplexityBasedSelector(llm_client)
        self.label_balanced_selector = LabelBalancedSelector()
        self.order_optimizer = ExampleOrderOptimizer(llm_client)
    
    def select(self, query: str, example_pool: list, 
               k: int = 3, task_type: str = None) -> list:
        """自适应选择示例"""
        # 1. 分析输入特征
        features = self._analyze_query(query)
        
        # 2. 选择策略
        strategy = self._select_strategy(features, task_type)
        
        # 3. 执行选择
        if strategy == 'similarity':
            examples = self.similarity_selector.select(query, example_pool, k)
        elif strategy == 'diversity':
            examples = self.diversity_selector.select(query, example_pool, k)
        elif strategy == 'complexity':
            examples = self.complexity_selector.select(query, example_pool, k)
        elif strategy == 'balanced':
            examples = self.label_balanced_selector.select(query, example_pool, k)
        elif strategy == 'hybrid':
            examples = self._hybrid_select(query, example_pool, k)
        
        # 4. 优化顺序
        examples = self.order_optimizer.optimize_order(examples, query)
        
        return examples
    
    def _analyze_query(self, query: str) -> dict:
        """分析输入特征"""
        return {
            'length': len(query.split()),
            'complexity_signals': sum(1 for w in query.split() 
                                     if len(w) > 8),
            'has_numbers': any(c.isdigit() for c in query),
            'language': self._detect_language(query),
        }
    
    def _select_strategy(self, features: dict, task_type: str = None) -> str:
        """根据特征选择策略"""
        if task_type == 'classification':
            return 'balanced'
        elif task_type == 'generation':
            return 'diversity'
        elif features['complexity_signals'] > 3:
            return 'complexity'
        elif features['length'] > 50:
            return 'similarity'
        else:
            return 'hybrid'
    
    def _hybrid_select(self, query: str, pool: list, k: int) -> list:
        """混合策略选择"""
        # 先用相似度选 2k 个候选
        candidates = self.similarity_selector.select(query, pool, k * 2)
        # 再用多样性从候选中选 k 个
        return self.diversity_selector.select(query, candidates, k)

class ExampleSelectionEvaluator:
    """示例选择效果评估"""
    
    def evaluate(self, selector, test_set: list, 
                 baseline_selector=None) -> dict:
        results = {
            'selector': selector.__class__.__name__,
            'metrics': {
                'accuracy': [],
                'consistency': [],
                'latency_ms': [],
            }
        }
        
        for case in test_set:
            import time
            start = time.time()
            
            examples = selector.select(case['input'], case['pool'], k=3)
            prompt = self._build_prompt(examples, case['input'])
            response = self.llm.generate(prompt)
            
            results['metrics']['accuracy'].append(
                self._score(response, case['expected'])
            )
            results['metrics']['latency_ms'].append(
                (time.time() - start) * 1000
            )
        
        # 汇总
        summary = {
            'mean_accuracy': np.mean(results['metrics']['accuracy']),
            'std_accuracy': np.std(results['metrics']['accuracy']),
            'mean_latency': np.mean(results['metrics']['latency_ms']),
        }
        
        return summary

class ExampleOptimizationLoop:
    """示例池持续优化"""
    
    def __init__(self, selector, llm_client):
        self.selector = selector
        self.llm = llm_client
        self.performance_log = []
    
    def step(self, query: str, response: str, 
             expected: str, examples_used: list):
        """记录每次使用并优化"""
        correct = self._is_correct(response, expected)
        
        self.performance_log.append({
            'query': query,
            'examples': examples_used,
            'correct': correct,
            'response': response,
        })
        
        # 每积累100条记录，优化示例池
        if len(self.performance_log) % 100 == 0:
            self._optimize_pool()
    
    def _optimize_pool(self):
        """基于历史表现优化示例池"""
        # 找出好示例（使用后正确率高）和坏示例
        example_stats = {}
        for log in self.performance_log:
            for ex in log['examples']:
                ex_id = ex['id']
                if ex_id not in example_stats:
                    example_stats[ex_id] = {'correct': 0, 'total': 0}
                example_stats[ex_id]['total'] += 1
                if log['correct']:
                    example_stats[ex_id]['correct'] += 1
        
        # 标记低效示例
        for ex_id, stats in example_stats.items():
            success_rate = stats['correct'] / stats['total']
            if success_rate < 0.5 and stats['total'] > 5:
                print(f"示例 {ex_id} 成功率低 ({success_rate:.0%})，建议替换")