Move to user message vs system message

libs/core/kiln_ai/adapters/base_adapter.py

@@ -42,5 +42,9 @@ def __init__(self, task: Task, adapter: BaseAdapter | None = None):
         self.adapter = adapter
 
     @abstractmethod
-    def build_prompt(self, input: str) -> str:
+    def build_prompt(self) -> str:
+        pass
+
+    @abstractmethod
+    def build_user_message(self, input: str) -> str:
         pass

libs/core/kiln_ai/adapters/langchain_adapters.py

@@ -3,6 +3,7 @@
 import kiln_ai.datamodel.models as models
 from kiln_ai.adapters.prompt_builders import SimplePromptBuilder
 from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_core.messages.base import BaseMessage
 
 from .base_adapter import BaseAdapter, BasePromptBuilder
@@ -58,9 +59,13 @@ def adapter_specific_instructions(self) -> str | None:
         return None
 
     async def _run(self, input: str) -> Dict | str:
-        # TODO cleanup
-        prompt = self.prompt_builder.build_prompt(input)
-        response = self.model.invoke(prompt)
+        prompt = self.prompt_builder.build_prompt()
+        user_msg = self.prompt_builder.build_user_message(input)
+        messages = [
+            SystemMessage(content=prompt),
+            HumanMessage(content=user_msg),
+        ]
+        response = self.model.invoke(messages)
         if self._is_structured:
             if (
                 not isinstance(response, dict)

libs/core/kiln_ai/adapters/prompt_builders.py

@@ -2,7 +2,7 @@
 
 
 class SimplePromptBuilder(BasePromptBuilder):
-    def build_prompt(self, input: str) -> str:
+    def build_prompt(self) -> str:
         base_prompt = self.task.instruction
 
         # TODO: this is just a quick version. Formatting and best practices TBD
@@ -19,6 +19,7 @@ def build_prompt(self, input: str) -> str:
             if adapter_instructions is not None:
                 base_prompt += f"\n\n{adapter_instructions}\n\n"
 
-        # TODO: should be another message, not just appended to prompt
-        base_prompt += f"\n\nThe input is:\n{input}"
         return base_prompt
+
+    def build_user_message(self, input: str) -> str:
+        return f"The input is:\n{input}"

libs/core/kiln_ai/adapters/test_prompt_builders.py

@@ -8,19 +8,20 @@ def test_simple_prompt_builder(tmp_path):
     task = build_test_task(tmp_path)
     builder = SimplePromptBuilder(task=task)
     input = "two plus two"
-    prompt = builder.build_prompt(input)
+    prompt = builder.build_prompt()
     assert (
         "You are an assistant which performs math tasks provided in plain text."
         in prompt
     )
 
-    # TODO this should be a user message later
-    assert input in prompt
-
     assert "1) " + task.requirements()[0].instruction in prompt
     assert "2) " + task.requirements()[1].instruction in prompt
     assert "3) " + task.requirements()[2].instruction in prompt
 
+    user_msg = builder.build_user_message(input)
+    assert input in user_msg
+    assert input not in prompt
+
 
 class MockAdapter(BaseAdapter):
     def adapter_specific_instructions(self) -> str | None:
@@ -35,11 +36,12 @@ def test_simple_prompt_builder_structured_output(tmp_path):
     builder = SimplePromptBuilder(task=task)
     builder.adapter = MockAdapter(task)
     input = "Cows"
-    prompt = builder.build_prompt(input)
+    prompt = builder.build_prompt()
     assert "You are an assistant which tells a joke, given a subject." in prompt
 
-    # TODO this should be a user message later
-    assert input in prompt
-
     # check adapter instructions are included
     assert "You are a mock, send me the response!" in prompt
+
+    user_msg = builder.build_user_message(input)
+    assert input in user_msg
+    assert input not in prompt

libs/core/kiln_ai/adapters/test_structured_output.py

@@ -24,7 +24,7 @@ async def test_structured_output_openrouter(tmp_path):
 
 @pytest.mark.paid
 async def test_structured_output_bedrock(tmp_path):
-    await run_structured_output_test(tmp_path, "llama_3_1_8b", "amazon_bedrock")
+    await run_structured_output_test(tmp_path, "llama_3_1_70b", "amazon_bedrock")
 
 
 @pytest.mark.ollama