diff --git a/core/agents/convo.py b/core/agents/convo.py index 0eb58f9d8..f4f6f200b 100644 --- a/core/agents/convo.py +++ b/core/agents/convo.py @@ -105,3 +105,7 @@ def remove_defs(d): f"YOU MUST NEVER add any additional fields to your response, and NEVER add additional preamble like 'Here is your JSON'." ) return self + + def remove_last_x_messages(self, x: int) -> "AgentConvo": + self.messages = self.messages[:-x] + return self diff --git a/core/agents/developer.py b/core/agents/developer.py index aaeed50f9..02f596c3d 100644 --- a/core/agents/developer.py +++ b/core/agents/developer.py @@ -1,11 +1,12 @@ -from typing import Optional +from enum import Enum +from typing import Annotated, Literal, Optional, Union from uuid import uuid4 from pydantic import BaseModel, Field from core.agents.base import BaseAgent from core.agents.convo import AgentConvo -from core.agents.mixins import TaskSteps +from core.agents.mixins import RelevantFilesMixin from core.agents.response import AgentResponse, ResponseType from core.config import TASK_BREAKDOWN_AGENT_NAME from core.db.models.project_state import IterationStatus, TaskStatus @@ -17,11 +18,48 @@ log = get_logger(__name__) -class RelevantFiles(BaseModel): - relevant_files: list[str] = Field(description="List of relevant files for the current task.") +class StepType(str, Enum): + COMMAND = "command" + SAVE_FILE = "save_file" + HUMAN_INTERVENTION = "human_intervention" -class Developer(BaseAgent): +class CommandOptions(BaseModel): + command: str = Field(description="Command to run") + timeout: int = Field(description="Timeout in seconds") + success_message: str = "" + + +class SaveFileOptions(BaseModel): + path: str + + +class SaveFileStep(BaseModel): + type: Literal[StepType.SAVE_FILE] = StepType.SAVE_FILE + save_file: SaveFileOptions + + +class CommandStep(BaseModel): + type: Literal[StepType.COMMAND] = StepType.COMMAND + command: CommandOptions + + +class HumanInterventionStep(BaseModel): + type: Literal[StepType.HUMAN_INTERVENTION] = StepType.HUMAN_INTERVENTION + human_intervention_description: str + + +Step = Annotated[ + Union[SaveFileStep, CommandStep, HumanInterventionStep], + Field(discriminator="type"), +] + + +class TaskSteps(BaseModel): + steps: list[Step] + + +class Developer(RelevantFilesMixin, BaseAgent): agent_type = "developer" display_name = "Developer" @@ -96,7 +134,8 @@ async def breakdown_current_iteration(self, task_review_feedback: Optional[str] log.debug(f"Breaking down the iteration {description}") await self.send_message("Breaking down the current task iteration ...") - await self.get_relevant_files(user_feedback, description) + if self.current_state.files and self.current_state.relevant_files is None: + return await self.get_relevant_files(user_feedback, description) await self.ui.send_task_progress( n_tasks, # iterations and reviews can be created only one at a time, so we are always on last one @@ -114,7 +153,6 @@ async def breakdown_current_iteration(self, task_review_feedback: Optional[str] AgentConvo(self) .template( "iteration", - current_task=current_task, user_feedback=user_feedback, user_feedback_qa=None, next_solution_to_try=None, @@ -175,7 +213,7 @@ async def breakdown_current_task(self) -> AgentResponse: log.debug(f"Current state files: {len(self.current_state.files)}, relevant {self.current_state.relevant_files}") # Check which files are relevant to the current task if self.current_state.files and self.current_state.relevant_files is None: - await self.get_relevant_files() + return await self.get_relevant_files() current_task_index = self.current_state.tasks.index(current_task) @@ -189,6 +227,8 @@ async def breakdown_current_task(self) -> AgentResponse: ) response: str = await llm(convo) + await self.get_relevant_files(None, response) + self.next_state.tasks[current_task_index] = { **current_task, "instructions": response, @@ -214,31 +254,6 @@ async def breakdown_current_task(self) -> AgentResponse: ) return AgentResponse.done(self) - async def get_relevant_files( - self, user_feedback: Optional[str] = None, solution_description: Optional[str] = None - ) -> AgentResponse: - log.debug("Getting relevant files for the current task") - await self.send_message("Figuring out which project files are relevant for the next task ...") - - llm = self.get_llm() - convo = ( - AgentConvo(self) - .template( - "filter_files", - current_task=self.current_state.current_task, - user_feedback=user_feedback, - solution_description=solution_description, - ) - .require_schema(RelevantFiles) - ) - - llm_response: list[str] = await llm(convo, parser=JSONParser(RelevantFiles), temperature=0) - - existing_files = {file.path for file in self.current_state.files} - self.next_state.relevant_files = [path for path in llm_response.relevant_files if path in existing_files] - - return AgentResponse.done(self) - def set_next_steps(self, response: TaskSteps, source: str): # For logging/debugging purposes, we don't want to remove the finished steps # until we're done with the task. diff --git a/core/agents/mixins.py b/core/agents/mixins.py index 4533afdfe..cc6ba97d6 100644 --- a/core/agents/mixins.py +++ b/core/agents/mixins.py @@ -1,50 +1,21 @@ -from enum import Enum -from typing import Annotated, Literal, Optional, Union +from typing import Optional from pydantic import BaseModel, Field from core.agents.convo import AgentConvo +from core.agents.response import AgentResponse +from core.config import GET_RELEVANT_FILES_AGENT_NAME +from core.llm.parser import JSONParser +from core.log import get_logger +log = get_logger(__name__) -class StepType(str, Enum): - COMMAND = "command" - SAVE_FILE = "save_file" - HUMAN_INTERVENTION = "human_intervention" - -class CommandOptions(BaseModel): - command: str = Field(description="Command to run") - timeout: int = Field(description="Timeout in seconds") - success_message: str = "" - - -class SaveFileOptions(BaseModel): - path: str - - -class SaveFileStep(BaseModel): - type: Literal[StepType.SAVE_FILE] = StepType.SAVE_FILE - save_file: SaveFileOptions - - -class CommandStep(BaseModel): - type: Literal[StepType.COMMAND] = StepType.COMMAND - command: CommandOptions - - -class HumanInterventionStep(BaseModel): - type: Literal[StepType.HUMAN_INTERVENTION] = StepType.HUMAN_INTERVENTION - human_intervention_description: str - - -Step = Annotated[ - Union[SaveFileStep, CommandStep, HumanInterventionStep], - Field(discriminator="type"), -] - - -class TaskSteps(BaseModel): - steps: list[Step] +class RelevantFiles(BaseModel): + read_files: list[str] = Field(description="List of files you want to read.") + add_files: list[str] = Field(description="List of files you want to add to the list of relevant files.") + remove_files: list[str] = Field(description="List of files you want to remove from the list of relevant files.") + done: bool = Field(description="Boolean flag to indicate that you are done selecting relevant files.") class IterationPromptMixin: @@ -74,7 +45,6 @@ async def find_solution( llm = self.get_llm() convo = AgentConvo(self).template( "iteration", - current_task=self.current_state.current_task, user_feedback=user_feedback, user_feedback_qa=user_feedback_qa, next_solution_to_try=next_solution_to_try, @@ -82,3 +52,57 @@ async def find_solution( ) llm_solution: str = await llm(convo) return llm_solution + + +class RelevantFilesMixin: + """ + Provides a method to get relevant files for the current task. + """ + + async def get_relevant_files( + self, user_feedback: Optional[str] = None, solution_description: Optional[str] = None + ) -> AgentResponse: + log.debug("Getting relevant files for the current task") + await self.send_message("Figuring out which project files are relevant for the next task ...") + + done = False + relevant_files = set() + llm = self.get_llm(GET_RELEVANT_FILES_AGENT_NAME) + convo = ( + AgentConvo(self) + .template( + "filter_files", + user_feedback=user_feedback, + solution_description=solution_description, + relevant_files=relevant_files, + ) + .require_schema(RelevantFiles) + ) + + while not done and len(convo.messages) < 13: + llm_response: RelevantFiles = await llm(convo, parser=JSONParser(RelevantFiles), temperature=0) + + # Check if there are files to add to the list + if llm_response.add_files: + # Add only the files from add_files that are not already in relevant_files + relevant_files.update(file for file in llm_response.add_files if file not in relevant_files) + + # Check if there are files to remove from the list + if llm_response.remove_files: + # Remove files from relevant_files that are in remove_files + relevant_files.difference_update(llm_response.remove_files) + + read_files = [file for file in self.current_state.files if file.path in llm_response.read_files] + + convo.remove_last_x_messages(1) + convo.assistant(llm_response.original_response) + convo.template("filter_files_loop", read_files=read_files, relevant_files=relevant_files).require_schema( + RelevantFiles + ) + done = llm_response.done + + existing_files = {file.path for file in self.current_state.files} + relevant_files = [path for path in relevant_files if path in existing_files] + self.next_state.relevant_files = relevant_files + + return AgentResponse.done(self) diff --git a/core/agents/spec_writer.py b/core/agents/spec_writer.py index 69b5a7655..baac7b590 100644 --- a/core/agents/spec_writer.py +++ b/core/agents/spec_writer.py @@ -73,6 +73,7 @@ async def initialize_spec(self) -> AgentResponse: }, ) + reviewed_spec = user_description if len(user_description) < ANALYZE_THRESHOLD and complexity != Complexity.SIMPLE: initial_spec = await self.analyze_spec(user_description) reviewed_spec = await self.review_spec(desc=user_description, spec=initial_spec) diff --git a/core/agents/task_reviewer.py b/core/agents/task_reviewer.py index f6b3999a4..5cdd4f366 100644 --- a/core/agents/task_reviewer.py +++ b/core/agents/task_reviewer.py @@ -28,6 +28,11 @@ async def review_code_changes(self) -> AgentResponse: # Some iterations are created by the task reviewer and have no user feedback if iteration["user_feedback"] ] + bug_hunter_instructions = [ + iteration["bug_hunting_cycles"][-1]["human_readable_instructions"].replace("```", "").strip() + for iteration in self.current_state.iterations + if iteration["bug_hunting_cycles"] + ] files_before_modification = self.current_state.modified_files files_after_modification = [ @@ -40,10 +45,10 @@ async def review_code_changes(self) -> AgentResponse: # TODO instead of sending files before and after maybe add nice way to show diff for multiple files convo = AgentConvo(self).template( "review_task", - current_task=self.current_state.current_task, all_feedbacks=all_feedbacks, files_before_modification=files_before_modification, files_after_modification=files_after_modification, + bug_hunter_instructions=bug_hunter_instructions, ) llm_response: str = await llm(convo, temperature=0.7) diff --git a/core/agents/troubleshooter.py b/core/agents/troubleshooter.py index 081f7812b..497eb7805 100644 --- a/core/agents/troubleshooter.py +++ b/core/agents/troubleshooter.py @@ -5,7 +5,7 @@ from core.agents.base import BaseAgent from core.agents.convo import AgentConvo -from core.agents.mixins import IterationPromptMixin +from core.agents.mixins import IterationPromptMixin, RelevantFilesMixin from core.agents.response import AgentResponse from core.db.models.file import File from core.db.models.project_state import IterationStatus, TaskStatus @@ -28,7 +28,7 @@ class RouteFilePaths(BaseModel): files: list[str] = Field(description="List of paths for files that contain routes") -class Troubleshooter(IterationPromptMixin, BaseAgent): +class Troubleshooter(IterationPromptMixin, RelevantFilesMixin, BaseAgent): agent_type = "troubleshooter" display_name = "Troubleshooter" @@ -102,6 +102,7 @@ async def create_iteration(self) -> AgentResponse: else: # should be - elif change_description is not None: - but to prevent bugs with the extension # this might be caused if we show the input field instead of buttons + await self.get_relevant_files(user_feedback) iteration_status = IterationStatus.NEW_FEATURE_REQUESTED self.next_state.iterations = self.current_state.iterations + [ diff --git a/core/config/__init__.py b/core/config/__init__.py index a0ae1945a..ea8b56da8 100644 --- a/core/config/__init__.py +++ b/core/config/__init__.py @@ -39,6 +39,7 @@ CHECK_LOGS_AGENT_NAME = "BugHunter.check_logs" TASK_BREAKDOWN_AGENT_NAME = "Developer.breakdown_current_task" SPEC_WRITER_AGENT_NAME = "SpecWriter" +GET_RELEVANT_FILES_AGENT_NAME = "get_relevant_files" # Endpoint for the external documentation EXTERNAL_DOCUMENTATION_API = "http://docs-pythagora-io-439719575.us-east-1.elb.amazonaws.com" @@ -330,6 +331,7 @@ class Config(_StrictModel): temperature=0.5, ), SPEC_WRITER_AGENT_NAME: AgentLLMConfig(model="gpt-4-0125-preview", temperature=0.0), + GET_RELEVANT_FILES_AGENT_NAME: AgentLLMConfig(model="claude-3-5-sonnet-20240620", temperature=0.0), } ) prompt: PromptConfig = PromptConfig() diff --git a/core/db/models/project_state.py b/core/db/models/project_state.py index 67ad685ad..3d520cc10 100644 --- a/core/db/models/project_state.py +++ b/core/db/models/project_state.py @@ -303,6 +303,7 @@ def complete_iteration(self): log.debug(f"Completing iteration {self.unfinished_iterations[0]}") self.unfinished_iterations[0]["status"] = IterationStatus.DONE + self.relevant_files = None self.flag_iterations_as_modified() def flag_iterations_as_modified(self): diff --git a/core/llm/parser.py b/core/llm/parser.py index 4d49b3d73..8cd026366 100644 --- a/core/llm/parser.py +++ b/core/llm/parser.py @@ -3,7 +3,7 @@ from enum import Enum from typing import Optional, Union -from pydantic import BaseModel, ValidationError +from pydantic import BaseModel, ValidationError, create_model class MultiCodeBlockParser: @@ -86,6 +86,7 @@ class JSONParser: def __init__(self, spec: Optional[BaseModel] = None, strict: bool = True): self.spec = spec self.strict = strict or (spec is not None) + self.original_response = None @property def schema(self): @@ -102,7 +103,8 @@ def errors_to_markdown(errors: list) -> str: return "\n".join(error_txt) def __call__(self, text: str) -> Union[BaseModel, dict, None]: - text = text.strip() + self.original_response = text.strip() # Store the original text + text = self.original_response if text.startswith("```"): try: text = CodeBlockParser()(text) @@ -130,7 +132,17 @@ def __call__(self, text: str) -> Union[BaseModel, dict, None]: except Exception as err: raise ValueError(f"Error parsing JSON: {err}") from err - return model + # Create a new model that includes the original model fields and the original text + ExtendedModel = create_model( + f"Extended{self.spec.__name__}", + original_response=(str, ...), + **{field_name: (field.annotation, field.default) for field_name, field in self.spec.__fields__.items()}, + ) + + # Instantiate the extended model + extended_model = ExtendedModel(original_response=self.original_response, **model.dict()) + + return extended_model class EnumParser: diff --git a/core/prompts/developer/filter_files.prompt b/core/prompts/developer/filter_files.prompt index b9c154cd7..264a244a7 100644 --- a/core/prompts/developer/filter_files.prompt +++ b/core/prompts/developer/filter_files.prompt @@ -2,7 +2,6 @@ We're starting work on a new task for a project we're working on. {% include "partials/project_details.prompt" %} {% include "partials/features_list.prompt" %} -{% include "partials/files_list.prompt" %} We've broken the development of the project down to these tasks: ``` @@ -14,7 +13,7 @@ We've broken the development of the project down to these tasks: The next task we need to work on, and have to focus on, is this task: ``` -{{ current_task.description }} +{{ state.current_task.description }} ``` {% if user_feedback %}User who was using the app sent you this feedback: @@ -28,8 +27,12 @@ Focus on solving this issue in the following way: ``` {% endif %} +{% include "partials/files_descriptions.prompt" %} + **IMPORTANT** The files necessary for a developer to understand, modify, implement, and test the current task are considered to be relevant files. -Your job is select which of existing files are relevant for the current task. From the above list of files that app currently contains, you have to select ALL files that are relevant to the current task. Think step by step of everything that has to be done in this task and which files contain needed information. If you are unsure if a file is relevant or not, it is always better to include it in the list of relevant files. +Your job is select which of existing files are relevant for the current task. From the above list of files that app currently contains, you have to select ALL files that are relevant to the current task. Think step by step of everything that has to be done in this task and which files contain needed information. + +{% include "partials/filter_files_actions.prompt" %} {% include "partials/relative_paths.prompt" %} diff --git a/core/prompts/developer/filter_files_loop.prompt b/core/prompts/developer/filter_files_loop.prompt new file mode 100644 index 000000000..a2da3c5be --- /dev/null +++ b/core/prompts/developer/filter_files_loop.prompt @@ -0,0 +1,13 @@ +{% if read_files %} +Here are the files that you wanted to read: +---START_OF_FILES--- +{% for file in read_files %} +File **`{{ file.path }}`** ({{file.content.content.splitlines()|length}} lines of code): +``` +{{ file.content.content }}``` + +{% endfor %} +---END_OF_FILES--- +{% endif %} + +{% include "partials/filter_files_actions.prompt" %} diff --git a/core/prompts/partials/files_descriptions.prompt b/core/prompts/partials/files_descriptions.prompt new file mode 100644 index 000000000..c2ac2c4c2 --- /dev/null +++ b/core/prompts/partials/files_descriptions.prompt @@ -0,0 +1,4 @@ +These files are currently implemented in the project: +{% for file in state.files %} +* `{{ file.path }}{% if file.meta.get("description") %}: {{file.meta.description}}{% endif %}` +{% endfor %} \ No newline at end of file diff --git a/core/prompts/partials/files_list.prompt b/core/prompts/partials/files_list.prompt index 3dbc70932..9c7037b61 100644 --- a/core/prompts/partials/files_list.prompt +++ b/core/prompts/partials/files_list.prompt @@ -1,18 +1,7 @@ {% if state.relevant_files %} -These files are currently implemented in the project: -{% for file in state.files %} -* `{{ file.path }}{% if file.meta.get("description") %}: {{file.meta.description}}{% endif %}` -{% endfor %} +{% include "partials/files_descriptions.prompt" %} -Here are the complete contents of files relevant to this task: ----START_OF_FILES--- -{% for file in state.relevant_file_objects %} -File **`{{ file.path }}`** ({{file.content.content.splitlines()|length}} lines of code): -``` -{{ file.content.content }}``` - -{% endfor %} ----END_OF_FILES--- +{% include "partials/files_list_relevant.prompt" %} {% elif state.files %} These files are currently implemented in the project: ---START_OF_FILES--- diff --git a/core/prompts/partials/files_list_relevant.prompt b/core/prompts/partials/files_list_relevant.prompt new file mode 100644 index 000000000..3aa7834fd --- /dev/null +++ b/core/prompts/partials/files_list_relevant.prompt @@ -0,0 +1,9 @@ +Here are the complete contents of files relevant to this task: +---START_OF_FILES--- +{% for file in state.relevant_file_objects %} +File **`{{ file.path }}`** ({{file.content.content.splitlines()|length}} lines of code): +``` +{{ file.content.content }}``` + +{% endfor %} +---END_OF_FILES--- \ No newline at end of file diff --git a/core/prompts/partials/filter_files_actions.prompt b/core/prompts/partials/filter_files_actions.prompt new file mode 100644 index 000000000..5c67a2544 --- /dev/null +++ b/core/prompts/partials/filter_files_actions.prompt @@ -0,0 +1,14 @@ +Here is the current relevant files list: +{% if relevant_files %}{{ relevant_files }}{% else %}[]{% endif %} + +Now, with multiple iterations you have to find relevant files for the current task. Here are commands that you can use: +- `read_files` - List of files that you want to read. +- `add_files` - Add file to the list of relevant files. +- `remove_files` - Remove file from the list of relevant files. +- `finished` - Boolean command that you will use when you finish with finding relevant files. + +Make sure to follow these rules: +- All files that you want to read or add to the list of relevant files, must exist in the project. Do not ask to read or add file that does not exist! In the first message you have list of all files that currently exist in the project. +- Do not repeat actions that you have already done. For example if you already added "index.js" to the list of relevant files you must not add it again. +- You must read the file before adding it to the list of relevant files. Do not `add_files` that you didn't read and see the content of the file. +- Focus only on your current task `{{ state.current_task.description }}` when selecting relevant files. diff --git a/core/prompts/task-reviewer/review_task.prompt b/core/prompts/task-reviewer/review_task.prompt index 27aa1236a..f19e6b910 100644 --- a/core/prompts/task-reviewer/review_task.prompt +++ b/core/prompts/task-reviewer/review_task.prompt @@ -12,7 +12,7 @@ Development process of this app was split into smaller tasks. Here is the list o You are currently working on, and have to focus only on, this task: ``` -{{ current_task.description }} +{{ state.current_task.description }} ``` A part of the app is already finished. @@ -26,7 +26,17 @@ While working on this task, your colleague who is testing the app "{{ state.bran {% endfor %} ``` -After you got each of these additional inputs, you tried to fix it as part of this task. {% endif %}Files that were modified during implementation of the task are: +After you got each of these additional inputs, you tried to fix it as part of this task. {% endif %} +{% if bug_hunter_instructions -%}Here are the last implementation instructions that were given while fixing a bug: +{% for instructions in bug_hunter_instructions %} +Instructions #{{ loop.index }} +``` +{{ instructions }} +``` +{% endfor %} +{% endif %} + +Files that were modified during implementation of the task are: {% for path, content in files_after_modification %} * `{{ path }}` {% endfor %} @@ -42,7 +52,6 @@ Now I will show you how those files looked before this task implementation start {% endif %}{% endfor %} ---end_of_files_at_start_of_task--- - **IMPORTANT** You have to review this task implementation. You are known to be very strict with your reviews and very good at noticing bugs but you don't mind minor changes like refactoring, adding or removing logs and so on. You think twice through all information given before giving any conclusions. diff --git a/core/prompts/troubleshooter/filter_files.prompt b/core/prompts/troubleshooter/filter_files.prompt new file mode 100644 index 000000000..99c9ebfc7 --- /dev/null +++ b/core/prompts/troubleshooter/filter_files.prompt @@ -0,0 +1,2 @@ +{# This is the same template as for Developer's filter files because Troubleshooter is reusing it in a conversation #} +{% extends "developer/filter_files.prompt" %} \ No newline at end of file diff --git a/core/prompts/troubleshooter/filter_files_loop.prompt b/core/prompts/troubleshooter/filter_files_loop.prompt new file mode 100644 index 000000000..74602d7de --- /dev/null +++ b/core/prompts/troubleshooter/filter_files_loop.prompt @@ -0,0 +1,2 @@ +{# This is the same template as for Developer's filter files because Troubleshooter is reusing it in a conversation #} +{% extends "developer/filter_files_loop.prompt" %} \ No newline at end of file diff --git a/core/prompts/troubleshooter/iteration.prompt b/core/prompts/troubleshooter/iteration.prompt index 8825992e6..d0560dd44 100644 --- a/core/prompts/troubleshooter/iteration.prompt +++ b/core/prompts/troubleshooter/iteration.prompt @@ -12,7 +12,7 @@ Development process of this app was split into smaller tasks. Here is the list o You are currently working on, and have to focus only on, this task: ``` -{{ current_task.description }} +{{ state.current_task.description }} ``` {% endif %} diff --git a/core/ui/base.py b/core/ui/base.py index 0cb8dbdd2..3f92c9546 100644 --- a/core/ui/base.py +++ b/core/ui/base.py @@ -269,6 +269,15 @@ async def send_project_stats(self, stats: dict): """ raise NotImplementedError() + async def generate_diff(self, file_old: str, file_new: str): + """ + Generate a diff between two files. + + :param file_old: Old file content. + :param file_new: New file content. + """ + raise NotImplementedError() + async def loading_finished(self): """ Notify the UI that loading has finished. diff --git a/core/ui/console.py b/core/ui/console.py index ed3128120..0716fc797 100644 --- a/core/ui/console.py +++ b/core/ui/console.py @@ -130,6 +130,9 @@ async def send_project_root(self, path: str): async def send_project_stats(self, stats: dict): pass + async def generate_diff(self, file_old: str, file_new: str): + pass + async def loading_finished(self): pass diff --git a/core/ui/virtual.py b/core/ui/virtual.py index 0d07a58fc..146ca440a 100644 --- a/core/ui/virtual.py +++ b/core/ui/virtual.py @@ -123,6 +123,9 @@ async def send_project_root(self, path: str): async def send_project_stats(self, stats: dict): pass + async def generate_diff(self, file_old: str, file_new: str): + pass + async def loading_finished(self): pass diff --git a/tests/llm/test_parser.py b/tests/llm/test_parser.py index fcc6ce9ab..16f8aadc7 100644 --- a/tests/llm/test_parser.py +++ b/tests/llm/test_parser.py @@ -141,7 +141,8 @@ class ParentModel(BaseModel): with pytest.raises(ValueError): parser(input) else: - assert parser(input).model_dump() == expected + result = parser(input) + assert result.model_dump() == {**expected, "original_response": input.strip()} def test_parse_json_schema():