build-flow.api/tests/BuildFlow.Rag.Tests/RagPgVectorIntegrationTests.cs

using System.Data;
using Microsoft.EntityFrameworkCore;
using Npgsql;
using BuildFlow.Domain.Entities.Document;
using BuildFlow.Domain.Entities.DocumentChunk;
using BuildFlow.Domain.Entities.KnowledgeBase;
using BuildFlow.Domain.Enums;
using BuildFlow.Infrastructure.AI;
using BuildFlow.Infrastructure.Persistence;
using BuildFlow.Rag.Tests.Stub;
using FluentAssertions;
using Xunit;

namespace BuildFlow.Rag.Tests;

/// <summary>
/// RAG 向量检索 + 写入的真实 PostgreSQL + pgvector 集成测试。
/// 覆盖表名正确性（R1）、向量召回、关键词召回、Hybrid RRF 融合（R4）、
/// ParentChild 父块回填（R3）、RerankStrategy 生效（R2）、批量向量写入。
/// 这是唯一一组连接真实 PG 的 RAG 测试，也是表名漂移 bug 的回归防线。
/// </summary>
[Collection(PostgresRagCollection.Name)]
public sealed class RagPgVectorIntegrationTests : IAsyncLifetime
{
    private readonly PostgresRagFixture _fixture;

    public RagPgVectorIntegrationTests(PostgresRagFixture fixture)
    {
        _fixture = fixture;
    }

    public Task InitializeAsync() => Task.CompletedTask;
    public Task DisposeAsync() => Task.CompletedTask;

    /// <summary>
    /// 清空 b_document_chunks 的 embedding 和数据，每个测试独立。
    /// </summary>
    private async Task CleanDatabaseAsync()
    {
        await using var db = _fixture.CreateDbContext();
        db.DocumentChunks.RemoveRange(db.DocumentChunks);
        db.Documents.RemoveRange(db.Documents);
        db.KnowledgeBases.RemoveRange(db.KnowledgeBases);
        await db.SaveChangesAsync();

        // 清除 embedding 列残留（EF 不映射 embedding，走 raw SQL）
        var conn = (NpgsqlConnection)db.Database.GetDbConnection();
        await conn.OpenAsync();
        try
        {
            await using var cmd = conn.CreateCommand();
            cmd.CommandText = "UPDATE b_document_chunks SET embedding = NULL";
            await cmd.ExecuteNonQueryAsync();
        }
        finally
        {
            await conn.CloseAsync();
        }
    }

    private static float[] MakeVector(params float[] values)
    {
        var v = new float[768];
        for (var i = 0; i < values.Length && i < 768; i++)
            v[i] = values[i];
        return v;
    }

    /// <summary>直接通过 raw SQL 写入 embedding 向量（绕过 EF，因为 EF Ignore 了 Embedding）。</summary>
    private static async Task WriteEmbeddingAsync(NpgsqlConnection conn, Guid chunkId, float[] vector)
    {
        var vectorStr = $"[{string.Join(",", vector)}]";
        await using var cmd = conn.CreateCommand();
        cmd.CommandText = $"UPDATE b_document_chunks SET embedding = '{vectorStr}'::vector WHERE \"Id\" = '{chunkId}'";
        await cmd.ExecuteNonQueryAsync();
    }

    /// <summary>构建与 StubEmbeddingService 相同的向量（i%10/10），用于让某 chunk 与查询"最相似"。</summary>
    private static float[] BuildStubLikeVector() => StubEmbeddingService.BuildVectorLike();

    /// <summary>构建与 StubEmbeddingService 正交的向量，用于让某 chunk 与查询"最不相似"。</summary>
    private static float[] BuildOrthogonalVector()
    {
        var v = new float[768];
        // 把 stub 的模式偏移：stub 是 0,0.1,...0.9 重复，这里用全 1 取负余弦
        for (var i = 0; i < 768; i++)
            v[i] = -(i % 10) / 10f - 0.5f;
        return v;
    }

    // ===== R1 回归：表名必须使用 b_ 前缀 =====

