Fixes to last_user_block (#1077)

* Improvements to get_last_user_message_block

* Fix test_messages_block

Author: jhrozek

src/codegate/pipeline/base.py

@@ -208,14 +208,12 @@ def get_last_user_message(
     @staticmethod
     def get_last_user_message_block(
         request: ChatCompletionRequest,
-        client: ClientType = ClientType.GENERIC,
     ) -> Optional[tuple[str, int]]:
         """
         Get the last block of consecutive 'user' messages from the request.
 
         Args:
             request (ChatCompletionRequest): The chat completion request to process
-            client (ClientType): The client type to consider when processing the request
 
         Returns:
             Optional[str, int]: A string containing all consecutive user messages in the

src/codegate/pipeline/codegate_context_retriever/codegate.py

@@ -75,7 +75,7 @@ async def process(  # noqa: C901
         Use RAG DB to add context to the user request
         """
         # Get the latest user message
-        last_message = self.get_last_user_message_block(request, context.client)
+        last_message = self.get_last_user_message_block(request)
         if not last_message:
             return PipelineResult(request=request)
         user_message, last_user_idx = last_message

src/codegate/pipeline/secrets/secrets.py

@@ -295,7 +295,7 @@ async def process(
         total_matches = []
 
         # get last user message block to get index for the first relevant user message
-        last_user_message = self.get_last_user_message_block(request, context.client)
+        last_user_message = self.get_last_user_message_block(request)
         last_assistant_idx = last_user_message[1] - 1 if last_user_message else -1
 
         # Process all messages

src/codegate/types/anthropic/_request_models.py

@@ -213,7 +213,6 @@ def last_user_block(self) -> Iterable[tuple[Message, int]]:
         for idx, msg in enumerate(reversed(self.messages)):
             if isinstance(msg, UserMessage):
                 yield  msg, len(self.messages) - 1 - idx
-                break
 
     def get_system_prompt(self) -> Iterable[str]:
         if isinstance(self.system, str):

src/codegate/types/ollama/_request_models.py

@@ -144,6 +144,11 @@ def last_user_block(self) -> Iterable[tuple[Message, int]]:
         for idx, msg in enumerate(reversed(self.messages)):
             if isinstance(msg, UserMessage):
                 yield msg, len(self.messages) - 1 - idx
+            elif isinstance(msg, (SystemMessage, ToolMessage)):
+                # these can occur in the middle of a user block
+                continue
+            elif isinstance(msg, AssistantMessage):
+                # these are LLM responses, end of user input, break on them
                 break
 
     def get_system_prompt(self) -> Iterable[str]:

src/codegate/types/openai/_request_models.py

@@ -341,6 +341,11 @@ def last_user_block(self) -> Iterable[tuple[Message, int]]:
         for idx, msg in enumerate(reversed(self.messages)):
             if isinstance(msg, UserMessage):
                 yield msg, len(self.messages) - 1 - idx
+            elif isinstance(msg, (SystemMessage, DeveloperMessage, ToolMessage)):
+                # these can occur in the middle of a user block
+                continue
+            elif isinstance(msg, (AssistantMessage, FunctionMessage)):
+                # these are LLM responses, end of user input, break on them
                 break
 
     def get_system_prompt(self) -> Iterable[str]: