kiln-ai 0.5.5__py3-none-any.whl → 0.6.1__py3-none-any.whl

kiln_ai/adapters/data_gen/test_data_gen_task.py

@@ -0,0 +1,293 @@
+import json
+
+import pytest
+
+from kiln_ai.adapters.data_gen.data_gen_task import (
+    DataGenCategoriesTask,
+    DataGenCategoriesTaskInput,
+    DataGenCategoriesTaskOutput,
+    DataGenSampleTask,
+    DataGenSampleTaskInput,
+    list_json_schema_for_task,
+)
+from kiln_ai.adapters.langchain_adapters import LangChainPromptAdapter
+from kiln_ai.adapters.ml_model_list import get_model_and_provider
+from kiln_ai.adapters.test_prompt_adaptors import get_all_models_and_providers
+from kiln_ai.datamodel import Project, Task
+
+
+@pytest.fixture
+def base_task():
+    project = Project(name="TestProject")
+    return Task(
+        name="Cowboy Speaker",
+        parent=project,
+        description="Reply like a cowboy",
+        instruction="Reply like a cowboy",
+        requirements=[],
+    )
+
+
+def test_data_gen_categories_task_input_initialization(base_task):
+    # Arrange
+    node_path = ["root", "branch", "leaf"]
+    num_subtopics = 4
+    human_guidance = "Test guidance"
+
+    # Act
+    input_model = DataGenCategoriesTaskInput.from_task(
+        task=base_task,
+        node_path=node_path,
+        num_subtopics=num_subtopics,
+        human_guidance=human_guidance,
+    )
+
+    # Assert
+    assert input_model.node_path == node_path
+    assert input_model.num_subtopics == num_subtopics
+    assert input_model.human_guidance == human_guidance
+    assert isinstance(input_model.system_prompt, str)
+    assert "Reply like a cowboy" in input_model.system_prompt
+
+
+def test_data_gen_categories_task_input_default_values(base_task):
+    # Act
+    input_model = DataGenCategoriesTaskInput.from_task(task=base_task)
+
+    # Assert
+    assert input_model.num_subtopics == 6
+    assert input_model.human_guidance is None
+    assert input_model.node_path == []
+
+
+def test_data_gen_categories_task_initialization():
+    # Act
+    task = DataGenCategoriesTask()
+
+    # Assert
+    assert task.name == "DataGen"
+    assert isinstance(task.parent, Project)
+    assert task.description is not None
+    assert task.instruction is not None
+    assert isinstance(task.input_json_schema, str)
+    assert isinstance(task.output_json_schema, str)
+
+
+def test_data_gen_categories_task_schemas():
+    # Act
+    task = DataGenCategoriesTask()
+
+    # Assert
+    input_schema = json.loads(task.input_json_schema)
+    output_schema = json.loads(task.output_json_schema)
+
+    assert isinstance(input_schema, dict)
+    assert isinstance(output_schema, dict)
+    assert output_schema["type"] == "object"
+    assert output_schema["properties"]["subtopics"]["type"] == "array"
+    assert input_schema["properties"]["node_path"]["type"] == "array"
+    assert input_schema["properties"]["num_subtopics"]["type"] == "integer"
+    assert set(input_schema["required"]) == {
+        "node_path",
+        "num_subtopics",
+        "system_prompt",
+    }
+
+
+@pytest.mark.paid
+@pytest.mark.ollama
+@pytest.mark.parametrize("model_name,provider_name", get_all_models_and_providers())
+async def test_data_gen_all_models_providers(
+    tmp_path, model_name, provider_name, base_task
+):
+    _, provider = get_model_and_provider(model_name, provider_name)
+    if not provider.supports_data_gen:
+        # pass if the model doesn't support data gen (testing the support flag is part of this)
+        return
+
+    data_gen_task = DataGenCategoriesTask()
+    data_gen_input = DataGenCategoriesTaskInput.from_task(base_task, num_subtopics=6)
+
+    adapter = LangChainPromptAdapter(
+        data_gen_task,
+        model_name=model_name,
+        provider=provider_name,
+    )
+
+    input_dict = data_gen_input.model_dump()
+    run = await adapter.invoke(input_dict)
+    parsed_output = DataGenCategoriesTaskOutput.model_validate_json(run.output.output)
+    assert len(parsed_output.subtopics) == 6
+    for subtopic in parsed_output.subtopics:
+        assert isinstance(subtopic, str)
+
+
+def test_data_gen_sample_task_input_initialization(base_task):
+    # Arrange
+    topic = ["cowboys", "hats"]
+    num_samples = 4
+    human_guidance = "Test guidance"
+
+    # Act
+    input_model = DataGenSampleTaskInput.from_task(
+        task=base_task,
+        topic=topic,
+        num_samples=num_samples,
+        human_guidance=human_guidance,
+    )
+
+    # Assert
+    assert input_model.topic == topic
+    assert input_model.num_samples == num_samples
+    assert input_model.human_guidance == human_guidance
+    assert isinstance(input_model.system_prompt, str)
+    assert "Reply like a cowboy" in input_model.system_prompt
+
+
+def test_data_gen_sample_task_input_default_values(base_task):
+    # Act
+    input_model = DataGenSampleTaskInput.from_task(task=base_task)
+
+    # Assert
+    assert input_model.num_samples == 8
+    assert input_model.human_guidance is None
+    assert input_model.topic == []
+
+
+def test_data_gen_sample_task_initialization(base_task):
+    # Act
+    task = DataGenSampleTask(target_task=base_task)
+
+    # Assert
+    assert task.name == "DataGenSample"
+    assert isinstance(task.parent, Project)
+    assert task.description is not None
+    assert task.instruction is not None
+
+    input_schema = json.loads(task.input_json_schema)
+    output_schema = json.loads(task.output_json_schema)
+
+    assert isinstance(input_schema, dict)
+    assert isinstance(output_schema, dict)
+    assert output_schema["type"] == "object"
+    assert output_schema["properties"]["generated_samples"]["type"] == "array"
+    assert input_schema["properties"]["topic"]["type"] == "array"
+    assert input_schema["properties"]["num_samples"]["type"] == "integer"
+    assert set(input_schema["required"]) == {
+        "topic",
+        "num_samples",
+        "system_prompt",
+    }
+
+
+def test_list_json_schema_for_task_with_output_schema(base_task):
+    # Arrange
+    base_task.input_json_schema = json.dumps(
+        {
+            "type": "object",
+            "properties": {"name": {"type": "string"}, "age": {"type": "integer"}},
+        }
+    )
+
+    # Act
+    schema = list_json_schema_for_task(base_task)
+    parsed_schema = json.loads(schema)
+
+    # Assert
+    assert parsed_schema["type"] == "object"
+    generated_samples_schema = parsed_schema["properties"]["generated_samples"]
+    assert generated_samples_schema["type"] == "array"
+    assert generated_samples_schema["items"]["type"] == "object"
+    assert generated_samples_schema["items"]["properties"]["name"]["type"] == "string"
+    assert generated_samples_schema["items"]["properties"]["age"]["type"] == "integer"
+
+
+def test_list_json_schema_for_task_without_output_schema(base_task):
+    # Arrange
+    base_task.output_json_schema = None
+
+    # Act
+    schema = list_json_schema_for_task(base_task)
+    parsed_schema = json.loads(schema)
+
+    # Assert
+    assert parsed_schema["type"] == "object"
+    assert parsed_schema["properties"]["generated_samples"]["type"] == "array"
+    assert parsed_schema["properties"]["generated_samples"]["items"]["type"] == "string"
+
+
+@pytest.mark.paid
+@pytest.mark.ollama
+@pytest.mark.parametrize("model_name,provider_name", get_all_models_and_providers())
+async def test_data_gen_sample_all_models_providers(
+    tmp_path, model_name, provider_name, base_task
+):
+    _, provider = get_model_and_provider(model_name, provider_name)
+    if not provider.supports_data_gen:
+        # pass if the model doesn't support data gen (testing the support flag is part of this)
+        return
+
+    data_gen_task = DataGenSampleTask(target_task=base_task)
+    data_gen_input = DataGenSampleTaskInput.from_task(
+        base_task, topic=["riding horses"], num_samples=4
+    )
+
+    adapter = LangChainPromptAdapter(
+        data_gen_task,
+        model_name=model_name,
+        provider=provider_name,
+    )
+
+    input_dict = data_gen_input.model_dump()
+    run = await adapter.invoke(input_dict)
+    parsed_output = json.loads(run.output.output)
+    samples = parsed_output["generated_samples"]
+    assert len(samples) == 4
+    for sample in samples:
+        assert isinstance(sample, str)
+
+
+@pytest.mark.paid
+@pytest.mark.ollama
+@pytest.mark.parametrize("model_name,provider_name", get_all_models_and_providers())
+async def test_data_gen_sample_all_models_providers_with_structured_output(
+    tmp_path, model_name, provider_name, base_task
+):
+    base_task.output_json_schema = json.dumps(
+        {
+            "type": "object",
+            "properties": {
+                "opening": {"type": "string"},
+                "closing": {"type": "string"},
+            },
+            "required": ["opening", "closing"],
+        }
+    )
+
+    _, provider = get_model_and_provider(model_name, provider_name)
+    if not provider.supports_data_gen:
+        # pass if the model doesn't support data gen (testing the support flag is part of this)
+        return
+
+    data_gen_task = DataGenSampleTask(target_task=base_task)
+    data_gen_input = DataGenSampleTaskInput.from_task(
+        base_task, topic=["riding horses"], num_samples=4
+    )
+
+    adapter = LangChainPromptAdapter(
+        data_gen_task,
+        model_name=model_name,
+        provider=provider_name,
+    )
+
+    input_dict = data_gen_input.model_dump()
+    run = await adapter.invoke(input_dict)
+    parsed_output = json.loads(run.output.output)
+    samples = parsed_output["generated_samples"]
+    assert len(samples) == 4
+    for sample in samples:
+        assert isinstance(sample, dict)
+        assert "opening" in sample
+        assert "closing" in sample
+        assert isinstance(sample["opening"], str)
+        assert isinstance(sample["closing"], str)

