from fastapi import FastAPI, Header, HTTPException
from pydantic import BaseModel
import time
import uuid
from collections import defaultdict
# 初始化FastAPI应用
app = FastAPI(title="大模型多租户隔离实战Demo")
# ===================== 1. 多租户核心配置：权限、资源、限流 =====================
# 租户配置库：存储TenantID、权限、限流配额、资源配额
# 格式：{租户ID: {"name": "租户名称", "model_perm": ["模型接口列表"], "qps": 每秒请求数, "quota": 剩余调用额度}}
TENANT_CONFIG = {
    "tenant_001": {
        "name": "电商企业租户",
        "model_perm": ["llm_chat", "llm_summary"],
        "qps": 5,
        "quota": 1000,
        "gpu_quota": 2  # 分配2GB显存资源
    },
    "tenant_002": {
        "name": "教育小型租户",
        "model_perm": ["llm_chat"],
        "qps": 2,
        "quota": 500,
        "gpu_quota": 1  # 分配1GB显存资源
    }
}
# 租户独立缓存：模拟Redis命名空间隔离，存储对话历史
TENANT_CACHE = defaultdict(dict)
# 租户请求计数器：用于限流（性能隔离）
TENANT_REQUEST_COUNTER = defaultdict(int)
# ===================== 2. 模拟大模型推理函数 =====================
def mock_llm_infer(tenant_id: str, prompt: str) -> str:
    """模拟大模型推理，仅使用租户自身缓存/资源"""
    # 数据隔离：仅读取当前租户的对话历史
    chat_history = TENANT_CACHE[tenant_id].get("chat_history", [])
    chat_history.append(prompt)
    TENANT_CACHE[tenant_id]["chat_history"] = chat_history[-5:]  # 保留最近5轮对话
    
    # 模拟推理耗时
    time.sleep(0.5)
    return f"【租户{tenant_id}】模型回答：已处理你的问题「{prompt}」，对话历史：{len(chat_history)}轮"
# ===================== 3. 多租户隔离核心中间件 =====================
def check_tenant_isolation(tenant_id: str, token: str, model_name: str):
    """
    多租户三合一校验：
    1. 权限管控（身份+接口权限）
    2. 性能隔离（限流）
    3. 资源隔离（调用额度）
    """
    # -------------- 权限管控：身份校验 --------------
    if tenant_id not in TENANT_CONFIG:
        raise HTTPException(status_code=403, detail=f"租户[{tenant_id}]不存在，无访问权限")
    
    tenant_name = TENANT_CONFIG[tenant_id]["name"]
    
    # 权限管控：接口访问权限校验
    if model_name not in TENANT_CONFIG[tenant_id]["model_perm"]:
        raise HTTPException(status_code=403, detail=f"租户[{tenant_name}]无权限访问{model_name}接口")
    
    # -------------- 性能隔离：限流管控 --------------
    if TENANT_REQUEST_COUNTER[tenant_id] >= TENANT_CONFIG[tenant_id]["qps"]:
        raise HTTPException(status_code=429, detail=f"租户[{tenant_name}]请求超限，性能隔离触发")
    TENANT_REQUEST_COUNTER[tenant_id] += 1
    
    # -------------- 资源隔离：调用额度管控 --------------
    if TENANT_CONFIG[tenant_id]["quota"] <= 0:
        raise HTTPException(status_code=402, detail=f"租户[{tenant_name}]资源额度耗尽，资源隔离触发")
    TENANT_CONFIG[tenant_id]["quota"] -= 1
# ===================== 4. 多租户大模型API =====================
class LLMRequest(BaseModel):
    prompt: str
    model_name: str = "llm_chat"
@app.post("/llm/chat")
def llm_chat(request: LLMRequest, TenantID: str = Header(...), Authorization: str = Header(...)):
    """
    多租户隔离大模型对话接口
    必须携带请求头：TenantID、Authorization
    """
    # 执行三合一隔离校验
    check_tenant_isolation(TenantID, Authorization, request.model_name)
    
    # 执行模型推理（数据隔离）
    response = mock_llm_infer(TenantID, request.prompt)
    
    # 释放计数器
    TENANT_REQUEST_COUNTER[TenantID] -= 1
    
    return {
        "code": 200,
        "tenant_id": TenantID,
        "response": response,
        "remaining_quota": TENANT_CONFIG[TenantID]["quota"]
    }
# ===================== 5. 直接启动（无需命令行）=====================
if __name__ == "__main__":
    import uvicorn
    print("🚀 启动多租户大模型服务...")
    print("📍 服务地址: http://0.0.0.0:8010")
    print("📖 API文档: http://0.0.0.0:8010/docs")
    uvicorn.run(app, host="0.0.0.0", port=8010)

