03 - 推理服务部署（全面版）¶

# 基础推理示例
from vllm import LLM, SamplingParams

# 加载模型
llm = LLM(
    model="meta-llama/Llama-2-7b-hf",
    tensor_parallel_size=1,  # 单卡
    gpu_memory_utilization=0.9,  # GPU内存使用率
)

# 设置采样参数
sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.9,
    max_tokens=256,
)

# 生成
prompts = [
    "The future of AI is",
    "In the beginning,",
    "Once upon a time,",
]

outputs = llm.generate(prompts, sampling_params)

# 打印结果
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}")
    print(f"Generated: {generated_text!r}")
    print("-" * 50)

from vllm import LLM, SamplingParams

class VLLMDeployment:
    """
    vLLM高级部署配置
    """

    def __init__(self, model_path, config=None):
        self.config = config or {}

        # 初始化LLM
        self.llm = LLM(
            model=model_path,

            # 并行配置
            tensor_parallel_size=self.config.get('tensor_parallel_size', 1),
            pipeline_parallel_size=self.config.get('pipeline_parallel_size', 1),

            # 内存配置
            gpu_memory_utilization=self.config.get('gpu_memory_utilization', 0.9),
            max_num_seqs=self.config.get('max_num_seqs', 256),
            max_model_len=self.config.get('max_model_len', 4096),

            # 量化配置
            quantization=self.config.get('quantization', None),
            # 可选: 'awq', 'gptq', 'squeezellm'

            # 其他配置
            dtype=self.config.get('dtype', 'auto'),
            # 可选: 'float16', 'bfloat16', 'float32'

            trust_remote_code=self.config.get('trust_remote_code', True),
        )

    def generate(self, prompts, **sampling_kwargs):
        """
        批量生成
        """
        sampling_params = SamplingParams(
            temperature=sampling_kwargs.get('temperature', 0.7),
            top_p=sampling_kwargs.get('top_p', 0.9),
            top_k=sampling_kwargs.get('top_k', -1),
            max_tokens=sampling_kwargs.get('max_tokens', 256),
            presence_penalty=sampling_kwargs.get('presence_penalty', 0.0),
            frequency_penalty=sampling_kwargs.get('frequency_penalty', 0.0),
            stop=sampling_kwargs.get('stop', None),
        )

        outputs = self.llm.generate(prompts, sampling_params)

        # 格式化输出
        results = []
        for output in outputs:
            results.append({
                'prompt': output.prompt,
                'text': output.outputs[0].text,
                'tokens': len(output.outputs[0].token_ids),
            })

        return results

    def chat(self, messages, **kwargs):  # **kwargs收集关键字参数
        """
        对话模式
        """
        # 构建prompt
        prompt = self._build_chat_prompt(messages)

        # 生成
        result = self.generate([prompt], **kwargs)

        return result[0]['text']

    def _build_chat_prompt(self, messages):
        """
        构建对话prompt
        """
        # 使用模型特定的chat template
        if hasattr(self.llm.get_tokenizer(), 'apply_chat_template'):  # hasattr检查对象是否有某属性
            prompt = self.llm.get_tokenizer().apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True
            )
        else:
            # 简单的对话格式
            prompt = ""
            for msg in messages:
                role = msg['role']
                content = msg['content']
                prompt += f"{role}: {content}\n"
            prompt += "assistant:"

        return prompt

# 配置示例
VLLM_CONFIG = {
    'tensor_parallel_size': 2,  # 2卡并行
    'gpu_memory_utilization': 0.85,
    'max_num_seqs': 128,
    'max_model_len': 8192,
    'quantization': None,
    'dtype': 'bfloat16',
}

# 使用示例
deployment = VLLMDeployment("meta-llama/Llama-2-7b-hf", VLLM_CONFIG)

# 批量生成
results = deployment.generate(
    ["Hello, how are you?", "What is machine learning?"],
    temperature=0.7,
    max_tokens=100
)

# 对话模式
response = deployment.chat([
    {'role': 'user', 'content': 'Explain quantum computing'}
])

# 启动vLLM OpenAI兼容API服务
# 命令行方式（vLLM 0.6+推荐使用 vllm serve）
"""
vllm serve meta-llama/Llama-2-7b-hf \
    --tensor-parallel-size 2 \
    --gpu-memory-utilization 0.85 \
    --max-num-seqs 256 \
    --port 8000

# 旧版本也可使用：
# python -m vllm.entrypoints.openai.api_server \
#     --model meta-llama/Llama-2-7b-hf ...
"""