kiln_ai/adapters/langchain_adapters.py

@@ -2,14 +2,14 @@ from typing import Dict
 
 from langchain_core.language_models import LanguageModelInput
 from langchain_core.language_models.chat_models import BaseChatModel
-from langchain_core.messages import HumanMessage, SystemMessage
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 from langchain_core.messages.base import BaseMessage
 from langchain_core.runnables import Runnable
 from pydantic import BaseModel
 
 import kiln_ai.datamodel as datamodel
 
-from .base_adapter import AdapterInfo, BaseAdapter, BasePromptBuilder
+from .base_adapter import AdapterInfo, BaseAdapter, BasePromptBuilder, RunOutput
 from .ml_model_list import langchain_model_from
 
 LangChainModelType = BaseChatModel | Runnable[LanguageModelInput, Dict | BaseModel]
@@ -84,15 +84,41 @@ class LangChainPromptAdapter(BaseAdapter):
             )
         return self._model
 
-    async def _run(self, input: Dict | str) -> Dict | str:
+    async def _run(self, input: Dict | str) -> RunOutput:
+        model = await self.model()
+        chain = model
+        intermediate_outputs = {}
+
         prompt = self.build_prompt()
         user_msg = self.prompt_builder.build_user_message(input)
         messages = [
             SystemMessage(content=prompt),
             HumanMessage(content=user_msg),
         ]
