fix: separate LLM and embedding token limits

.env.example

@@ -2,8 +2,9 @@ AZURE_OPENAI_API_KEY=YOURKEY
 AZURE_OPENAI_ENDPOINT=https://YOURENDPOINT.openai.azure.com
 AZURE_OPENAI_API_VERSION=2024-02-01
 MODEL=gpt-4o-mini
+LLM_MODEL_MAX_TOKENS=126000
 EMBEDDING=text-embedding-3-small
 EMBEDDING_MODEL_MAX_TOKENS=8192
 GITHUB_TOKEN=YOURTOKEN
 SUPABASE_URL=https://YOURLINK.supabase.co
-SUPABASE_KEY=your_supabase_key
+SUPABASE_KEY=your_supabase_key

helpers/devcontainer_helpers.py

@@ -4,21 +4,23 @@
 import logging
 import os
 import jsonschema
-import tiktoken
 from helpers.jinja_helper import process_template
+from helpers.token_helpers import (
+    DEFAULT_LLM_MAX_TOKENS,
+    encoding_for_model,
+    max_tokens_from_env,
+)
 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}")
+def truncate_context(context, model_name=None, max_tokens=DEFAULT_LLM_MAX_TOKENS):
+    model_name = model_name or os.getenv("MODEL", "gpt-4o-mini")
+    logging.info(f"Starting truncate_context with model={model_name}, max_tokens={max_tokens}")
     logging.debug(f"Initial context length: {len(context)} characters")
 
-    encoding = tiktoken.encoding_for_model("gpt-4o-mini")
+    encoding = encoding_for_model(model_name)
     tokens = encoding.encode(context)
 
     logging.info(f"Initial token count: {len(tokens)}")
