diff --git a/src/codemodder/codemodder.py b/src/codemodder/codemodder.py index 9ae4d37e..d646e863 100644 --- a/src/codemodder/codemodder.py +++ b/src/codemodder/codemodder.py @@ -14,7 +14,7 @@ from codemodder.codetf import CodeTF from codemodder.context import CodemodExecutionContext from codemodder.dependency import Dependency -from codemodder.llm import MisconfiguredAIClient, TokenUsage, log_token_usage +from codemodder.llm import TokenUsage, log_token_usage from codemodder.logging import configure_logger, log_list, log_section, logger from codemodder.project_analysis.file_parsers.package_store import PackageStore from codemodder.project_analysis.python_repo_manager import PythonRepoManager @@ -134,7 +134,6 @@ def run( original_cli_args: list[str] | None = None, codemod_registry: registry.CodemodRegistry | None = None, sast_only: bool = False, - ai_client: bool = True, log_matched_files: bool = False, remediation: bool = False, ) -> tuple[CodeTF | None, int, TokenUsage]: @@ -162,24 +161,18 @@ def run( repo_manager = PythonRepoManager(Path(directory)) - try: - context = CodemodExecutionContext( - Path(directory), - dry_run, - verbose, - codemod_registry, - provider_registry, - repo_manager, - path_include, - path_exclude, - tool_result_files_map, - max_workers, - ai_client, - ) - except MisconfiguredAIClient as e: - logger.error(e) - # Codemodder instructions conflicted (according to spec) - return None, 3, token_usage + context = CodemodExecutionContext( + Path(directory), + dry_run, + verbose, + codemod_registry, + provider_registry, + repo_manager, + path_include, + path_exclude, + tool_result_files_map, + max_workers, + ) context.repo_manager.parse_project() diff --git a/src/codemodder/context.py b/src/codemodder/context.py index 4fe48541..5dda1db0 100644 --- a/src/codemodder/context.py +++ b/src/codemodder/context.py @@ -17,7 +17,6 @@ build_failed_dependency_notification, ) from codemodder.file_context import FileContext -from codemodder.llm import setup_azure_llama_llm_client, setup_openai_llm_client from codemodder.logging import log_list, logger from codemodder.project_analysis.file_parsers.package_store import PackageStore from codemodder.project_analysis.python_repo_manager import PythonRepoManager @@ -28,9 +27,6 @@ from codemodder.utils.update_finding_metadata import update_finding_metadata if TYPE_CHECKING: - from azure.ai.inference import ChatCompletionsClient - from openai import OpenAI - from codemodder.codemods.base_codemod import BaseCodemod @@ -51,8 +47,6 @@ class CodemodExecutionContext: max_workers: int = 1 tool_result_files_map: dict[str, list[Path]] semgrep_prefilter_results: ResultSet | None = None - openai_llm_client: OpenAI | None = None - azure_llama_llm_client: ChatCompletionsClient | None = None def __init__( self, @@ -66,7 +60,6 @@ def __init__( path_exclude: list[str] | None = None, tool_result_files_map: dict[str, list[Path]] | None = None, max_workers: int = 1, - ai_client: bool = True, ): self.directory = directory self.dry_run = dry_run @@ -85,10 +78,6 @@ def __init__( self.max_workers = max_workers self.tool_result_files_map = tool_result_files_map or {} self.semgrep_prefilter_results = None - self.openai_llm_client = setup_openai_llm_client() if ai_client else None - self.azure_llama_llm_client = ( - setup_azure_llama_llm_client() if ai_client else None - ) def add_changesets(self, codemod_name: str, change_sets: List[ChangeSet]): self._changesets_by_codemod.setdefault(codemod_name, []).extend(change_sets) diff --git a/src/codemodder/llm.py b/src/codemodder/llm.py index f9d66a28..2ec60590 100644 --- a/src/codemodder/llm.py +++ b/src/codemodder/llm.py @@ -2,35 +2,13 @@ import os from dataclasses import dataclass -from typing import TYPE_CHECKING from typing_extensions import Self -try: - from openai import AzureOpenAI, OpenAI -except ImportError: - OpenAI = None - AzureOpenAI = None - -try: - from azure.ai.inference import ChatCompletionsClient - from azure.core.credentials import AzureKeyCredential -except ImportError: - ChatCompletionsClient = None - AzureKeyCredential = None - -if TYPE_CHECKING: - from openai import OpenAI - from azure.ai.inference import ChatCompletionsClient - from azure.core.credentials import AzureKeyCredential - from codemodder.logging import logger __all__ = [ "MODELS", - "setup_openai_llm_client", - "setup_azure_llama_llm_client", - "MisconfiguredAIClient", "TokenUsage", "log_token_usage", ] @@ -42,7 +20,6 @@ "o1-mini", "o1", ] -DEFAULT_AZURE_OPENAI_API_VERSION = "2024-02-01" class ModelRegistry(dict): @@ -66,68 +43,6 @@ def __getattr__(self, name): MODELS = ModelRegistry(models) -def setup_openai_llm_client() -> OpenAI | None: - """Configure either the Azure OpenAI LLM client or the OpenAI client, in that order.""" - if not AzureOpenAI: - logger.info("Azure OpenAI API client not available") - return None - - azure_openapi_key = os.getenv("CODEMODDER_AZURE_OPENAI_API_KEY") - azure_openapi_endpoint = os.getenv("CODEMODDER_AZURE_OPENAI_ENDPOINT") - if bool(azure_openapi_key) ^ bool(azure_openapi_endpoint): - raise MisconfiguredAIClient( - "Azure OpenAI API key and endpoint must both be set or unset" - ) - - if azure_openapi_key and azure_openapi_endpoint: - logger.info("Using Azure OpenAI API client") - return AzureOpenAI( - api_key=azure_openapi_key, - api_version=os.getenv( - "CODEMODDER_AZURE_OPENAI_API_VERSION", - DEFAULT_AZURE_OPENAI_API_VERSION, - ), - azure_endpoint=azure_openapi_endpoint, - ) - - if not OpenAI: - logger.info("OpenAI API client not available") - return None - - if not (api_key := os.getenv("CODEMODDER_OPENAI_API_KEY")): - logger.info("OpenAI API key not found") - return None - - logger.info("Using OpenAI API client") - return OpenAI(api_key=api_key) - - -def setup_azure_llama_llm_client() -> ChatCompletionsClient | None: - """Configure the Azure Llama LLM client.""" - if not ChatCompletionsClient: - logger.info("Azure Llama client not available") - return None - - azure_llama_key = os.getenv("CODEMODDER_AZURE_LLAMA_API_KEY") - azure_llama_endpoint = os.getenv("CODEMODDER_AZURE_LLAMA_ENDPOINT") - if bool(azure_llama_key) ^ bool(azure_llama_endpoint): - raise MisconfiguredAIClient( - "Azure Llama API key and endpoint must both be set or unset" - ) - - if azure_llama_key and azure_llama_endpoint: - logger.info("Using Azure Llama API client") - return ChatCompletionsClient( - credential=AzureKeyCredential(azure_llama_key), - endpoint=azure_llama_endpoint, - ) - return None - - -class MisconfiguredAIClient(ValueError): - pass - - @dataclass class TokenUsage: completion_tokens: int = 0 diff --git a/tests/test_context.py b/tests/test_context.py index 8c7ef197..0c4fc85c 100644 --- a/tests/test_context.py +++ b/tests/test_context.py @@ -1,8 +1,4 @@ -import os - import pytest -from azure.ai.inference import ChatCompletionsClient -from openai import AzureOpenAI, OpenAI from codemodder.codetf import ( DetectionTool, @@ -22,7 +18,6 @@ ) from codemodder.context import CodemodExecutionContext as Context from codemodder.dependency import Security -from codemodder.llm import DEFAULT_AZURE_OPENAI_API_VERSION, MisconfiguredAIClient from codemodder.project_analysis.python_repo_manager import PythonRepoManager from codemodder.registry import load_registered_codemods @@ -107,214 +102,6 @@ def test_failed_dependency_description(self, mocker): in description ) - def test_setup_llm_clients_no_env_vars(self, mocker): - mocker.patch.dict(os.environ, clear=True) - context = Context( - mocker.Mock(), - True, - False, - load_registered_codemods(), - None, - PythonRepoManager(mocker.Mock()), - [], - [], - ) - assert context.openai_llm_client is None - assert context.azure_llama_llm_client is None - - def test_setup_openai_llm_client(self, mocker): - mocker.patch.dict(os.environ, {"CODEMODDER_OPENAI_API_KEY": "test"}) - context = Context( - mocker.Mock(), - True, - False, - load_registered_codemods(), - None, - PythonRepoManager(mocker.Mock()), - [], - [], - ) - assert isinstance(context.openai_llm_client, OpenAI) - - def test_setup_both_llm_clients(self, mocker): - mocker.patch.dict( - os.environ, - { - "CODEMODDER_OPENAI_API_KEY": "test", - "CODEMODDER_AZURE_LLAMA_API_KEY": "test", - "CODEMODDER_AZURE_LLAMA_ENDPOINT": "test", - }, - ) - context = Context( - mocker.Mock(), - True, - False, - load_registered_codemods(), - None, - PythonRepoManager(mocker.Mock()), - [], - [], - ) - assert isinstance(context.openai_llm_client, OpenAI) - assert isinstance(context.azure_llama_llm_client, ChatCompletionsClient) - - def test_setup_azure_llm_client(self, mocker): - mocker.patch.dict( - os.environ, - { - "CODEMODDER_AZURE_OPENAI_API_KEY": "test", - "CODEMODDER_AZURE_OPENAI_ENDPOINT": "test", - }, - ) - context = Context( - mocker.Mock(), - True, - False, - load_registered_codemods(), - None, - PythonRepoManager(mocker.Mock()), - [], - [], - ) - assert isinstance(context.openai_llm_client, AzureOpenAI) - assert ( - context.openai_llm_client._api_version == DEFAULT_AZURE_OPENAI_API_VERSION - ) - - @pytest.mark.parametrize( - "env_var", - ["CODEMODDER_AZURE_OPENAI_API_KEY", "CODEMODDER_AZURE_OPENAI_ENDPOINT"], - ) - def test_setup_azure_llm_client_missing_one(self, mocker, env_var): - mocker.patch.dict(os.environ, {env_var: "test"}) - with pytest.raises(MisconfiguredAIClient): - Context( - mocker.Mock(), - True, - False, - load_registered_codemods(), - None, - PythonRepoManager(mocker.Mock()), - [], - [], - ) - - def test_setup_azure_llama_llm_client(self, mocker): - mocker.patch.dict( - os.environ, - { - "CODEMODDER_AZURE_LLAMA_API_KEY": "test", - "CODEMODDER_AZURE_LLAMA_ENDPOINT": "test", - }, - ) - context = Context( - mocker.Mock(), - True, - False, - load_registered_codemods(), - None, - PythonRepoManager(mocker.Mock()), - [], - [], - ) - assert isinstance(context.azure_llama_llm_client, ChatCompletionsClient) - - @pytest.mark.parametrize( - "env_var", - ["CODEMODDER_AZURE_LLAMA_API_KEY", "CODEMODDER_AZURE_LLAMA_ENDPOINT"], - ) - def test_setup_azure_llama_llm_client_missing_one(self, mocker, env_var): - mocker.patch.dict(os.environ, {env_var: "test"}) - with pytest.raises(MisconfiguredAIClient): - Context( - mocker.Mock(), - True, - False, - load_registered_codemods(), - None, - PythonRepoManager(mocker.Mock()), - [], - [], - ) - - def test_get_api_version_from_env(self, mocker): - version = "fake-version" - mocker.patch.dict( - os.environ, - { - "CODEMODDER_AZURE_OPENAI_API_KEY": "test", - "CODEMODDER_AZURE_OPENAI_ENDPOINT": "test", - "CODEMODDER_AZURE_OPENAI_API_VERSION": version, - }, - ) - context = Context( - mocker.Mock(), - True, - False, - load_registered_codemods(), - None, - PythonRepoManager(mocker.Mock()), - [], - [], - ) - assert isinstance(context.openai_llm_client, AzureOpenAI) - assert context.openai_llm_client._api_version == version - - def test_disable_ai_client_openai(self, mocker): - mocker.patch.dict(os.environ, {"CODEMODDER_OPENAI_API_KEY": "test"}) - context = Context( - mocker.Mock(), - True, - False, - load_registered_codemods(), - None, - PythonRepoManager(mocker.Mock()), - [], - [], - ai_client=False, - ) - assert context.openai_llm_client is None - - def test_disable_ai_client_azure(self, mocker): - mocker.patch.dict( - os.environ, - { - "CODEMODDER_AZURE_OPENAI_API_KEY": "test", - "CODEMODDER_AZURE_OPENAI_ENDPOINT": "test", - }, - ) - context = Context( - mocker.Mock(), - True, - False, - load_registered_codemods(), - None, - PythonRepoManager(mocker.Mock()), - [], - [], - ai_client=False, - ) - assert context.openai_llm_client is None - - @pytest.mark.parametrize( - "env_var", - ["CODEMODDER_AZURE_OPENAI_API_KEY", "CODEMODDER_AZURE_OPENAI_ENDPOINT"], - ) - def test_no_misconfiguration_ai_client_disabled(self, mocker, env_var): - mocker.patch.dict(os.environ, {env_var: "test"}) - context = Context( - mocker.Mock(), - True, - False, - load_registered_codemods(), - None, - PythonRepoManager(mocker.Mock()), - [], - [], - ai_client=False, - ) - assert context.openai_llm_client is None - def test_compile_results(self, mocker): rule = rule = Rule( id="roslyn.sonaranalyzer.security.cs:S5131",