From 62653d877f864498f31f2a258ad78525f5b890fe Mon Sep 17 00:00:00 2001 From: caydyan Date: Sun, 14 Jun 2026 22:15:39 +0800 Subject: [PATCH] Optimize context window content --- helpers/context_window_helpers.py | 144 ++++++++++++++++++++++ helpers/devcontainer_helpers.py | 51 ++------ tests/test_context_window_optimization.py | 72 +++++++++++ 3 files changed, 223 insertions(+), 44 deletions(-) create mode 100644 helpers/context_window_helpers.py create mode 100644 tests/test_context_window_optimization.py diff --git a/helpers/context_window_helpers.py b/helpers/context_window_helpers.py new file mode 100644 index 0000000..ba02159 --- /dev/null +++ b/helpers/context_window_helpers.py @@ -0,0 +1,144 @@ +from dataclasses import dataclass +import re + + +CONTEXT_SECTION_PATTERN = re.compile( + r"<.*?) >>\n(?P.*?)(?:\n)?<>", + re.DOTALL, +) + +HIGH_PRIORITY_FILES = { + "package.json", + "requirements.txt", + "pyproject.toml", + "Cargo.toml", + "Cargo.lock", + "go.mod", + "go.sum", + "pom.xml", + "build.gradle", + "Dockerfile", + "docker-compose.yml", + "compose.yml", + "Gemfile", + "Pipfile", + "Pipfile.lock", + "setup.py", + "Makefile", + "CMakeLists.txt", +} + +TRUNCATION_NOTICE = "\n\n<>" + + +@dataclass +class ContextSection: + title: str + text: str + order: int + priority: int + + +def _encoding_for_context(): + import tiktoken + + return tiktoken.encoding_for_model("gpt-4o-mini") + + +def count_context_tokens(text, encoding=None): + selected_encoding = encoding or _encoding_for_context() + return len(selected_encoding.encode(text)) + + +def _section_priority(title): + normalized_title = title.strip() + if normalized_title == "Repository Structure": + return 0 + if normalized_title == "Repository Languages": + return 1 + if normalized_title == "Existing devcontainer.json": + return 2 + if normalized_title.startswith("Content of "): + file_name = normalized_title.removeprefix("Content of ").strip() + if file_name in HIGH_PRIORITY_FILES: + return 3 + if file_name.lower() == "readme.md": + return 4 + return 5 + + +def _parse_context_sections(context): + sections = [] + for order, match in enumerate(CONTEXT_SECTION_PATTERN.finditer(context)): + title = match.group("title") + sections.append( + ContextSection( + title=title, + text=match.group(0), + order=order, + priority=_section_priority(title), + ) + ) + return sections + + +def _truncate_text_to_budget(text, max_tokens, encoding): + if max_tokens <= 0: + return "" + + tokens = encoding.encode(text) + if len(tokens) <= max_tokens: + return text + return encoding.decode(tokens[:max_tokens]) + + +def _truncate_section_to_budget(section, max_tokens, encoding): + notice_tokens = count_context_tokens(TRUNCATION_NOTICE, encoding) + if max_tokens <= notice_tokens: + return _truncate_text_to_budget(section.text, max_tokens, encoding) + + truncated = _truncate_text_to_budget(section.text, max_tokens - notice_tokens, encoding) + if not truncated: + return "" + return truncated.rstrip() + TRUNCATION_NOTICE + + +def optimize_context_window(context, max_tokens=120000, encoding=None): + selected_encoding = encoding or _encoding_for_context() + initial_tokens = count_context_tokens(context, selected_encoding) + + if initial_tokens <= max_tokens: + return context + + sections = _parse_context_sections(context) + if not sections: + return _truncate_text_to_budget(context, max_tokens, selected_encoding) + + selected_sections = [] + used_tokens = 0 + + for section in sorted(sections, key=lambda item: (item.priority, item.order)): + separator = "\n\n" if selected_sections else "" + separator_tokens = count_context_tokens(separator, selected_encoding) + section_tokens = count_context_tokens(section.text, selected_encoding) + available_tokens = max_tokens - used_tokens - separator_tokens + + if available_tokens <= 0: + break + + if section_tokens <= available_tokens: + selected_sections.append(section.text) + used_tokens += separator_tokens + section_tokens + continue + + truncated_section = _truncate_section_to_budget(section, available_tokens, selected_encoding) + if truncated_section: + selected_sections.append(truncated_section) + break + + optimized_context = "\n\n".join(selected_sections) + optimized_tokens = count_context_tokens(optimized_context, selected_encoding) + if optimized_tokens > max_tokens: + return _truncate_text_to_budget(optimized_context, max_tokens, selected_encoding) + + return optimized_context diff --git a/helpers/devcontainer_helpers.py b/helpers/devcontainer_helpers.py index ef4f34c..f83e6ba 100644 --- a/helpers/devcontainer_helpers.py +++ b/helpers/devcontainer_helpers.py @@ -4,61 +4,24 @@ import logging import os import jsonschema -import tiktoken +from helpers.context_window_helpers import count_context_tokens, optimize_context_window from helpers.jinja_helper import process_template from schemas import DevContainerModel from supabase_client import supabase from models import DevContainer -import logging -import tiktoken - def truncate_context(context, max_tokens=120000): logging.info(f"Starting truncate_context with max_tokens={max_tokens}") logging.debug(f"Initial context length: {len(context)} characters") - encoding = tiktoken.encoding_for_model("gpt-4o-mini") - tokens = encoding.encode(context) - - logging.info(f"Initial token count: {len(tokens)}") - - if len(tokens) <= max_tokens: - logging.info("Context is already within token limit. No truncation needed.") - return context - - logging.info(f"Context size is {len(tokens)} tokens. Truncation needed.") - - # Prioritize keeping the repository structure and languages - structure_end = context.find("<>") - languages_end = context.find("<>") - - logging.debug(f"Structure end position: {structure_end}") - logging.debug(f"Languages end position: {languages_end}") - - important_content = context[:languages_end] + "<>\n\n" - remaining_content = context[languages_end + len("<>\n\n"):] - - important_tokens = encoding.encode(important_content) - logging.debug(f"Important content token count: {len(important_tokens)}") - - if len(important_tokens) > max_tokens: - logging.warning("Important content alone exceeds max_tokens. Truncating important content.") - important_content = encoding.decode(important_tokens[:max_tokens]) - return important_content - - remaining_tokens = max_tokens - len(important_tokens) - logging.info(f"Tokens available for remaining content: {remaining_tokens}") - - truncated_remaining = encoding.decode(encoding.encode(remaining_content)[:remaining_tokens]) - - final_context = important_content + truncated_remaining - final_tokens = encoding.encode(final_context) + optimized_context = optimize_context_window(context, max_tokens=max_tokens) + final_tokens = count_context_tokens(optimized_context) - logging.info(f"Final token count: {len(final_tokens)}") - logging.debug(f"Final context length: {len(final_context)} characters") + logging.info(f"Final token count: {final_tokens}") + logging.debug(f"Final context length: {len(optimized_context)} characters") - return final_context + return optimized_context def generate_devcontainer_json(instructor_client, repo_url, repo_context, devcontainer_url=None, max_retries=2, regenerate=False): existing_devcontainer = None @@ -136,4 +99,4 @@ def save_devcontainer(new_devcontainer): return result.data[0] if result.data else None except Exception as e: logging.error(f"Error saving devcontainer to Supabase: {str(e)}") - raise \ No newline at end of file + raise diff --git a/tests/test_context_window_optimization.py b/tests/test_context_window_optimization.py new file mode 100644 index 0000000..5fa1674 --- /dev/null +++ b/tests/test_context_window_optimization.py @@ -0,0 +1,72 @@ +import unittest + +from helpers.context_window_helpers import optimize_context_window + + +class FakeEncoding: + def encode(self, text): + return [ord(character) for character in text] + + def decode(self, tokens): + return "".join(chr(token) for token in tokens) + + +ENCODING = FakeEncoding() + + +def section(title, body): + return f"<>\n{body}\n<>" + + +def count_tokens(text): + return len(ENCODING.encode(text)) + + +class ContextWindowOptimizationTest(unittest.TestCase): + def test_keeps_high_priority_sections_before_large_lower_priority_content(self): + structure = section("Repository Structure", "package.json\nREADME.md\nsrc/app.py") + languages = section("Repository Languages", "TypeScript: 120 lines\nCSS: 10 lines") + package_json = section("Content of package.json", '{"scripts":{"dev":"vite"}}') + readme = section("Content of README.md", "README_START " + ("details " * 300) + "README_TAIL") + changelog = section("Content of CHANGELOG.md", "CHANGELOG_START " + ("noise " * 300) + "CHANGELOG_TAIL") + context = "\n\n".join([readme, changelog, structure, package_json, languages]) + budget = count_tokens("\n\n".join([structure, languages, package_json])) + 10 + + result = optimize_context_window(context, max_tokens=budget, encoding=ENCODING) + + self.assertLessEqual(count_tokens(result), budget) + self.assertIn("Repository Structure", result) + self.assertIn("Repository Languages", result) + self.assertIn('"scripts":{"dev":"vite"}', result) + self.assertNotIn("README_TAIL", result) + self.assertNotIn("CHANGELOG_TAIL", result) + + def test_existing_devcontainer_is_preserved_before_readme_content(self): + structure = section("Repository Structure", ".devcontainer/devcontainer.json\nREADME.md") + languages = section("Repository Languages", "Python: 20 lines") + existing_devcontainer = section( + "Existing devcontainer.json", + '{"image":"mcr.microsoft.com/devcontainers/python:3.12"}', + ) + readme = section("Content of README.md", "README_START " + ("details " * 300) + "README_TAIL") + context = "\n\n".join([readme, structure, existing_devcontainer, languages]) + budget = count_tokens("\n\n".join([structure, languages, existing_devcontainer])) + 10 + + result = optimize_context_window(context, max_tokens=budget, encoding=ENCODING) + + self.assertLessEqual(count_tokens(result), budget) + self.assertIn("mcr.microsoft.com/devcontainers/python:3.12", result) + self.assertNotIn("README_TAIL", result) + + def test_plain_context_falls_back_to_token_prefix(self): + context = "alpha " * 300 + budget = 50 + + result = optimize_context_window(context, max_tokens=budget, encoding=ENCODING) + + self.assertLessEqual(count_tokens(result), budget) + self.assertTrue(result.startswith("alpha")) + + +if __name__ == "__main__": + unittest.main()