npm install @tidbcloud/serverless

// workers.ts / pages/api/xxx.ts
import { createClient } from '@tidbcloud/serverless';

// 1. 在 Cloudflare Dashboard -> Workers -> Settings -> Variables and Secrets 添加：
// TIDB_HOST = "gateway01.xxx.ap-northeast-1.prod.aws.tidbcloud.com" (不带端口)
// TIDB_PORT = "4000" (可选, 默认4000)
// TIDB_USER = "xxx.root"
// TIDB_PASSWORD = "xxx"
// TIDB_DATABASE = "test"

const client = createClient({
  host: env.TIDB_HOST,
  port: Number(env.TIDB_PORT || 4000),
  username: env.TIDB_USER,
  password: env.TIDB_PASSWORD,
  database: env.TIDB_DATABASE,
  // ssl: true // 默认开启
});

export default {
  async fetch(request, env, ctx): Promise<Response> {
    try {
      // 简单查询
      const [rows] = await client.query('SELECT VERSION() as version, NOW() as now');

      // 参数化查询 (防注入)
      const [users] = await client.query('SELECT id, name FROM users WHERE status = ? AND age > ?', ['active', 18]);

      // 事务支持
      // await client.execute('BEGIN');
      // await client.execute('INSERT INTO logs (msg) VALUES (?)', ['test']);
      // await client.execute('COMMIT');

      return Response.json({ success: true, data: { version: rows[0].version, users } });
    } catch (e: any) {
      return Response.json({ success: false, error: e.message }, { status: 500 });
    }
  },
} satisfies ExportedHandler<Env>;

npm install mysql2

// wrangler.toml 需配置兼容性
// compatibility_flags = ["nodejs_compat", "global_navigator"]
// compatibility_date = "2024-09-01" (或更新)

import { createConnection } from 'mysql2/promise';
import { connect } from 'cloudflare:sockets';

export default {
  async fetch(request, env, ctx): Promise<Response> {
    // 1. 建立原始 TCP 连接
    const socket = connect(`${env.TIDB_HOST}:${env.TIDB_PORT || 4000}`, {
      tls: true, // TiDB 强制 TLS
      // servername: env.TIDB_HOST // SNI, 通常自动推断
    });

    // 2. 将 Socket 适配为 Node.js net.Socket 兼容对象 (需要 nodejs_compat)
    // mysql2 需要一个类似 net.Socket 的对象
    // 这里利用 Cloudflare 提供的工具或自行封装
    // 简易封装示例 (实际生产建议用成熟封装库如 @mysql-js/client 或自行完善):

    // 注意：直接用 mysql2 在 Workers 里跑 TCP 还比较繁琐，需处理 readable/writable stream 转 duplex
    // 社区成熟方案：使用 `mysql2/promise` 配合 `cloudflare:sockets` 的封装库
    // 例如：<https://github.com/cloudflare/workers-sdk/tree/main/packages/cloudflare-sockets>

    // 伪代码逻辑：
    // const conn = await createConnection({
    //   stream: socket, // 需要实现 Duplex 接口
    //   user: env.TIDB_USER,
    //   password: env.TIDB_PASSWORD,
    //   database: env.TIDB_DATABASE,
    //   ssl: { rejectUnauthorized: true } // 在 socket 层已做 TLS
    // });

    // const [rows] = await conn.execute('SELECT 1');
    // await conn.end();

    return Response.json({ msg: "TCP Socket 方案需付费计划且代码较复杂，建议方案1或方案3" });
  }
}

// Hono + mysql2 (连接池复用)
import { Hono } from 'hono';
import { createPool } from 'mysql2/promise';

const app = new Hono();
const pool = createPool({
  host: process.env.TIDB_HOST,
  port: 4000,
  user: process.env.TIDB_USER,
  password: process.env.TIDB_PASSWORD,
  database: process.env.TIDB_DATABASE,
  ssl: { rejectUnauthorized: true },
  waitForConnections: true,
  connectionLimit: 10, // 连接池
  queueLimit: 0
});