-        model = await self.model()
-        response = model.invoke(messages)
+
+        # COT with structured output
+        cot_prompt = self.prompt_builder.chain_of_thought_prompt()
+        if cot_prompt and self.has_structured_output():
+            # Base model (without structured output) used for COT message
+            base_model = await langchain_model_from(
+                self.model_name, self.model_provider
+            )
+            messages.append(
+                SystemMessage(content=cot_prompt),
+            )
+
+            cot_messages = [*messages]
+            cot_response = base_model.invoke(cot_messages)
+            intermediate_outputs["chain_of_thought"] = cot_response.content
+            messages.append(AIMessage(content=cot_response.content))
+            messages.append(
+                SystemMessage(content="Considering the above, return a final result.")
+            )
+        elif cot_prompt:
+            # for plaintext output, we just add COT instructions. We still only make one call.
+            messages.append(SystemMessage(content=cot_prompt))
+
+        response = chain.invoke(messages)
 
         if self.has_structured_output():
             if (
@@ -102,14 +128,20 @@ class LangChainPromptAdapter(BaseAdapter):
             ):
                 raise RuntimeError(f"structured response not returned: {response}")
             structured_response = response["parsed"]
-            return self._munge_response(structured_response)
+            return RunOutput(
+                output=self._munge_response(structured_response),
+                intermediate_outputs=intermediate_outputs,
+            )
         else:
             if not isinstance(response, BaseMessage):
                 raise RuntimeError(f"response is not a BaseMessage: {response}")
             text_content = response.content
             if not isinstance(text_content, str):
                 raise RuntimeError(f"response is not a string: {text_content}")
