96 lines
3.9 KiB
C#
96 lines
3.9 KiB
C#
using System.Collections.Generic;
|
|
using System.Threading.Tasks;
|
|
using CodeBase.Models;
|
|
using Neo4j.Driver;
|
|
|
|
namespace CodeBase.Repositories
|
|
{
|
|
public class ScipRepository
|
|
{
|
|
private readonly IDriver _neo4jDriver;
|
|
|
|
public ScipRepository(IDriver neo4jDriver)
|
|
{
|
|
_neo4jDriver = neo4jDriver;
|
|
}
|
|
public async Task SaveChunksAsync(List<CodeChunk> chunks)
|
|
{
|
|
if (chunks == null || chunks.Count == 0) return;
|
|
|
|
await using var session = _neo4jDriver.AsyncSession();
|
|
|
|
await session.ExecuteWriteAsync(async tx =>
|
|
{
|
|
foreach (var chunk in chunks)
|
|
{
|
|
var uniqueId = $"{chunk.ProjectName}::{chunk.Id}";
|
|
|
|
var query = @"
|
|
MERGE (p:Project {name: $projectName})
|
|
|
|
MERGE (m:CodeEntity {id: $id})
|
|
SET m.projectName = $projectName,
|
|
m.name = $name,
|
|
m.filePath = $filePath,
|
|
m.code = $code,
|
|
m.language = $language,
|
|
m.embedding = $embedding
|
|
|
|
MERGE (p)-[:CONTAINS]->(m)";
|
|
|
|
await tx.RunAsync(query, new
|
|
{
|
|
id = uniqueId,
|
|
projectName = chunk.ProjectName,
|
|
name = chunk.EntityName,
|
|
filePath = chunk.FilePath,
|
|
code = chunk.Content,
|
|
language = chunk.Language,
|
|
embedding = chunk.Embedding
|
|
});
|
|
}
|
|
});
|
|
}
|
|
|
|
/// <summary>
|
|
/// Ищет в графе узлы, наиболее близкие к переданному вектору.
|
|
/// </summary>
|
|
/// <param name="queryVector">Вектор вопроса пользователя</param>
|
|
/// <param name="projectName">Имя проекта для фильтрации (опционально)</param>
|
|
/// <param name="topK">Сколько кусков кода вернуть</param>
|
|
public async Task<List<RetrievedContext>> FindSimilarNodesAsync(float[] queryVector, string projectName = null, int topK = 5)
|
|
{
|
|
await using var session = _neo4jDriver.AsyncSession();
|
|
|
|
var result = await session.ExecuteReadAsync(async tx =>
|
|
{
|
|
// Базовый запрос к векторному индексу
|
|
string cypherQuery = @"
|
|
CALL db.index.vector.queryNodes('code_embeddings', $topK, $queryVector)
|
|
YIELD node AS method, score
|
|
WHERE $projectName IS NULL OR method.projectName = $projectName
|
|
RETURN method.projectName AS Project, method.name AS Name, method.filePath AS Path, method.code AS Code, score
|
|
ORDER BY score DESC";
|
|
|
|
var cursor = await tx.RunAsync(cypherQuery, new { topK, queryVector, projectName });
|
|
var contexts = new List<RetrievedContext>();
|
|
|
|
while (await cursor.FetchAsync())
|
|
{
|
|
contexts.Add(new RetrievedContext
|
|
{
|
|
ProjectName = cursor.Current["Project"].As<string>(), // Читаем имя проекта
|
|
EntityName = cursor.Current["Name"].As<string>(),
|
|
FilePath = cursor.Current["Path"].As<string>(),
|
|
Content = cursor.Current["Code"].As<string>(),
|
|
SimilarityScore = cursor.Current["score"].As<double>()
|
|
});
|
|
}
|
|
|
|
return contexts;
|
|
});
|
|
|
|
return result;
|
|
}
|
|
}
|
|
} |