用 Rust 扩展 AI 网关：agentgateway 与 kgateway 的自定义转换实践

/// 请求转换：在请求被路由到上游模型之前执行
async fn transform_request(
    ctx: &RequestContext,
    request: LlmRequest,
) -> Result<LlmRequest, TransformError>;

/// 响应转换：在上游模型返回结果之后、返回给客户端之前执行
async fn transform_response(
    ctx: &RequestContext,
    response: LlmResponse,
) -> Result<LlmResponse, TransformError>;

// src/extensions/tenant_context.rs

use agentgateway::transform::{RequestContext, LlmRequest, TransformError};
use agentgateway::llm::Message;
use std::collections::HashMap;

/// 租户上下文配置——实际项目中可以从文件或数据库加载
#[derive(Debug, Clone)]
pub struct TenantContextConfig {
    pub prompts: HashMap<String, String>,
}

pub struct TenantContextTransform {
    config: TenantContextConfig,
}

impl TenantContextTransform {
    pub fn new(config: TenantContextConfig) -> Self {
        Self { config }
    }

    /// 核心转换逻辑：根据 RequestContext 中的 tenant_id，
    /// 查找对应的 system prompt，插入到 messages 最前面
    pub async fn transform_request(
        &self,
        ctx: &RequestContext,
        mut request: LlmRequest,
    ) -> Result<LlmRequest, TransformError> {
        let tenant_id = ctx.tenant_id();

        // 查找租户专属 prompt，找不到就用默认值
        let system_content = self.config.prompts
            .get(tenant_id)
            .cloned()
            .unwrap_or_else(|| "You are a helpful assistant.".to_string());

        // 构造 system message 并插入到 messages 列表头部
        let system_msg = Message {
            role: "system".to_string(),
            content: system_content,
            ..Default::default()
        };

        request.messages.insert(0, system_msg);

        // 记录日志，方便排查哪个租户用了哪段 prompt
        tracing::info!(
            tenant_id = %tenant_id,
            injected_prompt_len = system_content.len(),
            "Injected tenant-specific system prompt"
        );

        Ok(request)
    }
}

// src/main.rs

use agentgateway::GatewayBuilder;
use extensions::tenant_context::{TenantContextTransform, TenantContextConfig};
use std::collections::HashMap;

fn main() {
    // 构造租户配置——生产环境建议从 YAML / 环境变量 / DB 加载
    let mut prompts = HashMap::new();
    prompts.insert(
        "tenant-acme".to_string(),
        "You are Acme Corp's internal assistant. Always reference Acme policies. Answer in concise bullet points.".to_string(),
    );
    prompts.insert(
        "tenant-beta".to_string(),
        "You are Beta Inc's customer support bot. Be friendly, use emojis sparingly, never disclose internal pricing.".to_string(),
    );

    let config = TenantContextConfig { prompts };
    let transform = TenantContextTransform::new(config);

    // 把扩展注册到 GatewayBuilder
    let gateway = GatewayBuilder::new()
        .add_request_transform(transform)
        // 其他内置策略照常配置
        .add_rate_limiting(/* ... */)
        .add_auth(/* ... */)
        .build();

    // 启动网关
    gateway.run().await.expect("gateway failed");
}

# gateway-config.yaml
routes:
  - name: chat-completions
    path: /v1/chat/completions
    upstream:
      model: gpt-4o
      endpoint: https://api.openai.com/v1/chat/completions
    # 引用上面注册的 request_transform
    request_transforms:
      - tenant_context

  - name: chat-completions-anthropic
    path: /v1/anthropic/chat
    upstream:
      model: claude-sonnet-4-20250514
      endpoint: https://api.anthropic.com/v1/messages
    request_transforms:
      - tenant_context

# 假设网关监听 localhost:8080，认证层把 API Key 映射成 tenant_id
curl -X POST http://localhost:8080/v1/chat/completions \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer sk-acme-xxxx" \
  -d '{
    "model": "gpt-4o",
    "messages": [
      {"role": "user", "content": "What is our refund policy?"}
    ]
  }'

{
  "model": "gpt-4o",
  "messages": [
    {"role": "system", "content": "You are Acme Corp's internal assistant. Always reference Acme policies. Answer in concise bullet points."},
    {"role": "user", "content": "What is our refund policy?"}
  ]
}

// src/extensions/cost_obfuscation.rs

use agentgateway::transform::{RequestContext, LlmResponse, TransformError};

pub async fn transform_response(
    _ctx: &RequestContext,
    mut response: LlmResponse,
) -> Result<LlmResponse, TransformError> {
    if let Some(usage) = &response.usage {
        // 把精确 token 数替换成模糊费用区间
        let estimated_cost = estimate_cost(usage.prompt_tokens, usage.completion_tokens);
        // 移除原始 usage，注入自定义字段
        response.usage = None;
        response.metadata.insert(
            "estimated_cost_range".to_string(),
            estimated_cost,
        );
    }
    Ok(response)
}

fn estimate_cost(prompt_tokens: u32, completion_tokens: u32) -> String {
    let total = prompt_tokens + completion_tokens;
    // 粗粒度区间，不暴露精确数字
    if total < 1000 {
        "< $0.01".to_string()
    } else if total < 5000 {
        "$0.01 – $0.05".to_string()
    } else {
        "> $0.05".to_string()
    }
}