fix tests

Author: yrobla

src/codegate/pipeline/pii/analyzer.py

@@ -70,7 +70,7 @@ def __init__(self):
         PiiAnalyzer._instance = self
 
     def analyze(
-        self, text: str, session_id: str, context: Optional[PipelineContext] = None
+        self, session_id: str, text: str, context: Optional[PipelineContext] = None
     ) -> Tuple[str, List[Dict[str, Any]]]:
         # Prioritize credit card detection first
         entities = [
@@ -155,7 +155,7 @@ def analyze(
         # If no PII found, return original text, empty list, and session store
         return text, []
 
-    def restore_pii(self, anonymized_text: str, session_id: str) -> str:
+    def restore_pii(self, session_id: str, anonymized_text: str) -> str:
         """
         Restore the original PII (Personally Identifiable Information) in the given anonymized text.
 
@@ -164,7 +164,7 @@ def restore_pii(self, anonymized_text: str, session_id: str) -> str:
 
         Args:
             anonymized_text (str): The text containing placeholders for PII.
-            session_store (SessionStore): The session store containing mappings of placeholders
+            session_id (str): The session id containing mappings of placeholders
             to original PII.
 
         Returns:

src/codegate/pipeline/pii/manager.py

@@ -24,17 +24,19 @@ class PiiManager:
             Initializes the PiiManager with the singleton PiiAnalyzer instance and sets the
             session store.
 
-        analyze(text: str, session_id: str) -> Tuple[str, List[Dict[str, Any]]]:
+        analyze(session_id: str, text: str) -> Tuple[str, List[Dict[str, Any]]]:
             Analyzes the given text for PII, anonymizes it, and logs the detected PII details.
             Args:
+                session_id (str): The session id to store the PII.
                 text (str): The text to be analyzed for PII.
             Returns:
                 Tuple[str, List[Dict[str, Any]]]: A tuple containing the anonymized text and
                 a list of found PII details.
 
-        restore_pii(anonymized_text: str, session_id: str ) -> str:
+        restore_pii(session_id: str, anonymized_text: st ) -> str:
             Restores the PII in the given anonymized text using the current session.
             Args:
+                session_id (str): The session id for the PII to be restored.
                 anonymized_text (str): The text with anonymized PII to be restored.
             Returns:
                 str: The text with restored PII.
@@ -55,10 +57,10 @@ def session_store(self) -> SessionStore:
         return self.analyzer.session_store
 
     def analyze(
-        self, text: str, session_id: str, context: Optional[PipelineContext] = None
+        self, session_id: str, text: str, context: Optional[PipelineContext] = None
     ) -> Tuple[str, List[Dict[str, Any]]]:
         # Call analyzer and get results
-        anonymized_text, found_pii = self.analyzer.analyze(text, session_id, context=context)
+        anonymized_text, found_pii = self.analyzer.analyze(session_id, text, context=context)
 
         # Log found PII details (without modifying the found_pii list)
         if found_pii:
@@ -73,9 +75,11 @@ def analyze(
         # Return the exact same objects we got from the analyzer
         return anonymized_text, found_pii
 
-    def restore_pii(self, anonymized_text: str, session_id: str) -> str:
+    def restore_pii(self, session_id: str, anonymized_text: str) -> str:
         """
         Restore PII in the given anonymized text using the current session.
         """
+        if not session_id:
+            return anonymized_text
         # Use the analyzer's restore_pii method with the current session store
-        return self.analyzer.restore_pii(anonymized_text, session_id)
+        return self.analyzer.restore_pii(session_id, anonymized_text)

src/codegate/pipeline/pii/pii.py

@@ -38,7 +38,7 @@ class CodegatePii(PipelineStep):
             Processes the chat completion request to detect and redact PII. Updates the request with
             anonymized text and stores PII details in the context metadata.
 
-        restore_pii(anonymized_text: str, session_id: str) -> str:
+        restore_pii(session_id: str, anonymized_text: str) -> str:
             Restores the original PII from the anonymized text using the PiiManager.
     """
 
@@ -83,7 +83,7 @@ async def process(
                 # This is where analyze and anonymize the text
                 original_text = str(message["content"])
                 anonymized_text, pii_details = self.pii_manager.analyze(
-                    original_text, session_id, context
+                    session_id, original_text, context
                 )
 
                 if pii_details:
@@ -118,8 +118,8 @@ async def process(
 
         return PipelineResult(request=new_request, context=context)
 
-    def restore_pii(self, anonymized_text: str, session_id: str) -> str:
-        return self.pii_manager.restore_pii(anonymized_text, session_id)
+    def restore_pii(self, session_id: str, anonymized_text: str) -> str:
+        return self.pii_manager.restore_pii(session_id, anonymized_text)
 
 
 class PiiUnRedactionStep(OutputPipelineStep):
@@ -141,12 +141,12 @@ class PiiUnRedactionStep(OutputPipelineStep):
     """
 
     def __init__(self):
-        self.redacted_pattern = re.compile(r"<([0-9a-f-]{0,36})>")
+        self.redacted_pattern = re.compile(r"#([0-9a-f-]{0,36})#")
         self.complete_uuid_pattern = re.compile(
             r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$"
         )  # noqa: E501
-        self.marker_start = "<"
-        self.marker_end = ">"
+        self.marker_start = "#"
+        self.marker_end = "#"
 
     @property
     def name(self) -> str:
@@ -181,13 +181,13 @@ async def process_chunk(  # noqa: C901
         current_pos = 0
         result = []
         while current_pos < len(content):
-            start_idx = content.find("<", current_pos)
+            start_idx = content.find("#", current_pos)
             if start_idx == -1:
                 # No more markers!, add remaining content
                 result.append(content[current_pos:])
                 break
 
-            end_idx = content.find(">", start_idx)
+            end_idx = content.find("#", start_idx + 1)
             if end_idx == -1:
                 # Incomplete marker, buffer the rest
                 context.prefix_buffer = content[current_pos:]
@@ -199,7 +199,7 @@ async def process_chunk(  # noqa: C901
 
             # Extract potential UUID if it's a valid format!
             uuid_marker = content[start_idx : end_idx + 1]
-            uuid_value = uuid_marker[1:-1]  # Remove < >
+            uuid_value = uuid_marker[1:-1]  # Remove # #
 
             if self._is_complete_uuid(uuid_value):
                 # Get the PII manager from context metadata

src/codegate/pipeline/secrets/manager.py

@@ -1,5 +1,5 @@
 import json
-from typing import Optional
+from typing import Dict, Optional
 
 import structlog
 
@@ -19,6 +19,9 @@ def store_secret(self, session_id: str, value: str, service: str, secret_type: s
         Encrypts and stores a secret value.
         Returns the encrypted value.
         """
+        if not session_id:
+            raise ValueError("Session id must be provided")
+
         if not value:
             raise ValueError("Value must be provided")
         if not service:
@@ -35,6 +38,16 @@ def store_secret(self, session_id: str, value: str, service: str, secret_type: s
         )
         return uuid_placeholder
 
+    def get_by_session_id(self, session_id: str) -> Optional[Dict]:
+        session_data = self.session_store.get_by_session_id(session_id)
+        if not session_data:
+            return None
+        # Convert all string values to dictionary objects using json.loads
+        return {
+            key: json.loads(value) if isinstance(value, str) else value
+            for key, value in session_data.items()
+        }
+
     def get_original_value(self, session_id: str, uuid_placeholder: str) -> Optional[str]:
         """Retrieve original value for an encrypted value"""
         secret_entry_json = self.session_store.get_mapping(session_id, uuid_placeholder)
@@ -45,3 +58,6 @@ def get_original_value(self, session_id: str, uuid_placeholder: str) -> Optional
 
     def cleanup_session(self, session_id):
         self.session_store.cleanup_session(session_id)
+
+    def cleanup(self):
+        self.session_store.cleanup()

src/codegate/pipeline/secrets/secrets.py

@@ -430,14 +430,13 @@ async def process_chunk(
             if encrypted_value.startswith("$"):
                 encrypted_value = encrypted_value[1:]
 
-            session_id = context.sensitive.session_id
+            session_id = input_context.sensitive.session_id
             if not session_id:
                 raise ValueError("Session ID not found in context")
 
             original_value = input_context.sensitive.manager.get_original_value(
                 session_id,
                 encrypted_value,
-                input_context.sensitive.session_id,
             )
 
             if original_value is None:

src/codegate/session/session_store.py

@@ -17,6 +17,9 @@ def add_mapping(self, session_id: str, data: str) -> str:
         self.sessions[session_id][uuid_placeholder] = data
         return uuid_placeholder
 
+    def get_by_session_id(self, session_id: str) -> Optional[Dict]:
+        return self.sessions.get(session_id, None)
+
     def get_mapping(self, session_id: str, uuid_placeholder: str) -> Optional[str]:
         return self.sessions.get(session_id, {}).get(uuid_placeholder)
 

tests/pipeline/pii/test_analyzer.py

@@ -70,16 +70,17 @@ def test_singleton_pattern(self):
 
     def test_analyze_no_pii(self, analyzer, mock_analyzer_engine):
         text = "Hello world"
+        session_id = "session-id"
         mock_analyzer_engine.analyze.return_value = []
 
-        result_text, found_pii, session_store = analyzer.analyze(text)
+        result_text, found_pii = analyzer.analyze(session_id, text)
 
         assert result_text == text
         assert found_pii == []
-        assert isinstance(session_store, SessionStore)
 
     def test_analyze_with_pii(self, analyzer, mock_analyzer_engine):
         text = "My email is [email protected]"
+        session_id = "session-id"
         email_pii = RecognizerResult(
             entity_type="EMAIL_ADDRESS",
             start=12,
@@ -88,7 +89,7 @@ def test_analyze_with_pii(self, analyzer, mock_analyzer_engine):
         )
         mock_analyzer_engine.analyze.return_value = [email_pii]
 
-        result_text, found_pii, session_store = analyzer.analyze(text)
+        result_text, found_pii = analyzer.analyze(session_id, text)
 
         assert len(found_pii) == 1
         pii_info = found_pii[0]
@@ -101,36 +102,32 @@ def test_analyze_with_pii(self, analyzer, mock_analyzer_engine):
         # Verify the placeholder was used to replace the PII
         placeholder = pii_info["uuid_placeholder"]
         assert result_text == f"My email is {placeholder}"
-        # Verify the mapping was stored
-        assert session_store.get_pii(placeholder) == "[email protected]"
 
     def test_restore_pii(self, analyzer):
-        session_store = SessionStore()
         original_text = "[email protected]"
         session_id = "session-id"
-        placeholder = session_store.add_mapping(session_id, original_text)
-        anonymized_text = f"My email is {placeholder}"
 
-        restored_text = analyzer.restore_pii(anonymized_text, session_id)
+        placeholder = analyzer.session_store.add_mapping(session_id, original_text)
+        anonymized_text = f"My email is {placeholder}"
+        restored_text = analyzer.restore_pii(session_id, anonymized_text)
 
         assert restored_text == f"My email is {original_text}"
 
     def test_restore_pii_multiple(self, analyzer):
-        session_store = SessionStore()
         email = "[email protected]"
         phone = "123-456-7890"
         session_id = "session-id"
-        email_placeholder = session_store.add_mapping(session_id, email)
-        phone_placeholder = session_store.add_mapping(session_id, phone)
+        email_placeholder = analyzer.session_store.add_mapping(session_id, email)
+        phone_placeholder = analyzer.session_store.add_mapping(session_id, phone)
         anonymized_text = f"Email: {email_placeholder}, Phone: {phone_placeholder}"
 
-        restored_text = analyzer.restore_pii(anonymized_text, session_id)
+        restored_text = analyzer.restore_pii(session_id, anonymized_text)
 
         assert restored_text == f"Email: {email}, Phone: {phone}"
 
     def test_restore_pii_no_placeholders(self, analyzer):
         text = "No PII here"
         session_id = "session-id"
-        restored_text = analyzer.restore_pii(text, session_id)
+        restored_text = analyzer.restore_pii(session_id, text)
 
         assert restored_text == text

tests/pipeline/pii/test_pi.py

@@ -96,9 +96,10 @@ async def test_process_with_pii(self, pii_step):
     def test_restore_pii(self, pii_step):
         anonymized_text = "My email is <test-uuid>"
         original_text = "My email is [email protected]"
+        session_id = "session-id"
         pii_step.pii_manager.restore_pii = MagicMock(return_value=original_text)
 
-        restored = pii_step.restore_pii(anonymized_text)
+        restored = pii_step.restore_pii(session_id, anonymized_text)
 
         assert restored == original_text
 
@@ -148,7 +149,7 @@ async def test_process_chunk_with_uuid(self, unredaction_step):
                 StreamingChoices(
                     finish_reason=None,
                     index=0,
-                    delta=Delta(content=f"Text with <{uuid}>"),
+                    delta=Delta(content=f"Text with #{uuid}#"),
                     logprobs=None,
                 )
             ],
@@ -157,17 +158,16 @@ async def test_process_chunk_with_uuid(self, unredaction_step):
             object="chat.completion.chunk",
         )
         context = OutputPipelineContext()
-        input_context = PipelineContext()
+        input_context = PipelineContext(metadata={"session_id": "session-id"})
 
         # Mock PII manager in input context
         mock_pii_manager = MagicMock()
         mock_session = MagicMock()
-        mock_session.get_pii = MagicMock(return_value="[email protected]")
+        mock_session.get_mapping = MagicMock(return_value="[email protected]")
         mock_pii_manager.session_store = mock_session
         input_context.metadata["pii_manager"] = mock_pii_manager
 
         result = await unredaction_step.process_chunk(chunk, context, input_context)
-
         assert result[0].choices[0].delta.content == "Text with [email protected]"