From 4c1b4be32d172a42a15642fdab36e1b55f3f6b08 Mon Sep 17 00:00:00 2001 From: Luke Hinds Date: Sun, 15 Dec 2024 22:00:55 +0000 Subject: [PATCH 1/3] Optimize signtures --- src/codegate/cli.py | 9 +++++---- src/codegate/pipeline/secrets/secrets.py | 3 +++ src/codegate/pipeline/secrets/signatures.py | 8 +++++--- src/codegate/server.py | 4 ---- 4 files changed, 13 insertions(+), 11 deletions(-) diff --git a/src/codegate/cli.py b/src/codegate/cli.py index b230eba9..a23c3177 100644 --- a/src/codegate/cli.py +++ b/src/codegate/cli.py @@ -305,16 +305,16 @@ def serve( ca = CertificateAuthority.get_instance() ca.ensure_certificates_exist() - # Set up event loop - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - # Initialize secrets manager and pipeline factory secrets_manager = SecretsManager() pipeline_factory = PipelineFactory(secrets_manager) app = init_app(pipeline_factory) + # Set up event loop + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + # Run the server try: loop.run_until_complete(run_servers(cfg, app)) @@ -331,6 +331,7 @@ def serve( logger.exception("Unexpected error occurred") click.echo(f"Error: {e}", err=True) sys.exit(1) + async def run_servers(cfg: Config, app) -> None: diff --git a/src/codegate/pipeline/secrets/secrets.py b/src/codegate/pipeline/secrets/secrets.py index 67bc0386..3f1f13f2 100644 --- a/src/codegate/pipeline/secrets/secrets.py +++ b/src/codegate/pipeline/secrets/secrets.py @@ -26,6 +26,9 @@ class CodegateSecrets(PipelineStep): def __init__(self): """Initialize the CodegateSecrets pipeline step.""" super().__init__() + # Initialize signatures eagerly + CodegateSignatures.initialize("signatures.yaml") + CodegateSignatures._ensure_signatures_loaded() @property def name(self) -> str: diff --git a/src/codegate/pipeline/secrets/signatures.py b/src/codegate/pipeline/secrets/signatures.py index 1f1a7ddd..5fd2b6b2 100644 --- a/src/codegate/pipeline/secrets/signatures.py +++ b/src/codegate/pipeline/secrets/signatures.py @@ -61,9 +61,11 @@ def initialize(cls, yaml_path: str) -> None: raise FileNotFoundError(f"Signatures file not found: {yaml_path}") with cls._instance_lock: - cls._yaml_path = yaml_path - cls._signatures_loaded = False - logger.debug(f"SecretFinder initialized with {yaml_path}") + # Only initialize if not already initialized with this path + if cls._yaml_path != yaml_path: + cls._yaml_path = yaml_path + cls._signatures_loaded = False + logger.debug(f"SecretFinder initialized with {yaml_path}") @classmethod def _preprocess_yaml(cls, content: str) -> str: diff --git a/src/codegate/server.py b/src/codegate/server.py index 34da0b8c..d2d7801e 100644 --- a/src/codegate/server.py +++ b/src/codegate/server.py @@ -4,7 +4,6 @@ from codegate import __description__, __version__ from codegate.dashboard.dashboard import dashboard_router from codegate.pipeline.factory import PipelineFactory -from codegate.pipeline.secrets.signatures import CodegateSignatures from codegate.providers.anthropic.provider import AnthropicProvider from codegate.providers.llamacpp.provider import LlamaCppProvider from codegate.providers.ollama.provider import OllamaProvider @@ -31,9 +30,6 @@ def init_app(pipeline_factory: PipelineFactory) -> FastAPI: # Create provider registry registry = ProviderRegistry(app) - # Initialize SignaturesFinder - CodegateSignatures.initialize("signatures.yaml") - # Register all known providers registry.add_provider( "openai", From 02e3efdfca46e3f065c0dd12c1c7fc0e7f8d4816 Mon Sep 17 00:00:00 2001 From: Luke Hinds Date: Sun, 15 Dec 2024 22:23:45 +0000 Subject: [PATCH 2/3] simplify signature initialization for optimization Remove unnecessary lazy loading mechanism from signatures since they are always loaded during server startup when CodegateSecrets is created. This simplifies the code and makes the initialization flow singular, which will reduce uneeded lazy loading on every call --- src/codegate/cli.py | 1 - src/codegate/pipeline/secrets/secrets.py | 3 +- src/codegate/pipeline/secrets/signatures.py | 32 ++++++--------------- tests/pipeline/secrets/test_signatures.py | 5 ---- tests/test_cli.py | 7 ----- tests/test_server.py | 7 ----- 6 files changed, 9 insertions(+), 46 deletions(-) diff --git a/src/codegate/cli.py b/src/codegate/cli.py index a23c3177..ea871b5e 100644 --- a/src/codegate/cli.py +++ b/src/codegate/cli.py @@ -331,7 +331,6 @@ def serve( logger.exception("Unexpected error occurred") click.echo(f"Error: {e}", err=True) sys.exit(1) - async def run_servers(cfg: Config, app) -> None: diff --git a/src/codegate/pipeline/secrets/secrets.py b/src/codegate/pipeline/secrets/secrets.py index 3f1f13f2..0f3d61c2 100644 --- a/src/codegate/pipeline/secrets/secrets.py +++ b/src/codegate/pipeline/secrets/secrets.py @@ -26,9 +26,8 @@ class CodegateSecrets(PipelineStep): def __init__(self): """Initialize the CodegateSecrets pipeline step.""" super().__init__() - # Initialize signatures eagerly + # Initialize and load signatures immediately CodegateSignatures.initialize("signatures.yaml") - CodegateSignatures._ensure_signatures_loaded() @property def name(self) -> str: diff --git a/src/codegate/pipeline/secrets/signatures.py b/src/codegate/pipeline/secrets/signatures.py index 5fd2b6b2..f2df00ab 100644 --- a/src/codegate/pipeline/secrets/signatures.py +++ b/src/codegate/pipeline/secrets/signatures.py @@ -39,7 +39,6 @@ class CodegateSignatures: """Main class for detecting secrets in text using regex patterns.""" _instance_lock: ClassVar[Lock] = Lock() - _signatures_loaded: ClassVar[bool] = False _signature_groups: ClassVar[List[SignatureGroup]] = [] _compiled_regexes: ClassVar[Dict[str, re.Pattern]] = {} _yaml_path: ClassVar[Optional[str]] = None @@ -48,7 +47,6 @@ class CodegateSignatures: def reset(cls) -> None: """Reset the cached patterns.""" with cls._instance_lock: - cls._signatures_loaded = False cls._signature_groups = [] cls._compiled_regexes = {} cls._yaml_path = None @@ -56,7 +54,7 @@ def reset(cls) -> None: @classmethod def initialize(cls, yaml_path: str) -> None: - """Initialize the SecretFinder with a YAML file path.""" + """Initialize the SecretFinder with a YAML file path and load signatures.""" if not Path(yaml_path).exists(): raise FileNotFoundError(f"Signatures file not found: {yaml_path}") @@ -64,7 +62,7 @@ def initialize(cls, yaml_path: str) -> None: # Only initialize if not already initialized with this path if cls._yaml_path != yaml_path: cls._yaml_path = yaml_path - cls._signatures_loaded = False + cls._load_signatures() logger.debug(f"SecretFinder initialized with {yaml_path}") @classmethod @@ -174,6 +172,10 @@ def _add_signature_group(cls, name: str, patterns: Dict[str, str]) -> None: def _load_signatures(cls) -> None: """Load signature patterns from the YAML file.""" try: + # Clear existing signatures before loading new ones + cls._signature_groups = [] + cls._compiled_regexes = {} + yaml_data = cls._load_yaml(cls._yaml_path) # Add custom GitHub token patterns @@ -207,32 +209,14 @@ def _load_signatures(cls) -> None: logger.error(f"Error loading signatures: {e}") raise - @classmethod - def _ensure_signatures_loaded(cls) -> None: - """Ensure signatures are loaded before use.""" - if not cls._signatures_loaded: - with cls._instance_lock: - if not cls._signatures_loaded: - if not cls._yaml_path: - raise RuntimeError("SecretFinder not initialized. Call initialize() first.") - try: - cls._load_signatures() - cls._signatures_loaded = True - except Exception as e: - logger.error(f"Failed to load signatures: {e}") - raise - @classmethod def find_in_string(cls, text: str) -> List[Match]: """Search for secrets in the provided string.""" if not text: return [] - try: - cls._ensure_signatures_loaded() - except Exception as e: - logger.error(f"Failed to load signatures: {e}") - return [] + if not cls._yaml_path: + raise RuntimeError("SecretFinder not initialized. Call initialize() first.") matches = [] lines = text.splitlines() diff --git a/tests/pipeline/secrets/test_signatures.py b/tests/pipeline/secrets/test_signatures.py index c7e0b9e6..42f6470f 100644 --- a/tests/pipeline/secrets/test_signatures.py +++ b/tests/pipeline/secrets/test_signatures.py @@ -69,14 +69,12 @@ def test_initialize_and_reset(temp_yaml_file): CodegateSignatures.reset() assert CodegateSignatures._yaml_path is None - assert not CodegateSignatures._signatures_loaded assert not CodegateSignatures._signature_groups assert not CodegateSignatures._compiled_regexes def test_find_in_string_with_aws_credentials(temp_yaml_file): CodegateSignatures.initialize(temp_yaml_file) - CodegateSignatures._signatures_loaded = False # Force reload of signatures test_string = """ aws_access_key = 'AKIAIOSFODNN7EXAMPLE' @@ -100,7 +98,6 @@ def test_find_in_string_with_github_token(): try: CodegateSignatures.initialize(f.name) - CodegateSignatures._signatures_loaded = False # Force reload of signatures test_string = "github_token = 'ghp_1234567890abcdef1234567890abcdef123456'" matches = CodegateSignatures.find_in_string(test_string) @@ -121,7 +118,6 @@ def test_find_in_string_with_github_token(): def test_find_in_string_with_no_matches(temp_yaml_file): CodegateSignatures.initialize(temp_yaml_file) - CodegateSignatures._signatures_loaded = False # Force reload of signatures test_string = "No secrets here!" matches = CodegateSignatures.find_in_string(test_string) @@ -158,7 +154,6 @@ def test_duplicate_patterns(): try: CodegateSignatures.initialize(f.name) - CodegateSignatures._signatures_loaded = False # Force reload of signatures test_string = "aws_key = 'AKIAIOSFODNN7EXAMPLE'" matches = CodegateSignatures.find_in_string(test_string) diff --git a/tests/test_cli.py b/tests/test_cli.py index 8557a63c..caa3a215 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -105,13 +105,6 @@ def test_provider_registration(mock_registry, mock_secrets_mgr, mock_pipeline_fa assert "ollama" in provider_names -@patch("codegate.server.CodegateSignatures") -def test_signatures_initialization(mock_signatures, mock_pipeline_factory) -> None: - """Test that signatures are initialized correctly.""" - init_app(mock_pipeline_factory) - mock_signatures.initialize.assert_called_once_with("signatures.yaml") - - def test_pipeline_initialization(mock_pipeline_factory) -> None: """Test that pipelines are initialized correctly.""" app = init_app(mock_pipeline_factory) diff --git a/tests/test_server.py b/tests/test_server.py index 9170f36e..5420ad70 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -99,13 +99,6 @@ def test_provider_registration(mock_registry, mock_secrets_mgr, mock_pipeline_fa assert "ollama" in provider_names -@patch("codegate.server.CodegateSignatures") -def test_signatures_initialization(mock_signatures, mock_pipeline_factory) -> None: - """Test that signatures are initialized correctly.""" - init_app(mock_pipeline_factory) - mock_signatures.initialize.assert_called_once_with("signatures.yaml") - - def test_pipeline_initialization(mock_pipeline_factory) -> None: """Test that pipelines are initialized correctly.""" app = init_app(mock_pipeline_factory) From d2a1315f06625304e85935d440086e92a59955ae Mon Sep 17 00:00:00 2001 From: Luke Hinds Date: Sun, 15 Dec 2024 22:34:41 +0000 Subject: [PATCH 3/3] Old debug message --- src/codegate/pipeline/secrets/signatures.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/codegate/pipeline/secrets/signatures.py b/src/codegate/pipeline/secrets/signatures.py index f2df00ab..95d306e1 100644 --- a/src/codegate/pipeline/secrets/signatures.py +++ b/src/codegate/pipeline/secrets/signatures.py @@ -216,7 +216,7 @@ def find_in_string(cls, text: str) -> List[Match]: return [] if not cls._yaml_path: - raise RuntimeError("SecretFinder not initialized. Call initialize() first.") + raise RuntimeError("SecretFinder not initialized.") matches = [] lines = text.splitlines()