refactor: break up config models (#448)

Author: caroljung-cg

codegen-examples/examples/pr_review_bot/run.py

@@ -4,7 +4,7 @@
 from codegen.sdk.codebase.config import CodebaseConfig
 import json
 from codegen.sdk.secrets import Secrets
-from codegen.shared.configs.models import CodebaseFeatureFlags
+from codegen.shared.configs.models.feature_flags import CodebaseFeatureFlags
 
 github_token = "Your github token"
 open_ai_key = "your open ai key"

codegen-examples/examples/sqlalchemy_soft_delete/run.py

@@ -1,7 +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.configs.models.feature_flags import CodebaseFeatureFlags
 from codegen.shared.enums.programming_language import ProgrammingLanguage
 import shutil
 import subprocess

src/codegen/cli/auth/session.py

@@ -3,25 +3,24 @@
 from pygit2.repository import Repository
 
 from codegen.cli.git.repo import get_git_repo
-from codegen.shared.configs.config import load
 from codegen.shared.configs.constants import CODEGEN_DIR_NAME, CONFIG_FILENAME
-from codegen.shared.configs.global_config import config as global_config
-from codegen.shared.configs.models import Config
+from codegen.shared.configs.models.session import SessionConfig
+from codegen.shared.configs.session_configs import global_config, load_session_config
 
 
 class CodegenSession:
     """Represents an authenticated codegen session with user and repository context"""
 
     repo_path: Path  # TODO: rename to root_path
     codegen_dir: Path
-    config: Config
+    config: SessionConfig
     existing: bool
 
     def __init__(self, repo_path: Path):
         self.repo_path = repo_path
         self.codegen_dir = repo_path / CODEGEN_DIR_NAME
         self.existing = global_config.get_session(repo_path) is not None
-        self.config = load(self.codegen_dir / CONFIG_FILENAME)
+        self.config = load_session_config(self.codegen_dir / CONFIG_FILENAME)
         global_config.set_active_session(repo_path)
 
     @classmethod

src/codegen/extensions/lsp/protocol.py

@@ -10,7 +10,7 @@
 from codegen.extensions.lsp.utils import get_path
 from codegen.sdk.codebase.config import CodebaseConfig
 from codegen.sdk.core.codebase import Codebase
-from codegen.shared.configs.models import CodebaseFeatureFlags
+from codegen.shared.configs.models.feature_flags import CodebaseFeatureFlags
 
 if TYPE_CHECKING:
     from codegen.extensions.lsp.server import CodegenLanguageServer

src/codegen/runner/sandbox/runner.py

@@ -12,7 +12,7 @@
 from codegen.sdk.core.codebase import Codebase
 from codegen.sdk.secrets import Secrets
 from codegen.shared.compilation.string_to_code import create_execute_function_from_codeblock
-from codegen.shared.configs.config import config
+from codegen.shared.configs.session_configs import config
 from codegen.shared.performance.stopwatch_utils import stopwatch
 
 logger = logging.getLogger(__name__)

src/codegen/runner/sandbox/server.py

@@ -23,7 +23,7 @@
 )
 from codegen.runner.sandbox.middlewares import CodemodRunMiddleware
 from codegen.runner.sandbox.runner import SandboxRunner
-from codegen.shared.configs.config import config
+from codegen.shared.configs.session_configs import config
 from codegen.shared.enums.programming_language import ProgrammingLanguage
 from codegen.shared.performance.memory_utils import get_memory_stats
 

src/codegen/sdk/codebase/config.py

@@ -11,7 +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.configs.models.feature_flags import CodebaseFeatureFlags
 from codegen.shared.enums.programming_language import ProgrammingLanguage
 
 HARD_MAX_AI_LIMIT = 500  # Global limit for AI requests

src/codegen/sdk/codebase/factory/get_session.py

@@ -11,7 +11,7 @@
 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.configs.models.feature_flags import CodebaseFeatureFlags
 from codegen.shared.enums.programming_language import ProgrammingLanguage
 
 

src/codegen/shared/configs/models/feature_flags.py

@@ -0,0 +1,34 @@
+from pydantic import BaseModel, Field
+from pydantic_settings import BaseSettings
+
+from codegen.shared.configs.models.utils import get_setting_config
+
+prefix = "FEATURE_FLAGS"
+
+
+class TypescriptConfig(BaseSettings):
+    model_config = get_setting_config(f"{prefix}_TYPESCRIPT")
+
+    ts_dependency_manager: bool = False
+    ts_language_engine: bool = False
+    v8_ts_engine: bool = False
+
+
+class CodebaseFeatureFlags(BaseSettings):
+    model_config = get_setting_config(f"{prefix}")
+
+    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)
+
+
+class FeatureFlagsConfig(BaseModel):
+    codebase: CodebaseFeatureFlags = Field(default_factory=CodebaseFeatureFlags)