-            return text_content
+            return RunOutput(
+                output=text_content,
+                intermediate_outputs=intermediate_outputs,
+            )
 
     def adapter_info(self) -> AdapterInfo:
         return AdapterInfo(

kiln_ai/adapters/ml_model_list.py

@@ -89,6 +89,7 @@ class KilnModelProvider(BaseModel):
 
     name: ModelProviderName
     supports_structured_output: bool = True
+    supports_data_gen: bool = True
     provider_options: Dict = {}
 
 
@@ -176,6 +177,8 @@ built_in_models: List[KilnModel] = [
         providers=[
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
+                supports_structured_output=False,  # it should, but doesn't work on openrouter
+                supports_data_gen=False,  # doesn't work on openrouter
                 provider_options={"model": "google/gemini-pro-1.5"},
             ),
         ],
@@ -188,6 +191,7 @@ built_in_models: List[KilnModel] = [
         providers=[
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
+                supports_data_gen=False,
                 provider_options={"model": "google/gemini-flash-1.5"},
             ),
         ],
@@ -200,6 +204,8 @@ built_in_models: List[KilnModel] = [
         providers=[
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
+                supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={"model": "google/gemini-flash-1.5-8b"},
             ),
         ],
@@ -213,6 +219,7 @@ built_in_models: List[KilnModel] = [
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
                 supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={"model": "nvidia/llama-3.1-nemotron-70b-instruct"},
             ),
         ],
