89 lines
3.3 KiB
C#
89 lines
3.3 KiB
C#
using CodeBase.Models;
|
||
using Qdrant.Client;
|
||
using Qdrant.Client.Grpc;
|
||
using System;
|
||
using System.Collections.Generic;
|
||
using System.Linq;
|
||
using System.Threading.Tasks;
|
||
|
||
public class QdrantRepository
|
||
{
|
||
private readonly QdrantClient _client;
|
||
private const string CollectionName = "CodeBaseCollection";
|
||
|
||
public QdrantRepository(QdrantClient client)
|
||
{
|
||
_client = client;
|
||
}
|
||
|
||
// ==========================================
|
||
// 1. СОХРАНЕНИЕ ВЕКТОРОВ С ПРИВЯЗКОЙ К ПРОЕКТУ
|
||
// ==========================================
|
||
public async Task SaveChunksAsync(string projectName, List<CodeChunk> chunks)
|
||
{
|
||
var points = new List<PointStruct>();
|
||
|
||
foreach (var chunk in chunks)
|
||
{
|
||
if (chunk.Vector == null) continue;
|
||
|
||
// Формируем точку для БД
|
||
var point = new PointStruct
|
||
{
|
||
Id = Guid.NewGuid(), // Уникальный ID чанка
|
||
Vectors = chunk.Vector, // Тот самый массив float[]
|
||
|
||
// Вот здесь мы навешиваем все ярлыки!
|
||
Payload =
|
||
{
|
||
["ProjectName"] = projectName,
|
||
["FilePath"] = chunk.FilePath,
|
||
["MethodName"] = chunk.MethodName,
|
||
["Content"] = chunk.Content // Сам текст кода тоже храним прямо в БД
|
||
}
|
||
};
|
||
points.Add(point);
|
||
}
|
||
|
||
// Отправляем пачку векторов в базу
|
||
await _client.UpsertAsync(CollectionName, points);
|
||
}
|
||
|
||
// ==========================================
|
||
// 2. ПОИСК ТОЛЬКО В РАМКАХ ОДНОГО ПРОЕКТА
|
||
// ==========================================
|
||
public async Task<List<CodeChunk>> SearchAsync(string projectName, float[] queryVector, int topK = 15)
|
||
{
|
||
// Создаем фильтр: искать только там, где ProjectName совпадает с запрошенным
|
||
var projectFilter = new Filter
|
||
{
|
||
Must = {
|
||
new Condition
|
||
{
|
||
Field = new FieldCondition
|
||
{
|
||
Key = "ProjectName",
|
||
Match = new Match { Keyword = projectName }
|
||
}
|
||
}
|
||
}
|
||
};
|
||
|
||
// База делает всю тяжелую математику за нас!
|
||
var searchResults = await _client.SearchAsync(
|
||
collectionName: CollectionName,
|
||
vector: queryVector,
|
||
limit: (ulong)topK,
|
||
filter: projectFilter, // Передаем наш фильтр
|
||
payloadSelector: true // Просим БД вернуть Payload вместе с результатами
|
||
);
|
||
|
||
// Превращаем ответы базы обратно в наши C#-объекты CodeChunk
|
||
return searchResults.Select(hit => new CodeChunk
|
||
{
|
||
FilePath = hit.Payload["FilePath"].StringValue,
|
||
MethodName = hit.Payload["MethodName"].StringValue,
|
||
Content = hit.Payload["Content"].StringValue
|
||
}).ToList();
|
||
}
|
||
} |