Added specification update to new feature add on app load and diff di…

…splay on ui

core/agents/orchestrator.py

@@ -182,6 +182,9 @@ def create_agent(self, prev_response: Optional[AgentResponse]) -> BaseAgent:
                 return Importer(self.state_manager, self.ui, prev_response=prev_response)
             if prev_response.type == ResponseType.EXTERNAL_DOCS_REQUIRED:
                 return ExternalDocumentation(self.state_manager, self.ui, prev_response=prev_response)
+            if prev_response.type == ResponseType.UPDATE_SPECIFICATION:
+                return SpecWriter(self.state_manager, self.ui, prev_response=prev_response)
+
         if not state.specification.description:
             if state.files:
                 # The project has been imported, but not analyzed yet

core/agents/response.py

@@ -45,6 +45,9 @@ class ResponseType(str, Enum):
     EXTERNAL_DOCS_REQUIRED = "external-docs-required"
     """We need to fetch external docs for a task."""
 
+    UPDATE_SPECIFICATION = "update-specification"
+    """We need to update the project specification."""
+
 
 class AgentResponse:
     type: ResponseType = ResponseType.DONE
@@ -144,3 +147,13 @@ def import_project(agent: "BaseAgent") -> "AgentResponse":
     @staticmethod
     def external_docs_required(agent: "BaseAgent") -> "AgentResponse":
         return AgentResponse(type=ResponseType.EXTERNAL_DOCS_REQUIRED, agent=agent)
+
+    @staticmethod
+    def update_specification(agent: "BaseAgent", description: str) -> "AgentResponse":
+        return AgentResponse(
+            type=ResponseType.UPDATE_SPECIFICATION,
+            agent=agent,
+            data={
+                "description": description,
+            },
+        )

core/agents/spec_writer.py

@@ -1,6 +1,6 @@
 from core.agents.base import BaseAgent
 from core.agents.convo import AgentConvo
-from core.agents.response import AgentResponse
+from core.agents.response import AgentResponse, ResponseType
 from core.db.models import Complexity
 from core.db.models.project_state import IterationStatus
 from core.llm.parser import StringParser
@@ -29,7 +29,9 @@ class SpecWriter(BaseAgent):
     async def run(self) -> AgentResponse:
         current_iteration = self.current_state.current_iteration
         if current_iteration is not None and current_iteration.get("status") == IterationStatus.NEW_FEATURE_REQUESTED:
-            return await self.update_spec()
+            return await self.update_spec(iteration_mode=True)
+        elif self.prev_response and self.prev_response.type == ResponseType.UPDATE_SPECIFICATION:
+            return await self.update_spec(iteration_mode=False)
         else:
             return await self.initialize_spec()
 
@@ -84,23 +86,40 @@ async def initialize_spec(self) -> AgentResponse:
         self.next_state.action = SPEC_STEP_NAME
         return AgentResponse.done(self)
 
-    async def update_spec(self) -> AgentResponse:
-        feature_description = self.current_state.current_iteration["user_feedback"]
+    async def update_spec(self, iteration_mode) -> AgentResponse:
+        if iteration_mode:
+            feature_description = self.current_state.current_iteration["user_feedback"]
+        else:
+            feature_description = self.prev_response.data["description"]
+
         await self.send_message(
-            f"Adding feature with the following description to project specification:\n\n{feature_description}"
+            f"Making the following changes to project specification:\n\n{feature_description}\n\nUpdated project specification:"
         )
         llm = self.get_llm()
         convo = AgentConvo(self).template("add_new_feature", feature_description=feature_description)
         llm_response: str = await llm(convo, temperature=0, parser=StringParser())
         updated_spec = llm_response.strip()
+        await self.ui.generate_diff(self.current_state.specification.description, updated_spec)
+        user_response = await self.ask_question(
+            "Do you accept these changes to the project specification?",
+            buttons={"yes": "Yes", "no": "No"},
+            default="yes",
+            buttons_only=True,
+        )
+        await self.ui.close_diff()
 
-        self.next_state.specification = self.current_state.specification.clone()
-        self.next_state.specification.description = updated_spec
+        if user_response.button == "yes":
+            self.next_state.specification = self.current_state.specification.clone()
+            self.next_state.specification.description = updated_spec
+            telemetry.set("updated_prompt", updated_spec)
 
-        telemetry.set("updated_prompt", updated_spec)
+        if iteration_mode:
+            self.next_state.current_iteration["status"] = IterationStatus.FIND_SOLUTION
+            self.next_state.flag_iterations_as_modified()
+        else:
+            complexity = await self.check_prompt_complexity(user_response.text)
+            self.next_state.current_epic["complexity"] = complexity
 
