Step 424: add context window optimizer

2026-02-16 16:58:41 -07:00
parent 01951bac03
commit 337f8b09b2
4 changed files with 328 additions and 0 deletions
--- a/editor/src/ContextWindowOptimizer.h
+++ b/editor/src/ContextWindowOptimizer.h
@@ -0,0 +1,115 @@
+#pragma once
+
+#include "ModelProfileRegistry.h"
+
+#include <algorithm>
+#include <string>
+#include <vector>
+
+#include <nlohmann/json.hpp>
+using json = nlohmann::json;
+
+struct ContextPack {
+    std::string instructions;
+    std::string localContext;
+    std::string fileContext;
+    std::string projectContext;
+};
+
+struct OptimizedContext {
+    std::string scope;  // local|file|project
+    std::string context;
+    int tokenEstimate = 0;
+    int budgetTokens = 0;
+    bool truncated = false;
+
+    json toJson() const {
+        return {
+            {"scope", scope},
+            {"context", context},
+            {"tokenEstimate", tokenEstimate},
+            {"budgetTokens", budgetTokens},
+            {"truncated", truncated}
+        };
+    }
+};
+
+class ContextWindowOptimizer {
+public:
+    static int estimateTokens(const std::string& text) {
+        if (text.empty()) return 0;
+        int chars = (int)text.size();
+        int estimate = chars / 4;
+        return estimate > 0 ? estimate : 1;
+    }
+
+    static int budgetForContextWindow(int contextWindow) {
+        if (contextWindow <= 0) return 512;
+        // Reserve ~40% for output/system/tool wrappers.
+        int budget = (int)(contextWindow * 0.6);
+        if (budget < 512) budget = 512;
+        return budget;
+    }
+
+    static OptimizedContext optimize(const ContextPack& pack,
+                                     const ModelProfile& profile) {
+        OptimizedContext out;
+        out.budgetTokens = budgetForContextWindow(profile.contextWindow);
+
+        std::string chosenScope = chooseScope(profile.contextWindow);
+        std::string base = pack.instructions;
+        if (!base.empty()) base += "\n\n";
+
+        if (chosenScope == "project") {
+            base += "Project Context:\n" + pack.projectContext + "\n\n";
+            base += "File Context:\n" + pack.fileContext + "\n\n";
+            base += "Local Context:\n" + pack.localContext;
+        } else if (chosenScope == "file") {
+            base += "File Context:\n" + pack.fileContext + "\n\n";
+            base += "Local Context:\n" + pack.localContext;
+        } else {
+            base += "Local Context:\n" + pack.localContext;
+        }
+
+        out.scope = chosenScope;
+        out.context = trimToBudget(base, out.budgetTokens, out.truncated);
+        out.tokenEstimate = estimateTokens(out.context);
+        return out;
+    }
+
+    static OptimizedContext optimizeForWorker(const ContextPack& pack,
+                                              const std::string& workerType,
+                                              const ModelProfileRegistry& registry) {
+        auto profile = registry.resolveForWorker(workerType);
+        if (profile.name.empty()) {
+            // Fall back to local-only conservative budget.
+            profile.name = "fallback";
+            profile.contextWindow = 8192;
+        }
+        return optimize(pack, profile);
+    }
+
+private:
+    static std::string chooseScope(int contextWindow) {
+        if (contextWindow >= 120000) return "project";
+        if (contextWindow >= 24000) return "file";
+        return "local";
+    }
+
+    static std::string trimToBudget(const std::string& text,
+                                    int budgetTokens,
+                                    bool& truncated) {
+        int est = estimateTokens(text);
+        if (est <= budgetTokens) {
+            truncated = false;
+            return text;
+        }
+        truncated = true;
+        int maxChars = budgetTokens * 4;
+        if (maxChars <= 0) return "";
+        if ((int)text.size() <= maxChars) return text;
+        std::string out = text.substr(0, (size_t)maxChars);
+        out += "\n\n[truncated to fit model context window]";
+        return out;
+    }
+};