92 lines
2.7 KiB
C
92 lines
2.7 KiB
C
|
#ifndef HUMANUS_MEMORY_MEM0_VECTOR_STORE_BASE_H
|
||
|
|
#define HUMANUS_MEMORY_MEM0_VECTOR_STORE_BASE_H
|
||
|
|
|
||
|
|
#include "hnswlib/hnswlib.h"
|
||
|
|
|
||
|
|
namespace humanus {
|
||
|
|
|
||
|
|
namespace mem0 {
|
||
|
|
|
||
|
|
struct VectorStoreConfig {
|
||
|
|
int dim = 16; // Dimension of the elements
|
||
|
|
int max_elements = 10000; // Maximum number of elements, should be known beforehand
|
||
|
|
int M = 16; // Tightly connected with internal dimensionality of the data
|
||
|
|
// strongly affects the memory consumption
|
||
|
|
int ef_construction = 200; // Controls index search speed/build speed tradeoff
|
||
|
|
enum class Metric {
|
||
|
|
L2,
|
||
|
|
IP
|
||
|
|
};
|
||
|
|
Metric metric = Metric::L2;
|
||
|
|
};
|
||
|
|
|
||
|
|
struct VectorStoreBase {
|
||
|
|
VectorStoreConfig config;
|
||
|
|
|
||
|
|
VectorStoreBase(const VectorStoreConfig& config) : config(config) {
|
||
|
|
reset();
|
||
|
|
}
|
||
|
|
|
||
|
|
virtual void reset() = 0;
|
||
|
|
|
||
|
|
/**
|
||
|
|
* @brief 插入向量到集合中
|
||
|
|
* @param vectors 向量数据
|
||
|
|
* @param payloads 可选的负载数据
|
||
|
|
* @param ids 可选的ID列表
|
||
|
|
* @return 插入的向量ID列表
|
||
|
|
*/
|
||
|
|
virtual std::vector<size_t> insert(const std::vector<std::vector<float>>& vectors,
|
||
|
|
const std::vector<std::string>& payloads = {},
|
||
|
|
const std::vector<size_t>& ids = {}) = 0;
|
||
|
|
|
||
|
|
/**
|
||
|
|
* @brief 搜索相似向量
|
||
|
|
* @param query 查询向量
|
||
|
|
* @param limit 返回结果数量限制
|
||
|
|
* @param filters 可选的过滤条件
|
||
|
|
* @return 相似向量的ID和距离
|
||
|
|
*/
|
||
|
|
std::vector<std::pair<size_t, std::vector<float>>> search(const std::vector<float>& query,
|
||
|
|
int limit = 5,
|
||
|
|
const std::string& filters = "") = 0;
|
||
|
|
|
||
|
|
/**
|
||
|
|
* @brief 通过ID删除向量
|
||
|
|
* @param vector_id 向量ID
|
||
|
|
*/
|
||
|
|
virtual void delete_vector(size_t vector_id) = 0;
|
||
|
|
|
||
|
|
/**
|
||
|
|
* @brief 更新向量及其负载
|
||
|
|
* @param vector_id 向量ID
|
||
|
|
* @param vector 可选的新向量数据
|
||
|
|
* @param payload 可选的新负载数据
|
||
|
|
*/
|
||
|
|
virtual void update(size_t vector_id,
|
||
|
|
const std::vector<float>* vector = nullptr,
|
||
|
|
const std::string* payload = nullptr) = 0;
|
||
|
|
|
||
|
|
/**
|
||
|
|
* @brief 通过ID获取向量
|
||
|
|
* @param vector_id 向量ID
|
||
|
|
* @return 向量数据
|
||
|
|
*/
|
||
|
|
virtual std::vector<float> get(size_t vector_id) = 0;
|
||
|
|
|
||
|
|
/**
|
||
|
|
* @brief 列出所有记忆
|
||
|
|
* @param filters 可选的过滤条件
|
||
|
|
* @param limit 可选的结果数量限制
|
||
|
|
* @return 记忆ID列表
|
||
|
|
*/
|
||
|
|
virtual std::vector<size_t> list(const std::string& filters = "", int limit = 0) = 0;
|
||
|
|
};
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
#endif // HUMANUS_MEMORY_MEM0_VECTOR_STORE_BASE_H
|