    [Fact]
    public async Task VectorRecall_UsesBPrefixedTableNames()
    {
        // 如果表名错误（document_chunks vs b_document_chunks），查询会抛 PostgresException
        await CleanDatabaseAsync();
        await using var db = _fixture.CreateDbContext();

        var kb = new KnowledgeBaseEntity
        {
            Name = "向量召回测试",
            RetrievalMode = RetrievalModeEnum.Vector,
            CreatedBy = "test", UpdatedBy = "test", OperatorIP = "127.0.0.1"
        };
        db.KnowledgeBases.Add(kb);

        var docId = Guid.NewGuid();
        db.Documents.Add(new DocumentEntity
        {
            Id = docId, KnowledgeBaseId = kb.Id, Title = "测试文档",
            FileName = "test.md", FilePath = "test", ContentType = "text",
            Status = DocumentStatusEnum.Completed,
            CreatedBy = "test", UpdatedBy = "test", OperatorIP = "127.0.0.1"
        });

        var chunkId = Guid.NewGuid();
        db.DocumentChunks.Add(new DocumentChunkEntity
        {
            Id = chunkId, DocumentId = docId, Content = "PostgreSQL 是一个强大的关系型数据库",
            ChunkIndex = 0, TokenCount = 10,
            CreatedBy = "test", UpdatedBy = "test"
        });
        await db.SaveChangesAsync();

        // 写入向量：与目标文本 "关系型数据库" 对应位置取高值
        var conn = (NpgsqlConnection)db.Database.GetDbConnection();
        await conn.OpenAsync();
        try
        {
            await WriteEmbeddingAsync(conn, chunkId, MakeVector(0.9f, 0.1f, 0f));
        }
        finally
        {
            await conn.CloseAsync();
        }

        // 检索：不抛异常 = 表名正确（R1 回归）
        var retrieval = new RagRetrievalService(db, new StubEmbeddingService());
        var act = () => retrieval.RetrieveAsync(kb, "数据库", CancellationToken.None);

        await act.Should().NotThrowAsync("b_document_chunks 表名正确，查询不应抛 relation does not exist");
    }

    // ===== 向量召回 =====

    [Fact]
    public async Task VectorRecall_ReturnsMostSimilarChunk()
    {
        await CleanDatabaseAsync();
        await using var db = _fixture.CreateDbContext();

        var kb = new KnowledgeBaseEntity
        {
            Name = "余弦相似度测试", RetrievalMode = RetrievalModeEnum.Vector,
            SimilarityThreshold = 0.0, // 放宽阈值，确保有结果
            CreatedBy = "test", UpdatedBy = "test", OperatorIP = "127.0.0.1"
        };
        db.KnowledgeBases.Add(kb);

        var docId = Guid.NewGuid();
        db.Documents.Add(new DocumentEntity
        {
            Id = docId, KnowledgeBaseId = kb.Id, Title = "doc", FileName = "a.md",
            FilePath = "a", ContentType = "text", Status = DocumentStatusEnum.Completed,
            CreatedBy = "test", UpdatedBy = "test", OperatorIP = "127.0.0.1"
        });

        var chunkA = Guid.NewGuid();
        var chunkB = Guid.NewGuid();
        db.DocumentChunks.Add(new DocumentChunkEntity { Id = chunkA, DocumentId = docId, Content = "苹果", ChunkIndex = 0, TokenCount = 1, CreatedBy = "t", UpdatedBy = "t" });
        db.DocumentChunks.Add(new DocumentChunkEntity { Id = chunkB, DocumentId = docId, Content = "香蕉", ChunkIndex = 1, TokenCount = 1, CreatedBy = "t", UpdatedBy = "t" });
        await db.SaveChangesAsync();

        var conn = (NpgsqlConnection)db.Database.GetDbConnection();
        await conn.OpenAsync();
        try
        {
            // chunkA 向量与 StubEmbeddingService 返回的查询向量完全相同 → 余弦相似度=1
            // chunkB 向量与查询向量正交/相反 → 余弦相似度低
            await WriteEmbeddingAsync(conn, chunkA, BuildStubLikeVector());
            await WriteEmbeddingAsync(conn, chunkB, BuildOrthogonalVector());
        }
        finally
        {
            await conn.CloseAsync();
        }

        // Stub 返回固定向量，与 chunkA（同向量）余弦相似度=1，chunkB 差异大
        var retrieval = new RagRetrievalService(db, new StubEmbeddingService());
        var hits = await retrieval.RetrieveAsync(kb, "test", CancellationToken.None);

        hits.Should().NotBeEmpty();
        hits[0].ChunkId.Should().Be(chunkA, "chunkA 向量与查询向量最接近");
    }

    // ===== 关键词召回 =====