app.post('/query', async (c) => {
  const { sql, params } = await c.req.json();
  if (!sql) return c.json({ error: 'sql required' }, 400);
  try {
    const [rows] = await pool.execute(sql, params || []);
    return c.json({ data: rows });
  } catch (e: any) {
    return c.json({ error: e.message }, 500);
  }
});

export default app;

export default {
  async fetch(request, env) {
    const res = await fetch(`${env.GATEWAY_URL}/query`, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json', 'Authorization': `Bearer ${env.GATEWAY_TOKEN}` },
      body: JSON.stringify({ sql: 'SELECT * FROM users WHERE id = ?', params: [1] })
    });
    return res; // 直接透传
  }
}

// src/index.ts
import { createClient } from '@tidbcloud/serverless';
import { Redis } from '@upstash/redis';
import { Hono } from 'hono'; // 推荐用 Hono 做路由/中间件，极轻量
import { stream } from 'hono/streaming';
import { cors } from 'hono/cors';

type Bindings = {
  // TiDB
  TIDB_HOST: string;
  TIDB_USER: string;
  TIDB_PASSWORD: string;
  TIDB_DATABASE: string;
  // Upstash
  UPSTASH_REDIS_REST_URL: string;
  UPSTASH_REDIS_REST_TOKEN: string;
  // R2
  MY_BUCKET: R2Bucket;
  // API Keys
  SILICONFLOW_API_KEY: string;
  // Env
  ENVIRONMENT: string;
};

const app = new Hono<{ Bindings: Bindings }>();
app.use('*', cors()); // 允许前端跨域

// 初始化客户端 (全局复用，但注意 Workers 无状态，每次请求新建开销极小)
const getTidb = (env: Bindings) => createClient({
  host: env.TIDB_HOST, username: env.TIDB_USER, password: env.TIDB_PASSWORD, database: env.TIDB_DATABASE
});
const getRedis = (env: Bindings) => new Redis({ url: env.UPSTASH_REDIS_REST_URL, token: env.UPSTASH_REDIS_REST_TOKEN });

// --- 嵌入向量 ---
async function embed(texts: string[], env: Bindings): Promise<number[][]> {
  const res = await fetch('<https://api.siliconflow.cn/v1/embeddings>', {
    method: 'POST',
    headers: { 'Content-Type': 'application/json', 'Authorization': `Bearer ${env.SILICONFLOW_API_KEY}` },
    body: JSON.stringify({ model: 'BAAI/bge-m3', input: texts, encoding_format: 'float' })
  });
  const data = await res.json();
  return data.data.map((d: any) => d.embedding);
}

// --- 重排 ---
async function rerank(query: string, docs: string[], env: Bindings): Promise<number[]> {
  const res = await fetch('<https://api.siliconflow.cn/v1/rerank>', {
    method: 'POST',
    headers: { 'Content-Type': 'application/json', 'Authorization': `Bearer ${env.SILICONFLOW_API_KEY}` },
    body: JSON.stringify({ model: 'BAAI/bge-reranker-v2-m3', query, documents: docs, top_n: 3 })
  });
  const data = await res.json();
  return data.results.map((r: any) => r.index); // 返回排序后的原始索引
}

// --- LLM 流式生成 ---
async function* generateStream(messages: any[], env: Bindings): AsyncGenerator<string> {
  const res = await fetch('<https://api.siliconflow.cn/v1/chat/completions>', {
    method: 'POST',
    headers: { 'Content-Type': 'application/json', 'Authorization': `Bearer ${env.SILICONFLOW_API_KEY}` },
    body: JSON.stringify({ model: 'Qwen/Qwen2.5-72B-Instruct', messages, stream: true, temperature: 0.3 })
  });
  const reader = res.body?.getReader();
  const decoder = new TextDecoder();
  if (!reader) return;

  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    const chunk = decoder.decode(value);
    const lines = chunk.split('\\n').filter(l => l.startsWith('data: '));
    for (const line of lines) {
      if (line === 'data: [DONE]') return;
      try {
        const json = JSON.parse(line.slice(6));
        const content = json.choices[0]?.delta?.content;
        if (content) yield content;
      } catch (e) { /* ignore parse error */ }
    }
  }
}