src/codegen/shared/configs/global_config.py renamed to src/codegen/shared/configs/models/global_session.py

@@ -1,6 +1,5 @@
 """Global config to manage different codegen sessions, as well as user auth."""
 
-# TODO: rename this file to global.py
 import json
 from pathlib import Path
 
@@ -39,18 +38,3 @@ def save(self) -> None:
 
         with open(SESSION_FILE, "w") as f:
             json.dump(self.model_dump(), f)
-
-
-def _load_global_config() -> GlobalSessionConfig:
-    """Load configuration from the JSON file."""
-    if SESSION_FILE.exists():
-        with open(SESSION_FILE) as f:
-            json_config = json.load(f)
-            return GlobalSessionConfig.model_validate(json_config, strict=False)
-
-    new_config = GlobalSessionConfig(sessions=[])
-    new_config.save()
-    return new_config
-
-
-config = _load_global_config()

src/codegen/shared/configs/models/repository.py

@@ -0,0 +1,20 @@
+from pydantic_settings import BaseSettings
+
+from codegen.shared.configs.models.utils import get_setting_config
+
+prefix = "REPOSITORY"
+
+
+class RepositoryConfig(BaseSettings):
+    """Configuration for the repository context to run codegen.
+    To populate this config, call `codegen init` from within a git repository.
+    """
+
+    model_config = get_setting_config(prefix)
+
+    repo_path: str | None = None  # replace with base_dir
+    repo_name: str | None = None
+    full_name: str | None = None  # replace with org_name
+    language: str | None = None
+    user_name: str | None = None
+    user_email: str | None = None

src/codegen/shared/configs/models/secrets.py

@@ -0,0 +1,12 @@
+from pydantic_settings import BaseSettings
+
+from codegen.shared.configs.models.utils import get_setting_config
+
+prefix = "SECRETS"
+
+
+class SecretsConfig(BaseSettings):
+    model_config = get_setting_config(prefix)
+
+    github_token: str | None = None
+    openai_api_key: str | None = None

src/codegen/shared/configs/models.py renamed to src/codegen/shared/configs/models/session.py

@@ -5,69 +5,12 @@
 from pydantic import BaseModel, Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