    [Fact]
    public async Task KeywordRecall_MatchesExpectedChunks()
    {
        await CleanDatabaseAsync();
        await using var db = _fixture.CreateDbContext();

        var kb = new KnowledgeBaseEntity
        {
            Name = "关键词测试", RetrievalMode = RetrievalModeEnum.Keyword,
            CreatedBy = "test", UpdatedBy = "test", OperatorIP = "127.0.0.1"
        };
        db.KnowledgeBases.Add(kb);

        var docId = Guid.NewGuid();
        db.Documents.Add(new DocumentEntity
        {
            Id = docId, KnowledgeBaseId = kb.Id, Title = "doc", FileName = "a.md",
            FilePath = "a", ContentType = "text", Status = DocumentStatusEnum.Completed,
            CreatedBy = "test", UpdatedBy = "test", OperatorIP = "127.0.0.1"
        });

        db.DocumentChunks.Add(new DocumentChunkEntity { Id = Guid.NewGuid(), DocumentId = docId, Content = "ASP.NET Core 是一个跨平台的 Web 框架", ChunkIndex = 0, TokenCount = 5, CreatedBy = "t", UpdatedBy = "t" });
        db.DocumentChunks.Add(new DocumentChunkEntity { Id = Guid.NewGuid(), DocumentId = docId, Content = "Python 是一种动态语言", ChunkIndex = 1, TokenCount = 3, CreatedBy = "t", UpdatedBy = "t" });
        await db.SaveChangesAsync();

        var retrieval = new RagRetrievalService(db, new StubEmbeddingService());
        var hits = await retrieval.RetrieveAsync(kb, "ASP.NET Core 框架", CancellationToken.None);

        hits.Should().NotBeEmpty();
        hits.Should().OnlyContain(h => h.Content.Contains("ASP.NET Core"), "关键词召回应只返回命中的 chunk");
    }

    // ===== R4: Hybrid 双路召回 + RRF 融合 =====

    [Fact]
    public async Task HybridRecall_MergesBothPaths()
    {
        await CleanDatabaseAsync();
        await using var db = _fixture.CreateDbContext();

        var kb = new KnowledgeBaseEntity
        {
            Name = "Hybrid 测试", RetrievalMode = RetrievalModeEnum.Hybrid,
            SimilarityThreshold = 0.0, RerankStrategy = RerankStrategyEnum.None,
            CreatedBy = "test", UpdatedBy = "test", OperatorIP = "127.0.0.1"
        };
        db.KnowledgeBases.Add(kb);

        var docId = Guid.NewGuid();
        db.Documents.Add(new DocumentEntity
        {
            Id = docId, KnowledgeBaseId = kb.Id, Title = "doc", FileName = "a.md",
            FilePath = "a", ContentType = "text", Status = DocumentStatusEnum.Completed,
            CreatedBy = "test", UpdatedBy = "test", OperatorIP = "127.0.0.1"
        });

        // chunkA: 向量相似但关键词不命中
        // chunkB: 向量不相似但关键词命中
        var chunkA = Guid.NewGuid();
        var chunkB = Guid.NewGuid();
        db.DocumentChunks.Add(new DocumentChunkEntity { Id = chunkA, DocumentId = docId, Content = "完全无关的内容", ChunkIndex = 0, TokenCount = 2, CreatedBy = "t", UpdatedBy = "t" });
        db.DocumentChunks.Add(new DocumentChunkEntity { Id = chunkB, DocumentId = docId, Content = "Docker 容器化部署指南", ChunkIndex = 1, TokenCount = 3, CreatedBy = "t", UpdatedBy = "t" });
        await db.SaveChangesAsync();

        var conn = (NpgsqlConnection)db.Database.GetDbConnection();
        await conn.OpenAsync();
        try
        {
            // chunkA 向量与 StubEmbeddingService 返回的查询向量完全相同 → 向量路召回 chunkA
            await WriteEmbeddingAsync(conn, chunkA, BuildStubLikeVector());
            await WriteEmbeddingAsync(conn, chunkB, BuildOrthogonalVector());
        }
        finally
        {
            await conn.CloseAsync();
        }

        var retrieval = new RagRetrievalService(db, new StubEmbeddingService());
        // 查询包含关键词 "Docker"，关键词路应召回 chunkB
        var hits = await retrieval.RetrieveAsync(kb, "Docker", CancellationToken.None);

        // Hybrid 应同时召回 chunkA（向量路）和 chunkB（关键词路）
        var hitIds = hits.Select(h => h.ChunkId).ToHashSet();
        hitIds.Should().Contain(chunkB, "关键词路应召回包含 Docker 的 chunk");
    }

