Fixed mistral formatter

Maximilian-Winter · May 14, 2024 · 3a8234a · 3a8234a
1 parent a5c5c09
commit 3a8234a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/src/llama_cpp_agent/llm_agent.py b/src/llama_cpp_agent/llm_agent.py
@@ -396,7 +396,7 @@ def get_response_role_and_completion(
         additional_suffix = ""
         if self.add_tools_and_structures_documentation_to_system_prompt:
             if structured_output_settings.output_type != LlmStructuredOutputType.no_structured_output:
-                additional_suffix = "\n"
+                # additional_suffix = "\n"
                 if structured_output_settings.output_type == LlmStructuredOutputType.function_calling or structured_output_settings.output_type == LlmStructuredOutputType.parallel_function_calling:
                     if structured_output_settings.add_thoughts_and_reasoning_field and self.provider.is_using_json_schema_constraints():
                         messages[0][

diff --git a/src/llama_cpp_agent/messages_formatter.py b/src/llama_cpp_agent/messages_formatter.py
@@ -98,6 +98,7 @@ def _format_system_message(self, content: str) -> str:
 
     def _format_user_message(self, content: str) -> str:
         if self.include_sys_prompt_in_first_user_message and self.added_system_prompt:
+            self.added_system_prompt = False
             return content + self.prompt_markers[Roles.user].end
         return self.prompt_markers[Roles.user].start + content + self.prompt_markers[Roles.user].end
 
@@ -128,8 +129,8 @@ def _format_response(
 
 mixtral_prompt_markers = {
     Roles.system: PromptMarkers("", """\n\n"""),
-    Roles.user: PromptMarkers("""[INST] """, """ """),
-    Roles.assistant: PromptMarkers("""[/INST]""", """</s>"""),
+    Roles.user: PromptMarkers("""[INST] """, """ [/INST]"""),
+    Roles.assistant: PromptMarkers("""""", """</s>"""),
     Roles.tool: PromptMarkers("", ""),
 }