humanus.cpp/memory/mem0/base.h

#ifndef HUMANUS_MEMORY_MEM0_H
#define HUMANUS_MEMORY_MEM0_H

#include "memory/base.h"
#include "vector_store/base.h"
#include "embedding_model/base.h"
#include "schema.h"
#include "prompt.h"
#include "httplib.h"
#include "llm.h"
#include "utils.h"
#include "tool/fact_extract.h"

namespace humanus::mem0 {

struct Memory : BaseMemory {
    MemoryConfig config;

    std::string current_request;
    
    std::string fact_extraction_prompt;
    std::string update_memory_prompt;
    int max_messages;
    int retrieval_limit;
    FilterFunc filter;

    std::shared_ptr<EmbeddingModel> embedding_model;
    std::shared_ptr<VectorStore> vector_store;
    std::shared_ptr<LLM> llm;
    // std::shared_ptr<SQLiteManager> db;

    std::shared_ptr<FactExtract> fact_extract_tool;
    
    Memory(const MemoryConfig& config) : config(config) {
        fact_extraction_prompt = config.fact_extraction_prompt;
        update_memory_prompt = config.update_memory_prompt;
        max_messages = config.max_messages;
        retrieval_limit = config.retrieval_limit;
        filter = config.filter;

        size_t pos = fact_extraction_prompt.find("{current_date}");
        if (pos != std::string::npos) {
            // %Y-%d-%m
            auto current_date = std::chrono::system_clock::now();
            auto in_time_t = std::chrono::system_clock::to_time_t(current_date);
            std::stringstream ss;
            std::tm tm_info = *std::localtime(&in_time_t);
            ss << std::put_time(&tm_info, "%Y-%m-%d");
            std::string formatted_date = ss.str(); // YYYY-MM-DD
            fact_extraction_prompt.replace(pos, 14, formatted_date);
        }

        embedding_model = EmbeddingModel::get_instance("default", config.embedding_model_config);
        vector_store = VectorStore::get_instance("default", config.vector_store_config);
        
        llm = LLM::get_instance("default", config.llm_config);
        // db = std::make_shared<SQLiteManager>(config.history_db_path);

        fact_extract_tool = std::make_shared<FactExtract>();
    }

    void add_message(const Message& message) override {
        messages.push_back(message);
        // Message message_to_memory = message;
        // if (llm->enable_vision()) {
        //     message_to_memory = parse_vision_message(message_to_memory, llm, llm->vision_details());
        // } else {
        //     message_to_memory = parse_vision_message(message_to_memory);
        // }
        // _add_to_vector_store({message_to_memory}); 
        // while (!messages.empty() && (messages.size() > max_messages || messages.front().role == "assistant" || messages.front().role == "tool")) {
        //     // Ensure the first message is always a user or system message
        //     messages.pop_front();
        // }
        std::vector<Message> messages_to_memory;
        while (!messages.empty() && (messages.size() > max_messages || messages.front().role == "assistant" || messages.front().role == "tool")) {
            // Ensure the first message is always a user or system message
            messages_to_memory.push_back(messages.front());
            messages.pop_front();
        }
        if (!messages_to_memory.empty()) {
            if (llm->enable_vision()) {
                for (auto& m : messages_to_memory) {
                    m = parse_vision_message(m, llm, llm->vision_details());
                }
            } else {
                for (auto& m : messages_to_memory) {
                    m = parse_vision_message(m);
                }
            }
            _add_to_vector_store(messages_to_memory); 
        }
    }

    std::vector<Message> get_messages(const std::string& query = "") const override {
        std::vector<Message> messages_with_memory;

        if (!query.empty()) {
            auto embeddings = embedding_model->embed(query, EmbeddingType::SEARCH);
            std::vector<MemoryItem> memories;
            
            // 检查vector_store是否已初始化
            if (vector_store) {
                memories = vector_store->search(embeddings, retrieval_limit, filter);
            }

            if (!memories.empty()) {
                sort(memories.begin(), memories.end(), [](const MemoryItem& a, const MemoryItem& b) {
                    return a.updated_at < b.updated_at;
                });

                std::string memory_prompt;
                for (const auto& memory_item : memories) {
                    memory_prompt += "<memory>" + memory_item.memory + "</memory>";
                }

                messages_with_memory.push_back(Message::user_message(memory_prompt));

                logger->info("📤 Total retreived memories: " + std::to_string(memories.size()));
            }
        }

        messages_with_memory.insert(messages_with_memory.end(), messages.begin(), messages.end());

        return messages_with_memory;
    }

    void clear() override {
        if (messages.empty()) {
            return;
        }
        std::vector<Message> messages_to_memory(messages.begin(), messages.end());
        _add_to_vector_store(messages_to_memory);
        messages.clear();
    }

    void _add_to_vector_store(const std::vector<Message>& messages) {
        // 检查vector_store是否已初始化
        if (!vector_store) {
            logger->warn("Vector store is not initialized, skipping memory operation");
            return;
        }
        
        std::string parsed_message;
        
        for (const auto& message : messages) {
            parsed_message += message.role + ": " + (message.content.is_string() ? message.content.get<std::string>() : message.content.dump()) + "\n";

            for (const auto& tool_call : message.tool_calls) {
                parsed_message += "<tool_call>" + tool_call.to_json().dump() + "</tool_call>\n";
            }
        }
    
        std::string system_prompt = fact_extraction_prompt;

        size_t pos = system_prompt.find("{current_request}");
        if (pos != std::string::npos) {
            system_prompt = system_prompt.replace(pos, 17, current_request);
        }

        std::string user_prompt = "Input:\n" + parsed_message;

        Message user_message = Message::user_message(user_prompt);

        json response = llm->ask_tool(
            {user_message},
            system_prompt,
            "",
            json::array({fact_extract_tool->to_param()}),
            "required"
        );

        std::vector<std::string> new_facts; // ["fact1", "fact2", "fact3"]

        try {
            auto tool_calls = ToolCall::from_json_list(response["tool_calls"]);
            for (const auto& tool_call : tool_calls) {
                // Parse arguments
                json args = tool_call.function.arguments;

                if (args.is_string()) {
                    args = json::parse(args.get<std::string>());
                }

                auto facts = fact_extract_tool->execute(args).output.get<std::vector<std::string>>();
                new_facts.insert(new_facts.end(), facts.begin(), facts.end());
            }
        } catch (const std::exception& e) {
            logger->error("Error in new_facts: " + std::string(e.what()));
        }

        if (new_facts.empty()) {
            return;
        }

        logger->info("📫 New facts to remember: " + json(new_facts).dump());

        std::vector<json> old_memories;
        std::map<std::string, std::vector<float>> new_message_embeddings;

        for (const auto& fact : new_facts) {
            auto message_embedding = embedding_model->embed(fact, EmbeddingType::ADD);
            new_message_embeddings[fact] = message_embedding;
            auto existing_memories = vector_store->search(
                message_embedding,
                5
            );
            for (const auto& memory : existing_memories) {
                old_memories.push_back({
                    {"id", memory.id},
                    {"text", memory.memory}
                });
            }
        }
        // sort and unique by id
        std::sort(old_memories.begin(), old_memories.end(), [](const json& a, const json& b) {
            return a["id"] < b["id"];
        });
        old_memories.resize(std::unique(old_memories.begin(), old_memories.end(), [](const json& a, const json& b) {
            return a["id"] == b["id"];
        }) - old_memories.begin());
        logger->info("🧠 Existing memories about new facts: " + std::to_string(old_memories.size()));

        // mapping UUIDs with integers for handling ID hallucinations
        std::vector<size_t> temp_id_mapping;
        for (size_t idx = 0; idx < old_memories.size(); ++idx) {
            temp_id_mapping.push_back(old_memories[idx]["id"].get<size_t>());
            old_memories[idx]["id"] = idx;
        }

        std::string function_calling_prompt = get_update_memory_messages(old_memories, new_facts, update_memory_prompt);

        std::string new_memories_with_actions_str;
        json new_memories_with_actions = json::array();

        try {
            new_memories_with_actions_str = llm->ask(
                {Message::user_message(function_calling_prompt)}
            );
            new_memories_with_actions_str = remove_code_blocks(new_memories_with_actions_str);
        } catch (const std::exception& e) {
            logger->error("Error in new_memories_with_actions: " + std::string(e.what()));
        }
        
        try {
            new_memories_with_actions = json::parse(new_memories_with_actions_str);
        } catch (const std::exception& e) {
            logger->error("Invalid JSON response: " + std::string(e.what()));
        }

        try {
            for (const auto& resp : new_memories_with_actions["memory"]) {
                logger->debug("Processing memory: " + resp.dump(2));
                try {
                    if (!resp.contains("text")) {
                        logger->warn("Skipping memory entry because of empty `text` field.");
                        continue;
                    }
                    std::string event = resp.value("event", "NONE");
                    size_t memory_id;
                    try {
                        if (event != "ADD") {
                            memory_id = temp_id_mapping.at(resp["id"].get<size_t>());
                        } else {
                            memory_id = get_uuid_64();
                        }
                    } catch (...) {
                        memory_id = get_uuid_64();
                    }
                    if (event == "ADD") {
                        _create_memory(
                            memory_id,
                            resp["text"], // data
                            new_message_embeddings // existing_embeddings
                        );
                    } else if (event == "UPDATE") {
                        _update_memory(
                            memory_id,
                            resp["text"], // data
                            new_message_embeddings // existing_embeddings
                        );
                    } else if (event == "DELETE") {
                        _delete_memory(memory_id);
                    }
                } catch (const std::exception& e) {
                    logger->error("Error in new_memories_with_actions: " + std::string(e.what()));
                }
            }
        } catch (const std::exception& e) {
            logger->error("Error in new_memories_with_actions: " + std::string(e.what()));
        }
    }

