move chat with breakdown to mixin

core/agents/bug_hunter.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 TestSteps
+from core.agents.mixins import ChatWithBreakdownMixin, TestSteps
 from core.agents.response import AgentResponse
 from core.config import CHECK_LOGS_AGENT_NAME, magic_words
 from core.db.models.project_state import IterationStatus
@@ -43,7 +43,7 @@ class ImportantLogsForDebugging(BaseModel):
     logs: list[ImportantLog] = Field(description="Important logs that will help the human debug the current bug.")
 
 
-class BugHunter(BaseAgent):
+class BugHunter(ChatWithBreakdownMixin, BaseAgent):
     agent_type = "bug-hunter"
     display_name = "Bug Hunter"
 
@@ -93,22 +93,7 @@ class BugHunter(BaseAgent):
 
         convo.assistant(human_readable_instructions)
 
-        while True:
-            chat = await self.ui.ask_question(
-                "Are you happy with the breakdown? Now is a good time to ask questions or suggest changes.",
-                buttons={"yes": "Yes, looks good!"},
-                default="yes",
-                verbose=False,
-            )
-            if chat.button == "yes":
-                break
-
-            if len(convo.messages) > 11:
-                convo.trim(3, 2)
-
-            convo.user(chat.text)
-            human_readable_instructions: str = await llm(convo)
-            convo.assistant(human_readable_instructions)
+        human_readable_instructions = await self.chat_with_breakdown(convo, human_readable_instructions)
 
         convo = (
             AgentConvo(self)

core/agents/developer.py

@@ -7,7 +7,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.mixins import ChatWithBreakdownMixin, RelevantFilesMixin
 from core.agents.response import AgentResponse
 from core.config import PARSE_TASK_AGENT_NAME, TASK_BREAKDOWN_AGENT_NAME
 from core.db.models.project_state import IterationStatus, TaskStatus
@@ -72,7 +72,7 @@ class TaskSteps(BaseModel):
     steps: list[Step]
 
 
-class Developer(RelevantFilesMixin, BaseAgent):
+class Developer(ChatWithBreakdownMixin, RelevantFilesMixin, BaseAgent):
     agent_type = "developer"
     display_name = "Developer"
 
@@ -230,22 +230,7 @@ class Developer(RelevantFilesMixin, BaseAgent):
         response: str = await llm(convo)
         convo.assistant(response)
 
-        while True:
-            chat = await self.ask_question(
-                "Are you happy with the breakdown? Now is a good time to ask questions or suggest changes.",
-                buttons={"yes": "Yes, looks good!"},
-                default="yes",
-                verbose=False,
-            )
-            if chat.button == "yes":
-                break
-
-            if len(convo.messages) > 11:
-                convo.trim(3, 2)
-
-            convo.user(chat.text)
-            response: str = await llm(convo)
-            convo.assistant(response)
+        response = await self.chat_with_breakdown(convo, response)
 
         self.next_state.tasks[current_task_index] = {
             **current_task,

core/agents/mixins.py

@@ -6,7 +6,7 @@ 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, TROUBLESHOOTER_BUG_REPORT
+from core.config import GET_RELEVANT_FILES_AGENT_NAME, TASK_BREAKDOWN_AGENT_NAME, TROUBLESHOOTER_BUG_REPORT
 from core.llm.parser import JSONParser
 from core.log import get_logger
 
@@ -49,7 +49,42 @@ class TestSteps(BaseModel):
     steps: List[Test]
 
 
-class IterationPromptMixin:
+class ChatWithBreakdownMixin:
+    """
+    Provides a method to chat with the user and provide a breakdown of the conversation.
+    """
+
+    async def chat_with_breakdown(self, convo: AgentConvo, breakdown: str) -> AgentConvo:
+        """
+        Chat with the user and provide a breakdown of the conversation.
+
+        :param convo: The conversation object.
+        :param breakdown: The breakdown of the conversation.
+        :return: The breakdown.
+        """
+
+        llm = self.get_llm(TASK_BREAKDOWN_AGENT_NAME, stream_output=True)
+        while True:
+            chat = await self.ask_question(
+                "Are you happy with the breakdown? Now is a good time to ask questions or suggest changes.",
+                buttons={"yes": "Yes, looks good!"},
+                default="yes",
+                verbose=False,
+            )
+            if chat.button == "yes":
+                break
+
+            if len(convo.messages) > 11:
+                convo.trim(3, 2)
+
+            convo.user(chat.text)
+            breakdown: str = await llm(convo)
+            convo.assistant(breakdown)
+
+        return breakdown
+
+
+class IterationPromptMixin(ChatWithBreakdownMixin):
     """
     Provides a method to find a solution to a problem based on user feedback.
 
@@ -83,6 +118,9 @@ class IterationPromptMixin:
             test_instructions=json.loads(self.current_state.current_task.get("test_instructions") or "[]"),
         )
         llm_solution: str = await llm(convo)
+
+        llm_solution = await self.chat_with_breakdown(convo, llm_solution)
+
         return llm_solution
 
 

core/agents/troubleshooter.py

@@ -6,7 +6,7 @@ from pydantic import BaseModel, Field
 
 from core.agents.base import BaseAgent
 from core.agents.convo import AgentConvo
-from core.agents.mixins import IterationPromptMixin, RelevantFilesMixin, TestSteps
+from core.agents.mixins import ChatWithBreakdownMixin, IterationPromptMixin, RelevantFilesMixin, TestSteps
 from core.agents.response import AgentResponse
 from core.config import TROUBLESHOOTER_GET_RUN_COMMAND
 from core.db.models.file import File
@@ -31,7 +31,7 @@ class RouteFilePaths(BaseModel):
     files: list[str] = Field(description="List of paths for files that contain routes")
 
 
-class Troubleshooter(IterationPromptMixin, RelevantFilesMixin, BaseAgent):
+class Troubleshooter(ChatWithBreakdownMixin, IterationPromptMixin, RelevantFilesMixin, BaseAgent):
     agent_type = "troubleshooter"
     display_name = "Troubleshooter"