Files
whetstone_DSL/editor/src/ContextWindowOptimizer.h
2026-02-16 16:58:41 -07:00

116 lines
3.6 KiB
C++

#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;
}
};