Skip to content

Commit

Permalink
Fix handling of LLM model environment variables (#657)
Browse files Browse the repository at this point in the history
* Fix model environment variable handling

* Add GPT-3.5 turbo model
  • Loading branch information
drdavella authored Jun 20, 2024
1 parent 125b301 commit 318ef09
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 6 deletions.
4 changes: 2 additions & 2 deletions src/codemodder/llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
"MisconfiguredAIClient",
]

models = ["gpt-4-turbo-2024-04-09", "gpt-4o-2024-05-13"]
models = ["gpt-4-turbo-2024-04-09", "gpt-4o-2024-05-13", "gpt-35-turbo-0125"]
DEFAULT_AZURE_OPENAI_API_VERSION = "2024-02-01"


Expand All @@ -36,7 +36,7 @@ def __init__(self, models):
def __getattr__(self, name):
if name in self:
return os.getenv(
f"CODEMODDER_AZURE_OPENAI_{self[name].upper()}_DEPLOYMENT", self[name]
f"CODEMODDER_AZURE_OPENAI_{name.upper()}_DEPLOYMENT", self[name]
)
raise AttributeError(
f"'{self.__class__.__name__}' object has no attribute '{name}'"
Expand Down
9 changes: 5 additions & 4 deletions tests/test_llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,20 +2,21 @@

import pytest

from codemodder.llm import MODELS
from codemodder.llm import MODELS, models


class TestModels:
def test_get_model_name(self):
assert MODELS.gpt_4_turbo_2024_04_09 == "gpt-4-turbo-2024-04-09"

@pytest.mark.parametrize("model", ["gpt-4-turbo-2024-04-09", "gpt-4o-2024-05-13"])
@pytest.mark.parametrize("model", models)
def test_model_get_name_from_env(self, mocker, model):
name = "my-awesome-deployment"
attr_name = model.replace("-", "_")
mocker.patch.dict(
os.environ,
{
f"CODEMODDER_AZURE_OPENAI_{model.upper()}_DEPLOYMENT": name,
f"CODEMODDER_AZURE_OPENAI_{attr_name.upper()}_DEPLOYMENT": name,
},
)
assert getattr(MODELS, model.replace("-", "_")) == name
assert getattr(MODELS, attr_name) == name

0 comments on commit 318ef09

Please sign in to comment.