humanus.cpp/memory/mem0/storage.h

150 lines
5.6 KiB
C
Raw Normal View History

2025-03-23 14:35:54 +08:00
#ifndef HUMANUS_MEMORY_MEM0_STORAGE_H
#define HUMANUS_MEMORY_MEM0_STORAGE_H
#include <sqlite3.h>
namespace humanus {
namespace mem0 {
struct SQLiteManager {
std::shared_ptr<sqlite3> db;
std::mutex mutex;
SQLiteManager(const std::string& db_path) {
int rc = sqlite3_open(db_path.c_str(), &db);
if (rc) {
throw std::runtime_error("Failed to open database: " + std::string(sqlite3_errmsg(db)));
}
_migrate_history_table();
_create_history_table()
}
void _migrate_history_table() {
std::lock_guard<std::mutex> lock(mutex);
char* errmsg = nullptr;
sqlite3_stmt* stmt = nullptr;
// 检查历史表是否存在
int rc = sqlite3_prepare_v2(db.get(), "SELECT name FROM sqlite_master WHERE type='table' AND name='history'", -1, &stmt, nullptr);
if (rc != SQLITE_OK) {
throw std::runtime_error("Failed to prepare statement: " + std::string(sqlite3_errmsg(db.get())));
}
bool table_exists = false;
if (sqlite3_step(stmt) == SQLITE_ROW) {
table_exists = true;
}
sqlite3_finalize(stmt);
if (table_exists) {
// 获取当前表结构
std::map<std::string, std::string> current_schema;
rc = sqlite3_prepare_v2(db.get(), "PRAGMA table_info(history)", -1, &stmt, nullptr);
if (rc != SQLITE_OK) {
throw std::runtime_error("Failed to prepare statement: " + std::string(sqlite3_errmsg(db.get())));
}
while (sqlite3_step(stmt) == SQLITE_ROW) {
std::string column_name = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 1));
std::string column_type = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 2));
current_schema[column_name] = column_type;
}
sqlite3_finalize(stmt);
// 定义预期表结构
std::map<std::string, std::string> expected_schema = {
{"id", "TEXT"},
{"memory_id", "TEXT"},
{"old_memory", "TEXT"},
{"new_memory", "TEXT"},
{"new_value", "TEXT"},
{"event", "TEXT"},
{"created_at", "DATETIME"},
{"updated_at", "DATETIME"},
{"is_deleted", "INTEGER"}
};
// 检查表结构是否一致
if (current_schema != expected_schema) {
// 重命名旧表
rc = sqlite3_exec(db.get(), "ALTER TABLE history RENAME TO old_history", nullptr, nullptr, &errmsg);
if (rc != SQLITE_OK) {
std::string error = errmsg ? errmsg : "Unknown error";
sqlite3_free(errmsg);
throw std::runtime_error("Failed to rename table: " + error);
}
// 创建新表
rc = sqlite3_exec(db.get(),
"CREATE TABLE IF NOT EXISTS history ("
"id TEXT PRIMARY KEY,"
"memory_id TEXT,"
"old_memory TEXT,"
"new_memory TEXT,"
"new_value TEXT,"
"event TEXT,"
"created_at DATETIME,"
"updated_at DATETIME,"
"is_deleted INTEGER"
")", nullptr, nullptr, &errmsg);
if (rc != SQLITE_OK) {
std::string error = errmsg ? errmsg : "Unknown error";
sqlite3_free(errmsg);
throw std::runtime_error("Failed to create table: " + error);
}
// 复制数据
rc = sqlite3_exec(db.get(),
"INSERT INTO history (id, memory_id, old_memory, new_memory, new_value, event, created_at, updated_at, is_deleted) "
"SELECT id, memory_id, prev_value, new_value, new_value, event, timestamp, timestamp, is_deleted "
"FROM old_history", nullptr, nullptr, &errmsg);
if (rc != SQLITE_OK) {
std::string error = errmsg ? errmsg : "Unknown error";
sqlite3_free(errmsg);
throw std::runtime_error("Failed to copy data: " + error);
}
// 删除旧表
rc = sqlite3_exec(db.get(), "DROP TABLE old_history", nullptr, nullptr, &errmsg);
if (rc != SQLITE_OK) {
std::string error = errmsg ? errmsg : "Unknown error";
sqlite3_free(errmsg);
throw std::runtime_error("Failed to drop old table: " + error);
}
}
}
}
void _create_history_table() {
std::lock_guard<std::mutex> lock(mutex);
char* errmsg = nullptr;
int rc = sqlite3_exec(db.get(),
"CREATE TABLE IF NOT EXISTS history ("
"id TEXT PRIMARY KEY,"
"memory_id TEXT,"
"old_memory TEXT,"
"new_memory TEXT,"
"new_value TEXT,"
"event TEXT,"
"created_at DATETIME,"
"updated_at DATETIME,"
"is_deleted INTEGER"
")", nullptr, nullptr, &errmsg);
if (rc != SQLITE_OK) {
std::string error = errmsg ? errmsg : "Unknown error";
sqlite3_free(errmsg);
throw std::runtime_error("Failed to create history table: " + error);
}
}
};
} // namespace mem0
} // namespace humanus
#endif // HUMANUS_MEMORY_MEM0_STORAGE_H