chore: dedupe codebase feature flag (#429)

caroljung-cg · web-flow · commit ca4f995e8c0f · 2025-02-11T18:21:50.000-08:00
diff --git 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
@@ -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,15 +93,15 @@ 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"
 
     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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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":
