from llama_cpp_agent.messages_formatter import MessagesFormatter, PromptMarkers, Roles

mistral_v1_markers = {
    Roles.system: PromptMarkers(""" [INST]""", """ [/INST]"""),
    Roles.user: PromptMarkers(""" [INST]""", """ [/INST]"""),
    Roles.assistant: PromptMarkers(""" """, """</s>"""),
    Roles.tool: PromptMarkers("", ""),
}

mistral_v1_formatter = MessagesFormatter(
    "",
    mistral_v1_markers,
    False,
    ["</s>"],
)

mistral_v2_markers = {
    Roles.system: PromptMarkers("""[INST] """, """[/INST]"""),
    Roles.user: PromptMarkers("""[INST] """, """[/INST]"""),
    Roles.assistant: PromptMarkers(""" """, """</s>"""),
    Roles.tool: PromptMarkers("", ""),
}

mistral_v2_formatter = MessagesFormatter(
    "",
    mistral_v2_markers,
    False,
    ["</s>"],
)

mistral_v3_tekken_markers = {
    Roles.system: PromptMarkers("""[INST]""", """[/INST]"""),
    Roles.user: PromptMarkers("""[INST]""", """[/INST]"""),
    Roles.assistant: PromptMarkers("""""", """</s>"""),
    Roles.tool: PromptMarkers("", ""),
}

mistral_v3_tekken_formatter = MessagesFormatter(
    "",
    mistral_v3_tekken_markers,
    False,
    ["</s>"],
)