-        self.next_state.current_iteration["status"] = IterationStatus.FIND_SOLUTION
-        self.next_state.flag_iterations_as_modified()
         return AgentResponse.done(self)
 
     async def check_prompt_complexity(self, prompt: str) -> str:
@@ -188,6 +207,6 @@ async def review_spec(self, spec: str) -> str:
         llm = self.get_llm()
         llm_response: str = await llm(convo, temperature=0)
         additional_info = llm_response.strip()
-        if additional_info:
+        if additional_info and len(additional_info) > 6:
             spec += "\nAdditional info/examples:\n" + additional_info
         return spec

core/agents/tech_lead.py

@@ -5,7 +5,6 @@
 from core.agents.base import BaseAgent
 from core.agents.convo import AgentConvo
 from core.agents.response import AgentResponse
-from core.db.models import Complexity
 from core.db.models.project_state import TaskStatus
 from core.llm.parser import JSONParser
 from core.log import get_logger
@@ -136,12 +135,12 @@ async def ask_for_new_feature(self) -> AgentResponse:
                 "description": response.text,
                 "summary": None,
                 "completed": False,
-                "complexity": Complexity.HARD,
+                "complexity": None,  # Determined and defined in SpecWriter
             }
         ]
         # Orchestrator will rerun us to break down the new feature epic
         self.next_state.action = f"Start of feature #{len(self.current_state.epics)}"
-        return AgentResponse.done(self)
+        return AgentResponse.update_specification(self, response.text)
 
     async def plan_epic(self, epic) -> AgentResponse:
         log.debug(f"Planning tasks for the epic: {epic['name']}")

core/agents/troubleshooter.py

@@ -223,7 +223,7 @@ async def get_user_feedback(
         is_loop = False
         should_iterate = True
 
-        test_message = "Can you check if the app works please?"
+        test_message = "Please check if the app is working"
         if user_instructions:
             hint = " Here is a description of what should be working:\n\n" + user_instructions
 
@@ -259,13 +259,11 @@ async def get_user_feedback(
             is_loop = True
 
         elif user_response.button == "change":
-            user_description = await self.ask_question(
-                "Please describe the change you want to make (one at the time please)"
-            )
+            user_description = await self.ask_question("Please describe the change you want to make (one at a time)")
             change_description = user_description.text
 
         elif user_response.button == "bug":
-            user_description = await self.ask_question("Please describe the issue you found (one at the time please)")
+            user_description = await self.ask_question("Please describe the issue you found (one at a time)")
             bug_report = user_description.text
 
         return should_iterate, is_loop, bug_report, change_description

core/prompts/spec-writer/add_new_feature.prompt

@@ -7,15 +7,13 @@ This might include:
 * information which 3rd party packages or APIs to use or avoid
 * concrete examples of API requests/responses, library usage, or other external documentation
 
-Here is the client brief:
----PROJECT-SPECIFICATION-START---
+Here is the original project specification:
 {{ state.specification.description }}
----PROJECT-SPECIFICATION-END---
 
-Here is the specification your team came up with:
+Here is the new feature description:
 ---FEATURE-DESCRIPTION-START---
 {{ feature_description }}
 ---FEATURE-DESCRIPTION-END---
 
-In your response, output the new updated project specification. 
-If there is no feature description, just output the original project specification.
+In your response, output only the new updated project specification, without any additional messages to the user. 
+If there is no feature description just output the original project specification.

core/ui/ipc_client.py

@@ -42,6 +42,8 @@ class MessageType(str, Enum):
     IMPORT_PROJECT = "importProject"
     APP_FINISHED = "appFinished"
     FEATURE_FINISHED = "featureFinished"
+    GENERATE_DIFF = "generateDiff"
+    CLOSE_DIFF = "closeDiff"
 
 
 class Message(BaseModel):
@@ -356,6 +358,19 @@ async def send_project_stats(self, stats: dict):
             content=stats,
         )
 
+    async def generate_diff(self, file_old: str, file_new: str):
+        await self._send(
+            MessageType.GENERATE_DIFF,
+            content={
+                "file_old": file_old,
+                "file_new": file_new,
+            },
+        )
+
+    async def close_diff(self):
+        log.debug("Sending signal to close the generated diff file")
+        await self._send(MessageType.CLOSE_DIFF)
+
     async def loading_finished(self):
         log.debug("Sending project loading finished signal to the extension")
         await self._send(MessageType.LOADING_FINISHED)

tests/agents/test_tech_lead.py

@@ -56,7 +56,7 @@ async def test_ask_for_feature(agentcontext):
 
     tl = TechLead(sm, ui)
     response = await tl.run()
-    assert response.type == ResponseType.DONE
+    assert response.type == ResponseType.UPDATE_SPECIFICATION
 
     await sm.commit()