TiDB for AI Memory: Vector Search, HTAP, and Horizontal Scaling in One Database

# Connection string format
mysql -h gateway01.us-west-2.prod.aws.tidbcloud.com -P 4000 \
  -u '<your-username>' -p'<your-password>' --ssl-mode=VERIFY_IDENTITY \
  --ssl-ca=/path/to/ca.pem

# Install TiUP
curl --proto '=https' --tlsv1.2 -sSf https://tiup-mirrors.pingcap.com/install.sh | sh

# Start a local cluster with TiFlash
tiup playground --db 1 --pd 1 --kv 1 --tiflash 1

# Connect using MySQL client
mysql --host 127.0.0.1 --port 4000 -u root

-- Create database
CREATE DATABASE ai_memory;
USE ai_memory;

-- Main memory table
CREATE TABLE memories (
    id BIGINT PRIMARY KEY AUTO_INCREMENT,
    user_id BIGINT NOT NULL,
    session_id VARCHAR(255),

    -- Content and embedding
    content TEXT NOT NULL,
    embedding VECTOR(1536) COMMENT "hnsw(distance=cosine)",

    -- Classification
    memory_type ENUM('episodic', 'semantic', 'working') NOT NULL,
    importance FLOAT DEFAULT 0.5,
    confidence FLOAT DEFAULT 1.0,

    -- Temporal tracking
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    last_accessed TIMESTAMP NULL,
    access_count INT DEFAULT 0,
    expires_at TIMESTAMP NULL,

    -- Metadata (JSON for flexibility)
    metadata JSON,

    -- Memory relationships
    source_memory_ids JSON COMMENT 'IDs of memories that support this one',

    -- Soft deletion
    is_deleted BOOLEAN DEFAULT FALSE,

    -- Indexes
    INDEX idx_user_type (user_id, memory_type),
    INDEX idx_user_created (user_id, created_at),
    INDEX idx_expires (expires_at),
    INDEX idx_session (session_id),

    -- Vector index with HNSW for fast approximate search
    VECTOR INDEX idx_embedding (embedding)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin;

-- Memory access log for analytics
CREATE TABLE memory_access_log (
    id BIGINT PRIMARY KEY AUTO_INCREMENT,
    memory_id BIGINT NOT NULL,
    user_id BIGINT NOT NULL,
    accessed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    retrieval_score FLOAT COMMENT 'Cosine similarity score',
    query_context TEXT COMMENT 'What query retrieved this memory',

    INDEX idx_memory (memory_id),
    INDEX idx_user_time (user_id, accessed_at),
    INDEX idx_accessed (accessed_at)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin;

-- Enable TiFlash replica for memory analytics
ALTER TABLE memories SET TIFLASH REPLICA 1;

-- Check replication progress
SELECT 
    table_schema, 
    table_name, 
    replica_count, 
    available, 
    progress 
FROM information_schema.tiflash_replica 
WHERE table_schema = 'ai_memory';

import mysql from 'mysql2/promise';
import OpenAI from 'openai';

const openai = new OpenAI({
  apiKey: process.env.OPENAI_API_KEY
});

// Create TiDB connection pool
const pool = mysql.createPool({
  host: process.env.TIDB_HOST,
  port: parseInt(process.env.TIDB_PORT || '4000'),
  user: process.env.TIDB_USER,
  password: process.env.TIDB_PASSWORD,
  database: 'ai_memory',
  ssl: {
    minVersion: 'TLSv1.2',
    rejectUnauthorized: true
  },
  connectionLimit: 10
});

/**
 * Generate embedding for text using OpenAI
 */
async function generateEmbedding(text) {
  const response = await openai.embeddings.create({
    model: 'text-embedding-3-small', // 1536 dimensions
    input: text
  });

  return response.data[0].embedding;
}

/**
 * Extract memorable facts from a message
 */
async function extractMemorableFacts(message, userId) {
  const extraction = await openai.chat.completions.create({
    model: 'gpt-4o-mini',
    messages: [
      {
        role: 'system',
        content: `Extract important, long-term facts from the user's message.

        Rules:
        - Ignore pleasantries, confirmations, and generic responses
        - Extract specific preferences, attributes, or requirements
        - Each fact should be self-contained and meaningful
        - Classify each as episodic, semantic, or working memory
        - Assign importance score (0.0 to 1.0)

        Return JSON array:
        {
          "facts": [
            {
              "fact": "User prefers dark mode interfaces",
              "type": "semantic",
              "importance": 0.85
            }
          ]
        }`
      },
      { role: 'user', content: message }
    ],
    response_format: { type: 'json_object' }
  });

  const result = JSON.parse(extraction.choices[0].message.content);
  return result.facts || [];
}

/**
 * Store a memory in TiDB with duplicate detection
 */
async function storeMemory(userId, fact, sessionId = null) {
  const connection = await pool.getConnection();

  try {
    // Generate embedding
    const embedding = await generateEmbedding(fact.fact);
    const embeddingJson = JSON.stringify(embedding);

    // Check for similar existing memories (within 0.15 cosine distance)
    const [existingMemories] = await connection.query(`
      SELECT 
        id, 
        content, 
        confidence,
        VEC_COSINE_DISTANCE(embedding, ?) AS distance
      FROM memories
      WHERE user_id = ?
        AND memory_type = ?
        AND is_deleted = FALSE
        AND VEC_COSINE_DISTANCE(embedding, ?) < 0.15
      ORDER BY distance ASC
      LIMIT 1
    `, [embeddingJson, userId, fact.type, embeddingJson]);

    if (existingMemories.length > 0) {
      // Similar memory exists, reinforce it
      const existing = existingMemories[0];

      await connection.query(`
        UPDATE memories
        SET 
          last_accessed = CURRENT_TIMESTAMP,
          access_count = access_count + 1,
          confidence = LEAST(confidence + 0.1, 1.0),
          metadata = JSON_SET(
            COALESCE(metadata, '{}'),
            '$.last_reinforced',
            CURRENT_TIMESTAMP()
          )
        WHERE id = ?
      `, [existing.id]);

      return { 
        action: 'reinforced', 
        memoryId: existing.id,
        content: existing.content
      };
    }

    // No similar memory, create new one
    const [result] = await connection.query(`
      INSERT INTO memories (
        user_id,
        session_id,
        content,
        embedding,
        memory_type,
        importance,
        confidence,
        created_at,
        last_accessed
      ) VALUES (?, ?, ?, ?, ?, ?, 1.0, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
    `, [
      userId,
      sessionId,
      fact.fact,
      embeddingJson,
      fact.type,
      fact.importance
    ]);

    return {
      action: 'created',
      memoryId: result.insertId,
      content: fact.fact
    };

  } finally {
    connection.release();
  }
}

/**
 * Process a user message and store extracted memories
 */
async function processMessage(userId, message, sessionId = null) {
  // Extract facts worth remembering
  const facts = await extractMemorableFacts(message, userId);

  // Filter by importance threshold
  const importantFacts = facts.filter(f => f.importance > 0.6);

  // Store each fact
  const results = [];
  for (const fact of importantFacts) {
    const result = await storeMemory(userId, fact, sessionId);
    results.push(result);
  }

  return results;
}

// Example usage
const userId = 123;
const sessionId = 'sess_abc123';
const message = "I really prefer dark mode for all interfaces. I am currently working on a healthcare AI project focused on cardiology research.";

const stored = await processMessage(userId, message, sessionId);
console.log('Stored memories:', stored);

/**
 * Retrieve relevant memories with multi-factor scoring
 */
async function retrieveMemories(userId, queryText, options = {}) {
  const {
    memoryTypes = ['semantic', 'episodic'],
    limit = 10,
    recencyHalfLifeDays = 7, // Memories lose half their recency score every 7 days
    minImportance = 0.0
  } = options;

  const connection = await pool.getConnection();

  try {
    // Generate query embedding
    const queryEmbedding = await generateEmbedding(queryText);
    const embeddingJson = JSON.stringify(queryEmbedding);

    // Retrieve and score memories
    // Score formula:
    // - 60% semantic relevance (1 - cosine distance)
    // - 25% recency (exponential decay)
    // - 15% importance
    const [memories] = await connection.query(`
      SELECT 
        id,
        content,
        memory_type,
        importance,
        confidence,
        created_at,
        last_accessed,
        access_count,
        metadata,

        -- Calculate semantic similarity (1 - distance = similarity)
        (1 - VEC_COSINE_DISTANCE(embedding, ?)) AS semantic_similarity,

        -- Calculate recency score with exponential decay
        EXP(-TIMESTAMPDIFF(DAY, created_at, CURRENT_TIMESTAMP) / ?) AS recency_score,

        -- Combined weighted score
        (
          (1 - VEC_COSINE_DISTANCE(embedding, ?)) * 0.60 +
          EXP(-TIMESTAMPDIFF(DAY, created_at, CURRENT_TIMESTAMP) / ?) * 0.25 +
          importance * 0.15
        ) AS final_score

      FROM memories
      WHERE user_id = ?
        AND memory_type IN (?)
        AND is_deleted = FALSE
        AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
        AND importance >= ?
      ORDER BY final_score DESC
      LIMIT ?
    `, [
      embeddingJson,
      recencyHalfLifeDays,
      embeddingJson,
      recencyHalfLifeDays,
      userId,
      memoryTypes,
      minImportance,
      limit * 2 // Get extra for filtering
    ]);

    // Update access tracking for retrieved memories
    if (memories.length > 0) {
      const memoryIds = memories.map(m => m.id);

      // Batch update access tracking
      await connection.query(`
        UPDATE memories
        SET 
          last_accessed = CURRENT_TIMESTAMP,
          access_count = access_count + 1
        WHERE id IN (?)
      `, [memoryIds]);

      // Log access for analytics (optional)
      const accessLogs = memories.map(m => [
        m.id,
        userId,
        m.final_score,
        queryText.substring(0, 500) // Limit query context length
      ]);

      await connection.query(`
        INSERT INTO memory_access_log (memory_id, user_id, retrieval_score, query_context)
        VALUES ?
      `, [accessLogs]);
    }

    return memories.slice(0, limit);

  } finally {
    connection.release();
  }
}

// Example usage
const userId = 123;
const query = "What are my interface preferences?";

const relevantMemories = await retrieveMemories(userId, query, {
  memoryTypes: ['semantic'],
  limit: 5,
  recencyHalfLifeDays: 30
});

console.log('Retrieved memories:');
relevantMemories.forEach(m => {
  console.log(`- ${m.content} (score: ${m.final_score.toFixed(3)})`);
});

-- Force query to use TiFlash for analytics
-- TiDB usually auto-selects, but you can specify explicitly
SET SESSION tidb_isolation_read_engines = 'tiflash';

-- Memory growth over time
SELECT 
    DATE(created_at) AS date,
    memory_type,
    COUNT(*) AS memories_created,
    AVG(importance) AS avg_importance,
    AVG(confidence) AS avg_confidence
FROM memories
WHERE created_at > DATE_SUB(CURRENT_TIMESTAMP, INTERVAL 30 DAY)
    AND is_deleted = FALSE
GROUP BY DATE(created_at), memory_type
ORDER BY date DESC, memory_type;

-- Most accessed memories (cache candidate analysis)
SELECT 
    id,
    content,
    memory_type,
    access_count,
    DATEDIFF(CURRENT_TIMESTAMP, created_at) AS age_days,
    access_count / GREATEST(DATEDIFF(CURRENT_TIMESTAMP, created_at), 1) AS access_rate_per_day
FROM memories
WHERE user_id = 123
    AND is_deleted = FALSE
ORDER BY access_rate_per_day DESC
LIMIT 20;

-- Memory distribution by user
SELECT 
    user_id,
    COUNT(*) AS total_memories,
    SUM(CASE WHEN memory_type = 'episodic' THEN 1 ELSE 0 END) AS episodic_count,
    SUM(CASE WHEN memory_type = 'semantic' THEN 1 ELSE 0 END) AS semantic_count,
    SUM(CASE WHEN memory_type = 'working' THEN 1 ELSE 0 END) AS working_count,
    AVG(importance) AS avg_importance,
    MAX(created_at) AS last_memory_created
FROM memories
WHERE is_deleted = FALSE
GROUP BY user_id
ORDER BY total_memories DESC
LIMIT 100;

-- Memory consolidation candidates
-- Find clusters of similar episodic memories that should be consolidated into semantic facts
SELECT 
    m1.id AS memory_id_1,
    m2.id AS memory_id_2,
    m1.content AS content_1,
    m2.content AS content_2,
    VEC_COSINE_DISTANCE(m1.embedding, m2.embedding) AS distance,
    DATEDIFF(m2.created_at, m1.created_at) AS days_apart
FROM memories m1
JOIN memories m2 ON m1.user_id = m2.user_id AND m1.id < m2.id
WHERE m1.user_id = 123
    AND m1.memory_type = 'episodic'
    AND m2.memory_type = 'episodic'
    AND m1.is_deleted = FALSE
    AND m2.is_deleted = FALSE
    AND VEC_COSINE_DISTANCE(m1.embedding, m2.embedding) < 0.20
ORDER BY distance ASC
LIMIT 50;

-- Reset to auto-selection
SET SESSION tidb_isolation_read_engines = 'tikv,tidb,tiflash';

/**
 * Find clusters of similar episodic memories for consolidation
 */
async function findConsolidationCandidates(userId, similarityThreshold = 0.20) {
  const connection = await pool.getConnection();

  try {
    const [clusters] = await connection.query(`
      SELECT 
        m1.id AS id1,
        m2.id AS id2,
        m1.content AS content1,
        m2.content AS content2,
        m1.importance AS importance1,
        m2.importance AS importance2,
        VEC_COSINE_DISTANCE(m1.embedding, m2.embedding) AS distance
      FROM memories m1
      JOIN memories m2 ON m1.user_id = m2.user_id AND m1.id < m2.id
      WHERE m1.user_id = ?
        AND m1.memory_type = 'episodic'
        AND m2.memory_type = 'episodic'
        AND m1.is_deleted = FALSE
        AND m2.is_deleted = FALSE
        AND m1.created_at > DATE_SUB(CURRENT_TIMESTAMP, INTERVAL 90 DAY)
        AND VEC_COSINE_DISTANCE(m1.embedding, m2.embedding) < ?
      ORDER BY distance ASC
      LIMIT 100
    `, [userId, similarityThreshold]);

    // Group into consolidation clusters
    const processed = new Set();
    const consolidationGroups = [];

    for (const pair of clusters) {
      if (processed.has(pair.id1) || processed.has(pair.id2)) continue;

      consolidationGroups.push({
        memberIds: [pair.id1, pair.id2],
        contents: [pair.content1, pair.content2],
        avgImportance: (pair.importance1 + pair.importance2) / 2
      });

      processed.add(pair.id1);
      processed.add(pair.id2);
    }

    return consolidationGroups;

  } finally {
    connection.release();
  }
}

/**
 * Consolidate a group of episodic memories into semantic knowledge
 */
async function consolidateMemoryGroup(userId, group) {
  const connection = await pool.getConnection();

  try {
    await connection.beginTransaction();

    // Use LLM to synthesize semantic fact from episodes
    const synthesis = await openai.chat.completions.create({
      model: 'gpt-4o-mini',
      messages: [
        {
          role: 'system',
          content: `You are consolidating multiple related episodic memories into a single semantic fact.

          Given these related memories:
          ${group.contents.map((c, i) => `${i + 1}. ${c}`).join('\n')}

          Synthesize a single, concise semantic fact that captures the core information.
          Be specific but general enough to apply across contexts.

          Return JSON:
          {
            "fact": "...",
            "confidence": 0.9
          }`
        }
      ],
      response_format: { type: 'json_object' }
    });

    const result = JSON.parse(synthesis.choices[0].message.content);

    // Generate embedding for consolidated fact
    const embedding = await generateEmbedding(result.fact);
    const embeddingJson = JSON.stringify(embedding);

    // Create semantic memory
    const [insertResult] = await connection.query(`
      INSERT INTO memories (
        user_id,
        content,
        embedding,
        memory_type,
        importance,
        confidence,
        source_memory_ids,
        created_at,
        last_accessed
      ) VALUES (?, ?, ?, 'semantic', ?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
    `, [
      userId,
      result.fact,
      embeddingJson,
      group.avgImportance,
      result.confidence,
      JSON.stringify(group.memberIds)
    ]);

    // Mark source episodes as consolidated (but keep for audit trail)
    await connection.query(`
      UPDATE memories
      SET metadata = JSON_SET(
        COALESCE(metadata, '{}'),
        '$.consolidated',
        TRUE,
        '$.consolidated_into',
        ?
      )
      WHERE id IN (?)
    `, [insertResult.insertId, group.memberIds]);

    await connection.commit();

    return {
      semanticMemoryId: insertResult.insertId,
      fact: result.fact,
      sourceCount: group.memberIds.length
    };

  } catch (error) {
    await connection.rollback();
    throw error;
  } finally {
    connection.release();
  }
}

/**
 * Run full consolidation pipeline for a user
 */
async function consolidateUserMemories(userId) {
  console.log(`Starting consolidation for user ${userId}`);

  const candidates = await findConsolidationCandidates(userId);
  console.log(`Found ${candidates.length} consolidation candidates`);

  const results = [];
  for (const group of candidates) {
    const result = await consolidateMemoryGroup(userId, group);
    results.push(result);

    console.log(`Consolidated: ${result.fact}`);
  }

  return results;
}

// Run consolidation (could be scheduled daily/weekly)
await consolidateUserMemories(123);

# TiKV configuration for hot storage nodes
server_configs:
  tikv:
    server.labels:
      performance: "high"
      tier: "hot"

# TiKV configuration for cold storage nodes
server_configs:
  tikv:
    server.labels:
      performance: "standard"
      tier: "cold"

-- Policy for hot data (recent memories)
CREATE PLACEMENT POLICY hot_memory
  CONSTRAINTS='[+tier=hot]'
  FOLLOWERS=2;

-- Policy for cold data (archived memories)
CREATE PLACEMENT POLICY cold_memory
  CONSTRAINTS='[+tier=cold]'
  FOLLOWERS=1;

-- Apply to partitioned table
CREATE TABLE memories_partitioned (
    id BIGINT PRIMARY KEY AUTO_INCREMENT,
    user_id BIGINT NOT NULL,
    content TEXT NOT NULL,
    embedding VECTOR(1536) COMMENT "hnsw(distance=cosine)",
    memory_type ENUM('episodic', 'semantic', 'working') NOT NULL,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    -- ... other columns

    INDEX idx_user_created (user_id, created_at),
    VECTOR INDEX idx_embedding (embedding)
)
PARTITION BY RANGE (UNIX_TIMESTAMP(created_at)) (
    PARTITION p_recent VALUES LESS THAN (UNIX_TIMESTAMP('2026-04-01')) PLACEMENT POLICY hot_memory,
    PARTITION p_last_month VALUES LESS THAN (UNIX_TIMESTAMP('2026-03-01')) PLACEMENT POLICY cold_memory,
    PARTITION p_archive VALUES LESS THAN (MAXVALUE) PLACEMENT POLICY cold_memory
);

# Scale out TiKV for storage
tiup cluster scale-out my-cluster scale-out.yaml

# Scale out TiFlash for analytics
tiup cluster scale-out my-cluster tiflash-scale-out.yaml

-- Check Region distribution across TiKV nodes
SELECT 
    store_id,
    address,
    COUNT(*) AS region_count,
    SUM(approximate_size) / 1024 / 1024 AS size_gb
FROM information_schema.tikv_region_status
GROUP BY store_id, address
ORDER BY region_count DESC;

-- Monitor memory table size
SELECT 
    table_schema,
    table_name,
    ROUND((data_length + index_length) / 1024 / 1024 / 1024, 2) AS size_gb,
    table_rows
FROM information_schema.tables
WHERE table_schema = 'ai_memory'
    AND table_name = 'memories';

import mysql from 'mysql2/promise';
import OpenAI from 'openai';

const openai = new OpenAI();
const pool = mysql.createPool({
  host: process.env.TIDB_HOST,
  port: 4000,
  user: process.env.TIDB_USER,
  password: process.env.TIDB_PASSWORD,
  database: 'ai_memory'
});

class RAGSystem {
  constructor(userId) {
    this.userId = userId;
  }

  /**
   * Chat with memory-enhanced AI
   */
  async chat(userMessage, sessionId = null) {
    // 1. Retrieve relevant memories
    const memories = await this.retrieveMemories(userMessage);

    // 2. Build context from memories
    const memoryContext = memories.length > 0
      ? `Here is what I remember about you:\n${memories.map(m => `- ${m.content}`).join('\n')}\n\n`
      : '';

    // 3. Generate response with memory context
    const response = await openai.chat.completions.create({
      model: 'gpt-4o',
      messages: [
        {
          role: 'system',
          content: `You are a helpful AI assistant with memory.

          ${memoryContext}

          Use the provided memories to personalize your responses.
          Reference past interactions naturally when relevant.`
        },
        { role: 'user', content: userMessage }
      ]
    });

    const aiResponse = response.choices[0].message.content;

    // 4. Extract and store new memories (background)
    this.processMessage(userMessage, sessionId).catch(console.error);

    // 5. Store this interaction as episodic memory
    this.storeInteraction(userMessage, aiResponse, sessionId).catch(console.error);

    return aiResponse;
  }

  async retrieveMemories(queryText) {
    const connection = await pool.getConnection();

    try {
      const queryEmbedding = await this.generateEmbedding(queryText);
      const embeddingJson = JSON.stringify(queryEmbedding);

      const [memories] = await connection.query(`
        SELECT 
          content,
          memory_type,
          (1 - VEC_COSINE_DISTANCE(embedding, ?)) AS relevance,
          (
            (1 - VEC_COSINE_DISTANCE(embedding, ?)) * 0.60 +
            EXP(-TIMESTAMPDIFF(DAY, created_at, CURRENT_TIMESTAMP) / 7.0) * 0.25 +
            importance * 0.15
          ) AS score
        FROM memories
        WHERE user_id = ?
          AND memory_type IN ('semantic', 'episodic')
          AND is_deleted = FALSE
          AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
        ORDER BY score DESC
        LIMIT 10
      `, [embeddingJson, embeddingJson, this.userId]);

      if (memories.length > 0) {
        const memoryIds = memories.map(m => m.id);
        await connection.query(`
          UPDATE memories
          SET last_accessed = CURRENT_TIMESTAMP,
              access_count = access_count + 1
          WHERE id IN (?)
        `, [memoryIds]);
      }

      return memories;

    } finally {
      connection.release();
    }
  }

  async processMessage(message, sessionId) {
    const connection = await pool.getConnection();

    try {
      // Extract facts using LLM
      const extraction = await openai.chat.completions.create({
        model: 'gpt-4o-mini',
        messages: [
          {
            role: 'system',
            content: `Extract memorable facts from the user message.
            Return JSON: {"facts": [{"fact": "...", "type": "semantic|episodic", "importance": 0.8}]}`
          },
          { role: 'user', content: message }
        ],
        response_format: { type: 'json_object' }
      });

      const result = JSON.parse(extraction.choices[0].message.content);
      const facts = result.facts || [];

      // Store each important fact
      for (const fact of facts.filter(f => f.importance > 0.6)) {
        const embedding = await this.generateEmbedding(fact.fact);
        const embeddingJson = JSON.stringify(embedding);

        // Check for duplicates
        const [existing] = await connection.query(`
          SELECT id FROM memories
          WHERE user_id = ?
            AND memory_type = ?
            AND VEC_COSINE_DISTANCE(embedding, ?) < 0.15
          LIMIT 1
        `, [this.userId, fact.type, embeddingJson]);

        if (existing.length === 0) {
          await connection.query(`
            INSERT INTO memories (
              user_id, session_id, content, embedding,
              memory_type, importance
            ) VALUES (?, ?, ?, ?, ?, ?)
          `, [
            this.userId, sessionId, fact.fact, embeddingJson,
            fact.type, fact.importance
          ]);
        }
      }

    } finally {
      connection.release();
    }
  }

  async storeInteraction(userMessage, aiResponse, sessionId) {
    const connection = await pool.getConnection();

    try {
      const interaction = `User: ${userMessage}\nAssistant: ${aiResponse}`;
      const embedding = await this.generateEmbedding(interaction);

      await connection.query(`
        INSERT INTO memories (
          user_id, session_id, content, embedding,
          memory_type, importance
        ) VALUES (?, ?, ?, ?, 'episodic', 0.5)
      `, [this.userId, sessionId, interaction, JSON.stringify(embedding)]);

    } finally {
      connection.release();
    }
  }

  async generateEmbedding(text) {
    const response = await openai.embeddings.create({
      model: 'text-embedding-3-small',
      input: text
    });
    return response.data[0].embedding;
  }
}

// Usage
const rag = new RAGSystem(123);

console.log(await rag.chat("I am working on a healthcare AI project"));
// Later...
console.log(await rag.chat("What projects am I working on?"));
// AI remembers the healthcare project

-- Use EXPLAIN to check if vector index is used
EXPLAIN SELECT content
FROM memories
WHERE user_id = 123
ORDER BY VEC_COSINE_DISTANCE(embedding, '[0.1, 0.2, ...]')
LIMIT 10;

-- Output should show VectorIndexScan
-- If not, check that:
-- 1. Vector index exists
-- 2. Query uses same distance metric as index
-- 3. TiDB version is 8.4.0+

-- Drop and recreate vector index with custom HNSW parameters
ALTER TABLE memories DROP INDEX idx_embedding;

-- Create optimized HNSW index
ALTER TABLE memories ADD VECTOR INDEX idx_embedding_optimized (embedding)
USING HNSW
COMMENT 'hnsw(distance=cosine, m=32, efconstruction=400)';

-- Parameters:
-- m: connections per layer (default 16, higher = more accurate but slower build)
-- efconstruction: build-time search depth (default 200, higher = better recall)

// Bad: Individual inserts
for (const memory of memories) {
  await db.query('INSERT INTO memories VALUES (?, ...)', [memory]);
}

// Good: Batch insert
const values = memories.map(m => [m.userId, m.content, m.embedding]);
await db.query('INSERT INTO memories (user_id, content, embedding) VALUES ?', [values]);

const pool = mysql.createPool({
  host: process.env.TIDB_HOST,
  port: 4000,
  user: process.env.TIDB_USER,
  password: process.env.TIDB_PASSWORD,
  database: 'ai_memory',
  connectionLimit: 20, // Adjust based on load
  queueLimit: 0,
  waitForConnections: true
});