# 使用Python自定义服务
import uvicorn
from fastapi import FastAPI
from vllm import LLM, SamplingParams

app = FastAPI()

# 全局模型实例
llm_engine = None

@app.on_event("startup")
async def startup_event():  # async def定义协程函数
    global llm_engine
    llm_engine = LLM(
        model="meta-llama/Llama-2-7b-hf",
        tensor_parallel_size=2,
    )

@app.post("/v1/completions")
async def create_completion(request: CompletionRequest):
    """
    OpenAI兼容的completions API
    """
    sampling_params = SamplingParams(
        temperature=request.temperature,
        max_tokens=request.max_tokens,
        top_p=request.top_p,
    )

    outputs = llm_engine.generate(request.prompt, sampling_params)

    return {
        "id": "cmpl-" + str(uuid.uuid4()),
        "object": "text_completion",
        "created": int(time.time()),
        "model": request.model,
        "choices": [
            {
                "text": output.outputs[0].text,
                "index": i,
                "logprobs": None,
                "finish_reason": "stop"
            }
            for i, output in enumerate(outputs)  # enumerate同时获取索引和元素
        ]
    }

@app.post("/v1/chat/completions")
async def create_chat_completion(request: ChatCompletionRequest):
    """
    OpenAI兼容的chat completions API
    """
    # 构建prompt
    prompt = build_chat_prompt(request.messages)

    sampling_params = SamplingParams(
        temperature=request.temperature,
        max_tokens=request.max_tokens,
    )

    outputs = llm_engine.generate(prompt, sampling_params)

    return {
        "id": "chatcmpl-" + str(uuid.uuid4()),
        "object": "chat.completion",
        "created": int(time.time()),
        "model": request.model,
        "choices": [
            {
                "index": i,
                "message": {
                    "role": "assistant",
                    "content": output.outputs[0].text
                },
                "finish_reason": "stop"
            }
            for i, output in enumerate(outputs)
        ]
    }

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

# Python客户端
from text_generation import Client

client = Client("http://localhost:8080")

# 文本生成
text = client.generate(
    "The future of AI is",
    max_new_tokens=100,
    temperature=0.7,
    top_p=0.9,
).generated_text

print(text)

# 流式生成
for response in client.generate_stream(
    "Explain machine learning:",
    max_new_tokens=100
):
    print(response.token.text, end="", flush=True)

# 使用AutoGPTQ进行量化
from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig

# 量化配置
quantize_config = BaseQuantizeConfig(
    bits=4,  # 4-bit量化
    group_size=128,  # 分组大小
    desc_act=False,  # 是否降序激活
)

# 加载并量化模型
model = AutoGPTQForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantize_config,
)

# 准备校准数据
calib_data = [
    "auto-gptq is an easy-to-use model quantization library",
    "with user-friendly apis",
    # ... 更多校准数据
]

# 执行量化
model.quantize(calib_data)

# 保存量化模型
model.save_quantized("Llama-2-7b-4bit-gptq")

# 加载量化模型进行推理
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

model = AutoGPTQForCausalLM.from_quantized(
    "Llama-2-7b-4bit-gptq",
    device="cuda:0",
)

# 生成
inputs = tokenizer("Hello, how are you?", return_tensors="pt").to("cuda")
outputs = model.generate(**inputs, max_new_tokens=50)
print(tokenizer.decode(outputs[0]))

# AWQ量化（通常比GPTQ质量更好）
from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

# 加载模型
model_path = "meta-llama/Llama-2-7b-hf"
quant_path = "Llama-2-7b-awq"

# 量化配置
quant_config = {
    "zero_point": True,
    "q_group_size": 128,
    "w_bit": 4,
    "version": "GEMM"
}

# 加载模型
model = AutoAWQForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

# 准备校准数据
examples = [
    tokenizer("auto-gptq is an easy-to-use model quantization library")
    for _ in range(8)
]

# 量化
model.quantize(
    tokenizer,
    quant_config=quant_config,
    calib_data=examples,
)

# 保存
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)

# 加载量化模型
model = AutoAWQForCausalLM.from_quantized(
    quant_path,
    fuse_layers=True,  # 融合层以加速
)

# vLLM支持多种量化方式

# 1. AWQ量化
llm = LLM(
    model="TheBloke/Llama-2-7B-AWQ",
    quantization="awq",
    tensor_parallel_size=1,
)