import requests
response = requests.post(
    "http://localhost:8010/llm/chat",
    headers={
        "TenantID": "tenant_001",
        "Authorization": "Bearer xxx",
        "Content-Type": "application/json"
    },
    json={
        "prompt": "你好，请介绍一下自己",
        "model_name": "llm_chat"
    }
)
print(response.json())

"""
多租户大模型服务 - 测试脚本合集
包含：正常请求、限流测试、权限测试、额度耗尽测试
"""
import requests
import threading
import time
import json
BASE_URL = "http://localhost:8010"
# ===================== 测试配置 =====================
TENANT_001 = {"id": "tenant_001", "name": "电商企业租户", "auth": "Bearer token001"}
TENANT_002 = {"id": "tenant_002", "name": "教育小型租户", "auth": "Bearer token002"}
INVALID_TENANT = {"id": "tenant_999", "name": "非法租户", "auth": "Bearer xxx"}
def print_result(test_name, response):
    """打印测试结果"""
    print(f"\n{'='*60}")
    print(f"【{test_name}】")
    print(f"{'='*60}")
    print(f"状态码: {response.status_code}")
    try:
        print(f"响应内容: {json.dumps(response.json(), ensure_ascii=False, indent=2)}")
    except:
        print(f"响应内容: {response.text}")
# ===================== 1. 正常请求测试 =====================
def test_normal_request():
    """正常对话请求测试"""
    url = f"{BASE_URL}/llm/chat"
    headers = {
        "TenantID": TENANT_001["id"],
        "Authorization": TENANT_001["auth"],
        "Content-Type": "application/json"
    }
    data = {
        "prompt": "你好，请介绍一下自己",
        "model_name": "llm_chat"
    }
    response = requests.post(url, headers=headers, json=data)
    print_result("正常请求测试", response)
    return response
# ===================== 2. 限流测试（高频请求）=====================
def test_rate_limit():
    """限流测试 - 并发发送多个请求触发限流"""
    url = f"{BASE_URL}/llm/chat"
    headers = {
        "TenantID": TENANT_002["id"],  # 教育租户 QPS=2
        "Authorization": TENANT_002["auth"],
        "Content-Type": "application/json"
    }
    data = {"prompt": "测试限流", "model_name": "llm_chat"}
    
    print(f"\n{'='*60}")
    print("【限流测试】并发发送5个请求（租户QPS限制为2）")
    print(f"{'='*60}")
    
    results = [None] * 5
    
    def send_request(idx):
        response = requests.post(url, headers=headers, json=data)
        results[idx] = (idx, response)
    
    # 创建线程
    threads = []
    for i in range(5):
        t = threading.Thread(target=send_request, args=(i,))
        threads.append(t)
    
    # 几乎同时启动所有线程
    for t in threads:
        t.start()
    
    for t in threads:
        t.join()
    
    # 打印结果
    for idx, response in sorted(results, key=lambda x: x[0]):
        status = "✓ 成功" if response.status_code == 200 else f"✗ 被限流"
        detail = ""
        if response.status_code == 429:
            detail = f" - {response.json().get('detail', '')[:40]}"
        print(f"  请求{idx+1}: {status} (HTTP {response.status_code}){detail}")
    
    return [r[1] for r in results]
# ===================== 3. 并发限流测试 =====================
def test_concurrent_rate_limit():
    """并发请求测试"""
    url = f"{BASE_URL}/llm/chat"
    headers = {
        "TenantID": TENANT_001["id"],  # 电商租户 QPS=5
        "Authorization": TENANT_001["auth"],
        "Content-Type": "application/json"
    }
    data = {"prompt": "并发测试", "model_name": "llm_chat"}
    
    print(f"\n{'='*60}")
    print("【并发限流测试】同时启动10个线程请求（租户QPS限制为5）")
    print(f"{'='*60}")
    
    results = {"success": 0, "rate_limited": 0}
    
    def make_request():
        try:
            response = requests.post(url, headers=headers, json=data, timeout=5)
            if response.status_code == 200:
                results["success"] += 1
            elif response.status_code == 429:
                results["rate_limited"] += 1
        except Exception as e:
            print(f"  请求异常: {e}")
    
    threads = []
    for i in range(10):
        t = threading.Thread(target=make_request)
        threads.append(t)
    
    # 同时启动所有线程
    for t in threads:
        t.start()
    
    for t in threads:
        t.join()
    
    print(f"  成功: {results['success']} 个")
    print(f"  被限流: {results['rate_limited']} 个")
