From a66417128fb57dd365870faed4bcf33d2e71fc62 Mon Sep 17 00:00:00 2001 From: Jakub Hrozek Date: Thu, 12 Dec 2024 18:32:40 +0000 Subject: [PATCH] Notify about the number of secrets redacted since the last assistant message instead Copilot sends snippets as individual messages, which means that if a user asks about a snippet, there will be `[snippet, text]` and even though we had redacted the secret in snippet we wouldn't have counted it. Instead, let's count the secrets redacted since the last assistant (LLM) message. --- src/codegate/pipeline/secrets/secrets.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/codegate/pipeline/secrets/secrets.py b/src/codegate/pipeline/secrets/secrets.py index f12eda5f..fb874472 100644 --- a/src/codegate/pipeline/secrets/secrets.py +++ b/src/codegate/pipeline/secrets/secrets.py @@ -189,6 +189,12 @@ async def process( new_request = request.copy() total_redacted = 0 + # Process all messages + last_assistant_idx = -1 + for i, message in enumerate(new_request["messages"]): + if message.get("role", "") == "assistant": + last_assistant_idx = i + # Process all messages for i, message in enumerate(new_request["messages"]): if "content" in message and message["content"]: @@ -198,11 +204,11 @@ async def process( ) new_request["messages"][i]["content"] = protected_string - # only sum to the count if it is the last message - if i == len(new_request["messages"]) - 1: + # Sum redacted count for messages after the last assistant message + if i > last_assistant_idx: total_redacted += redacted_count - logger.info(f"Total secrets redacted: {total_redacted}") + logger.info(f"Total secrets redacted since last assistant message: {total_redacted}") # Store the count in context metadata context.metadata["redacted_secrets_count"] = total_redacted @@ -350,7 +356,7 @@ async def process_chunk( return [chunk] # Check if this is the first chunk (delta role will be present, others will not) - if chunk.choices[0].delta.role: + 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