# 2. GPTQ量化
llm = LLM(
    model="TheBloke/Llama-2-7B-GPTQ",
    quantization="gptq",
    tensor_parallel_size=1,
)

# 3. SqueezeLLM量化
llm = LLM(
    model="path/to/squeezellm/model",
    quantization="squeezellm",
)

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional
import uvicorn
from vllm import LLM, SamplingParams

app = FastAPI(title="LLM Inference API")

# 全局模型
llm = None

class CompletionRequest(BaseModel):  # Pydantic BaseModel：自动数据验证和序列化
    model: str
    prompt: str
    max_tokens: Optional[int] = 256  # Optional表示值可以为None
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 0.9
    top_k: Optional[int] = -1
    stop: Optional[List[str]] = None
    stream: Optional[bool] = False

class ChatCompletionRequest(BaseModel):
    model: str
    messages: List[dict]
    max_tokens: Optional[int] = 256
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 0.9
    stream: Optional[bool] = False

@app.on_event("startup")
async def load_model():
    global llm
    llm = LLM(
        model="meta-llama/Llama-2-7b-hf",
        tensor_parallel_size=1,
        gpu_memory_utilization=0.9,
    )
    print("Model loaded successfully")

@app.post("/v1/completions")
async def create_completion(request: CompletionRequest):
    try:  # try/except捕获异常，防止程序崩溃
        sampling_params = SamplingParams(
            temperature=request.temperature,
            top_p=request.top_p,
            top_k=request.top_k,
            max_tokens=request.max_tokens,
            stop=request.stop,
        )

        outputs = llm.generate(request.prompt, sampling_params)

        return {
            "id": f"cmpl-{uuid.uuid4()}",
            "object": "text_completion",
            "created": int(time.time()),
            "model": request.model,
            "choices": [
                {
                    "text": output.outputs[0].text,
                    "index": i,
                    "logprobs": None,
                    "finish_reason": "stop"
                }
                for i, output in enumerate(outputs)
            ],
            "usage": {
                "prompt_tokens": len(outputs[0].prompt_token_ids),
                "completion_tokens": len(outputs[0].outputs[0].token_ids),
                "total_tokens": len(outputs[0].prompt_token_ids) + len(outputs[0].outputs[0].token_ids)
            }
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/v1/chat/completions")
async def create_chat_completion(request: ChatCompletionRequest):
    try:
        # 构建chat prompt
        prompt = build_chat_prompt(request.messages)

        sampling_params = SamplingParams(
            temperature=request.temperature,
            top_p=request.top_p,
            max_tokens=request.max_tokens,
        )

        outputs = llm.generate(prompt, sampling_params)

        return {
            "id": f"chatcmpl-{uuid.uuid4()}",
            "object": "chat.completion",
            "created": int(time.time()),
            "model": request.model,
            "choices": [
                {
                    "index": i,
                    "message": {
                        "role": "assistant",
                        "content": output.outputs[0].text
                    },
                    "finish_reason": "stop"
                }
                for i, output in enumerate(outputs)
            ],
            "usage": {
                "prompt_tokens": len(outputs[0].prompt_token_ids),
                "completion_tokens": len(outputs[0].outputs[0].token_ids),
                "total_tokens": len(outputs[0].prompt_token_ids) + len(outputs[0].outputs[0].token_ids)
            }
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

def build_chat_prompt(messages):
    """构建chat prompt"""
    prompt = ""
    for msg in messages:
        role = msg.get('role', '')
        content = msg.get('content', '')
        if role == 'system':
            prompt += f"[INST] <<SYS>>\n{content}\n<</SYS>>\n\n"
        elif role == 'user':
            prompt += f"{content} [/INST]"
        elif role == 'assistant':
            prompt += f" {content} </s><s>[INST] "
    return prompt

@app.get("/health")
async def health_check():
    return {"status": "healthy", "model_loaded": llm is not None}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

import requests
import json

class LLMClient:
    """
    LLM API客户端
    """
    def __init__(self, base_url="http://localhost:8000"):
        self.base_url = base_url

    def complete(self, prompt, **kwargs):
        """
        文本补全
        """
        response = requests.post(
            f"{self.base_url}/v1/completions",
            json={
                "model": "llama-2-7b",
                "prompt": prompt,
                **kwargs
            }
        )
        return response.json()

    def chat(self, messages, **kwargs):
        """
        对话
        """
        response = requests.post(
            f"{self.base_url}/v1/chat/completions",
            json={
                "model": "llama-2-7b",
                "messages": messages,
                **kwargs
            }
        )
        return response.json()

# 使用示例
client = LLMClient()

# 文本补全
result = client.complete(
    "The future of AI is",
    max_tokens=100,
    temperature=0.7
)
print(result['choices'][0]['text'])

# 对话
result = client.chat([
    {"role": "user", "content": "Explain quantum computing"}
])
print(result['choices'][0]['message']['content'])

class PerformanceMonitor:
    """
    性能监控器
    """

    def __init__(self):
        self.metrics = {
            'latency': [],  # 延迟（秒）
            'throughput': [],  # 吞吐量（tokens/秒）
            'queue_size': [],  # 队列大小
        }

    def measure_latency(self, func, *args, **kwargs):  # func为传入的函数，*args/**kwargs透传任意参数，实现通用性能测量
        """
        测量延迟
        """
        import time

        start = time.time()
        result = func(*args, **kwargs)  # 将收集的参数解包后原样传给被测量的函数
        end = time.time()

        latency = end - start
        self.metrics['latency'].append(latency)

        return result, latency

    def calculate_throughput(self, num_tokens, latency):
        """
        计算吞吐量
        """
        throughput = num_tokens / latency
        self.metrics['throughput'].append(throughput)
        return throughput

    def get_statistics(self):
        """
        获取统计信息
        """
        import numpy as np

        stats = {}
        for key, values in self.metrics.items():
            if values:
                stats[key] = {
                    'mean': np.mean(values),
                    'median': np.median(values),
                    'p95': np.percentile(values, 95),
                    'p99': np.percentile(values, 99),
                    'min': np.min(values),
                    'max': np.max(values),
                }

        return stats

# 基准测试
def benchmark_inference(model, prompts, max_tokens=100):
    """
    推理基准测试
    """
    monitor = PerformanceMonitor()

    for prompt in prompts:
        # 测量延迟
        outputs, latency = monitor.measure_latency(
            model.generate,
            prompt,
            max_tokens=max_tokens
        )

        # 计算吞吐量
        num_tokens = len(outputs[0].outputs[0].token_ids)
        throughput = monitor.calculate_throughput(num_tokens, latency)

        print(f"Prompt: {prompt[:50]}...")
        print(f"Latency: {latency:.3f}s")
        print(f"Throughput: {throughput:.1f} tokens/s")
        print("-" * 50)

    # 打印统计
    stats = monitor.get_statistics()
    print("\n=== Benchmark Results ===")
    for metric, values in stats.items():
        print(f"\n{metric.upper()}:")
        for stat, value in values.items():
            print(f"  {stat}: {value:.3f}")

from vllm import LLM, SamplingParams

# 启用 Automatic Prefix Caching（APC）
llm = LLM(
    model="meta-llama/Llama-3-70B-Instruct",
    enable_prefix_caching=True,          # 关键参数
    gpu_memory_utilization=0.90,
    tensor_parallel_size=4,
    max_num_seqs=256,
    # 缓存预热：启动时预计算固定 System Prompt
    # preemption_mode="recompute",       # 显存不足时重算（默认）
    # preemption_mode="swap",            # 显存不足时换出到 CPU
)

# ---- 场景 1：多用户共享 System Prompt ----
SYSTEM_PROMPT = """
You are an expert AI assistant with deep knowledge in...
[此处填充 2000 tokens 的详细系统提示]
"""

def build_messages(user_question: str):
    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user",   "content": user_question},
    ]

# 第一批请求：System Prompt 被计算并缓存
prompts_batch1 = [
    llm.get_tokenizer().apply_chat_template(
        build_messages(q), tokenize=False, add_generation_prompt=True
    )
    for q in ["What is RAG?", "Explain Attention", "How does RLHF work?"]
]

sampling_params = SamplingParams(temperature=0.7, max_tokens=512)
# 首次生成：System Prompt 进行 Prefill 并缓存
outputs1 = llm.generate(prompts_batch1, sampling_params)

# 第二批请求：System Prompt KV Cache 命中，TTFT 大幅下降
prompts_batch2 = [
    llm.get_tokenizer().apply_chat_template(
        build_messages(q), tokenize=False, add_generation_prompt=True
    )
    for q in ["What is MoE?", "Explain vLLM", "How does LoRA work?"]
]
# 再次生成：System Prompt 命中缓存，只 Prefill 用户输入部分
outputs2 = llm.generate(prompts_batch2, sampling_params)


# ---- 场景 2：多轮对话的前缀复用 ----
class CachedMultiTurnChat:
    """多轮对话：历史轮次命中 Prefix Cache"""

    def __init__(self, llm_engine, system_prompt: str):
        self.llm = llm_engine
        self.system_prompt = system_prompt
        self.history = []

    def chat(self, user_input: str) -> str:
        self.history.append({"role": "user", "content": user_input})

        messages = [{"role": "system", "content": self.system_prompt}] + self.history
        prompt = self.llm.get_tokenizer().apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )

        # 每次请求，历史对话前缀命中缓存，只需计算新增 token
        output = self.llm.generate(
            [prompt],
            SamplingParams(temperature=0.7, max_tokens=256)
        )
        response = output[0].outputs[0].text
        self.history.append({"role": "assistant", "content": response})
        return response

