marsview-enterprise-node/service/ai/db.service.js

const { drizzle } = require('drizzle-orm/node-postgres');
const { Pool } = require('pg');
const { sql, gt, desc } = require('drizzle-orm');
const { vercelAiEmbeddings } = require('./schema.service');
const config = require('../../config');

const { cosineDistance } = require('drizzle-orm/sql');
const pool = new Pool({
  connectionString: config.DATABASE_URL,
  min: 2,
  max: 10,
});

const db = drizzle(pool);

// 实现相似度查询
async function findSimilarContent(userQueryEmbedded) {
  // 手动将数组转换为 PostgreSQL 支持的数组格式
  const embeddingArray = `{${userQueryEmbedded.join(',')}}`;

  // 使用 sql 模板字面量构建原始 SQL 查询
  //   const similarity = sql`1 - (cosine_distance(${vercelAiEmbeddings.embedding}, ${embeddingArray}::float[]))`;
  const similarity = sql`1 - (${cosineDistance(vercelAiEmbeddings.embedding, userQueryEmbedded)})`;
  return db
    .select({
      content: vercelAiEmbeddings.content,
      similarity,
    })
    .from(vercelAiEmbeddings)
    .where(gt(similarity, 0.5)) // 过滤相似度大于 0.5 的结果
    .orderBy((t) => desc(t.similarity)) // 按相似度降序排列
    .limit(4); // 限制返回 4 条结果
}

module.exports = {
  db,
  findSimilarContent,
};