From 88612632119ef394d1af05c9a07a4545024b0743 Mon Sep 17 00:00:00 2001
From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
Date: Sat, 5 Oct 2024 09:17:47 -0400
Subject: [PATCH] Add LLM rules for custom agents

---
 docs/concepts/agents.mdx        | 28 ++++++++++++++++++++++++++++
 src/controlflow/agents/agent.py |  9 ++++++++-
 src/controlflow/llm/rules.py    | 15 ++++++++++++---
 tests/agents/test_agents.py     | 19 ++++++++++++++++++-
 4 files changed, 66 insertions(+), 5 deletions(-)

diff --git a/docs/concepts/agents.mdx b/docs/concepts/agents.mdx
index 96240c7b..86e07a97 100644
--- a/docs/concepts/agents.mdx
+++ b/docs/concepts/agents.mdx
@@ -69,6 +69,34 @@ Each agent has a model, which is the LLM that powers the agent responses. This a
 
 ControlFlow supports any LangChain LLM that supports chat and function calling. For more details on how to configure models, see the [LLMs guide](/guides/configure-llms).
 
+```python
+import controlflow as cf
+
+
+agent1 = cf.Agent(model="openai/gpt-4o")
+agent2 = cf.Agent(model="anthropic/claude-3-5-sonnet-20240620")
+```
+
+### LLM rules
+
+Each LLM provider may have different requirements for how messages are formatted or presented. For example, OpenAI permits system messages to be interspersed between user messages, but Anthropic requires them to be at the beginning of the conversation. ControlFlow uses provider-specific rules to properly compile messages for each agent's API. 
+
+For common providers like OpenAI and Anthropic, LLM rules can be automatically inferred from the agent's model. However, you can use a custom `LLMRules` object to override these rules or provide rules for non-standard providers.
+
+Here is an example of how to tell the agent to use the Anthropic compilation rules with a custom model that can't be automatically inferred:
+
+```python
+import controlflow as cf
+
+# note: this is just an example
+llm_model = CustomAnthropicModel()
+
+agent = cf.Agent(
+    model=model,
+    llm_rules=cf.llm.rules.AnthropicRules(model=model)
+)
+```
+
 ### Interactivity
 
 By default, agents have no way to communicate with users. If you want to chat with an agent, set `interactive=True`. By default, this will let the agent communicate with users on the command line.
diff --git a/src/controlflow/agents/agent.py b/src/controlflow/agents/agent.py
index fb98d1b8..8d088eaf 100644
--- a/src/controlflow/agents/agent.py
+++ b/src/controlflow/agents/agent.py
@@ -81,6 +81,10 @@ class Agent(ControlFlowModel, abc.ABC):
         description="The LangChain BaseChatModel used by the agent. If not provided, the default model will be used. A compatible string can be passed to automatically retrieve the model.",
         exclude=True,
     )
+    llm_rules: Optional[LLMRules] = Field(
+        None,
+        description="The LLM rules for the agent. If not provided, the rules will be inferred from the model (if possible).",
+    )
 
     _cm_stack: list[contextmanager] = []
 
@@ -164,7 +168,10 @@ def get_llm_rules(self) -> LLMRules:
         """
         Retrieve the LLM rules for this agent's model
         """
-        return controlflow.llm.rules.rules_for_model(self.get_model())
+        if self.llm_rules is None:
+            return controlflow.llm.rules.rules_for_model(self.get_model())
+        else:
+            return self.llm_rules
 
     def get_tools(self) -> list["Tool"]:
         from controlflow.tools.input import cli_input
diff --git a/src/controlflow/llm/rules.py b/src/controlflow/llm/rules.py
index 3c54e1ff..f8626ff6 100644
--- a/src/controlflow/llm/rules.py
+++ b/src/controlflow/llm/rules.py
@@ -80,7 +80,16 @@ class AnthropicRules(LLMRules):
 def rules_for_model(model: BaseChatModel) -> LLMRules:
     if isinstance(model, (ChatOpenAI, AzureChatOpenAI)):
         return OpenAIRules(model=model)
-    elif isinstance(model, ChatAnthropic):
+    if isinstance(model, ChatAnthropic):
         return AnthropicRules(model=model)
-    else:
-        return LLMRules(model=model)
+
+    try:
+        from langchain_google_vertexai.model_garden import ChatAnthropicVertex
+
+        if isinstance(model, ChatAnthropicVertex):
+            return AnthropicRules(model=model)
+    except ImportError:
+        pass
+
+    # catchall
+    return LLMRules(model=model)
diff --git a/tests/agents/test_agents.py b/tests/agents/test_agents.py
index 274dfcd4..cafe7343 100644
--- a/tests/agents/test_agents.py
+++ b/tests/agents/test_agents.py
@@ -7,7 +7,7 @@
 from controlflow.events.base import Event
 from controlflow.events.events import AgentMessage
 from controlflow.instructions import instructions
-from controlflow.llm.rules import LLMRules
+from controlflow.llm.rules import AnthropicRules, LLMRules, OpenAIRules
 from controlflow.orchestration.handler import Handler
 from controlflow.tasks.task import Task
 
@@ -198,3 +198,20 @@ async def test_agent_run_async_with_handlers(self):
 
         assert len(handler.events) > 0
         assert len(handler.agent_messages) == 1
+
+
+class TestLLMRules:
+    def test_llm_rules_from_model_openai(self):
+        agent = Agent(model=ChatOpenAI(model="gpt-4o-mini"))
+        rules = agent.get_llm_rules()
+        assert isinstance(rules, OpenAIRules)
+
+    def test_llm_rules_from_model_anthropic(self):
+        agent = Agent(model=ChatAnthropic(model="claude-3-haiku-20240307"))
+        rules = agent.get_llm_rules()
+        assert isinstance(rules, AnthropicRules)
+
+    def test_custom_llm_rules(self):
+        rules = LLMRules(model=None)
+        agent = Agent(llm_rules=rules, model=ChatOpenAI(model="gpt-4o-mini"))
+        assert agent.get_llm_rules() is rules