fix: use the latest user messages block instead of single message - WIP

Depending on the tooling, the user request is splitted between several
user/assistant blocks. So use this logic instead of just picking the
latest user message one, to identify code snippets and secrets

Closes: #580

Author: yrobla

src/codegate/pipeline/base.py

@@ -231,8 +231,45 @@ def get_last_user_message(
             return None
         for i in reversed(range(len(request["messages"]))):
             if request["messages"][i]["role"] == "user":
-                content = request["messages"][i]["content"]
-                return content, i
+                content = request["messages"][i]["content"]  # type: ignore
+                return str(content), i
+
+        return None
+
+    @staticmethod
+    def get_last_user_message_block(
+        request: ChatCompletionRequest,
+    ) -> Optional[str]:
+        """
+        Get the last block of consecutive 'user' messages from the request.
+
+        Args:
+            request (ChatCompletionRequest): The chat completion request to process
+
+        Returns:
+            Optional[str]: A string containing all consecutive user messages in the
+                        last user message block, separated by newlines, or None if
+                        no user message block is found.
+        """
+        if request.get("messages") is None:
+            return None
+
+        user_messages = []
+        messages = request["messages"]
+
+        # Iterate in reverse to find the last block of consecutive 'user' messages
+        for i in reversed(range(len(messages))):
+            if messages[i]["role"] == "user" or messages[i]["role"] == "assistant":
+                if messages[i]["role"] == "user":
+                    user_messages.append(messages[i]["content"])  # type: ignore
+            else:
+                # Stop when a message with a different role is encountered
+                if user_messages:
+                    break
+
+        # Reverse the collected user messages to preserve the original order
+        if user_messages:
+            return "\n".join(reversed(user_messages))
 
         return None
 

src/codegate/pipeline/codegate_context_retriever/codegate.py

@@ -59,38 +59,36 @@ async def process(
         """
         Use RAG DB to add context to the user request
         """
-        # Get the latest user messages
-        user_messages = self.get_latest_user_messages(request)
-
-        # Nothing to do if the user_messages string is empty
-        if len(user_messages) == 0:
+        # Get the latest user message
+        user_message = self.get_last_user_message_block(request)
+        if not user_message:
             return PipelineResult(request=request)
 
         # Create storage engine object
         storage_engine = StorageEngine()
 
         # Extract any code snippets
-        snippets = extract_snippets(user_messages)
+        snippets = extract_snippets(user_message)
 
         bad_snippet_packages = []
         if len(snippets) > 0:
             # Collect all packages referenced in the snippets
             snippet_packages = []
             for snippet in snippets:
                 snippet_packages.extend(
-                    PackageExtractor.extract_packages(snippet.code, snippet.language)
+                    PackageExtractor.extract_packages(snippet.code, snippet.language)  # type: ignore
                 )
             logger.info(f"Found {len(snippet_packages)} packages in code snippets.")
 
             # Find bad packages in the snippets
             bad_snippet_packages = await storage_engine.search(
-                language=snippets[0].language, packages=snippet_packages
+                language=snippets[0].language, packages=snippet_packages  # 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
         # in the rest of the user query/messsages
-        user_messages = re.sub(r"```.*?```", "", user_messages, flags=re.DOTALL)
+        user_messages = re.sub(r"```.*?```", "", user_message, flags=re.DOTALL)
 
         # Vector search to find bad packages
         bad_packages = await storage_engine.search(query=user_messages, distance=0.5, limit=100)
@@ -119,7 +117,7 @@ async def process(
             # 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"]}'
+            context_msg = f'Context: {context_str} \n\n Query: {message["content"]}'  # type: ignore
             message["content"] = context_msg
 
             logger.debug("Final context message", context_message=context_msg)

src/codegate/pipeline/extract_snippets/extract_snippets.py

@@ -85,6 +85,7 @@ def extract_snippets(message: str) -> List[CodeSnippet]:
 
     # Find all code block matches
     for match in CODE_BLOCK_PATTERN.finditer(message):
+        print("i try to extract snippet")
         matched_language = match.group("language") if match.group("language") else None
         filename = match.group("filename") if match.group("filename") else None
         content = match.group("content")
@@ -94,18 +95,24 @@ def extract_snippets(message: str) -> List[CodeSnippet]:
         # format ` ```python ` in output snippets
         if filename and not matched_language and "." not in filename:
             lang = filename
+            print("lang is")
+            print(lang)
             filename = None
         else:
             # Determine language from the message, either by the short
             # language identifier or by the filename
             lang = None
             if matched_language:
+                print("i have a matched language")
                 lang = ecosystem_from_message(matched_language.strip())
             if lang is None and filename:
+                print("I try to get from filename")
                 filename = filename.strip()
                 # Determine language from the filename
                 lang = ecosystem_from_filepath(filename)
 
+        print("language is")
+        print(lang)
         snippets.append(CodeSnippet(filepath=filename, code=content, language=lang))
 
     return snippets
@@ -129,10 +136,9 @@ async def process(
         request: ChatCompletionRequest,
         context: PipelineContext,
     ) -> PipelineResult:
-        last_user_message = self.get_last_user_message(request)
-        if not last_user_message:
+        msg_content = self.get_last_user_message_block(request)
+        if not msg_content:
             return PipelineResult(request=request, context=context)
-        msg_content, _ = last_user_message
         snippets = extract_snippets(msg_content)
 
         logger.info(f"Extracted {len(snippets)} code snippets from the user message")

src/codegate/utils/package_extractor.py

@@ -74,16 +74,21 @@ class PackageExtractor:
 
     @staticmethod
     def extract_packages(code: str, language_name: str) -> list[str]:
+        print("packages are")
+        print(code)
+        print(language_name)
         if (code is None) or (language_name is None):
             return []
 
         language_name = language_name.lower()
 
         if language_name not in PackageExtractor.__languages.keys():
+            print("no langauge")
             return []
 
         language = PackageExtractor.__languages[language_name]
         parser = PackageExtractor.__parsers[language_name]
+        print("here")
 
         # Create tree
         tree = parser.parse(bytes(code, "utf8"))