@@ -230,6 +237,7 @@ built_in_models: List[KilnModel] = [
             KilnModelProvider(
                 name=ModelProviderName.amazon_bedrock,
                 supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={
                     "model": "meta.llama3-1-8b-instruct-v1:0",
                     "region_name": "us-west-2",  # Llama 3.1 only in west-2
@@ -237,6 +245,7 @@ built_in_models: List[KilnModel] = [
             ),
             KilnModelProvider(
                 name=ModelProviderName.ollama,
+                supports_data_gen=False,
                 provider_options={
                     "model": "llama3.1:8b",
                     "model_aliases": ["llama3.1"],  # 8b is default
@@ -245,6 +254,7 @@ built_in_models: List[KilnModel] = [
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
                 supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={"model": "meta-llama/llama-3.1-8b-instruct"},
             ),
         ],
@@ -261,7 +271,9 @@ built_in_models: List[KilnModel] = [
             ),
             KilnModelProvider(
                 name=ModelProviderName.amazon_bedrock,
+                # not sure how AWS manages to break this, but it's not working
                 supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={
                     "model": "meta.llama3-1-70b-instruct-v1:0",
                     "region_name": "us-west-2",  # Llama 3.1 only in west-2
@@ -285,6 +297,7 @@ built_in_models: List[KilnModel] = [
         providers=[
             KilnModelProvider(
                 name=ModelProviderName.amazon_bedrock,
+                supports_data_gen=False,
                 provider_options={
                     "model": "meta.llama3-1-405b-instruct-v1:0",
                     "region_name": "us-west-2",  # Llama 3.1 only in west-2
@@ -344,8 +357,15 @@ built_in_models: List[KilnModel] = [
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
                 supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={"model": "meta-llama/llama-3.2-3b-instruct"},
             ),
+            KilnModelProvider(
+                name=ModelProviderName.ollama,
+                supports_structured_output=False,
+                supports_data_gen=False,
+                provider_options={"model": "llama3.2"},
+            ),
         ],
     ),
     # Llama 3.2 11B
@@ -357,8 +377,15 @@ built_in_models: List[KilnModel] = [
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
                 supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={"model": "meta-llama/llama-3.2-11b-vision-instruct"},
             ),
+            KilnModelProvider(
+                name=ModelProviderName.ollama,
+                supports_structured_output=False,
+                supports_data_gen=False,
+                provider_options={"model": "llama3.2-vision"},
+            ),
         ],
     ),
     # Llama 3.2 90B
@@ -370,8 +397,15 @@ built_in_models: List[KilnModel] = [
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
                 supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={"model": "meta-llama/llama-3.2-90b-vision-instruct"},
             ),
+            KilnModelProvider(
+                name=ModelProviderName.ollama,
+                supports_structured_output=False,
+                supports_data_gen=False,
+                provider_options={"model": "llama3.2-vision:90b"},
+            ),
         ],
     ),
     # Phi 3.5
@@ -384,10 +418,13 @@ built_in_models: List[KilnModel] = [
             KilnModelProvider(
                 name=ModelProviderName.ollama,
                 supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={"model": "phi3.5"},
             ),
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
+                supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={"model": "microsoft/phi-3.5-mini-128k-instruct"},
             ),
         ],
@@ -402,6 +439,7 @@ built_in_models: List[KilnModel] = [
             KilnModelProvider(
                 name=ModelProviderName.ollama,
                 supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={
                     "model": "gemma2:2b",
                 },
@@ -417,12 +455,14 @@ built_in_models: List[KilnModel] = [
         providers=[
             KilnModelProvider(
                 name=ModelProviderName.ollama,
+                supports_data_gen=False,
                 provider_options={
                     "model": "gemma2:9b",
                 },
             ),
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
+                supports_data_gen=False,
                 provider_options={"model": "google/gemma-2-9b-it"},
             ),
         ],
@@ -436,12 +476,14 @@ built_in_models: List[KilnModel] = [
         providers=[
             KilnModelProvider(
                 name=ModelProviderName.ollama,
+                supports_data_gen=False,
                 provider_options={
                     "model": "gemma2:27b",
                 },
             ),
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
+                supports_data_gen=False,
                 provider_options={"model": "google/gemma-2-27b-it"},
             ),
         ],
@@ -449,6 +491,19 @@ built_in_models: List[KilnModel] = [
 ]
 
 
+def get_model_and_provider(
+    model_name: str, provider_name: str
+) -> tuple[KilnModel | None, KilnModelProvider | None]:
+    model = next(filter(lambda m: m.name == model_name, built_in_models), None)
+    if model is None:
+        return None, None
+    provider = next(filter(lambda p: p.name == provider_name, model.providers), None)
+    # all or nothing
+    if provider is None or model is None:
+        return None, None
+    return model, provider
+
+
 def provider_name_from_id(id: str) -> str:
     """
     Converts a provider ID to its human-readable name.
@@ -687,7 +742,6 @@ def parse_ollama_tags(tags: Any) -> OllamaConnection | None:
         models = tags["models"]
         if isinstance(models, list):
             model_names = [model["model"] for model in models]
-            print(f"model_names: {model_names}")
             available_supported_models = [
                 model
                 for model in model_names

kiln_ai/adapters/prompt_builders.py

@@ -54,6 +54,28 @@ class BasePromptBuilder(metaclass=ABCMeta):
 
         return f"The input is:\n{input}"
 
+    def chain_of_thought_prompt(self) -> str | None:
+        """Build and return the chain of thought prompt string.
+
+        Returns:
+            str: The constructed chain of thought prompt.
+        """
+        return None
+
+    def build_prompt_for_ui(self) -> str:
+        """Build a prompt for the UI. It includes additional instructions (like chain of thought), even if they are passed to the model in stages.
+
+        Designed for end-user consumption, not for model consumption.
+
+        Returns:
+            str: The constructed prompt string.
+        """
+        base_prompt = self.build_prompt()
+        cot_prompt = self.chain_of_thought_prompt()
+        if cot_prompt:
+            base_prompt += "\n# Thinking Instructions\n\n" + cot_prompt
+        return base_prompt
+
 
 class SimplePromptBuilder(BasePromptBuilder):
     """A basic prompt builder that combines task instruction with requirements."""
@@ -187,11 +209,49 @@ class RepairsPromptBuilder(MultiShotPromptBuilder):
         return prompt_section
 
 
+def chain_of_thought_prompt(task: Task) -> str | None:
+    """Standard implementation to build and return the chain of thought prompt string.
+
+    Returns:
+        str: The constructed chain of thought prompt.
+    """
+
+    cot_instruction = task.thinking_instruction
+    if not cot_instruction:
+        cot_instruction = "Think step by step, explaining your reasoning."
+
+    return cot_instruction
+
+
+class SimpleChainOfThoughtPromptBuilder(SimplePromptBuilder):
+    """A prompt builder that includes a chain of thought prompt on top of the simple prompt."""
+
+    def chain_of_thought_prompt(self) -> str | None:
+        return chain_of_thought_prompt(self.task)
+
+
+class FewShotChainOfThoughtPromptBuilder(FewShotPromptBuilder):
+    """A prompt builder that includes a chain of thought prompt on top of the few shot prompt."""
+
+    def chain_of_thought_prompt(self) -> str | None:
+        return chain_of_thought_prompt(self.task)
+
+
+class MultiShotChainOfThoughtPromptBuilder(MultiShotPromptBuilder):
+    """A prompt builder that includes a chain of thought prompt on top of the multi shot prompt."""
+
+    def chain_of_thought_prompt(self) -> str | None:
+        return chain_of_thought_prompt(self.task)
+
+
 prompt_builder_registry = {
     "simple_prompt_builder": SimplePromptBuilder,
     "multi_shot_prompt_builder": MultiShotPromptBuilder,
     "few_shot_prompt_builder": FewShotPromptBuilder,
     "repairs_prompt_builder": RepairsPromptBuilder,
+    "simple_chain_of_thought_prompt_builder": SimpleChainOfThoughtPromptBuilder,
+    "few_shot_chain_of_thought_prompt_builder": FewShotChainOfThoughtPromptBuilder,
+    "multi_shot_chain_of_thought_prompt_builder": MultiShotChainOfThoughtPromptBuilder,
 }
 
 
@@ -217,5 +277,11 @@ def prompt_builder_from_ui_name(ui_name: str) -> type[BasePromptBuilder]:
             return MultiShotPromptBuilder
         case "repairs":
             return RepairsPromptBuilder
+        case "simple_chain_of_thought":
+            return SimpleChainOfThoughtPromptBuilder
+        case "few_shot_chain_of_thought":
+            return FewShotChainOfThoughtPromptBuilder
+        case "multi_shot_chain_of_thought":
+            return MultiShotChainOfThoughtPromptBuilder
         case _:
             raise ValueError(f"Unknown prompt builder: {ui_name}")

kiln_ai/adapters/repair/test_repair_task.py

@@ -5,6 +5,7 @@ from unittest.mock import AsyncMock, patch
 import pytest
 from pydantic import ValidationError
 
+from kiln_ai.adapters.base_adapter import RunOutput
 from kiln_ai.adapters.langchain_adapters import (
     LangChainPromptAdapter,
 )
@@ -222,7 +223,9 @@ async def test_mocked_repair_task_run(sample_task, sample_task_run, sample_repai
     with patch.object(
         LangChainPromptAdapter, "_run", new_callable=AsyncMock
     ) as mock_run:
-        mock_run.return_value = mocked_output
+        mock_run.return_value = RunOutput(
+            output=mocked_output, intermediate_outputs=None
+        )
 
         adapter = LangChainPromptAdapter(
             repair_task, model_name="llama_3_1_8b", provider="groq"