    void _create_memory(const size_t& memory_id, const std::string& data, const std::map<std::string, std::vector<float>>& existing_embeddings) {
        if (!vector_store) {
            logger->warn("Vector store is not initialized, skipping create memory");
            return;
        }
        
        std::vector<float> embedding;
        if (existing_embeddings.find(data) != existing_embeddings.end()) {
            embedding = existing_embeddings.at(data);
        } else {
            embedding = embedding_model->embed(data, EmbeddingType::ADD);
        }

        MemoryItem metadata{
            memory_id,
            data,
            httplib::detail::MD5(data)
        };

        vector_store->insert(
            embedding,
            memory_id,
            metadata
        );
    }

    void _update_memory(const size_t& memory_id, const std::string& data, const std::map<std::string, std::vector<float>>& existing_embeddings) {
        if (!vector_store) {
            logger->warn("Vector store is not initialized, skipping update memory");
            return;
        }
        
        logger->info("Updating memory with " + data);

        MemoryItem existing_memory;

        try {
            existing_memory = vector_store->get(memory_id);
        } catch (const std::exception& e) {
            logger->error("Error fetching existing memory: " + std::string(e.what()));
            return;
        }

        std::vector<float> embedding;
        if (existing_embeddings.find(data) != existing_embeddings.end()) {
            embedding = existing_embeddings.at(data);
        } else {
            embedding = embedding_model->embed(data, EmbeddingType::ADD);
        }

        existing_memory.memory = data;
        existing_memory.hash = httplib::detail::MD5(data);
        existing_memory.updated_at = std::chrono::system_clock::now().time_since_epoch().count();

        vector_store->update(
            memory_id,
            embedding,
            existing_memory
        );
    }

    void _delete_memory(const size_t& memory_id) {
        if (!vector_store) {
            logger->warn("Vector store is not initialized, skipping delete memory");
            return;
        }
        
        logger->info("Deleting memory: " + std::to_string(memory_id));
        vector_store->delete_vector(memory_id);
    }
};

} // namespace humanus::mem0

#endif // HUMANUS_MEMORY_MEM0_H