fix: fix prompt builder

src/crllm/model/model.py

@@ -21,11 +21,11 @@ def generate(self, prompt_template, prompt_args):
         if config["rag"]["enabled"]:
             self.add_rag_context(prompt_args, rag_config)
 
-        logging.debug(
-            prompt_template.format(
-                context=prompt_args["context"], code=prompt_args["code"]
+            logging.debug(
+                prompt_template.format(
+                    context=prompt_args["context"], code=prompt_args["code"]
+                )
             )
-        )
 
         result = chain.invoke(prompt_args)
 

src/crllm/prompt_builder.py

@@ -9,20 +9,26 @@ def build(self):
 
         logging.debug(config["prompt"])
 
-        prompt = ChatPromptTemplate.from_messages(
-            [
-                ("system", config["prompt"]["template"]),
-                (
-                    "system",
-                    """
+        messages = [
+            ("system", config["prompt"]["template"]),
+            ("human", "Do code review for the following code changes: \n {code}"),
+        ]
+
+        if config["rag"]["enabled"]:
+            self._insert_rag_message(messages)
+
+        return ChatPromptTemplate.from_messages(messages)
+
+    def _insert_rag_message(self, messages):
+        messages.insert(
+            1,
+            (
+                "system",
+                """
 This code could be relevant to understand the changes:
 <context>
     {context}
 </context>
-                 """,
-                ),
-                ("human", "Do code review for the following code changes: \n {code}"),
-            ]
+                    """,
+            ),
         )
-
-        return prompt