// --- 主路由：RAG 问答 ---
app.post('/api/chat', async (c) => {
  const { question, history = [] } = await c.req.json();
  const env = c.env;
  const tidb = getTidb(env);
  const redis = getRedis(env);

  // 1. 限流 (IP 维度)
  const ip = c.req.header('CF-Connecting-IP') || 'unknown';
  const ratelimit = await redis.incr(`ratelimit:${ip}`);
  if (ratelimit === 1) await redis.expire(`ratelimit:${ip}`, 60); // 1分钟窗口
  if (ratelimit > 30) return c.json({ error: 'Rate limited' }, 429);

  // 2. 向量化问题
  const [qVec] = await embed([question], env);

  // 3. TiDB 向量检索 (SQL + 向量距离)
  // 表结构假设: chunks(id, doc_id, content, embedding VECTOR(1024), metadata JSON)
  const searchSql = `
    SELECT id, content, metadata, VEC_COSINE_DISTANCE(embedding, ?) AS dist
    FROM knowledge_chunks
    ORDER BY dist ASC
    LIMIT 10
  `;
  // @tidbcloud/serverless query 参数需手动序列化向量为字符串 '[0.1, 0.2...]'
  const vecStr = `[${qVec.join(',')}]`;
  const [candidates] = await tidb.query(searchSql, [vecStr]);

  if (candidates.length === 0) {
    return stream(c, async (stream) => {
      await stream.write('暂无相关知识库内容。');
    });
  }

  // 4. Rerank 重排
  const contents = candidates.map((c: any) => c.content);
  const rerankedIndices = await rerank(question, contents, env);
  const topChunks = rerankedIndices.map(i => candidates[i]).slice(0, 3);

  // 5. 构建 Prompt
  const context = topChunks.map((c: any, i: number) => `[来源 ${i+1}] ${c.content}`).join('\\n\\n');
  const systemPrompt = `你是专业助手。请仅根据以下参考资料回答问题，引用来源编号如[1]。若无法回答请说不知道。\\n\\n参考资料：\\n${context}`;

  const messages = [
    { role: 'system', content: systemPrompt },
    ...history.slice(-6), // 保留最近 3 轮对话
    { role: 'user', content: question }
  ];

  // 6. 流式返回
  return stream(c, async (stream) => {
    // 发送来源信息给前端 (可选，自定义协议)
    await stream.write(`__SOURCES__${JSON.stringify(topChunks.map(c => ({id: c.id, metadata: c.metadata})))}\\n`);

    for await (const chunk of generateStream(messages, env)) {
      await stream.write(chunk);
    }
  });
});

// --- 健康检查 ---
app.get('/health', (c) => c.json({ ok: true, ts: Date.now() }));

export default app;

name = "my-rag-worker"
main = "src/index.ts"
compatibility_date = "2024-10-01"
compatibility_flags = ["nodejs_compat"] # 部分库可能需要

# 绑定 R2
[[r2_buckets]]
binding = "MY_BUCKET"
bucket_name = "my-knowledge-files"

# 环境变量 (敏感信息请在 Dashboard -> Settings -> Variables/Secrets 设置，不要写这里)
[vars]
TIDB_HOST = "gateway01.xxx.ap-northeast-1.prod.aws.tidbcloud.com"
TIDB_USER = "xxx.root"
TIDB_DATABASE = "rag_db"
UPSTASH_REDIS_REST_URL = "<https://xxxxx.upstash.io>"
# TIDB_PASSWORD, UPSTASH_REDIS_REST_TOKEN, SILICONFLOW_API_KEY 放在 Secrets 里

# 1. 登录
npx wrangler login

# 2. 创建 R2 Bucket (首次)
npx wrangler r2 bucket create my-knowledge-files

# 3. 部署 (会提示输入 Secrets)
npx wrangler deploy

# 4. 开发调试 (本地跑，连远程资源)
npx wrangler dev --remote