diff --git a/codegen-examples/examples/pr_review_bot/run.py b/codegen-examples/examples/pr_review_bot/run.py index c3cd3206e..81982b72f 100644 --- a/codegen-examples/examples/pr_review_bot/run.py +++ b/codegen-examples/examples/pr_review_bot/run.py @@ -1,9 +1,12 @@ import codegen from codegen import Codebase from codegen.sdk.enums import ProgrammingLanguage -from codegen.sdk.codebase.config import CodebaseConfig, GSFeatureFlags, Secrets +from codegen.sdk.codebase.config import CodebaseConfig import json +from codegen.sdk.secrets import Secrets +from codegen.shared.configs.models import CodebaseFeatureFlags + github_token = "Your github token" open_ai_key = "your open ai key" pr_number = 0 # Your PR number must be an integer @@ -81,7 +84,7 @@ def run(codebase: Codebase): programming_language=ProgrammingLanguage.PYTHON, config=CodebaseConfig( secrets=Secrets(openai_key=open_ai_key, github_api_key=github_token), - feature_flags=GSFeatureFlags( + feature_flags=CodebaseFeatureFlags( sync_enabled=True, ), ), diff --git a/codegen-examples/examples/sqlalchemy_soft_delete/run.py b/codegen-examples/examples/sqlalchemy_soft_delete/run.py index d971d8520..3634342fb 100644 --- a/codegen-examples/examples/sqlalchemy_soft_delete/run.py +++ b/codegen-examples/examples/sqlalchemy_soft_delete/run.py @@ -1,6 +1,7 @@ import codegen from codegen import Codebase from codegen.sdk.core.detached_symbols.function_call import FunctionCall +from codegen.shared.configs.models import CodebaseFeatureFlags from codegen.shared.enums.programming_language import ProgrammingLanguage import shutil import subprocess @@ -92,7 +93,7 @@ def process_soft_deletes(codebase): if __name__ == "__main__": from codegen.sdk.core.codebase import Codebase - from codegen.sdk.codebase.config import CodebaseConfig, GSFeatureFlags + from codegen.sdk.codebase.config import CodebaseConfig repo_path = Path("/tmp/core") repo_url = "https://github.com/hasgeek/funnel.git" @@ -100,7 +101,7 @@ def process_soft_deletes(codebase): try: clone_repo(repo_url, repo_path) subprocess.run(["git", "-C", str(repo_path), "checkout", "8454e15"], check=True) - codebase = Codebase(str(repo_path), programming_language=ProgrammingLanguage.PYTHON, config=CodebaseConfig(feature_flags=GSFeatureFlags(disable_graph=True))) + codebase = Codebase(str(repo_path), programming_language=ProgrammingLanguage.PYTHON, config=CodebaseConfig(feature_flags=CodebaseFeatureFlags(disable_graph=True))) process_soft_deletes(codebase) finally: shutil.rmtree(repo_path) diff --git a/src/codegen/extensions/lsp/protocol.py b/src/codegen/extensions/lsp/protocol.py index 78e9bd945..d6326edd4 100644 --- a/src/codegen/extensions/lsp/protocol.py +++ b/src/codegen/extensions/lsp/protocol.py @@ -8,8 +8,9 @@ from codegen.extensions.lsp.io import LSPIO from codegen.extensions.lsp.utils import get_path -from codegen.sdk.codebase.config import CodebaseConfig, GSFeatureFlags +from codegen.sdk.codebase.config import CodebaseConfig from codegen.sdk.core.codebase import Codebase +from codegen.shared.configs.models import CodebaseFeatureFlags if TYPE_CHECKING: from codegen.extensions.lsp.server import CodegenLanguageServer @@ -25,7 +26,7 @@ def _init_codebase(self, params: InitializeParams) -> None: root = get_path(params.root_uri) else: root = os.getcwd() - config = CodebaseConfig(feature_flags=GSFeatureFlags(full_range_index=True)) + config = CodebaseConfig(feature_flags=CodebaseFeatureFlags(full_range_index=True)) io = LSPIO(self.workspace) self._server.codebase = Codebase(repo_path=str(root), config=config, io=io) self._server.io = io @@ -38,7 +39,7 @@ def lsp_initialize(self, params: InitializeParams) -> InitializeResult: root = get_path(params.root_uri) else: root = os.getcwd() - config = CodebaseConfig(feature_flags=GSFeatureFlags(full_range_index=True)) + config = CodebaseConfig(feature_flags=CodebaseFeatureFlags(full_range_index=True)) ret = super().lsp_initialize(params) self._worker = threading.Thread(target=self._init_codebase, args=(params,)) diff --git a/src/codegen/runner/sandbox/runner.py b/src/codegen/runner/sandbox/runner.py index 94242379c..f3fc8ba76 100644 --- a/src/codegen/runner/sandbox/runner.py +++ b/src/codegen/runner/sandbox/runner.py @@ -7,7 +7,7 @@ from codegen.git.schemas.repo_config import RepoConfig from codegen.runner.models.apis import CreateBranchRequest, CreateBranchResponse, GetDiffRequest, GetDiffResponse from codegen.runner.sandbox.executor import SandboxExecutor -from codegen.sdk.codebase.config import CodebaseConfig, GSFeatureFlags, ProjectConfig, SessionOptions +from codegen.sdk.codebase.config import CodebaseConfig, ProjectConfig, SessionOptions from codegen.sdk.codebase.factory.codebase_factory import CodebaseType from codegen.sdk.core.codebase import Codebase from codegen.sdk.secrets import Secrets @@ -46,9 +46,8 @@ async def warmup(self) -> None: async def _build_graph(self) -> Codebase: logger.info("> Building graph...") projects = [ProjectConfig(programming_language=self.repo.language, repo_operator=self.op, base_path=self.repo.base_path, subdirectories=self.repo.subdirectories)] - gs_ffs = GSFeatureFlags(**config.feature_flags.model_dump()) secrets = Secrets(openai_key=config.secrets.openai_api_key) - codebase_config = CodebaseConfig(secrets=secrets, feature_flags=gs_ffs) + codebase_config = CodebaseConfig(secrets=secrets, feature_flags=config.feature_flags.codebase) return Codebase(projects=projects, config=codebase_config) @stopwatch diff --git a/src/codegen/sdk/codebase/config.py b/src/codegen/sdk/codebase/config.py index 10d14cc00..44f73e731 100644 --- a/src/codegen/sdk/codebase/config.py +++ b/src/codegen/sdk/codebase/config.py @@ -11,6 +11,7 @@ from codegen.git.utils.file_utils import split_git_path from codegen.git.utils.language import determine_project_language from codegen.sdk.secrets import Secrets +from codegen.shared.configs.models import CodebaseFeatureFlags from codegen.shared.enums.programming_language import ProgrammingLanguage HARD_MAX_AI_LIMIT = 500 # Global limit for AI requests @@ -25,36 +26,10 @@ class SessionOptions(BaseModel): max_ai_requests: int = Field(default=150, le=HARD_MAX_AI_LIMIT) -class GSFeatureFlags(BaseModel): - """Config for building the graph sitter graph. These are non-repo specific options that are set per-usecase. - - Attributes: - debug: Warning if there are errors during parsing (such as unimplemented nodes) - verify_graph: Verify the accuracy of the graph between resets. Will result in lag - method_usages: Resolve . usages - """ - - model_config = ConfigDict(frozen=True) - debug: bool = False - verify_graph: bool = False - track_graph: bool = False # Track the initial graph state - method_usages: bool = True - sync_enabled: bool = True - ts_dependency_manager: bool = False # Enable Typescript Dependency Manager - ts_language_engine: bool = False # Enable Typescript Language Engine - v8_ts_engine: bool = False # Enable V8 Based Typescript Language Engine Instead of NodeJS - full_range_index: bool = False - ignore_process_errors: bool = True # Ignore errors from dependency manager and language engine - import_resolution_overrides: dict[str, str] = {} # Override import resolution for specific modules - disable_graph: bool = False # Turn of graph generation entirely. Speeds up parsing but disables usages and dependencies - generics: bool = True # Enable parsing of generic types - - -DefaultFlags = GSFeatureFlags(sync_enabled=False) - -TestFlags = GSFeatureFlags(debug=True, track_graph=True, verify_graph=True, full_range_index=True) -LintFlags = GSFeatureFlags(method_usages=False) -ParseTestFlags = GSFeatureFlags(debug=False, track_graph=False) +DefaultFlags = CodebaseFeatureFlags(sync_enabled=False) +TestFlags = CodebaseFeatureFlags(debug=True, track_graph=True, verify_graph=True, full_range_index=True) +LintFlags = CodebaseFeatureFlags(method_usages=False) +ParseTestFlags = CodebaseFeatureFlags(debug=False, track_graph=False) class ProjectConfig(BaseModel): @@ -103,7 +78,7 @@ class CodebaseConfig(BaseModel): model_config = ConfigDict(frozen=True) secrets: Secrets = Secrets() - feature_flags: GSFeatureFlags = DefaultFlags + feature_flags: CodebaseFeatureFlags = DefaultFlags DefaultConfig = CodebaseConfig() diff --git a/src/codegen/sdk/codebase/factory/get_session.py b/src/codegen/sdk/codebase/factory/get_session.py index 7444fd4c4..f9de00d2d 100644 --- a/src/codegen/sdk/codebase/factory/get_session.py +++ b/src/codegen/sdk/codebase/factory/get_session.py @@ -6,11 +6,12 @@ from codegen.git.repo_operator.local_repo_operator import LocalRepoOperator from codegen.sdk.codebase.codebase_context import CodebaseContext -from codegen.sdk.codebase.config import CodebaseConfig, GSFeatureFlags, ProjectConfig, SessionOptions, TestFlags +from codegen.sdk.codebase.config import CodebaseConfig, ProjectConfig, SessionOptions, TestFlags from codegen.sdk.codebase.factory.codebase_factory import CodebaseFactory from codegen.sdk.core.codebase import Codebase, PyCodebaseType, TSCodebaseType from codegen.sdk.secrets import Secrets from codegen.sdk.tree_sitter_parser import print_errors +from codegen.shared.configs.models import CodebaseFeatureFlags from codegen.shared.enums.programming_language import ProgrammingLanguage @@ -23,7 +24,7 @@ def get_codebase_session( sync_graph: bool = True, verify_input: bool = True, verify_output: bool = True, - feature_flags: GSFeatureFlags = TestFlags, + feature_flags: CodebaseFeatureFlags = TestFlags, session_options: SessionOptions = SessionOptions(), secrets: Secrets = Secrets(), ) -> AbstractContextManager[PyCodebaseType]: ... @@ -38,7 +39,7 @@ def get_codebase_session( sync_graph: bool = True, verify_input: bool = True, verify_output: bool = True, - feature_flags: GSFeatureFlags = TestFlags, + feature_flags: CodebaseFeatureFlags = TestFlags, session_options: SessionOptions = SessionOptions(), secrets: Secrets = Secrets(), ) -> AbstractContextManager[PyCodebaseType]: ... @@ -53,7 +54,7 @@ def get_codebase_session( sync_graph: bool = True, verify_input: bool = True, verify_output: bool = True, - feature_flags: GSFeatureFlags = TestFlags, + feature_flags: CodebaseFeatureFlags = TestFlags, session_options: SessionOptions = SessionOptions(), secrets: Secrets = Secrets(), ) -> AbstractContextManager[TSCodebaseType]: ... @@ -68,7 +69,7 @@ def get_codebase_session( sync_graph: bool = True, verify_input: bool = True, verify_output: bool = True, - feature_flags: GSFeatureFlags = TestFlags, + feature_flags: CodebaseFeatureFlags = TestFlags, session_options: SessionOptions = SessionOptions(), secrets: Secrets = Secrets(), ) -> Generator[Codebase, None, None]: diff --git a/src/codegen/sdk/core/external/dependency_manager.py b/src/codegen/sdk/core/external/dependency_manager.py index 8132bd028..578ea46b8 100644 --- a/src/codegen/sdk/core/external/dependency_manager.py +++ b/src/codegen/sdk/core/external/dependency_manager.py @@ -30,7 +30,7 @@ def remove_dependencies(self): def get_dependency_manager(language: ProgrammingLanguage, codebase_context: "CodebaseContext", enabled: bool = False) -> DependencyManager | None: from codegen.sdk.typescript.external.dependency_manager import TypescriptDependencyManager - ts_enabled = enabled or codebase_context.config.feature_flags.ts_dependency_manager + ts_enabled = enabled or codebase_context.config.feature_flags.typescript.ts_dependency_manager if language == ProgrammingLanguage.TYPESCRIPT: if ts_enabled: return TypescriptDependencyManager(repo_path=codebase_context.repo_path, base_path=codebase_context.projects[0].base_path) diff --git a/src/codegen/sdk/core/external/language_engine.py b/src/codegen/sdk/core/external/language_engine.py index 5dad68381..c5457e39e 100644 --- a/src/codegen/sdk/core/external/language_engine.py +++ b/src/codegen/sdk/core/external/language_engine.py @@ -24,8 +24,8 @@ def get_return_type(self, node: "Editable") -> str | None: def get_language_engine(language: ProgrammingLanguage, codebase_context: "CodebaseContext", use_ts: bool = False, use_v8: bool = False) -> LanguageEngine | None: from codegen.sdk.typescript.external.ts_analyzer_engine import NodeTypescriptEngine, V8TypescriptEngine - use_ts = use_ts or codebase_context.config.feature_flags.ts_language_engine - use_v8 = use_v8 or codebase_context.config.feature_flags.v8_ts_engine + use_ts = use_ts or codebase_context.config.feature_flags.typescript.ts_language_engine + use_v8 = use_v8 or codebase_context.config.feature_flags.typescript.v8_ts_engine if language == ProgrammingLanguage.TYPESCRIPT: if use_ts and use_v8: # Enables with both ts_language_engine and v8_ts_engine feature flags are on diff --git a/src/codegen/shared/configs/models.py b/src/codegen/shared/configs/models.py index e379c7058..6e8a5ec7c 100644 --- a/src/codegen/shared/configs/models.py +++ b/src/codegen/shared/configs/models.py @@ -20,22 +20,24 @@ def _get_setting_config(group_name: str) -> SettingsConfigDict: class TypescriptConfig(BaseSettings): model_config = _get_setting_config("FEATURE_FLAGS_TYPESCRIPT") - ts_dependency_manager: bool | None = None - ts_language_engine: bool | None = None - v8_ts_engine: bool | None = None + + ts_dependency_manager: bool = False + ts_language_engine: bool = False + v8_ts_engine: bool = False class CodebaseFeatureFlags(BaseSettings): model_config = _get_setting_config("FEATURE_FLAGS") - debug: bool | None = None - verify_graph: bool | None = None - track_graph: bool | None = None - method_usages: bool | None = None - sync_enabled: bool | None = None - full_range_index: bool | None = None - ignore_process_errors: bool | None = None - disable_graph: bool | None = None - generics: bool | None = None + + debug: bool = False + verify_graph: bool = False + track_graph: bool = False + method_usages: bool = True + sync_enabled: bool = True + full_range_index: bool = False + ignore_process_errors: bool = True + disable_graph: bool = False + generics: bool = True import_resolution_overrides: dict[str, str] = Field(default_factory=lambda: {}) typescript: TypescriptConfig = Field(default_factory=TypescriptConfig) @@ -57,6 +59,7 @@ class RepositoryConfig(BaseSettings): class SecretsConfig(BaseSettings): model_config = _get_setting_config("SECRETS") + github_token: str | None = None openai_api_key: str | None = None diff --git a/tests/integration/codemod/conftest.py b/tests/integration/codemod/conftest.py index ae71dd7cb..a5cdb11c4 100644 --- a/tests/integration/codemod/conftest.py +++ b/tests/integration/codemod/conftest.py @@ -12,8 +12,9 @@ from codegen.git.repo_operator.local_repo_operator import LocalRepoOperator from codegen.git.repo_operator.repo_operator import RepoOperator -from codegen.sdk.codebase.config import CodebaseConfig, GSFeatureFlags, ProjectConfig +from codegen.sdk.codebase.config import CodebaseConfig, ProjectConfig from codegen.sdk.core.codebase import Codebase +from codegen.shared.configs.models import CodebaseFeatureFlags from tests.shared.codemod.constants import DIFF_FILEPATH from tests.shared.codemod.models import BASE_PATH, BASE_TMP_DIR, VERIFIED_CODEMOD_DIFFS, CodemodMetadata, Repo, Size from tests.shared.codemod.test_discovery import find_codemod_test_cases, find_repos, find_verified_codemod_cases @@ -142,7 +143,7 @@ def op(repo: Repo, token: str | None) -> YieldFixture[LocalRepoOperator]: def _codebase(repo: Repo, op: RepoOperator, request) -> YieldFixture[Codebase]: sync = request.config.getoption("sync-graph").lower() == "true" log_parse = request.config.getoption("log-parse").lower() == "true" - feature_flags = GSFeatureFlags(verify_graph=sync, debug=log_parse) + feature_flags = CodebaseFeatureFlags(verify_graph=sync, debug=log_parse) if repo.name not in Codebases: projects = [ProjectConfig(repo_operator=op, programming_language=repo.language, subdirectories=repo.subdirectories, base_path=repo.base_path)] Codebases[repo.name] = Codebase(projects=projects, config=CodebaseConfig(feature_flags=feature_flags)) diff --git a/tests/shared/codemod/models.py b/tests/shared/codemod/models.py index 6aa1e3a43..79da07863 100644 --- a/tests/shared/codemod/models.py +++ b/tests/shared/codemod/models.py @@ -10,7 +10,7 @@ from pydantic import BaseModel, ConfigDict from codegen.git.repo_operator.local_repo_operator import LocalRepoOperator -from codegen.sdk.codebase.config import GSFeatureFlags +from codegen.shared.configs.models import CodebaseFeatureFlags from codegen.shared.enums.programming_language import ProgrammingLanguage from tests.shared.codemod.constants import DIFF_FILEPATH from tests.shared.codemod.verified_codemod_utils import CodemodAPI @@ -54,7 +54,7 @@ class Repo(BaseModel): repo_id: int | None = None priority: int = 0 base_path: str | None = None - feature_flags: GSFeatureFlags | None = None + feature_flags: CodebaseFeatureFlags | None = None @classmethod def from_json(cls, json_str: str) -> "Repo":