WIP

core/agents/convo.py

@@ -105,3 +105,7 @@ class AgentConvo(Convo):
             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

core/agents/developer.py

@@ -6,6 +6,7 @@ from pydantic import BaseModel, Field
 
 from core.agents.base import BaseAgent
 from core.agents.convo import AgentConvo
+from core.agents.mixins import RelevantFilesMixin
 from core.agents.response import AgentResponse, ResponseType
 from core.db.models.project_state import TaskStatus
 from core.db.models.specification import Complexity
@@ -57,11 +58,7 @@ class TaskSteps(BaseModel):
     steps: list[Step]
 
 
-class RelevantFiles(BaseModel):
-    relevant_files: list[str] = Field(description="List of relevant files for the current task.")
-
-
-class Developer(BaseAgent):
+class Developer(RelevantFilesMixin, BaseAgent):
     agent_type = "developer"
     display_name = "Developer"
 
@@ -141,7 +138,6 @@ class Developer(BaseAgent):
             AgentConvo(self)
             .template(
                 "iteration",
-                current_task=current_task,
                 user_feedback=user_feedback,
                 user_feedback_qa=None,
                 next_solution_to_try=None,
@@ -226,31 +222,6 @@ class Developer(BaseAgent):
         )
         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.

core/agents/mixins.py

@@ -1,6 +1,19 @@
-from typing import Optional
+from typing import Any, Optional
+
+from pydantic import BaseModel, Field
 
 from core.agents.convo import AgentConvo
+from core.llm.parser import JSONParser
+from core.log import get_logger
+
+log = get_logger(__name__)
+
+
+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:
@@ -28,10 +41,65 @@ class IterationPromptMixin:
         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,
         )
         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
+    ) -> list[str | None | Any]:
+        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()
+        read_files = None
+        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,
+                relevant_files=relevant_files,
+            )
+            .require_schema(RelevantFiles)
+        )
+
+        while not done:
+            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 relevant_files  # todo fix this

core/agents/troubleshooter.py

@@ -5,7 +5,7 @@ from pydantic import BaseModel, Field
 
 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 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"
 
@@ -73,6 +73,7 @@ class Troubleshooter(IterationPromptMixin, BaseAgent):
                 llm_solution = ""
                 await self.trace_loop("loop-feedback")
         else:
+            await self.get_relevant_files(user_feedback)
             llm_solution = await self.find_solution(user_feedback, user_feedback_qa=user_feedback_qa)
 
         self.next_state.iterations = self.current_state.iterations + [

core/llm/parser.py

@@ -3,7 +3,7 @@ import re
 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 @@ class JSONParser:
         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 @@ class JSONParser:
         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:

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:
 ```
@@ -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" %}

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" %}

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 %}

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---

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---

core/prompts/partials/filter_files_actions.prompt

@@ -0,0 +1,8 @@
+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_file` - list of files that you want to read
+- `add_file` - add file to the list of relevant files
+- `remove_file` - remove file from the list of relevant files
+- `finished` - boolean command that you will use when you finish with adding files

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" %}

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" %}

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 %}

tests/llm/test_parser.py

@@ -141,7 +141,8 @@ def test_parse_json_with_spec(input, expected):
         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():