unify interface in sensitive data

Author: yrobla

src/codegate/cli.py

@@ -16,7 +16,7 @@
 from codegate.config import Config, ConfigurationError
 from codegate.db.connection import init_db_sync, init_session_if_not_exists
 from codegate.pipeline.factory import PipelineFactory
-from codegate.pipeline.secrets.manager import SecretsManager
+from codegate.pipeline.sensitive_data.manager import SensitiveDataManager
 from codegate.providers import crud as provendcrud
 from codegate.providers.copilot.provider import CopilotProvider
 from codegate.server import init_app
@@ -331,8 +331,8 @@ def serve(  # noqa: C901
             click.echo("Existing Certificates are already present.")
 
         # Initialize secrets manager and pipeline factory
-        secrets_manager = SecretsManager()
-        pipeline_factory = PipelineFactory(secrets_manager)
+        sensitive_data_manager = SensitiveDataManager()
+        pipeline_factory = PipelineFactory(sensitive_data_manager)
 
         app = init_app(pipeline_factory)
 

src/codegate/pipeline/base.py

@@ -12,14 +12,14 @@
 from codegate.clients.clients import ClientType
 from codegate.db.models import Alert, AlertSeverity, Output, Prompt
 from codegate.extract_snippets.message_extractor import CodeSnippet
-from codegate.pipeline.secrets.manager import SecretsManager
+from codegate.pipeline.sensitive_data.manager import SensitiveDataManager
 
 logger = structlog.get_logger("codegate")
 
 
 @dataclass
 class PipelineSensitiveData:
-    manager: SecretsManager
+    manager: SensitiveDataManager
     session_id: str
     api_key: Optional[str] = None
     model: Optional[str] = None
@@ -266,19 +266,19 @@ class InputPipelineInstance:
     def __init__(
         self,
         pipeline_steps: List[PipelineStep],
-        secret_manager: SecretsManager,
+        sensitive_data_manager: SensitiveDataManager,
         is_fim: bool,
         client: ClientType = ClientType.GENERIC,
     ):
         self.pipeline_steps = pipeline_steps
-        self.secret_manager = secret_manager
+        self.sensitive_data_manager = sensitive_data_manager
         self.is_fim = is_fim
         self.context = PipelineContext(client=client)
 
         # we create the sesitive context here so that it is not shared between individual requests
         # TODO: could we get away with just generating the session ID for an instance?
         self.context.sensitive = PipelineSensitiveData(
-            manager=self.secret_manager,
+            manager=self.sensitive_data_manager,
             session_id=str(uuid.uuid4()),
         )
         self.context.metadata["is_fim"] = is_fim
@@ -335,20 +335,20 @@ class SequentialPipelineProcessor:
     def __init__(
         self,
         pipeline_steps: List[PipelineStep],
-        secret_manager: SecretsManager,
+        sensitive_data_manager: SensitiveDataManager,
         client_type: ClientType,
         is_fim: bool,
     ):
         self.pipeline_steps = pipeline_steps
-        self.secret_manager = secret_manager
+        self.sensitive_data_manager = sensitive_data_manager
         self.is_fim = is_fim
         self.instance = self._create_instance(client_type)
 
     def _create_instance(self, client_type: ClientType) -> InputPipelineInstance:
         """Create a new pipeline instance for processing a request"""
         return InputPipelineInstance(
             self.pipeline_steps,
-            self.secret_manager,
+            self.sensitive_data_manager,
             self.is_fim,
             client_type,
         )

src/codegate/pipeline/factory.py

@@ -12,18 +12,18 @@
     PiiRedactionNotifier,
     PiiUnRedactionStep,
 )
-from codegate.pipeline.secrets.manager import SecretsManager
 from codegate.pipeline.secrets.secrets import (
     CodegateSecrets,
     SecretRedactionNotifier,
     SecretUnredactionStep,
 )
+from codegate.pipeline.sensitive_data.manager import SensitiveDataManager
 from codegate.pipeline.system_prompt.codegate import SystemPrompt
 
 
 class PipelineFactory:
-    def __init__(self, secrets_manager: SecretsManager):
-        self.secrets_manager = secrets_manager
+    def __init__(self, sensitive_data_manager: SensitiveDataManager):
+        self.sensitive_data_manager = sensitive_data_manager
 
     def create_input_pipeline(self, client_type: ClientType) -> SequentialPipelineProcessor:
         input_steps: List[PipelineStep] = [
@@ -41,7 +41,7 @@ def create_input_pipeline(self, client_type: ClientType) -> SequentialPipelinePr
         ]
         return SequentialPipelineProcessor(
             input_steps,
-            self.secrets_manager,
+            self.sensitive_data_manager,
             client_type,
             is_fim=False,
         )
@@ -53,7 +53,7 @@ def create_fim_pipeline(self, client_type: ClientType) -> SequentialPipelineProc
         ]
         return SequentialPipelineProcessor(
             fim_steps,
-            self.secrets_manager,
+            self.sensitive_data_manager,
             client_type,
             is_fim=True,
         )

src/codegate/pipeline/pii/analyzer.py

@@ -1,13 +1,12 @@
-import uuid
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, List, Optional
 
 import structlog
 from presidio_analyzer import AnalyzerEngine
 from presidio_anonymizer import AnonymizerEngine
 
 from codegate.db.models import AlertSeverity
 from codegate.pipeline.base import PipelineContext
-from codegate.session.session_store import SessionStore
+from codegate.pipeline.sensitive_data.session_store import SessionStore
 
 logger = structlog.get_logger("codegate.pii.analyzer")
 
@@ -69,9 +68,7 @@ def __init__(self):
 
         PiiAnalyzer._instance = self
 
-    def analyze(
-        self, session_id: str, text: str, context: Optional[PipelineContext] = None
-    ) -> Tuple[str, List[Dict[str, Any]]]:
+    def analyze(self, text: str, context: Optional[PipelineContext] = None) -> List:
         # Prioritize credit card detection first
         entities = [
             "PHONE_NUMBER",
@@ -95,65 +92,7 @@ def analyze(
             language="en",
             score_threshold=0.3,  # Lower threshold to catch more potential matches
         )
-
-        # Track found PII
-        found_pii = []
-
-        # Only anonymize if PII was found
-        if analyzer_results:
-            # Log each found PII instance and anonymize
-            anonymized_text = text
-            for result in analyzer_results:
-                pii_value = text[result.start : result.end]
-                uuid_placeholder = self.session_store.add_mapping(session_id, pii_value)
-                pii_info = {
-                    "type": result.entity_type,
-                    "value": pii_value,
-                    "score": result.score,
-                    "start": result.start,
-                    "end": result.end,
-                    "uuid_placeholder": uuid_placeholder,
-                }
-                found_pii.append(pii_info)
-                anonymized_text = anonymized_text.replace(pii_value, uuid_placeholder)
-
-                # Log each PII detection with its UUID mapping
-                logger.info(
-                    "PII detected and mapped",
-                    pii_type=result.entity_type,
-                    score=f"{result.score:.2f}",
-                    uuid=uuid_placeholder,
-                    # Don't log the actual PII value for security
-                    value_length=len(pii_value),
-                    session_id=session_id,
-                )
-
-            # Log summary of all PII found in this analysis
-            if found_pii and context:
-                # Create notification string for alert
-                notify_string = (
-                    f"**PII Detected** 🔒\n"
-                    f"- Total PII Found: {len(found_pii)}\n"
-                    f"- Types Found: {', '.join(set(p['type'] for p in found_pii))}\n"
-                )
-                context.add_alert(
-                    self._name,
-                    trigger_string=notify_string,
-                    severity_category=AlertSeverity.CRITICAL,
-                )
-
-                logger.info(
-                    "PII analysis complete",
-                    total_pii_found=len(found_pii),
-                    pii_types=[p["type"] for p in found_pii],
-                    session_id=session_id,
-                )
-
-            # Return the anonymized text, PII details, and session store
-            return anonymized_text, found_pii
-
-        # If no PII found, return original text, empty list, and session store
-        return text, []
+        return analyzer_results
 
     def restore_pii(self, session_id: str, anonymized_text: str) -> str:
         """