reuse logic for run mode, add tests

Kiln-AI · Feb 1, 2025 · 6d5980e · 6d5980e
1 parent ca6672d
commit 6d5980e
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/libs/core/kiln_ai/adapters/model_adapters/langchain_adapters.py b/libs/core/kiln_ai/adapters/model_adapters/langchain_adapters.py
@@ -144,16 +144,17 @@ async def _run(self, input: Dict | str) -> RunOutput:
         if run_strategy == "cot_as_message":
             if not cot_prompt:
                 raise ValueError("cot_prompt is required for cot_as_message strategy")
-            messages.append({"role": "system", "content": cot_prompt})
+            messages.append(SystemMessage(content=cot_prompt))
         elif run_strategy == "cot_two_call":
             if not cot_prompt:
                 raise ValueError("cot_prompt is required for cot_two_call strategy")
-            # Base model (without structured output) used for COT message
-            base_model = await self.langchain_model_from()
             messages.append(
                 SystemMessage(content=cot_prompt),
             )
 
+            # Base model (without structured output) used for COT message
+            base_model = await self.langchain_model_from()
+
             cot_messages = [*messages]
             cot_response = await base_model.ainvoke(cot_messages)
             intermediate_outputs["chain_of_thought"] = cot_response.content