import time
import threading
from collections import defaultdict
from typing import Dict, List, Optional

class PrefixCacheMonitor:
    """
    生产级 Prefix Cache 命中率监控
    支持实时统计、滑动窗口、告警
    """

    def __init__(self, window_size: int = 1000):
        self.window_size = window_size
        self._lock = threading.Lock()

        # 总体统计
        self.total_prefill_tokens = 0
        self.cache_hit_tokens = 0
        self.total_requests = 0
        self.hit_requests = 0

        # 滑动窗口（最近 N 次请求）
        self.window: List[Dict] = []

        # 按场景分类统计
        self.scene_stats: Dict[str, Dict] = defaultdict(lambda: {
            "total_tokens": 0,
            "hit_tokens": 0,
            "requests": 0,
        })

    def record(self,
               prompt_tokens: int,
               cached_tokens: int,
               scene: str = "default"):
        """
        记录一次请求的缓存情况

        Args:
            prompt_tokens: 本次 Prefill 的总 token 数
            cached_tokens: 其中命中缓存的 token 数（无需重算）
            scene: 业务场景标记（用于分场景分析）
        """
        with self._lock:
            self.total_prefill_tokens += prompt_tokens
            self.cache_hit_tokens += cached_tokens
            self.total_requests += 1
            if cached_tokens > 0:
                self.hit_requests += 1

            # 场景统计
            self.scene_stats[scene]["total_tokens"] += prompt_tokens
            self.scene_stats[scene]["hit_tokens"] += cached_tokens
            self.scene_stats[scene]["requests"] += 1

            # 滑动窗口
            entry = {
                "timestamp": time.time(),
                "prompt_tokens": prompt_tokens,
                "cached_tokens": cached_tokens,
                "hit_rate": cached_tokens / max(prompt_tokens, 1),
                "scene": scene,
            }
            self.window.append(entry)
            if len(self.window) > self.window_size:
                self.window.pop(0)

    @property
    def overall_token_hit_rate(self) -> float:
        """整体 Token 级命中率"""
        if self.total_prefill_tokens == 0:
            return 0.0
        return self.cache_hit_tokens / self.total_prefill_tokens

    @property
    def overall_request_hit_rate(self) -> float:
        """整体请求级命中率"""
        if self.total_requests == 0:
            return 0.0
        return self.hit_requests / self.total_requests

    @property
    def recent_token_hit_rate(self) -> float:
        """最近滑动窗口内的 Token 命中率"""
        if not self.window:
            return 0.0
        total = sum(e["prompt_tokens"] for e in self.window)
        hit = sum(e["cached_tokens"] for e in self.window)
        return hit / max(total, 1)

    def compute_cost_saving(self,
                            price_per_1k_tokens: float = 0.003) -> Dict:
        """
        计算节省的 API 费用（适用于 OpenAI / Anthropic 等按 Token 计费场景）

        Args:
            price_per_1k_tokens: 每千输入 token 的价格（美元）
        """
        saved_tokens = self.cache_hit_tokens
        original_cost = self.total_prefill_tokens / 1000 * price_per_1k_tokens
        saved_cost = saved_tokens / 1000 * price_per_1k_tokens
        return {
            "total_prefill_tokens": self.total_prefill_tokens,
            "saved_tokens": saved_tokens,
            "original_cost_usd": round(original_cost, 4),
            "saved_cost_usd": round(saved_cost, 4),
            "saving_rate": round(saved_cost / max(original_cost, 1e-9), 4),
        }

    def get_report(self) -> str:
        """输出可读性报告"""
        lines = [
            "=" * 60,
            "  Prefix Cache 命中率报告",
            "=" * 60,
            f"  总请求数        : {self.total_requests:,}",
            f"  总 Prefill Token: {self.total_prefill_tokens:,}",
            f"  命中 Token 数   : {self.cache_hit_tokens:,}",
            f"  Token 命中率    : {self.overall_token_hit_rate:.1%}",
            f"  请求命中率      : {self.overall_request_hit_rate:.1%}",
            f"  近期命中率(滑窗): {self.recent_token_hit_rate:.1%}",
            "-" * 60,
            "  分场景统计：",
        ]
        for scene, stat in self.scene_stats.items():
            scene_rate = stat["hit_tokens"] / max(stat["total_tokens"], 1)
            lines.append(
                f"    [{scene}] 命中率={scene_rate:.1%}  "
                f"请求={stat['requests']}  "
                f"hit_tokens={stat['hit_tokens']:,}"
            )
        lines.append("=" * 60)
        return "\n".join(lines)


