From 7cc7ce880d32d143cfeef9e05301b0a6afa85548 Mon Sep 17 00:00:00 2001 From: Senko Rasic Date: Sun, 26 May 2024 10:44:30 +0200 Subject: [PATCH] use correct sources for cli output and success messages --- .gitignore | 1 + core/agents/executor.py | 11 ++++++++--- core/agents/tech_lead.py | 9 +++++++-- core/agents/tech_writer.py | 26 +++++++++++++++++++++++++- core/ui/base.py | 11 ++++++++++- 5 files changed, 51 insertions(+), 7 deletions(-) diff --git a/.gitignore b/.gitignore index 5c9918f7..e7a1d6e6 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,7 @@ htmlcov/ dist/ workspace/ pilot-env/ +venv/ .coverage *.code-workspace diff --git a/core/agents/executor.py b/core/agents/executor.py index d759da50..38748ebb 100644 --- a/core/agents/executor.py +++ b/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.stream_handler(err) + await self.ui.send_stream_chunk(out, source=self.cmd_ui_source) + await self.ui.send_stream_chunk(err, source=self.cmd_ui_source) async def exit_handler(self, process): pass diff --git a/core/agents/tech_lead.py b/core/agents/tech_lead.py index e268e417..0de4dee4 100644 --- a/core/agents/tech_lead.py +++ b/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 = ( diff --git a/core/agents/tech_writer.py b/core/agents/tech_writer.py index d6bef712..3ab625e7 100644 --- a/core/agents/tech_writer.py +++ b/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") diff --git a/core/ui/base.py b/core/ui/base.py index 5ccbea63..ebc48757 100644 --- a/core/ui/base.py +++ b/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", +]