use correct sources for cli output and success messages

.gitignore

@@ -7,6 +7,7 @@ htmlcov/
 dist/
 workspace/
 pilot-env/
+venv/
 
 .coverage
 *.code-workspace

core/agents/executor.py

@@ -11,10 +11,13 @@ from core.log import get_logger
 from core.proc.exec_log import ExecLog
 from core.proc.process_manager import ProcessManager
 from core.state.state_manager import StateManager
-from core.ui.base import AgentSource, UIBase
+from core.ui.base import AgentSource, UIBase, UISource
 
 log = get_logger(__name__)
 
+CMD_OUTPUT_SOURCE_NAME = "Command output"
+CMD_OUTPUT_SOURCE_TYPE = "cli-output"
+
 
 class CommandResult(BaseModel):
     """
@@ -42,6 +45,8 @@ class Executor(BaseAgent):
         Create a new Executor agent
         """
         self.ui_source = AgentSource(self.display_name, self.agent_type)
+        self.cmd_ui_source = UISource(CMD_OUTPUT_SOURCE_NAME, CMD_OUTPUT_SOURCE_TYPE)
+
         self.ui = ui
         self.state_manager = state_manager
         self.process_manager = ProcessManager(
@@ -58,8 +63,8 @@ class Executor(BaseAgent):
         return self
 
     async def output_handler(self, out, err):
-        await self.stream_handler(out)
+        await self.ui.send_stream_chunk(out, source=self.cmd_ui_source)
-        await self.stream_handler(err)
+        await self.ui.send_stream_chunk(err, source=self.cmd_ui_source)
 
     async def exit_handler(self, process):
         pass

core/agents/tech_lead.py

@@ -11,7 +11,7 @@ from core.db.models.project_state import TaskStatus
 from core.llm.parser import JSONParser
 from core.log import get_logger
 from core.templates.registry import apply_project_template, get_template_description, get_template_summary
-from core.ui.base import ProjectStage
+from core.ui.base import ProjectStage, success_source
 
 log = get_logger(__name__)
 
@@ -94,6 +94,11 @@ class TechLead(BaseAgent):
         return summary
 
     async def ask_for_new_feature(self) -> AgentResponse:
+        if len(self.current_state.epics) > 2:
+            await self.ui.send_message("Your new feature is complete!", source=success_source)
+        else:
+            await self.ui.send_message("Your app is DONE!!! You can start using it right now!", source=success_source)
+
         log.debug("Asking for new feature")
         response = await self.ask_question(
             "Do you have a new feature to add to the project? Just write it here",
@@ -163,7 +168,7 @@ class TechLead(BaseAgent):
         finished_tasks.append(self.next_state.current_task)
 
         log.debug(f"Updating development plan for {epic['name']}")
-        await self.ui.send_message("Updating development plan ...")
+        await self.send_message("Updating development plan ...")
 
         llm = self.get_llm()
         convo = (

core/agents/tech_writer.py

@@ -3,6 +3,7 @@ from core.agents.convo import AgentConvo
 from core.agents.response import AgentResponse
 from core.db.models.project_state import TaskStatus
 from core.log import get_logger
+from core.ui.base import success_source
 
 log = get_logger(__name__)
 
@@ -19,14 +20,37 @@ class TechnicalWriter(BaseAgent):
 
         if n_unfinished in [n_tasks // 2, 1]:
             # Halfway through the initial project, and at the last task
+            await self.send_congratulations()
             await self.create_readme()
 
         self.next_state.action = "Create README.md"
         self.next_state.set_current_task_status(TaskStatus.DOCUMENTED)
         return AgentResponse.done(self)
 
+    async def send_congratulations(self):
+        n_tasks = len(self.current_state.tasks)
+        if not n_tasks:
+            log.warning("No tasks found in the project")
+            return
+
+        n_unfinished = len(self.current_state.unfinished_tasks) - 1
+        n_finished = n_tasks - n_unfinished
+        pct_finished = int(n_finished / n_tasks * 100)
+        n_files = len(self.current_state.files)
+        n_lines = sum(len(f.content.content.splitlines()) for f in self.current_state.files)
+        await self.ui.send_message(
+            "\n\n".join(
+                [
+                    f"CONGRATULATIONS! You reached {pct_finished}% of your project generation!",
+                    f"For now, you have created {n_files} files with a total of {n_lines} lines of code.",
+                    "Before continuing, Pythagora will create some documentation for the project...",
+                ]
+            ),
+            source=success_source,
+        )
+
     async def create_readme(self):
-        await self.ui.send_message("Creating README ...")
+        await self.send_message("Creating README ...")
 
         llm = self.get_llm()
         convo = AgentConvo(self).template("create_readme")

core/ui/base.py

@@ -259,5 +259,14 @@ class UIBase:
 
 
 pythagora_source = UISource("Pythagora", "pythagora")
+success_source = UISource("Congratulations", "success")
 
-__all__ = ["UISource", "AgentSource", "UserInput", "UIBase", "pythagora_source"]
+
+__all__ = [
+    "UISource",
+    "AgentSource",
+    "UserInput",
+    "UIBase",
+    "pythagora_source",
+    "success_source",
+]