108 lines
3.3 KiB
C#
108 lines
3.3 KiB
C#
using BotSharp.Abstraction.Agents;
|
|
using BotSharp.Abstraction.VectorStorage;
|
|
using Microsoft.Extensions.DependencyInjection;
|
|
using Qdrant.Client;
|
|
using Qdrant.Client.Grpc;
|
|
using System;
|
|
using System.Collections.Generic;
|
|
using System.IO;
|
|
using System.Linq;
|
|
using System.Threading.Tasks;
|
|
|
|
namespace BotSharp.Plugin.Qdrant;
|
|
|
|
public class QdrantDb : IVectorDb
|
|
{
|
|
private QdrantClient _client;
|
|
private readonly QdrantSetting _setting;
|
|
private readonly IServiceProvider _services;
|
|
|
|
public QdrantDb(QdrantSetting setting,
|
|
IServiceProvider services)
|
|
{
|
|
_setting = setting;
|
|
_services = services;
|
|
|
|
}
|
|
|
|
private QdrantClient GetClient()
|
|
{
|
|
if (_client == null)
|
|
{
|
|
_client = new QdrantClient
|
|
(
|
|
host: _setting.Url,
|
|
apiKey: _setting.ApiKey
|
|
);
|
|
}
|
|
return _client;
|
|
}
|
|
|
|
public async Task<List<string>> GetCollections()
|
|
{
|
|
// List all the collections
|
|
var collections = await _client.ListCollectionsAsync();
|
|
return collections.ToList();
|
|
}
|
|
|
|
public async Task CreateCollection(string collectionName, int dim)
|
|
{
|
|
var collections = await GetCollections();
|
|
if (!collections.Contains(collectionName))
|
|
{
|
|
// Create a new collection
|
|
await GetClient().CreateCollectionAsync(collectionName, new VectorParams()
|
|
{
|
|
Size = (ulong)dim,
|
|
Distance = Distance.Cosine
|
|
});
|
|
|
|
var agentService = _services.GetRequiredService<IAgentService>();
|
|
var agentDataDir = agentService.GetAgentDataDir(collectionName);
|
|
var knowledgePath = Path.Combine(agentDataDir, "knowledge.txt");
|
|
File.WriteAllLines(knowledgePath, new string[0]);
|
|
}
|
|
|
|
// Get collection info
|
|
var collectionInfo = await _client.GetCollectionInfoAsync(collectionName);
|
|
if (collectionInfo == null)
|
|
{
|
|
throw new Exception($"Create {collectionName} failed.");
|
|
}
|
|
}
|
|
|
|
public async Task Upsert(string collectionName, int id, float[] vector, string text)
|
|
{
|
|
// Insert vectors
|
|
await GetClient().UpsertAsync(collectionName, points: new List<PointStruct>
|
|
{
|
|
new PointStruct()
|
|
{
|
|
Id = new PointId()
|
|
{
|
|
Num = (ulong)id,
|
|
},
|
|
Vectors = vector
|
|
}
|
|
});
|
|
|
|
// Store chunks in local file system
|
|
var agentService = _services.GetRequiredService<IAgentService>();
|
|
var agentDataDir = agentService.GetAgentDataDir(collectionName);
|
|
var knowledgePath = Path.Combine(agentDataDir, "knowledge.txt");
|
|
File.AppendAllLines(knowledgePath, new[] { text });
|
|
}
|
|
|
|
public async Task<List<string>> Search(string collectionName, float[] vector, int limit = 5)
|
|
{
|
|
var result = await GetClient().SearchAsync(collectionName, vector, limit: (ulong)limit);
|
|
|
|
var agentService = _services.GetRequiredService<IAgentService>();
|
|
var agentDataDir = agentService.GetAgentDataDir(collectionName);
|
|
var knowledgePath = Path.Combine(agentDataDir, "knowledge.txt");
|
|
var texts = File.ReadAllLines(knowledgePath);
|
|
|
|
return result.Select(x => texts[x.Id.Num]).ToList();
|
|
}
|
|
}
|