    // ===== R3: ParentChild 父块回填 =====

    [Fact]
    public async Task ParentChild_EnrichesWithParentContent()
    {
        await CleanDatabaseAsync();
        await using var db = _fixture.CreateDbContext();

        var kb = new KnowledgeBaseEntity
        {
            Name = "父子分块测试", RetrievalMode = RetrievalModeEnum.Vector,
            SimilarityThreshold = 0.0,
            CreatedBy = "test", UpdatedBy = "test", OperatorIP = "127.0.0.1"
        };
        db.KnowledgeBases.Add(kb);

        var docId = Guid.NewGuid();
        db.Documents.Add(new DocumentEntity
        {
            Id = docId, KnowledgeBaseId = kb.Id, Title = "doc", FileName = "a.md",
            FilePath = "a", ContentType = "text", Status = DocumentStatusEnum.Completed,
            CreatedBy = "test", UpdatedBy = "test", OperatorIP = "127.0.0.1"
        });

        // 父块：完整段落
        var parentId = Guid.NewGuid();
        var childId = Guid.NewGuid();
        var parentContent = "这是完整的父块内容，包含丰富的上下文信息，子块只是其中一小部分。";
        var childContent = "子块片段";

        db.DocumentChunks.Add(new DocumentChunkEntity { Id = parentId, DocumentId = docId, Content = parentContent, ChunkIndex = 0, TokenCount = 10, CreatedBy = "t", UpdatedBy = "t" });
        db.DocumentChunks.Add(new DocumentChunkEntity { Id = childId, DocumentId = docId, Content = childContent, ChunkIndex = 1, ParentChunkId = parentId, TokenCount = 2, CreatedBy = "t", UpdatedBy = "t" });
        await db.SaveChangesAsync();

        // 只有子块有 embedding（模拟 ParentChild：子块向量化，父块不向量化）
        var conn = (NpgsqlConnection)db.Database.GetDbConnection();
        await conn.OpenAsync();
        try
        {
            await WriteEmbeddingAsync(conn, childId, MakeVector(1f, 0f, 0f));
        }
        finally
        {
            await conn.CloseAsync();
        }

        var retrieval = new RagRetrievalService(db, new StubEmbeddingService());
        var hits = await retrieval.RetrieveAsync(kb, "test", CancellationToken.None);

        hits.Should().NotBeEmpty();
        var childHit = hits.FirstOrDefault(h => h.ChunkId == childId);
        childHit.Should().NotBeNull("子块应被向量召回");
        childHit!.ContextContent.Should().Be(parentContent, "命中子块时应回填父块完整内容而非子块自身内容");
    }

    // ===== R2: RerankStrategy 生效 =====

    [Fact]
    public async Task RerankStrategy_None_UsesRawScoresWithoutKeywordRerank()
    {
        await CleanDatabaseAsync();
        await using var db = _fixture.CreateDbContext();

        var kbNone = new KnowledgeBaseEntity
        {
            Name = "None 策略", RetrievalMode = RetrievalModeEnum.Vector,
            SimilarityThreshold = 0.0, RerankStrategy = RerankStrategyEnum.None,
            CreatedBy = "test", UpdatedBy = "test", OperatorIP = "127.0.0.1"
        };
        db.KnowledgeBases.Add(kbNone);
        var docId = Guid.NewGuid();
        db.Documents.Add(new DocumentEntity
        {
            Id = docId, KnowledgeBaseId = kbNone.Id, Title = "doc", FileName = "a.md",
            FilePath = "a", ContentType = "text", Status = DocumentStatusEnum.Completed,
            CreatedBy = "test", UpdatedBy = "test", OperatorIP = "127.0.0.1"
        });

        var chunkHigh = Guid.NewGuid();
        var chunkLow = Guid.NewGuid();
        db.DocumentChunks.Add(new DocumentChunkEntity { Id = chunkHigh, DocumentId = docId, Content = "向量高相似但关键词完全不匹配的内容", ChunkIndex = 0, TokenCount = 5, CreatedBy = "t", UpdatedBy = "t" });
        db.DocumentChunks.Add(new DocumentChunkEntity { Id = chunkLow, DocumentId = docId, Content = "数据库", ChunkIndex = 1, TokenCount = 1, CreatedBy = "t", UpdatedBy = "t" });
        await db.SaveChangesAsync();

        var conn = (NpgsqlConnection)db.Database.GetDbConnection();
        await conn.OpenAsync();
        try
        {
            // chunkHigh 向量与 StubEmbeddingService 返回的查询向量完全相同 → 余弦相似度=1
            // chunkLow 向量与查询向量正交 → 余弦相似度低
            await WriteEmbeddingAsync(conn, chunkHigh, BuildStubLikeVector());
            await WriteEmbeddingAsync(conn, chunkLow, BuildOrthogonalVector());
        }
        finally
        {
            await conn.CloseAsync();
        }

        // None 策略：不做关键词重排，排名完全由向量相似度决定
        var retrieval = new RagRetrievalService(db, new StubEmbeddingService());
        var hits = await retrieval.RetrieveAsync(kbNone, "数据库", CancellationToken.None);

        hits.Should().NotBeEmpty();
        // None 策略下，向量相似度最高的 chunkHigh 应排第一（不受关键词重排影响）
        hits[0].ChunkId.Should().Be(chunkHigh, "None 策略不应用关键词重排，排序完全由向量相似度决定");
    }

