Update the integration_tests.py to extract correctly FIM repsonses

Signed-off-by: Radoslav Dimitrov <[email protected]>

Author: rdimitrov

src/codegate/providers/llamacpp/completion_handler.py

@@ -60,7 +60,8 @@ async def execute_completion(
         model_path = f"{Config.get_config().model_base_path}/{request['model']}.gguf"
 
         # Create a copy of the request dict and remove stream_options
-        # Reason - Request error as JSON: {'error': "Llama.create_completion() got an unexpected keyword argument 'stream_options'"}
+        # Reason - Request error as JSON:
+        # {'error': "Llama.create_completion() got an unexpected keyword argument 'stream_options'"}
         request_dict = dict(request)
         request_dict.pop("stream_options", None)
 

tests/integration/integration_tests.py

@@ -67,17 +67,25 @@ def parse_response_message(response, streaming=True):
                     if "DONE" in decoded_line or "message_stop" in decoded_line:
                         break
 
-                    decoded_line = decoded_line.replace("data:", "")
+                    decoded_line = decoded_line.replace("data:", "").strip()
                     json_line = json.loads(decoded_line)
-
                     message_content = None
+                    # Handle both chat and FIM responses
                     if "choices" in json_line:
-                        if "finish_reason" in json_line["choices"][0]:
+                        choice = json_line["choices"][0]
+                        # Break if the conversation is over
+                        if choice.get("finish_reason") == "stop":
                             break
-                        if "delta" in json_line["choices"][0]:
-                            message_content = json_line["choices"][0]["delta"].get("content", "")
-                        elif "text" in json_line["choices"][0]:
-                            message_content = json_line["choices"][0].get("text", "")
+                        # Handle chat responses
+                        if "delta" in choice:
+                            delta = choice["delta"]
+                            if "content" in delta and delta["content"] is not None:
+                                message_content = delta["content"]
+                        # Handle FIM responses
+                        elif "text" in choice:
+                            text = choice["text"]
+                            if text is not None:
+                                message_content = text
                     elif "delta" in json_line:
                         message_content = json_line["delta"].get("text", "")
                     elif "message" in json_line:
@@ -87,7 +95,6 @@ def parse_response_message(response, streaming=True):
 
                     if message_content is not None:
                         response_message += message_content
-
             else:
                 if "choices" in response.json():
                     response_message = response.json()["choices"][0]["message"].get("content", "")
@@ -97,7 +104,8 @@ def parse_response_message(response, streaming=True):
         except Exception as e:
             logger.exception("An error occurred: %s", e)
 
-        return response_message
+        # Remove any trailing newlines and return
+        return response_message.strip()
 
     @staticmethod
     def replace_env_variables(input_string, env):