fix secret redaction

Author: yrobla

src/codegate/config.py

@@ -20,6 +20,7 @@
     "anthropic": "https://api.anthropic.com/v1",
     "vllm": "http://localhost:8000",  # Base URL without /v1 path
     "ollama": "http://localhost:11434",  # Default Ollama server URL
+    "lm_studio": "http://localhost:1234"
 }
 
 

src/codegate/pipeline/base.py

@@ -321,8 +321,7 @@ async def process(
 
 class InputPipelineInstance:
     def __init__(
-        self, pipeline_steps: List[PipelineStep], secret_manager: SecretsManager, is_fim: bool
-    ):
+        self, pipeline_steps: List[PipelineStep], secret_manager: SecretsManager, is_fim: bool):
         self.pipeline_steps = pipeline_steps
         self.secret_manager = secret_manager
         self.is_fim = is_fim
@@ -385,8 +384,7 @@ async def process_request(
 
 class SequentialPipelineProcessor:
     def __init__(
-        self, pipeline_steps: List[PipelineStep], secret_manager: SecretsManager, is_fim: bool
-    ):
+        self, pipeline_steps: List[PipelineStep], secret_manager: SecretsManager, is_fim: bool):
         self.pipeline_steps = pipeline_steps
         self.secret_manager = secret_manager
         self.is_fim = is_fim

src/codegate/pipeline/codegate_context_retriever/codegate.py

@@ -91,10 +91,11 @@ async def process(
             )  # type: ignore
             logger.info(f"Found {len(bad_snippet_packages)} bad packages in code snippets.")
 
-        # Remove code snippets from the user messages and search for bad packages
+        # Remove code snippets and file listing from the user messages and search for bad packages
         # in the rest of the user query/messsages
         user_messages = re.sub(r"```.*?```", "", user_message, flags=re.DOTALL)
         user_messages = re.sub(r"⋮...*?⋮...\n\n", "", user_messages, flags=re.DOTALL)
+        user_messages = re.sub(r"<environment_details>.*?</environment_details>", "", user_messages, flags=re.DOTALL)
 
         # split messages into double newlines, to avoid passing so many content in the search
         split_messages = re.split(r'</?task>|(\n\n)', user_messages)
@@ -126,10 +127,23 @@ async def process(
             # Make a copy of the request
             new_request = request.copy()
 
-            # Add the context to the last user message
             # Format: "Context: {context_str} \n Query: {last user message content}"
             message = new_request["messages"][last_user_idx]
-            context_msg = f'Context: {context_str} \n\n Query: {message["content"]}'  # type: ignore
+            message_str = str(message["content"])  # type: ignore
+            # Add the context to the last user message
+            if message_str.strip().startswith("<task>"):
+                # formatting of cline
+                match = re.match(r"(<task>)(.*?)(</task>)(.*)", message_str, re.DOTALL)
+                if match:
+                    task_start, task_content, task_end, rest_of_message = match.groups()
+
+                # Embed the context into the task block
+                updated_task_content = f"{task_start}Context: {context_str}\nQuery: {task_content.strip()}</details>{task_end}"
+
+                # Combine the updated task block with the rest of the message
+                context_msg = updated_task_content + rest_of_message
+            else:
+                context_msg = f'Context: {context_str} \n\n Query: {message_str}'  # type: ignore
             message["content"] = context_msg
 
             logger.debug("Final context message", context_message=context_msg)

src/codegate/pipeline/secrets/secrets.py

@@ -451,17 +451,30 @@ async def process_chunk(
         ):
             return [chunk]
 
+        is_cline_client = any(
+            "Cline" in str(message.trigger_string or "") for message in input_context.alerts_raised or []
+        )
+
         # Check if this is the first chunk (delta role will be present, others will not)
         if len(chunk.choices) > 0 and chunk.choices[0].delta.role:
             redacted_count = input_context.metadata["redacted_secrets_count"]
             secret_text = "secret" if redacted_count == 1 else "secrets"
             # Create notification chunk
-            notification_chunk = self._create_chunk(
-                chunk,
-                f"\n🛡️ [CodeGate prevented {redacted_count} {secret_text}]"
-                f"(http://localhost:9090/?search=codegate-secrets) from being leaked "
-                f"by redacting them.\n\n",
-            )
+            if is_cline_client:
+                notification_chunk = self._create_chunk(
+                    chunk,
+                    f"<thinking>\n🛡️ [CodeGate prevented {redacted_count} {secret_text}]"
+                    f"(http://localhost:9090/?search=codegate-secrets) from being leaked "
+                    f"by redacting them.</thinking>\n\n",
+                )
+                notification_chunk.choices[0].delta.role = "assistant"
+            else:
+                notification_chunk = self._create_chunk(
+                    chunk,
+                    f"\n🛡️ [CodeGate prevented {redacted_count} {secret_text}]"
+                    f"(http://localhost:9090/?search=codegate-secrets) from being leaked "
+                    f"by redacting them.\n\n",
+                )
 
             # Reset the counter
             input_context.metadata["redacted_secrets_count"] = 0

src/codegate/providers/base.py

@@ -199,8 +199,7 @@ async def _cleanup_after_streaming(
                     context.sensitive.secure_cleanup()
 
     async def complete(
-        self, data: Dict, api_key: Optional[str], is_fim_request: bool
-    ) -> Union[ModelResponse, AsyncIterator[ModelResponse]]:
+        self, data: Dict, api_key: Optional[str], is_fim_request: bool) -> Union[ModelResponse, AsyncIterator[ModelResponse]]:
         """
         Main completion flow with pipeline integration
 
@@ -220,20 +219,21 @@ async def complete(
             data.get("base_url"),
             is_fim_request,
         )
-        if input_pipeline_result.response:
+        if input_pipeline_result.response and input_pipeline_result.context:
             return await self._pipeline_response_formatter.handle_pipeline_response(
                 input_pipeline_result.response, streaming, context=input_pipeline_result.context
             )
 
-        provider_request = self._input_normalizer.denormalize(input_pipeline_result.request)
+        if input_pipeline_result.request:
+            provider_request = self._input_normalizer.denormalize(input_pipeline_result.request)
         if is_fim_request:
-            provider_request = self._fim_normalizer.denormalize(provider_request)
+            provider_request = self._fim_normalizer.denormalize(provider_request)  # type: ignore
 
         # Execute the completion and translate the response
         # This gives us either a single response or a stream of responses
         # based on the streaming flag
         model_response = await self._completion_handler.execute_completion(
-            provider_request, api_key=api_key, stream=streaming, is_fim_request=is_fim_request
+            provider_request, api_key=api_key, stream=streaming, is_fim_request=is_fim_request  # type: ignore
         )
         if not streaming:
             normalized_response = self._output_normalizer.normalize(model_response)
@@ -242,9 +242,9 @@ async def complete(
             return self._output_normalizer.denormalize(pipeline_output)
 
         pipeline_output_stream = await self._run_output_stream_pipeline(
-            input_pipeline_result.context, model_response, is_fim_request=is_fim_request
+            input_pipeline_result.context, model_response, is_fim_request=is_fim_request  # type: ignore
         )
-        return self._cleanup_after_streaming(pipeline_output_stream, input_pipeline_result.context)
+        return self._cleanup_after_streaming(pipeline_output_stream, input_pipeline_result.context)  # type: ignore
 
     def get_routes(self) -> APIRouter:
         return self.router

src/codegate/providers/ollama/completion_handler.py

@@ -63,7 +63,6 @@ async def ollama_stream_generator(
                     for field in optional_fields:
                         if field in chunk_dict:
                             response[field] = chunk_dict[field]
-
                     yield f"data: {json.dumps(response)}\n"
             except Exception as e:
                 logger.error(f"Error in stream generator: {str(e)}")

src/codegate/providers/ollama/provider.py

@@ -80,6 +80,8 @@ async def show_model(request: Request):
         @self.router.post(f"/{self.provider_route_name}/v1/generate")
         async def create_completion(request: Request):
             body = await request.body()
+            print("i request")
+            print(body)
             data = json.loads(body)
             # `base_url` is used in the providers pipeline to do the packages lookup.
             # Force it to be the one that comes in the configuration.

src/codegate/providers/openai/provider.py

@@ -4,6 +4,7 @@
 import structlog
 from fastapi import Header, HTTPException, Request
 
+from codegate.config import Config
 from codegate.pipeline.factory import PipelineFactory
 from codegate.providers.base import BaseProvider
 from codegate.providers.litellmshim import LiteLLmShim, sse_stream_generator
@@ -16,6 +17,11 @@ def __init__(
         pipeline_factory: PipelineFactory,
     ):
         completion_handler = LiteLLmShim(stream_generator=sse_stream_generator)
+        config = Config.get_config()
+        if config is not None:
+            provided_urls = config.provider_urls
+            self.lm_studio_url = provided_urls.get("lm_studio", "http://localhost:11434/")
+
         super().__init__(
             OpenAIInputNormalizer(),
             OpenAIOutputNormalizer(),
@@ -47,6 +53,10 @@ async def create_completion(
             api_key = authorization.split(" ")[1]
             body = await request.body()
             data = json.loads(body)
+
+            # if model starts with lm_studio, propagate it
+            if data.get("model", "").startswith("lm_studio"):
+                data["base_url"] = self.lm_studio_url+"/v1/"
             is_fim_request = self._is_fim_request(request, data)
             try:
                 stream = await self.complete(data, api_key, is_fim_request=is_fim_request)