# ---- 与 vLLM 集成的封装示例 ----
monitor = PrefixCacheMonitor(window_size=2000)

from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-3-8B-Instruct",
    enable_prefix_caching=True,
    gpu_memory_utilization=0.90,
)

def monitored_generate(prompts: List[str],
                       scene: str = "default",
                       **kwargs) -> List:
    """
    带监控的推理封装：自动记录缓存命中情况
    """
    params = SamplingParams(**kwargs)
    outputs = llm.generate(prompts, params)

    for output in outputs:
        prompt_tokens = len(output.prompt_token_ids)
        # vLLM RequestOutput 包含 num_cached_tokens 字段（v0.5+）
        cached_tokens = getattr(output, "num_cached_tokens", 0)
        monitor.record(prompt_tokens, cached_tokens, scene=scene)

    return outputs


# 使用示例
results = monitored_generate(
    ["What is RAG?", "Explain Attention Mechanism"],
    scene="qa",
    temperature=0.7,
    max_tokens=256,
)
print(monitor.get_report())

# SGLang 以激进的 KV Cache 复用著称
# pip install sglang[all]

import sglang as sgl

@sgl.function
def multi_turn_qa(s, system_prompt, questions):
    """
    SGLang 原生多轮对话：自动识别可复用前缀
    """
    s += sgl.system(system_prompt)
    answers = []
    for q in questions:
        s += sgl.user(q)
        s += sgl.assistant(sgl.gen("answer", max_new_tokens=256))
        answers.append(s["answer"])
    return answers

