Switch usage of re package for regex which is slightly more performant (#1127)

`regex` is a drop-in replacement of `re` and provides better
performance.

Signed-off-by: Juan Antonio Osorio <ozz@stacklok.com>

Author: JAORMX

poetry.lock

@@ -40,6 +40,7 @@ onnxruntime = "==1.20.1"
 onnx = "==1.17.0"
 spacy = "<3.8.0"
 en-core-web-sm = {url = "https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl"}
+regex = "==2024.11.6"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "==8.3.4"

pyproject.toml

@@ -1,10 +1,10 @@
 import asyncio
 import json
-import re
 from collections import defaultdict
 from typing import AsyncGenerator, Dict, List, Optional, Tuple
 
 import cachetools.func
+import regex as re
 import requests
 import structlog
 

src/codegate/api/v1_processing.py

@@ -1,8 +1,8 @@
-import re
 from abc import ABC, abstractmethod
 from functools import wraps
 from typing import List, Optional
 
+import regex as re
 import structlog
 from fastapi import Request
 

src/codegate/clients/detector.py

@@ -1,9 +1,9 @@
 import datetime
 import hashlib
 import json
-import re
 from typing import Dict, List, Optional
 
+import regex as re
 import structlog
 from pydantic import BaseModel
 
@@ -21,6 +21,11 @@ class CachedFim(BaseModel):
     initial_id: str
 
 
+# Regular expression to match file paths in FIM messages.
+# Compiled regex to improve performance.
+filepath_matcher = re.compile(r"^(#|//|<!--|--|%|;).*?\b([a-zA-Z0-9_\-\/]+\.\w+)\b", re.MULTILINE)
+
+
 class FimCache:
 
     def __init__(self):
@@ -55,8 +60,8 @@ def _match_filepath(self, message: str, provider: str) -> Optional[str]:
         # folder/testing_file.py
         # Path: file3.py
         # // Path: file3.js <-- Javascript
-        pattern = r"^(#|//|<!--|--|%|;).*?\b([a-zA-Z0-9_\-\/]+\.\w+)\b"
-        matches = re.findall(pattern, message, re.MULTILINE)
+        matches = filepath_matcher.findall(message)
+
         # If no path is found, hash the entire prompt message.
         if not matches:
             return None

src/codegate/db/fim_cache.py

@@ -1,8 +1,8 @@
-import re
 from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Dict, List, Optional, Self
 
+import regex as re
 import structlog
 from pydantic import BaseModel, field_validator, model_validator
 from pygments.lexers import guess_lexer

src/codegate/extract_snippets/message_extractor.py

@@ -1,7 +1,7 @@
-import re
 import shlex
 from typing import Optional
 
+import regex as re
 from litellm import ChatCompletionRequest
 
 from codegate.clients.clients import ClientType

src/codegate/pipeline/cli/cli.py

@@ -1,6 +1,6 @@
 import json
-import re
 
+import regex as re
 import structlog
 from litellm import ChatCompletionRequest
 
@@ -19,6 +19,12 @@
 logger = structlog.get_logger("codegate")
 
 
+# Pre-compiled regex patterns for performance
+markdown_code_block = re.compile(r"```.*?```", flags=re.DOTALL)
+markdown_file_listing = re.compile(r"⋮...*?⋮...\n\n", flags=re.DOTALL)
+environment_details = re.compile(r"<environment_details>.*?</environment_details>", flags=re.DOTALL)
+
+
 class CodegateContextRetriever(PipelineStep):
     """
     Pipeline step that adds a context message to the completion request when it detects
@@ -95,11 +101,9 @@ async def process(  # noqa: C901
 
         # 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
-        )
+        user_messages = markdown_code_block.sub("", user_message)
+        user_messages = markdown_file_listing.sub("", user_messages)
+        user_messages = environment_details.sub("", user_messages)
 
         # split messages into double newlines, to avoid passing so many content in the search
         split_messages = re.split(r"</?task>|\n|\\n", user_messages)

src/codegate/pipeline/codegate_context_retriever/codegate.py

@@ -1,6 +1,6 @@
-import re
 from typing import Any, Dict, List, Optional
 
+import regex as re
 import structlog
 from litellm import ChatCompletionRequest, ChatCompletionSystemMessage, ModelResponse
 from litellm.types.utils import Delta, StreamingChoices

src/codegate/pipeline/pii/pii.py

@@ -1,7 +1,7 @@
-import re
 from abc import abstractmethod
 from typing import List, Optional, Tuple
 
+import regex as re
 import structlog
 from litellm import ChatCompletionRequest, ChatCompletionSystemMessage, ModelResponse
 from litellm.types.utils import Delta, StreamingChoices