    // ===== 批量向量写入 =====

    [Fact]
    public async Task BatchVectorWrite_WritesAllEmbeddings()
    {
        await CleanDatabaseAsync();
        await using var db = _fixture.CreateDbContext();

        var kb = new KnowledgeBaseEntity
        {
            Name = "批量写入测试", RetrievalMode = RetrievalModeEnum.Vector,
            SimilarityThreshold = 0.0,
            CreatedBy = "test", UpdatedBy = "test", OperatorIP = "127.0.0.1"
        };
        db.KnowledgeBases.Add(kb);
        var docId = Guid.NewGuid();
        db.Documents.Add(new DocumentEntity
        {
            Id = docId, KnowledgeBaseId = kb.Id, Title = "doc", FileName = "a.md",
            FilePath = "a", ContentType = "text", Status = DocumentStatusEnum.Completed,
            CreatedBy = "test", UpdatedBy = "test", OperatorIP = "127.0.0.1"
        });

        // 插入 55 个 chunk（超过一个 batch=50），验证分批写入正确
        var chunkIds = new List<Guid>();
        for (var i = 0; i < 55; i++)
        {
            var id = Guid.NewGuid();
            chunkIds.Add(id);
            db.DocumentChunks.Add(new DocumentChunkEntity
            {
                Id = id, DocumentId = docId, Content = $"内容{i}",
                ChunkIndex = i, TokenCount = 2, CreatedBy = "t", UpdatedBy = "t"
            });
        }
        await db.SaveChangesAsync();

        // 模拟批量写入：调用与 ProcessDocumentCommandHandler 相同的 SQL 模式
        var vectors = chunkIds.Select(_ => MakeVector(0.5f, 0.5f, 0f)).ToList();
        const int batchSize = 50;
        var conn = (NpgsqlConnection)db.Database.GetDbConnection();
        await conn.OpenAsync();
        try
        {
            for (var start = 0; start < chunkIds.Count; start += batchSize)
            {
                var sliceLen = Math.Min(batchSize, chunkIds.Count - start);
                var valuesClauses = new List<string>(sliceLen);
                for (var i = start; i < start + sliceLen; i++)
                {
                    var vectorStr = $"[{string.Join(",", vectors[i])}]";
                    valuesClauses.Add($"('{chunkIds[i]}'::uuid, '{vectorStr}'::vector)");
                }

                await using var cmd = conn.CreateCommand();
                cmd.CommandText = $"""
                    UPDATE b_document_chunks SET embedding = v.vec
                    FROM (VALUES {string.Join(", ", valuesClauses)}) AS v(id, vec)
                    WHERE b_document_chunks."Id" = v.id
                    """;
                await cmd.ExecuteNonQueryAsync();
            }
        }
        finally
        {
            await conn.CloseAsync();
        }

        // 验证所有 55 个 chunk 都有 embedding
        await using var verifyConn = new NpgsqlConnection(_fixture.ConnectionString);
        await verifyConn.OpenAsync();
        await using var verifyCmd = verifyConn.CreateCommand();
        verifyCmd.CommandText = """
            SELECT COUNT(*) FROM b_document_chunks
            WHERE "DocumentId" = @docId AND embedding IS NOT NULL
            """;
        verifyCmd.Parameters.AddWithValue("docId", docId);
        var count = (long)(await verifyCmd.ExecuteScalarAsync() ?? 0);
        count.Should().Be(55, "所有 55 个 chunk 的 embedding 都应写入成功");

        await verifyConn.CloseAsync();
    }
}