# 初始化运行时（启用 Radix Cache）
runtime = sgl.Runtime(
    model_path="meta-llama/Llama-3-8B-Instruct",
    tp_size=1,
    # RadixAttention 默认开启，无需额外参数
)
sgl.set_default_backend(runtime)

# 批处理：相同 System Prompt → 自动共享前缀 KV Cache
SYSTEM = "You are a helpful coding assistant specializing in Python."
queries = [
    ["What is a decorator?", "How does yield work?"],
    ["Explain asyncio", "What is GIL?"],
]

for q_list in queries:
    results = multi_turn_qa.run(
        system_prompt=SYSTEM,
        questions=q_list,
    )
    print(results["answer"])

runtime.shutdown()

import anthropic
import time
from typing import List, Dict

client = anthropic.Anthropic()

# =====================================================================
# Anthropic Prompt Caching 实战
# =====================================================================

LARGE_SYSTEM_PROMPT = """
You are an expert AI assistant with the following detailed knowledge base:
[此处插入 2000+ tokens 的长系统提示、RAG 结果、Few-shot 示例等]
"""

def call_with_cache(user_message: str) -> Dict:
    """
    使用 Prompt Caching 的 Anthropic API 调用
    """
    response = client.messages.create(
        model="claude-opus-4-5",
        max_tokens=1024,
        system=[
            {
                "type": "text",
                "text": LARGE_SYSTEM_PROMPT,
                "cache_control": {"type": "ephemeral"},  # 显式标记缓存点
            }
        ],
        messages=[{"role": "user", "content": user_message}],
    )

    usage = response.usage
    return {
        "content": response.content[0].text,
        "input_tokens": usage.input_tokens,
        "cache_creation_tokens": getattr(usage, "cache_creation_input_tokens", 0),
        "cache_read_tokens": getattr(usage, "cache_read_input_tokens", 0),
    }

