Merge pull request #1057 from Pythagora-io/deep_focus_logging

Implemented the first step in deep focus initiative. Now, when the us…

core/agents/bug_hunter.py

@@ -0,0 +1,172 @@
+from enum import Enum
+
+from pydantic import BaseModel, Field
+
+from core.agents.base import BaseAgent
+from core.agents.convo import AgentConvo
+from core.agents.response import AgentResponse
+from core.config import CHECK_LOGS_AGENT_NAME, magic_words
+from core.db.models.project_state import IterationStatus
+from core.llm.parser import JSONParser
+from core.log import get_logger
+
+log = get_logger(__name__)
+
+
+class StepType(str, Enum):
+    ADD_LOG = "add_log"
+    EXPLAIN_PROBLEM = "explain_problem"
+    GET_ADDITIONAL_FILES = "get_additional_files"
+
+
+class HuntConclusionType(str, Enum):
+    ADD_LOGS = magic_words.ADD_LOGS
+    PROBLEM_IDENTIFIED = magic_words.PROBLEM_IDENTIFIED
+
+
+class HuntConclusionOptions(BaseModel):
+    conclusion: HuntConclusionType = Field(
+        description=f"If more logs are needed to identify the problem, respond with '{magic_words.ADD_LOGS}'. If the problem is identified, respond with '{magic_words.PROBLEM_IDENTIFIED}'."
+    )
+
+
+class BugHunter(BaseAgent):
+    agent_type = "bug-hunter"
+    display_name = "Bug Hunter"
+
+    async def run(self) -> AgentResponse:
+        current_iteration = self.current_state.current_iteration
+
+        if "bug_reproduction_description" not in current_iteration:
+            await self.get_bug_reproduction_instructions()
+        if current_iteration["status"] == IterationStatus.HUNTING_FOR_BUG:
+            # TODO determine how to find a bug (eg. check in db, ask user a question, etc.)
+            return await self.check_logs()
+        elif current_iteration["status"] == IterationStatus.AWAITING_USER_TEST:
+            return await self.ask_user_to_test(False, True)
+        elif current_iteration["status"] == IterationStatus.AWAITING_BUG_REPRODUCTION:
+            return await self.ask_user_to_test(True, False)
+
+    async def get_bug_reproduction_instructions(self):
+        llm = self.get_llm()
+        convo = AgentConvo(self).template(
+            "get_bug_reproduction_instructions",
+            current_task=self.current_state.current_task,
+            user_feedback=self.current_state.current_iteration["user_feedback"],
+            user_feedback_qa=self.current_state.current_iteration["user_feedback_qa"],
+            docs=self.current_state.docs,
+            next_solution_to_try=None,
+        )
+        bug_reproduction_instructions = await llm(convo, temperature=0)
+        self.next_state.current_iteration["bug_reproduction_description"] = bug_reproduction_instructions
+
+    async def check_logs(self, logs_message: str = None):
+        llm = self.get_llm(CHECK_LOGS_AGENT_NAME)
+        convo = AgentConvo(self).template(
+            "iteration",
+            current_task=self.current_state.current_task,
+            user_feedback=self.current_state.current_iteration["user_feedback"],
+            user_feedback_qa=self.current_state.current_iteration["user_feedback_qa"],
+            docs=self.current_state.docs,
+            magic_words=magic_words,
+            next_solution_to_try=None,
+        )
+
+        for hunting_cycle in self.current_state.current_iteration.get("bug_hunting_cycles", []):
+            convo = convo.assistant(hunting_cycle["human_readable_instructions"]).template(
+                "log_data",
+                backend_logs=hunting_cycle["backend_logs"],
+                frontend_logs=hunting_cycle["frontend_logs"],
+                fix_attempted=hunting_cycle["fix_attempted"],
+            )
+
+        human_readable_instructions = await llm(convo, temperature=0.5)
+
+        convo = (
+            AgentConvo(self)
+            .template(
+                "bug_found_or_add_logs",
+                hunt_conclusion=human_readable_instructions,
+            )
+            .require_schema(HuntConclusionOptions)
+        )
+        llm = self.get_llm()
+        hunt_conclusion = await llm(convo, parser=JSONParser(HuntConclusionOptions), temperature=0)
+
+        self.next_state.current_iteration["description"] = human_readable_instructions
+        self.next_state.current_iteration["bug_hunting_cycles"] += [
+            {
+                "human_readable_instructions": human_readable_instructions,
+                "fix_attempted": any(
+                    c["fix_attempted"] for c in self.current_state.current_iteration["bug_hunting_cycles"]
+                ),
+            }
+        ]
+
+        if hunt_conclusion.conclusion == magic_words.PROBLEM_IDENTIFIED:
+            # if no need for logs, implement iteration same as before
+            self.next_state.current_iteration["status"] = IterationStatus.AWAITING_BUG_FIX
+            await self.send_message("The bug is found - I'm  attempting to fix it.")
+        else:
+            # if logs are needed, add logging steps
+            self.next_state.current_iteration["status"] = IterationStatus.AWAITING_LOGGING
+            await self.send_message("Adding more logs to identify the bug.")
+
+        self.next_state.flag_iterations_as_modified()
+        return AgentResponse.done(self)
+
+    async def ask_user_to_test(self, awaiting_bug_reproduction: bool = False, awaiting_user_test: bool = False):
+        await self.send_message(
+            "You can reproduce the bug like this:\n\n"
+            + self.current_state.current_iteration["bug_reproduction_description"]
+        )
+
+        if self.current_state.run_command:
+            await self.ui.send_run_command(self.current_state.run_command)
+
+        if awaiting_user_test:
+            user_feedback = await self.ask_question(
+                "Is the bug you reported fixed now?",
+                buttons={"yes": "Yes, the issue is fixed", "no": "No"},
+                default="continue",
+                buttons_only=True,
+                hint="Instructions for testing:\n\n"
+                + self.current_state.current_iteration["bug_reproduction_description"],
+            )
+            self.next_state.current_iteration["bug_hunting_cycles"][-1]["fix_attempted"] = True
+
+            if user_feedback.button == "yes":
+                self.next_state.complete_iteration()
+            else:
+                awaiting_bug_reproduction = True
+
+        if awaiting_bug_reproduction:
+            # TODO how can we get FE and BE logs automatically?
+            backend_logs = await self.ask_question(
+                "Please do exactly what you did in the last iteration, paste **BACKEND** logs here and click CONTINUE.",
+                buttons={"continue": "Continue", "done": "Bug is fixed"},
+                default="continue",
+                hint="Instructions for testing:\n\n"
+                + self.current_state.current_iteration["bug_reproduction_description"],
+            )
+
+            if backend_logs.button == "done":
+                self.next_state.complete_iteration()
+            else:
+                frontend_logs = await self.ask_question(
+                    "Please paste **frontend** logs here and click CONTINUE.",
+                    buttons={"continue": "Continue", "done": "Bug is fixed"},
+                    default="continue",
+                    hint="Instructions for testing:\n\n"
+                    + self.current_state.current_iteration["bug_reproduction_description"],
+                )
+
+                # TODO select only the logs that are new (with PYTHAGORA_DEBUGGING_LOG)
+                self.next_state.current_iteration["bug_hunting_cycles"][-1]["backend_logs"] = backend_logs.text
+                self.next_state.current_iteration["bug_hunting_cycles"][-1]["frontend_logs"] = frontend_logs.text
+                self.next_state.current_iteration["status"] = IterationStatus.HUNTING_FOR_BUG
+
+                if frontend_logs.button == "done":
+                    self.next_state.complete_iteration()
+
+        return AgentResponse.done(self)