-
-def _get_setting_config(group_name: str) -> SettingsConfigDict:
-    return SettingsConfigDict(
-        env_prefix=f"CODEGEN_{group_name}__",
-        case_sensitive=False,
-        extra="ignore",
-        exclude_defaults=False,
-    )
-
-
-# TODO: break up these models into separate files and nest in shared/configs/models
-class TypescriptConfig(BaseSettings):
-    model_config = _get_setting_config("FEATURE_FLAGS_TYPESCRIPT")
-
-    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 = 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)
-
-
-class RepositoryConfig(BaseSettings):
-    """Configuration for the repository context to run codegen.
-    To populate this config, call `codegen init` from within a git repository.
-    """
-
-    model_config = _get_setting_config("REPOSITORY")
-
-    repo_path: str | None = None  # replace with base_dir
-    repo_name: str | None = None
-    full_name: str | None = None  # replace with org_name
-    language: str | None = None
-    user_name: str | None = None
-    user_email: str | None = None
-
-
-class SecretsConfig(BaseSettings):
-    model_config = _get_setting_config("SECRETS")
-
-    github_token: str | None = None
-    openai_api_key: str | None = None
-
-
-class FeatureFlagsConfig(BaseModel):
-    codebase: CodebaseFeatureFlags = Field(default_factory=CodebaseFeatureFlags)
+from codegen.shared.configs.models.feature_flags import FeatureFlagsConfig
+from codegen.shared.configs.models.repository import RepositoryConfig
+from codegen.shared.configs.models.secrets import SecretsConfig
 
 
-# TODO: rename to SessionConfig
-class Config(BaseSettings):
+class SessionConfig(BaseSettings):
     model_config = SettingsConfigDict(
         extra="ignore",
         exclude_defaults=False,
@@ -124,7 +67,7 @@ def set(self, full_key: str, value: str) -> None:
         field_info = current_attr.model_fields[keys[-1]].annotation
         if isinstance(field_info, BaseModel):
             try:
-                Config.model_validate(value, strict=False)
+                SessionConfig.model_validate(value, strict=False)
             except Exception as e:
                 msg = f"Value does not match the expected type for key: {full_key}\n\nError:{e}"
                 raise ValueError(msg)

src/codegen/shared/configs/models/utils.py

@@ -0,0 +1,10 @@
+from pydantic_settings import SettingsConfigDict
+
+
+def get_setting_config(group_name: str) -> SettingsConfigDict:
+    return SettingsConfigDict(
+        env_prefix=f"CODEGEN_{group_name}__",
+        case_sensitive=False,
+        extra="ignore",
+        exclude_defaults=False,
+    )

src/codegen/shared/configs/config.py renamed to src/codegen/shared/configs/session_configs.py

@@ -1,13 +1,14 @@
-# TODO: rename this file to local.py
+import json
 from pathlib import Path
 
 import tomllib
 
-from codegen.shared.configs.constants import CONFIG_PATH
-from codegen.shared.configs.models import Config
+from codegen.shared.configs.constants import CONFIG_PATH, SESSION_FILE
+from codegen.shared.configs.models.global_session import GlobalSessionConfig
+from codegen.shared.configs.models.session import SessionConfig
 
 
-def load(config_path: Path) -> Config:
+def load_session_config(config_path: Path) -> SessionConfig:
     """Loads configuration from various sources."""
     # Load from .env file
     env_config = _load_from_env(config_path)
@@ -17,28 +18,28 @@ def load(config_path: Path) -> Config:
 
     # Merge configurations recursively
     config_dict = _merge_configs(env_config.model_dump(), toml_config.model_dump())
-    loaded_config = Config(**config_dict)
+    loaded_config = SessionConfig(**config_dict)
 
     # Save the configuration to file if it doesn't exist
     if not config_path.exists():
         loaded_config.save()
     return loaded_config
 
 
-def _load_from_env(config_path: Path) -> Config:
+def _load_from_env(config_path: Path) -> SessionConfig:
     """Load configuration from the environment variables."""
-    return Config(file_path=str(config_path))
+    return SessionConfig(file_path=str(config_path))
 
 
-def _load_from_toml(config_path: Path) -> Config:
+def _load_from_toml(config_path: Path) -> SessionConfig:
     """Load configuration from the TOML file."""
     if config_path.exists():
         with open(config_path, "rb") as f:
             toml_config = tomllib.load(f)
             toml_config["file_path"] = str(config_path)
-            return Config.model_validate(toml_config, strict=False)
+            return SessionConfig.model_validate(toml_config, strict=False)
 
-    return Config(file_path=str(config_path))
+    return SessionConfig(file_path=str(config_path))
 
 
 def _merge_configs(base: dict, override: dict) -> dict:
@@ -54,7 +55,22 @@ def _merge_configs(base: dict, override: dict) -> dict:
     return merged
 
 
-config = load(CONFIG_PATH)
+def _load_global_config() -> GlobalSessionConfig:
+    """Load configuration from the JSON file."""
+    if SESSION_FILE.exists():
+        with open(SESSION_FILE) as f:
+            json_config = json.load(f)
+            return GlobalSessionConfig.model_validate(json_config, strict=False)
+
+    new_config = GlobalSessionConfig(sessions=[])
+    new_config.save()
+    return new_config
+
+
+config = load_session_config(CONFIG_PATH)
+global_config = _load_global_config()
+
 
 if __name__ == "__main__":
     print(config)
+    print(global_config)

tests/integration/codegen/git/codebase/conftest.py

@@ -8,15 +8,14 @@
 from codegen.git.utils.clone_url import get_authenticated_clone_url_for_repo_config
 from codegen.sdk.codebase.config import ProjectConfig
 from codegen.sdk.core.codebase import Codebase
-from codegen.shared.configs.config import config
+from codegen.shared.configs.session_configs import config
 
 
 @pytest.fixture
 def repo_config(tmpdir):
     repo_config = RepoConfig(
         name="Kevin-s-Adventure-Game",
         full_name="codegen-sh/Kevin-s-Adventure-Game",
-        organization_name="codegen-sh",
         base_dir=str(tmpdir),
     )
     yield repo_config

tests/integration/codegen/git/conftest.py

@@ -19,7 +19,6 @@ def repo_config(tmpdir):
     repo_config = RepoConfig(
         name="Kevin-s-Adventure-Game",
         full_name="codegen-sh/Kevin-s-Adventure-Game",
-        organization_name="codegen-sh",
         base_dir=str(tmpdir),
     )
     yield repo_config

tests/integration/codegen/runner/conftest.py

@@ -9,7 +9,7 @@
 from codegen.git.repo_operator.remote_repo_operator import RemoteRepoOperator
 from codegen.git.schemas.repo_config import RepoConfig
 from codegen.runner.clients.sandbox_client import SandboxClient
-from codegen.shared.configs.config import config
+from codegen.shared.configs.session_configs import config
 from codegen.shared.enums.programming_language import ProgrammingLanguage