compaction.cpp

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
#include <entropic/core/compaction.h>
#include <entropic/core/engine_types.h>
#include <entropic/types/logging.h>
 
#include <sstream>
 
static auto logger = entropic::log::get("core.compaction");
 
namespace entropic {
 
// ── TokenCounter ─────────────────────────────────────────
 
TokenCounter::TokenCounter(int max_tokens)
    : max_tokens(max_tokens) {}
 
int TokenCounter::count_text(const std::string& text) {
    if (text.empty()) {
        return 0;
    }
    return static_cast<int>(text.size()) / 4 + 1;
}
 
int TokenCounter::count_message(const Message& msg) const {
    const auto* key = static_cast<const void*>(&msg);
    auto it = cache_.find(key);
    if (it != cache_.end()) {
        return it->second;
    }
    int count = count_text(msg.content) + 4; // +4 for role tokens
    cache_[key] = count;
    return count;
}
 
int TokenCounter::count_messages(
    const std::vector<Message>& messages) const {
    int total = 0;
    for (const auto& msg : messages) {
        total += count_message(msg);
    }
    return total;
}
 
float TokenCounter::usage_percent(
    const std::vector<Message>& messages) const {
    if (max_tokens == 0) {
        return 0.0f;
    }
    return static_cast<float>(count_messages(messages))
         / static_cast<float>(max_tokens);
}
 
void TokenCounter::clear_cache() {
    cache_.clear();
}
 
// ── CompactionManager ────────────────────────────────────
 
CompactionManager::CompactionManager(
    const CompactionConfig& config,
    TokenCounter& counter)
    : config(config), counter(counter) {}
 
CompactionResult CompactionManager::check_and_compact(
    std::vector<Message>& messages,
    bool force,
    const std::string& conversation_id) {
    int current = counter.count_messages(messages);
    int threshold = static_cast<int>(
        static_cast<float>(counter.max_tokens) * config.threshold_percent);
 
    if (!force && current < threshold) {
        return {false, current, current};
    }
 
    if (!config.enabled) {
        logger->warn("Context at {}/{} tokens, compaction disabled",
                     current, counter.max_tokens);
        return {false, current, current};
    }
 
    logger->info("Compacting conversation ({} tokens)", current);
 
    // Save full history before compacting (v1.8.8)
    if (config.save_full_history && !conversation_id.empty()) {
        save_snapshot(conversation_id, messages);
    }
 
    std::string summary;
    int stripped = 0;
    auto compacted = compact(messages, summary, stripped);
 
    counter.clear_cache();
    int new_count = counter.count_messages(compacted);
 
    if (new_count >= current) {
        logger->error("Compaction did not reduce tokens: {} -> {}",
                      current, new_count);
    }
 
    logger->info("Compacted {} -> {} tokens", current, new_count);
    messages = std::move(compacted);
 
    CompactionResult result;
    result.compacted = true;
    result.old_token_count = current;
    result.new_token_count = new_count;
    result.summary = summary;
    result.preserved_messages = static_cast<int>(messages.size()) - 1;
    result.messages_summarized = stripped;
    return result;
}
 
static void partition_messages(const std::vector<Message>& messages,
                               size_t start,
                               std::vector<const Message*>& user_msgs,
                               std::vector<const Message*>& assistant_msgs,
                               int& stripped_count) {
    stripped_count = 0;
    for (size_t i = start; i < messages.size(); ++i) {
        const auto& msg = messages[i];
        auto src = msg.metadata.find("source");
        if (src != msg.metadata.end() && src->second == "user") {
            user_msgs.push_back(&msg);
        } else if (msg.role == "assistant") {
            assistant_msgs.push_back(&msg);
        } else {
            ++stripped_count;
        }
    }
}
 
static std::vector<Message> assemble_compacted(
    const Message* system_msg, Message summary_msg,
    const std::vector<const Message*>& user_msgs,
    const std::vector<const Message*>& assistant_msgs) {
    std::vector<Message> result;
    if (system_msg != nullptr) { result.push_back(*system_msg); }
    result.push_back(std::move(summary_msg));
    for (const auto* m : user_msgs) { result.push_back(*m); }
    if (!assistant_msgs.empty()) {
        result.push_back(*assistant_msgs.back());
    }
    return result;
}
 
std::vector<Message> CompactionManager::compact(
    const std::vector<Message>& messages,
    std::string& summary,
    int& stripped_count) {
    Message const* system_msg = nullptr;
    size_t start = 0;
    if (!messages.empty() && messages[0].role == "system") {
        system_msg = &messages[0];
        start = 1;
    }
 
    std::vector<const Message*> user_msgs;
    std::vector<const Message*> assistant_msgs;
    partition_messages(messages, start, user_msgs, assistant_msgs,
                       stripped_count);
 
    std::vector<Message> working(messages.begin() + static_cast<long>(start),
                                 messages.end());
    logger->info("Compact: {} user, {} assistant, {} stripped",
                 user_msgs.size(), assistant_msgs.size(), stripped_count);
    summary = structured_summary(working);
    Message summary_msg;
    summary_msg.role = "user";
    summary_msg.content = format_summary(
        summary, static_cast<int>(working.size()));
 
    return assemble_compacted(system_msg, std::move(summary_msg),
                              user_msgs, assistant_msgs);
}
 
std::string CompactionManager::structured_summary(
    const std::vector<Message>& messages) {
    std::string lines = "Original task: "
                      + extract_original_task(messages);
 
    auto tool_log = extract_tool_log(messages);
    if (!tool_log.empty()) {
        lines += "\n\nTool calls made (oldest first):";
        for (const auto& [name, brief] : tool_log) {
            lines += "\n- " + name + ": " + brief;
        }
    }
    return lines;
}
 
static bool is_tool_result(const Message& msg) {
    if (msg.content.rfind("Tool `", 0) == 0) { return true; }
    auto src = msg.metadata.find("source");
    return src != msg.metadata.end() && src->second == "tool";
}
 
static std::string find_tagged(
    const std::vector<Message>& messages,
    const std::string& source) {
    for (const auto& msg : messages) {
        auto src = msg.metadata.find("source");
        if (src != msg.metadata.end() && src->second == source) {
            return msg.content;
        }
    }
    return {};
}
 
static std::string find_first_user_task(
    const std::vector<Message>& messages) {
    for (const auto& msg : messages) {
        if (msg.role != "user" || msg.content.empty()) { continue; }
        if (is_tool_result(msg)) { continue; }
        return msg.content;
    }
    return {};
}
 
std::string CompactionManager::extract_original_task(
    const std::vector<Message>& messages) {
    std::string task = find_tagged(messages, "user");
    if (task.empty()) { task = find_first_user_task(messages); }
    if (task.empty()) { return "(no user message found)"; }
    if (task.size() > 500) { return task.substr(0, 500) + "..."; }
    return task;
}
 
std::vector<std::pair<std::string, std::string>>
CompactionManager::extract_tool_log(
    const std::vector<Message>& messages) {
    std::vector<std::pair<std::string, std::string>> log;
    for (const auto& msg : messages) {
        auto it = msg.metadata.find("tool_name");
        if (it == msg.metadata.end()) {
            continue;
        }
        const auto& name = it->second;
        if (msg.content.rfind("[Previous:", 0) == 0) {
            log.emplace_back(name, "(pruned)");
            continue;
        }
        auto nl = msg.content.find('\n');
        std::string brief = msg.content.substr(0, std::min(nl, size_t{100}));
        log.emplace_back(name, brief);
    }
    return log;
}
 
std::string CompactionManager::format_summary(
    const std::string& summary,
    int message_count) {
    return "[CONVERSATION SUMMARY]\n"
           "The following summarizes "
         + std::to_string(message_count)
         + " previous messages that have been compacted"
           " to save context space.\n\n"
         + summary
         + "\n\n[END SUMMARY - Recent conversation continues below]";
}
 
CompactionResult CompactionManager::compact_messages(
    const std::vector<Message>& messages) {
    int old_count = counter.count_messages(messages);
 
    std::string summary;
    int stripped = 0;
    auto compacted = compact(messages, summary, stripped);
 
    counter.clear_cache();
    int new_count = counter.count_messages(compacted);
 
    CompactionResult result;
    result.compacted = true;
    result.old_token_count = old_count;
    result.new_token_count = new_count;
    result.summary = summary;
    result.preserved_messages =
        static_cast<int>(compacted.size()) - 1;
    result.messages_summarized = stripped;
    result.messages = compacted;
    result.compactor_source = "default";
    return result;
}
 
void CompactionManager::set_storage(const StorageInterface* storage) {
    storage_ = storage;
}
 
static std::string json_escape(const std::string& input) {
    std::ostringstream oss;
    for (char c : input) {
        if (c == '"') oss << "\\\"";
        else if (c == '\\') oss << "\\\\";
        else if (c == '\n') oss << "\\n";
        else oss << c;
    }
    return oss.str();
}
 
static std::string serialize_messages_json(
        const std::vector<Message>& messages) {
    std::ostringstream oss;
    oss << '[';
    for (size_t i = 0; i < messages.size(); ++i) {
        if (i > 0) oss << ',';
        oss << "{\"role\":\"" << messages[i].role
            << "\",\"content\":\"" << json_escape(messages[i].content)
            << "\"}";
    }
    oss << ']';
    return oss.str();
}
 
void CompactionManager::save_snapshot(
    const std::string& conversation_id,
    const std::vector<Message>& messages) {
    if (!storage_ || !storage_->save_snapshot) {
        return;
    }
 
    auto json_str = serialize_messages_json(messages);
    storage_->save_snapshot(
        conversation_id.c_str(), json_str.c_str(),
        storage_->user_data);
    logger->info("Saved compaction snapshot for {} ({} messages)",
                 conversation_id, messages.size());
}
 
} // namespace entropic