@@ -76,7 +78,9 @@ def generate_devcontainer_json(instructor_client, repo_url, repo_context, devcon
     logging.info("Generating devcontainer.json...")
 
     # Truncate the context to fit within token limits
-    truncated_context = truncate_context(repo_context, max_tokens=126000)
+    model_name = os.getenv("MODEL")
+    max_context_tokens = max_tokens_from_env("LLM_MODEL_MAX_TOKENS", DEFAULT_LLM_MAX_TOKENS)
+    truncated_context = truncate_context(repo_context, model_name=model_name, max_tokens=max_context_tokens)
 
     template_data = {
         "repo_url": repo_url,
@@ -90,7 +94,7 @@ def generate_devcontainer_json(instructor_client, repo_url, repo_context, devcon
         try:
             logging.debug(f"Attempt {attempt + 1} to generate devcontainer.json")
             response = instructor_client.chat.completions.create(
-                model=os.getenv("MODEL"),
+                model=model_name,
                 response_model=DevContainerModel,
                 messages=[
                     {"role": "system", "content": "You are a helpful assistant that generates devcontainer.json files."},
@@ -136,4 +140,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
+        raise

helpers/token_helpers.py

@@ -1,14 +1,46 @@
+import os
 import tiktoken
 
+DEFAULT_ENCODING = "cl100k_base"
+DEFAULT_LLM_MAX_TOKENS = 126000
+DEFAULT_EMBEDDING_MAX_TOKENS = 8192
+
+
+def encoding_for_model(model_name):
+    try:
+        return tiktoken.encoding_for_model(model_name)
+    except KeyError:
+        return tiktoken.get_encoding(DEFAULT_ENCODING)
+
+
 def count_tokens(text):
-    encoder = tiktoken.encoding_for_model("gpt-4o")
+    encoder = encoding_for_model("gpt-4o")
     tokens = encoder.encode(text)
     return len(tokens)
 
+
+def max_tokens_from_env(env_name, default):
+    try:
+        default = int(default)
+    except (TypeError, ValueError):
+        default = DEFAULT_LLM_MAX_TOKENS
+
+    raw_value = os.getenv(env_name)
+    if not raw_value:
+        return default
+
+    try:
+        max_tokens = int(raw_value)
+    except ValueError:
+        return default
+
+    return max_tokens if max_tokens > 0 else default
+
+
 def truncate_to_token_limit(text, model_name, max_tokens):
-    encoding = tiktoken.encoding_for_model(model_name)
+    encoding = encoding_for_model(model_name)
     tokens = encoding.encode(text)
     if len(tokens) > max_tokens:
         truncated_tokens = tokens[:max_tokens]
         return encoding.decode(truncated_tokens)
-    return text
+    return text

main.py

@@ -9,7 +9,12 @@
 from helpers.openai_helpers import setup_azure_openai, setup_instructor
 from helpers.github_helpers import fetch_repo_context, check_url_exists
 from helpers.devcontainer_helpers import generate_devcontainer_json, validate_devcontainer_json
-from helpers.token_helpers import count_tokens, truncate_to_token_limit
+from helpers.token_helpers import (
+    DEFAULT_EMBEDDING_MAX_TOKENS,
+    count_tokens,
+    max_tokens_from_env,
+    truncate_to_token_limit,
+)
 from models import DevContainer
 from schemas import DevContainerModel
 from content import *
@@ -133,7 +138,7 @@ async def post(repo_url: str, regenerate: bool = False):
             try:
                 if hasattr(openai_client.embeddings, "create"):
                     embedding_model = os.getenv("EMBEDDING", "text-embedding-ada-002")
-                    max_tokens = int(os.getenv("EMBEDDING_MODEL_MAX_TOKENS", 8192))
+                    max_tokens = max_tokens_from_env("EMBEDDING_MODEL_MAX_TOKENS", DEFAULT_EMBEDDING_MAX_TOKENS)
 
                     truncated_context = truncate_to_token_limit(repo_context, embedding_model, max_tokens)
 
@@ -207,4 +212,4 @@ async def get(fname:str, ext:str):
 
 if __name__ == "__main__":
     logging.info("Starting FastHTML app...")
-    serve()
+    serve()

test_devcontainer_context.py

@@ -0,0 +1,65 @@
+import importlib
+import os
+import sys
+import types
+import unittest
+from unittest.mock import patch
+
+
+def import_devcontainer_helpers():
+    fake_supabase_client = types.ModuleType("supabase_client")
+    fake_supabase_client.supabase = object()
+    sys.modules["supabase_client"] = fake_supabase_client
+    sys.modules.pop("helpers.devcontainer_helpers", None)
+    return importlib.import_module("helpers.devcontainer_helpers")
+
+
+class DevcontainerContextTest(unittest.TestCase):
+    def tearDown(self):
+        os.environ.pop("MODEL", None)
+        os.environ.pop("LLM_MODEL_MAX_TOKENS", None)
+
+    def test_generation_uses_llm_model_token_limit_from_env(self):
+        helpers = import_devcontainer_helpers()
+        os.environ["MODEL"] = "custom-chat-model"
+        os.environ["LLM_MODEL_MAX_TOKENS"] = "7"
+
+        captured = {}
+
+        def fake_truncate_context(repo_context, model_name=None, max_tokens=None):
+            captured["repo_context"] = repo_context
+            captured["model_name"] = model_name
+            captured["max_tokens"] = max_tokens
+            return "truncated context"
+
+        class Completion:
+            def create(self, **kwargs):
+                captured["model"] = kwargs["model"]
+                captured["prompt"] = kwargs["messages"][1]["content"]
+                return types.SimpleNamespace(
+                    dict=lambda exclude_none=True: {"name": "example", "image": "python:3.12"}
+                )
+
+        instructor_client = types.SimpleNamespace(
+            chat=types.SimpleNamespace(completions=Completion())
+        )
+
+        with patch.object(helpers, "truncate_context", side_effect=fake_truncate_context), \
+                patch.object(helpers, "validate_devcontainer_json", return_value=True):
+            devcontainer_json, devcontainer_url = helpers.generate_devcontainer_json(
+                instructor_client,
+                "https://github.com/example/repo",
+                "full repository context",
+            )
+
+        self.assertIn('"name": "example"', devcontainer_json)
+        self.assertIsNone(devcontainer_url)
+        self.assertEqual(captured["repo_context"], "full repository context")
+        self.assertEqual(captured["model_name"], "custom-chat-model")
+        self.assertEqual(captured["max_tokens"], 7)
+        self.assertEqual(captured["model"], "custom-chat-model")
+        self.assertIn("truncated context", captured["prompt"])
+
+
+if __name__ == "__main__":
+    unittest.main()

test_token_helpers.py

@@ -0,0 +1,43 @@
+import os
+import unittest
+
+from helpers.token_helpers import (
+    DEFAULT_EMBEDDING_MAX_TOKENS,
+    encoding_for_model,
+    max_tokens_from_env,
+    truncate_to_token_limit,
+)
+
+
+class TokenHelpersTest(unittest.TestCase):
+    def tearDown(self):
+        os.environ.pop("TEST_MAX_TOKENS", None)
+
+    def test_unknown_model_uses_fallback_encoding(self):
+        encoding = encoding_for_model("unknown-provider/new-chat-model")
+
+        self.assertGreater(len(encoding.encode("hello world")), 0)
+
+    def test_max_tokens_from_env_uses_positive_integer(self):
+        os.environ["TEST_MAX_TOKENS"] = "4096"
+
+        self.assertEqual(max_tokens_from_env("TEST_MAX_TOKENS", DEFAULT_EMBEDDING_MAX_TOKENS), 4096)
+
+    def test_max_tokens_from_env_uses_default_for_invalid_values(self):
+        os.environ["TEST_MAX_TOKENS"] = "not-a-number"
+
+        self.assertEqual(
+            max_tokens_from_env("TEST_MAX_TOKENS", DEFAULT_EMBEDDING_MAX_TOKENS),
+            DEFAULT_EMBEDDING_MAX_TOKENS,
+        )
+
+    def test_truncate_to_token_limit_supports_unknown_model_names(self):
+        text = "one two three four"
+
+        truncated = truncate_to_token_limit(text, "custom-embedding-model", 2)
+
+        self.assertLessEqual(len(encoding_for_model("custom-embedding-model").encode(truncated)), 2)
+
+
+if __name__ == "__main__":
+    unittest.main()