diff --git a/.env.example b/.env.example index ab3da74..a38be5e 100644 --- a/.env.example +++ b/.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 \ No newline at end of file +SUPABASE_KEY=your_supabase_key diff --git a/helpers/devcontainer_helpers.py b/helpers/devcontainer_helpers.py index ef4f34c..cea9da2 100644 --- a/helpers/devcontainer_helpers.py +++ b/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 \ No newline at end of file + raise diff --git a/helpers/token_helpers.py b/helpers/token_helpers.py index 27042b0..1285260 100644 --- a/helpers/token_helpers.py +++ b/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 \ No newline at end of file + return text diff --git a/main.py b/main.py index 2e6410d..28ed406 100644 --- a/main.py +++ b/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() \ No newline at end of file + serve() diff --git a/test_devcontainer_context.py b/test_devcontainer_context.py new file mode 100644 index 0000000..b595927 --- /dev/null +++ b/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() diff --git a/test_token_helpers.py b/test_token_helpers.py new file mode 100644 index 0000000..eafb5ab --- /dev/null +++ b/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()