def demo_cache_savings():
    """演示缓存带来的成本节约"""
    questions = [
        "Summarize the key points",
        "What are the main challenges?",
        "Suggest three improvements",
        "Compare with alternative approaches",
    ]

    total_input = 0
    total_cache_read = 0
    total_cache_create = 0

    for i, q in enumerate(questions):
        result = call_with_cache(q)
        total_input += result["input_tokens"]
        total_cache_read += result["cache_read_tokens"]
        total_cache_create += result["cache_creation_tokens"]

        status = "CACHE HIT" if result["cache_read_tokens"] > 0 else "CACHE MISS"
        print(f"请求 {i+1}: [{status}] "
              f"input={result['input_tokens']} "
              f"cache_read={result['cache_read_tokens']} "
              f"cache_create={result['cache_creation_tokens']}")

    # 成本计算（claude-opus-4-5 定价）
    PRICE_INPUT = 15.0 / 1_000_000   # $15/M tokens
    PRICE_CACHE_CREATE = 18.75 / 1_000_000  # $18.75/M tokens（创建缓存溢价）
    PRICE_CACHE_READ = 1.5 / 1_000_000    # $1.5/M tokens（缓存读取 -90%）

    actual_cost = (
        total_input * PRICE_INPUT
        + total_cache_create * PRICE_CACHE_CREATE
        + total_cache_read * PRICE_CACHE_READ
    )
    no_cache_cost = (total_input + total_cache_read + total_cache_create) * PRICE_INPUT

    print(f"\n总计：")
    print(f"  实际成本（有缓存）: ${actual_cost:.6f}")
    print(f"  理论成本（无缓存）: ${no_cache_cost:.6f}")
    print(f"  节省比例         : {1 - actual_cost/no_cache_cost:.1%}")

demo_cache_savings()


# =====================================================================
# OpenAI Prompt Caching（自动触发，无需显式标记）
# =====================================================================
from openai import OpenAI

oai_client = OpenAI()

def call_openai_with_cache(system_prompt: str, user_message: str) -> Dict:
    """
    OpenAI 自动 Prefix Caching：
    - Prompt 的前缀（1024 token 对齐）超过 1024 tokens 自动触发
    - 无需任何 API 参数修改
    """
    response = oai_client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user",   "content": user_message},
        ],
    )

    usage = response.usage
    prompt_details = getattr(usage, "prompt_tokens_details", None)
    cached_tokens = getattr(prompt_details, "cached_tokens", 0) if prompt_details else 0

    token_hit_rate = cached_tokens / max(usage.prompt_tokens, 1)
    return {
        "content": response.choices[0].message.content,
        "prompt_tokens": usage.prompt_tokens,
        "cached_tokens": cached_tokens,
        "token_hit_rate": token_hit_rate,
    }

import hashlib
import json
import time
from typing import Optional, Tuple
import numpy as np

# 依赖：pip install redis sentence-transformers numpy