core/agents/developer.py

@@ -1,13 +1,13 @@
-from enum import Enum
-from typing import Annotated, Literal, Optional, Union
+from typing import Optional
 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.response import AgentResponse, ResponseType
-from core.db.models.project_state import TaskStatus
+from core.db.models.project_state import IterationStatus, TaskStatus
 from core.db.models.specification import Complexity
 from core.llm.parser import JSONParser
 from core.log import get_logger
@@ -16,47 +16,6 @@
 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):
     relevant_files: list[str] = Field(description="List of relevant files for the current task.")
 
@@ -109,6 +68,19 @@ async def breakdown_current_iteration(self, task_review_feedback: Optional[str]
             n_tasks = 1
             log.debug(f"Breaking down the task review feedback {task_review_feedback}")
             await self.send_message("Breaking down the task review feedback...")
+        elif self.current_state.current_iteration["status"] in (
+            IterationStatus.AWAITING_BUG_FIX,
+            IterationStatus.AWAITING_LOGGING,
+        ):
+            iteration = self.current_state.current_iteration
+            current_task["task_review_feedback"] = None
+
+            description = iteration["bug_hunting_cycles"][-1]["human_readable_instructions"]
+            user_feedback = iteration["user_feedback"]
+            source = "bug_hunt"
+            n_tasks = len(self.next_state.iterations)
+            log.debug(f"Breaking down the logging cycle {description}")
+            await self.send_message("Breaking down the current iteration logging cycle ...")
         else:
             iteration = self.current_state.current_iteration
             current_task["task_review_feedback"] = None
@@ -156,8 +128,22 @@ async def breakdown_current_iteration(self, task_review_feedback: Optional[str]
         self.set_next_steps(response, source)
 
         if iteration:
-            self.next_state.complete_iteration()
-            self.next_state.action = f"Troubleshooting #{len(self.current_state.iterations)}"
+            if "status" not in iteration or (
+                iteration["status"] in (IterationStatus.AWAITING_USER_TEST, IterationStatus.AWAITING_BUG_REPRODUCTION)
+            ):
+                # This is just a support for old iterations that don't have status
+                self.next_state.complete_iteration()
+                self.next_state.action = f"Troubleshooting #{len(self.current_state.iterations)}"
+            elif iteration["status"] == IterationStatus.IMPLEMENT_SOLUTION:
+                # If the user requested a change, then, we'll implement it and go straight back to testing
+                self.next_state.complete_iteration()
+                self.next_state.action = f"Troubleshooting #{len(self.current_state.iterations)}"
+            elif iteration["status"] == IterationStatus.AWAITING_BUG_FIX:
+                # If bug fixing is done, ask user to test again
+                self.next_state.current_iteration["status"] = IterationStatus.AWAITING_USER_TEST
+            elif iteration["status"] == IterationStatus.AWAITING_LOGGING:
+                # If logging is done, ask user to reproduce the bug
+                self.next_state.current_iteration["status"] = IterationStatus.AWAITING_BUG_REPRODUCTION
         else:
             self.next_state.action = "Task review feedback"
 
@@ -265,7 +251,14 @@ def set_next_steps(self, response: TaskSteps, source: str):
             }
             for step in response.steps
         ]
