From 0c2ddc0e96128574adf0a329d16e5fb354298b5b Mon Sep 17 00:00:00 2001
From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
Date: Tue, 8 Oct 2024 21:29:23 -0400
Subject: [PATCH] Add flag for removing all system messages

---
 src/controlflow/events/message_compiler.py | 23 +++++++++++++---------
 src/controlflow/llm/rules.py               |  3 +++
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/src/controlflow/events/message_compiler.py b/src/controlflow/events/message_compiler.py
index 6ed4d54..4dc9417 100644
--- a/src/controlflow/events/message_compiler.py
+++ b/src/controlflow/events/message_compiler.py
@@ -113,19 +113,22 @@ def convert_system_messages(
     messages: list[BaseMessage], rules: LLMRules
 ) -> list[BaseMessage]:
     """
-    Converts system messages to human messages if the LLM doesnt support system messages.
+    Converts system messages to human messages if the LLM doesnt support system messages, either at all or in the first position.
     """
-    if not messages or not rules.require_system_message_first:
-        return messages
 
     new_messages = []
-    for message in messages:
+    for i, message in enumerate(messages):
         if isinstance(message, SystemMessage):
-            new_messages.append(
-                HumanMessage(
-                    content=f"ORCHESTRATOR: {message.content}", name=message.name
+            # if system messages are not supported OR if they must be first and this is not the first message
+            # THEN convert the message to a human message
+            if not rules.allow_system_messages or (
+                i > 0 and rules.require_system_message_first
+            ):
+                new_messages.append(
+                    HumanMessage(
+                        content=f"ORCHESTRATOR: {message.content}", name=message.name
+                    )
                 )
-            )
         else:
             new_messages.append(message)
     return new_messages
@@ -249,7 +252,9 @@ def compile_to_messages(self, agent: "Agent") -> list[BaseMessage]:
         messages = break_up_consecutive_ai_messages(messages, rules=context.llm_rules)
         messages = format_message_name(messages, rules=context.llm_rules)
 
+        messages = system_prompt + messages
+
         # this should go last
         messages = convert_system_messages(messages, rules=context.llm_rules)
 
-        return system_prompt + messages
+        return messages
diff --git a/src/controlflow/llm/rules.py b/src/controlflow/llm/rules.py
index 21f7596..fed84d9 100644
--- a/src/controlflow/llm/rules.py
+++ b/src/controlflow/llm/rules.py
@@ -22,6 +22,9 @@ class LLMRules(ControlFlowModel):
     # require at least one non-system message
     require_at_least_one_message: bool = False
 
+    # system messages are supported as a role
+    allow_system_messages: bool = True
+
     # system messages can only be provided as the very first message in a thread
     require_system_message_first: bool = False