class SemanticCache:
    """
    生产级语义缓存：
    - 精确哈希缓存（L1）+ 语义相似度缓存（L2）
    - L1 命中：无需 Embedding，亚毫秒响应
    - L2 命中：Embedding 相似度匹配，跳过 LLM 推理
    - 未命中：调用 LLM，结果写入双层缓存
    """

    def __init__(
        self,
        redis_url: str = "redis://localhost:6379",
        similarity_threshold: float = 0.93,
        ttl_seconds: int = 3600,
        embedding_model: str = "all-MiniLM-L6-v2",
        max_cache_size: int = 10_000,
    ):
        import redis
        from sentence_transformers import SentenceTransformer

        self.redis = redis.from_url(redis_url, decode_responses=True)
        self.threshold = similarity_threshold
        self.ttl = ttl_seconds
        self.max_size = max_cache_size

        self.embedder = SentenceTransformer(embedding_model)

        # 内存向量索引（小规模）
        # 生产建议换 Faiss / Milvus / Qdrant
        self.vectors: list = []   # [(embedding, cache_key)]

        # 监控计数器
        self.stats = {"l1_hits": 0, "l2_hits": 0, "misses": 0}

    def _exact_key(self, query: str) -> str:
        """L1 精确哈希键"""
        return "sem_cache:exact:" + hashlib.sha256(query.encode()).hexdigest()

    def _semantic_key(self, cache_key: str) -> str:
        """L2 语义缓存的 Redis 键"""
        return "sem_cache:semantic:" + cache_key

    def _embed(self, text: str) -> np.ndarray:
        """计算归一化 Embedding"""
        vec = self.embedder.encode(text, normalize_embeddings=True)
        return vec

    def _find_similar(self, query_vec: np.ndarray) -> Optional[Tuple[float, str]]:
        """
        在内存向量库中找最相似的缓存项
        返回 (similarity, cache_key) 或 None
        """
        if not self.vectors:
            return None

        vecs = np.array([v for v, _ in self.vectors])  # (N, dim)
        keys = [k for _, k in self.vectors]

        # 余弦相似度（已归一化 → 直接点积）
        sims = vecs @ query_vec
        best_idx = int(np.argmax(sims))
        best_sim = float(sims[best_idx])

        if best_sim >= self.threshold:
            return best_sim, keys[best_idx]
        return None

    def get(self, query: str) -> Optional[Dict]:
        """
        查询缓存
        返回 {"response": str, "cache_level": "L1"/"L2", "similarity": float}
        或 None（未命中）
        """
        # L1：精确哈希命中
        exact_key = self._exact_key(query)
        cached = self.redis.get(exact_key)
        if cached:
            self.stats["l1_hits"] += 1
            return {"response": cached, "cache_level": "L1", "similarity": 1.0}

        # L2：语义相似度命中
        query_vec = self._embed(query)
        result = self._find_similar(query_vec)
        if result:
            sim, sem_key = result
            response = self.redis.get(self._semantic_key(sem_key))
            if response:
                self.stats["l2_hits"] += 1
                return {"response": response, "cache_level": "L2", "similarity": sim}

        self.stats["misses"] += 1
        return None

    def set(self, query: str, response: str):
        """写入缓存（同时写 L1 和 L2）"""
        # L1 精确缓存
        exact_key = self._exact_key(query)
        self.redis.setex(exact_key, self.ttl, response)

        # L2 语义缓存
        query_vec = self._embed(query)
        cache_key = hashlib.sha256(query.encode()).hexdigest()[:16]
        self.redis.setex(self._semantic_key(cache_key), self.ttl, response)

        # 更新向量索引（LRU 管理大小）
        if len(self.vectors) >= self.max_size:
            self.vectors.pop(0)  # 简单 FIFO，生产用 LRU
        self.vectors.append((query_vec, cache_key))

    @property
    def hit_rate(self) -> float:
        total = sum(self.stats.values())
        if total == 0:
            return 0.0
        return (self.stats["l1_hits"] + self.stats["l2_hits"]) / total

    def get_stats(self) -> Dict:
        total = sum(self.stats.values())
        return {
            "l1_hits": self.stats["l1_hits"],
            "l2_hits": self.stats["l2_hits"],
            "misses": self.stats["misses"],
            "total_requests": total,
            "hit_rate": f"{self.hit_rate:.1%}",
            "l1_hit_rate": f"{self.stats['l1_hits']/max(total,1):.1%}",
            "l2_hit_rate": f"{self.stats['l2_hits']/max(total,1):.1%}",
        }


# ---- 集成 LLM 调用的完整示例 ----
from openai import OpenAI

sem_cache = SemanticCache(similarity_threshold=0.93, ttl_seconds=7200)
oai = OpenAI()

def cached_llm_call(question: str, system_prompt: str = "") -> Dict:
    """
    带语义缓存的 LLM 调用
    - 命中：直接返回缓存，延迟 < 10ms
    - 未命中：调用 LLM，写入缓存
    """
    t0 = time.time()
    cache_result = sem_cache.get(question)

    if cache_result:
        return {
            "answer": cache_result["response"],
            "source": f"cache({cache_result['cache_level']})",
            "similarity": cache_result["similarity"],
            "latency_ms": (time.time() - t0) * 1000,
        }

    # 未命中 → 调用 LLM
    response = oai.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user",   "content": question},
        ],
    )
    answer = response.choices[0].message.content

    # 写入缓存
    sem_cache.set(question, answer)

    return {
        "answer": answer,
        "source": "llm",
        "similarity": None,
        "latency_ms": (time.time() - t0) * 1000,
    }


# 测试
questions = [
    "What is machine learning?",
    "What is ML?",                   # 语义相似 → L2 命中
    "Explain machine learning",      # 语义相似 → L2 命中
    "What is machine learning?",     # 精确匹配 → L1 命中
    "How does deep learning work?",  # 新问题 → 未命中
]

for q in questions:
    result = cached_llm_call(q)
    print(f"[{result['source']:12}] {q[:40]:40} | {result['latency_ms']:.1f}ms")

print("\n缓存统计:", json.dumps(sem_cache.get_stats(), indent=2, ensure_ascii=False))