-        if len(self.next_state.unfinished_steps) > 0 and source != "review":
+        if (
+            len(self.next_state.unfinished_steps) > 0
+            and source != "review"
+            and (
+                self.next_state.current_iteration is None
+                or self.next_state.current_iteration["status"] != IterationStatus.AWAITING_LOGGING
+            )
+        ):
             self.next_state.steps += [
                 # TODO: add refactor step here once we have the refactor agent
                 {

core/agents/error_handler.py

@@ -3,6 +3,7 @@
 from core.agents.base import BaseAgent
 from core.agents.convo import AgentConvo
 from core.agents.response import AgentResponse
+from core.db.models.project_state import IterationStatus
 from core.log import get_logger
 
 log = get_logger(__name__)
@@ -110,7 +111,8 @@ async def handle_command_error(self, message: str, details: dict) -> AgentRespon
                 "description": llm_response,
                 "alternative_solutions": [],
                 "attempts": 1,
-                "completed": False,
+                "status": IterationStatus.HUNTING_FOR_BUG,
+                "bug_hunting_cycles": [],
             }
         ]
         # TODO: maybe have ProjectState.finished_steps as well? would make the debug/ran_command prompts nicer too

core/agents/mixins.py

@@ -1,8 +1,52 @@
-from typing import Optional
+from enum import Enum
+from typing import Annotated, Literal, Optional, Union
+
+from pydantic import BaseModel, Field
 
 from core.agents.convo import AgentConvo
 
 
+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 IterationPromptMixin:
     """
     Provides a method to find a solution to a problem based on user feedback.
@@ -16,13 +60,15 @@ async def find_solution(
         *,
         user_feedback_qa: Optional[list[str]] = None,
         next_solution_to_try: Optional[str] = None,
+        bug_hunting_cycles: Optional[dict] = None,
     ) -> str:
         """
         Generate a new solution for the problem the user reported.
 
         :param user_feedback: User feedback about the problem.
         :param user_feedback_qa: Additional q/a about the problem provided by the user (optional).
         :param next_solution_to_try: Hint from ProblemSolver on which solution to try (optional).
+        :param bug_hunting_cycles: Data about logs that need to be added to the code (optional).
         :return: The generated solution to the problem.
         """
         llm = self.get_llm()
@@ -32,6 +78,7 @@ async def find_solution(
             user_feedback=user_feedback,
             user_feedback_qa=user_feedback_qa,
             next_solution_to_try=next_solution_to_try,
+            bug_hunting_cycles=bug_hunting_cycles,
         )
         llm_solution: str = await llm(convo)
         return llm_solution