# ===================== 4. 权限测试（非法租户）=====================
def test_invalid_tenant():
    """非法租户测试"""
    url = f"{BASE_URL}/llm/chat"
    headers = {
        "TenantID": INVALID_TENANT["id"],
        "Authorization": INVALID_TENANT["auth"],
        "Content-Type": "application/json"
    }
    data = {"prompt": "测试", "model_name": "llm_chat"}
    response = requests.post(url, headers=headers, json=data)
    print_result("非法租户测试", response)
    return response
# ===================== 5. 权限测试（无接口权限）=====================
def test_no_permission():
    """无接口权限测试 - tenant_002 没有 llm_summary 权限"""
    url = f"{BASE_URL}/llm/chat"
    headers = {
        "TenantID": TENANT_002["id"],
        "Authorization": TENANT_002["auth"],
        "Content-Type": "application/json"
    }
    data = {"prompt": "总结这段文本", "model_name": "llm_summary"}
    response = requests.post(url, headers=headers, json=data)
    print_result("无接口权限测试", response)
    return response
# ===================== 6. 额度耗尽测试 =====================
def test_quota_exhausted():
    """额度耗尽测试 - 快速消耗完租户额度"""
    url = f"{BASE_URL}/llm/chat"
    headers = {
        "TenantID": TENANT_002["id"],  # 教育租户 quota=500
        "Authorization": TENANT_002["auth"],
        "Content-Type": "application/json"
    }
    data = {"prompt": "消耗额度", "model_name": "llm_chat"}
    
    print(f"\n{'='*60}")
    print("【额度耗尽测试】快速发送请求消耗额度")
    print(f"{'='*60}")
    
    # 先发送几个正常请求
    for i in range(3):
        response = requests.post(url, headers=headers, json=data)
        if response.status_code == 200:
            remaining = response.json().get("remaining_quota", "未知")
            print(f"  请求{i+1}: 成功，剩余额度: {remaining}")
    
    # 模拟额度耗尽（直接修改配置，实际测试用）
    print("  ...模拟额度耗尽...")
    # 这里只是演示，实际需要通过大量请求或修改服务端配置
# ===================== 7. 多租户隔离测试 =====================
def test_tenant_isolation():
    """多租户数据隔离测试"""
    url = f"{BASE_URL}/llm/chat"
    
    print(f"\n{'='*60}")
    print("【多租户隔离测试】两个租户分别发送请求")
    print(f"{'='*60}")
    
    for tenant in [TENANT_001, TENANT_002]:
        headers = {
            "TenantID": tenant["id"],
            "Authorization": tenant["auth"],
            "Content-Type": "application/json"
        }
        data = {"prompt": "这是第一条消息", "model_name": "llm_chat"}
        response = requests.post(url, headers=headers, json=data)
        
        if response.status_code == 200:
            resp_data = response.json()
            print(f"  {tenant['name']}: {resp_data.get('response', '')}")
        else:
            print(f"  {tenant['name']}: 请求失败 {response.status_code}")
# ===================== 8. 压力测试 =====================
def test_stress():
    """压力测试 - 持续发送请求"""
    url = f"{BASE_URL}/llm/chat"
    headers = {
        "TenantID": TENANT_001["id"],
        "Authorization": TENANT_001["auth"],
        "Content-Type": "application/json"
    }
    data = {"prompt": "压力测试", "model_name": "llm_chat"}
    
    print(f"\n{'='*60}")
    print("【压力测试】10秒内持续发送请求")
    print(f"{'='*60}")
    
    success = 0
    failed = 0
    start_time = time.time()
    
    while time.time() - start_time < 10:
        response = requests.post(url, headers=headers, json=data)
        if response.status_code == 200:
            success += 1
        else:
            failed += 1
        time.sleep(0.2)  # 200ms间隔
    
    print(f"  10秒内成功: {success} 个")
    print(f"  被限流/失败: {failed} 个")
    print(f"  平均每秒: {(success+failed)/10:.1f} 请求")
# ===================== 主程序 =====================
if __name__ == "__main__":
    print("🚀 多租户大模型服务测试脚本")
    print(f"服务端地址: {BASE_URL}")
    print("="*60)
    
    # 检查服务是否可用
    try:
        requests.get(f"{BASE_URL}/docs", timeout=3)
        print("✅ 服务连接正常\n")
    except:
        print("❌ 无法连接到服务，请确保服务端已启动")
        print(f"   启动命令: python 260421.（156）1-main.py")
        exit(1)
    
    # 执行测试
    test_normal_request()
    test_tenant_isolation()
    test_rate_limit()
    test_concurrent_rate_limit()
    test_invalid_tenant()
    test_no_permission()
    test_quota_exhausted()
    test_stress()
    
    print(f"\n{'='*60}")
    print("✅ 所有测试执行完成")
    print(f"{'='*60}")