Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions codegen-examples/examples/pr_review_bot/run.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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,
),
),
Expand Down
5 changes: 3 additions & 2 deletions codegen-examples/examples/sqlalchemy_soft_delete/run.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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)
7 changes: 4 additions & 3 deletions src/codegen/extensions/lsp/protocol.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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,))
Expand Down
5 changes: 2 additions & 3 deletions src/codegen/runner/sandbox/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
37 changes: 6 additions & 31 deletions src/codegen/sdk/codebase/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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):
Expand Down Expand Up @@ -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()
11 changes: 6 additions & 5 deletions src/codegen/sdk/codebase/factory/get_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -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]: ...
Expand All @@ -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]: ...
Expand All @@ -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]: ...
Expand All @@ -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]:
Expand Down
2 changes: 1 addition & 1 deletion src/codegen/sdk/core/external/dependency_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
4 changes: 2 additions & 2 deletions src/codegen/sdk/core/external/language_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
27 changes: 15 additions & 12 deletions src/codegen/shared/configs/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand All @@ -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

Expand Down
5 changes: 3 additions & 2 deletions tests/integration/codemod/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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))
Expand Down
4 changes: 2 additions & 2 deletions tests/shared/codemod/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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":
Expand Down
Loading