From 6416c781a03ff8a28c24003d04b0b3c026d2934f Mon Sep 17 00:00:00 2001 From: scosman Date: Mon, 9 Sep 2024 22:31:57 -0400 Subject: [PATCH] Move to user message vs system message --- libs/core/kiln_ai/adapters/base_adapter.py | 6 +++++- .../kiln_ai/adapters/langchain_adapters.py | 11 ++++++++--- libs/core/kiln_ai/adapters/prompt_builders.py | 7 ++++--- .../kiln_ai/adapters/test_prompt_builders.py | 18 ++++++++++-------- .../kiln_ai/adapters/test_structured_output.py | 2 +- 5 files changed, 28 insertions(+), 16 deletions(-) diff --git a/libs/core/kiln_ai/adapters/base_adapter.py b/libs/core/kiln_ai/adapters/base_adapter.py index 2e26e8e..bf47b97 100644 --- a/libs/core/kiln_ai/adapters/base_adapter.py +++ b/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 diff --git a/libs/core/kiln_ai/adapters/langchain_adapters.py b/libs/core/kiln_ai/adapters/langchain_adapters.py index 2383bca..0d57f11 100644 --- a/libs/core/kiln_ai/adapters/langchain_adapters.py +++ b/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) diff --git a/libs/core/kiln_ai/adapters/prompt_builders.py b/libs/core/kiln_ai/adapters/prompt_builders.py index 5d1ffb8..45e4fb8 100644 --- a/libs/core/kiln_ai/adapters/prompt_builders.py +++ b/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}" diff --git a/libs/core/kiln_ai/adapters/test_prompt_builders.py b/libs/core/kiln_ai/adapters/test_prompt_builders.py index 8ede68b..f151185 100644 --- a/libs/core/kiln_ai/adapters/test_prompt_builders.py +++ b/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 diff --git a/libs/core/kiln_ai/adapters/test_structured_output.py b/libs/core/kiln_ai/adapters/test_structured_output.py index ed00b12..87d012d 100644 --- a/libs/core/kiln_ai/adapters/test_structured_output.py +++ b/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