555 lines
24 KiB
C#
555 lines
24 KiB
C#
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(), new NgramKeywordExtractor(), new StubLlmRerankService());
|
||
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(), new NgramKeywordExtractor(), new StubLlmRerankService());
|
||
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(), new NgramKeywordExtractor(), new StubLlmRerankService());
|
||
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(), new NgramKeywordExtractor(), new StubLlmRerankService());
|
||
// 查询包含关键词 "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(), new NgramKeywordExtractor(), new StubLlmRerankService());
|
||
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(), new NgramKeywordExtractor(), new StubLlmRerankService());
|
||
var hits = await retrieval.RetrieveAsync(kbNone, "数据库", CancellationToken.None);
|
||
|
||
hits.Should().NotBeEmpty();
|
||
// None 策略下，向量相似度最高的 chunkHigh 应排第一（不受关键词重排影响）
|
||
hits[0].ChunkId.Should().Be(chunkHigh, "None 策略不应用关键词重排，排序完全由向量相似度决定");
|
||
}
|
||
|
||
// ===== LLM 重排（真实现 + 降级） =====
|
||
|
||
private async Task<(KnowledgeBaseEntity Kb, Guid DocId, Guid ChunkWeakVec, Guid ChunkMidVec)> SeedLlmRerankDataAsync(RagDbContext db, RerankStrategyEnum strategy)
|
||
{
|
||
var kb = new KnowledgeBaseEntity
|
||
{
|
||
Name = "LLM 重排测试", RetrievalMode = RetrievalModeEnum.Vector,
|
||
// 正交向量余弦相似度为负，阈值需低于负值避免候选被剔除
|
||
SimilarityThreshold = -1.0, RerankStrategy = strategy, ContextTopK = 2,
|
||
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 chunkWeakVec = Guid.NewGuid(); // 关键词命中但向量最不相似
|
||
var chunkMidVec = Guid.NewGuid(); // 关键词不中、向量相似度居中
|
||
db.DocumentChunks.Add(new DocumentChunkEntity { Id = chunkWeakVec, DocumentId = docId, Content = "数据库 数据库 数据库", ChunkIndex = 0, TokenCount = 5, CreatedBy = "t", UpdatedBy = "t" });
|
||
db.DocumentChunks.Add(new DocumentChunkEntity { Id = chunkMidVec, DocumentId = docId, Content = "完全不相关的正文内容", ChunkIndex = 1, TokenCount = 5, CreatedBy = "t", UpdatedBy = "t" });
|
||
await db.SaveChangesAsync();
|
||
|
||
var conn = (NpgsqlConnection)db.Database.GetDbConnection();
|
||
await conn.OpenAsync();
|
||
try
|
||
{
|
||
await WriteEmbeddingAsync(conn, chunkWeakVec, BuildOrthogonalVector());
|
||
await WriteEmbeddingAsync(conn, chunkMidVec, BuildStubLikeVector());
|
||
}
|
||
finally
|
||
{
|
||
await conn.CloseAsync();
|
||
}
|
||
return (kb, docId, chunkWeakVec, chunkMidVec);
|
||
}
|
||
|
||
[Fact]
|
||
public async Task Rerank_LlmSucceeds_AppliesModelOrderOverVectorOrder()
|
||
{
|
||
await CleanDatabaseAsync();
|
||
await using var db = _fixture.CreateDbContext();
|
||
var (kb, _, chunkWeakVec, chunkMidVec) = await SeedLlmRerankDataAsync(db, RerankStrategyEnum.LLM);
|
||
|
||
// 模型判定：关键词命中的 chunkWeakVec 相关度最高（覆盖"向量相似优先"的召回顺序）
|
||
var stub = new StubLlmRerankService([chunkWeakVec, chunkMidVec]);
|
||
var retrieval = new RagRetrievalService(db, new StubEmbeddingService(), new NgramKeywordExtractor(), stub);
|
||
var hits = await retrieval.RetrieveAsync(kb, "数据库", CancellationToken.None);
|
||
|
||
stub.CallCount.Should().Be(1, "LLM 策略必须真正调用重排服务");
|
||
hits.Should().HaveCount(2);
|
||
hits[0].ChunkId.Should().Be(chunkWeakVec, "LLM 给出的顺序必须生效（关键词相关 chunk 置顶）");
|
||
hits[1].ChunkId.Should().Be(chunkMidVec);
|
||
}
|
||
|
||
[Fact]
|
||
public async Task Rerank_LlmNull_FallsBackToKeywordRerank()
|
||
{
|
||
await CleanDatabaseAsync();
|
||
await using var db = _fixture.CreateDbContext();
|
||
var (kb, _, chunkWeakVec, _) = await SeedLlmRerankDataAsync(db, RerankStrategyEnum.LLM);
|
||
|
||
// LLM 失败（null）→ 降级 Keyword 重排：向量分 ×0.7 + 关键词分 ×0.3
|
||
var stub = new StubLlmRerankService((List<Guid>?)null);
|
||
var retrieval = new RagRetrievalService(db, new StubEmbeddingService(), new NgramKeywordExtractor(), stub);
|
||
var hits = await retrieval.RetrieveAsync(kb, "数据库", CancellationToken.None);
|
||
|
||
stub.CallCount.Should().Be(1, "降级前必须先尝试过 LLM");
|
||
hits.Should().NotBeEmpty();
|
||
// 降级路径等价于 Keyword 策略：两个 chunk 都召回，不因 LLM 失败而丢结果
|
||
hits.Should().HaveCount(2);
|
||
// Keyword 融合下向量满分的 chunkMidVec（1.0×0.7 + 0×0.3 = 0.7）
|
||
// vs 关键词满分的 chunkWeakVec（低向量分×0.7 + 1×0.3），前者更高 → chunkMidVec 第一
|
||
hits[0].Content.Should().Contain("完全不相关");
|
||
hits.Should().Contain(h => h.ChunkId == chunkWeakVec);
|
||
}
|
||
|
||
// ===== 批量向量写入 =====
|
||
|
||
[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();
|
||
}
|
||
}
|