From dac640b20460073b5643f5c002799f951bf2fedb Mon Sep 17 00:00:00 2001 From: mijauexe Date: Thu, 29 May 2025 08:26:29 +0200 Subject: [PATCH] Clean up code --- core/agents/orchestrator.py | 37 ++++++----- core/agents/spec_writer.py | 34 +++++----- core/agents/wizard.py | 1 - core/cli/main.py | 2 +- core/db/models/project.py | 3 - core/db/models/project_state.py | 4 ++ core/disk/vfs.py | 2 + core/prompts/spec-writer/need_auth.prompt | 4 +- core/prompts/spec-writer/project_name.prompt | 2 +- core/state/state_manager.py | 67 ++++++-------------- 10 files changed, 66 insertions(+), 90 deletions(-) diff --git a/core/agents/orchestrator.py b/core/agents/orchestrator.py index da1e1f60..0d20f78d 100644 --- a/core/agents/orchestrator.py +++ b/core/agents/orchestrator.py @@ -153,9 +153,6 @@ class Orchestrator(BaseAgent, GitMixin): await self.ui.send_project_root(self.state_manager.get_full_project_root()) continue - if response.type == ResponseType.CREATE_SPECIFICATION: - continue - # TODO: rollback changes to "next" so they aren't accidentally committed? return True @@ -412,23 +409,25 @@ class Orchestrator(BaseAgent, GitMixin): will trigger the HumanInput agent to ask the user to provide the required input. """ - n_epics = len(self.next_state.epics) - n_finished_epics = n_epics - len(self.next_state.unfinished_epics) - n_tasks = len(self.next_state.tasks) - n_finished_tasks = n_tasks - len(self.next_state.unfinished_tasks) - n_iterations = len(self.next_state.iterations) - n_finished_iterations = n_iterations - len(self.next_state.unfinished_iterations) - n_steps = len(self.next_state.steps) - n_finished_steps = n_steps - len(self.next_state.unfinished_steps) + if self.next_state and self.next_state.tasks: + n_epics = len(self.next_state.epics) + n_finished_epics = n_epics - len(self.next_state.unfinished_epics) + n_tasks = len(self.next_state.tasks) + n_finished_tasks = n_tasks - len(self.next_state.unfinished_tasks) + n_iterations = len(self.next_state.iterations) + n_finished_iterations = n_iterations - len(self.next_state.unfinished_iterations) + n_steps = len(self.next_state.steps) + n_finished_steps = n_steps - len(self.next_state.unfinished_steps) + + log.debug( + f"Agent {agent.__class__.__name__} is done, " + f"committing state for step {self.current_state.step_index}: " + f"{n_finished_epics}/{n_epics} epics, " + f"{n_finished_tasks}/{n_tasks} tasks, " + f"{n_finished_iterations}/{n_iterations} iterations, " + f"{n_finished_steps}/{n_steps} dev steps." + ) - log.debug( - f"Agent {agent.__class__.__name__} is done, " - f"committing state for step {self.current_state.step_index}: " - f"{n_finished_epics}/{n_epics} epics, " - f"{n_finished_tasks}/{n_tasks} tasks, " - f"{n_finished_iterations}/{n_iterations} iterations, " - f"{n_finished_steps}/{n_steps} dev steps." - ) await self.state_manager.commit() # If there are any new or modified files changed outside Pythagora, diff --git a/core/agents/spec_writer.py b/core/agents/spec_writer.py index f8b799a5..9d1285f7 100644 --- a/core/agents/spec_writer.py +++ b/core/agents/spec_writer.py @@ -3,7 +3,7 @@ import secrets from core.agents.base import BaseAgent from core.agents.convo import AgentConvo from core.agents.response import AgentResponse, ResponseType -from core.config import DESCRIBE_FILES_AGENT_NAME, SPEC_WRITER_AGENT_NAME +from core.config import DEFAULT_AGENT_NAME, SPEC_WRITER_AGENT_NAME from core.config.actions import SPEC_CHANGE_FEATURE_STEP_NAME, SPEC_CHANGE_STEP_NAME, SPEC_CREATE_STEP_NAME from core.db.models import Complexity from core.db.models.project_state import IterationStatus @@ -50,6 +50,8 @@ class SpecWriter(BaseAgent): self.state_manager, self.process_manager, ) + if not self.state_manager.template: + self.state_manager.template = {} self.state_manager.template["template"] = template log.info(f"Applying project template: {template.name}") summary = await template.apply() @@ -81,19 +83,19 @@ class SpecWriter(BaseAgent): await self.ui.send_project_stage({"stage": ProjectStage.PROJECT_NAME}) - llm = self.get_llm(DESCRIBE_FILES_AGENT_NAME) + llm = self.get_llm(DEFAULT_AGENT_NAME) convo = AgentConvo(self).template( "project_name", description=llm_assisted_description, ) llm_response: str = await llm(convo, temperature=0) - project_name = llm_response.strip().replace(" ", "_").replace("-", "_") + project_name = llm_response.strip() + + self.state_manager.project.name = project_name + self.state_manager.project.folder_name = project_name.replace(" ", "_").replace("-", "_") + + self.state_manager.file_system = await self.state_manager.init_file_system(load_existing=False) - await self.state_manager.rename_project(self.state_manager.project.id, project_name) - self.state_manager.file_system.root = ( - self.state_manager.get_full_parent_project_root() + "/" + self.state_manager.project.folder_name - ) - self.state_manager.file_system.ignore_matcher.root_path = self.state_manager.file_system.root self.process_manager.root_dir = self.state_manager.file_system.root self.next_state.knowledge_base.user_options["original_description"] = description @@ -106,6 +108,11 @@ class SpecWriter(BaseAgent): llm_assisted_description = self.current_state.knowledge_base.user_options["project_description"] description = self.current_state.knowledge_base.user_options["original_description"] + convo = AgentConvo(self).template( + "build_full_specification", + initial_prompt=llm_assisted_description.strip(), + ) + while True: user_done_with_description = await self.ask_question( "Are you satisfied with the project description?", @@ -125,12 +132,7 @@ class SpecWriter(BaseAgent): allow_empty=False, ) - convo = ( - AgentConvo(self).template( - "build_full_specification", - initial_prompt=llm_assisted_description.strip(), - ) - ).template("add_to_specification", user_message=user_add_to_spec.text.strip()) + convo = convo.template("add_to_specification", user_message=user_add_to_spec.text.strip()) if len(convo.messages) > 6: convo.slice(1, 4) @@ -138,6 +140,8 @@ class SpecWriter(BaseAgent): await self.ui.start_important_stream() llm_assisted_description = await llm(convo) + convo = convo.assistant(llm_assisted_description) + llm = self.get_llm(SPEC_WRITER_AGENT_NAME) convo = AgentConvo(self).template( "need_auth", @@ -162,7 +166,7 @@ class SpecWriter(BaseAgent): self.next_state.specification.original_description = description self.next_state.specification.description = llm_assisted_description - complexity = await self.check_prompt_complexity(description) + complexity = await self.check_prompt_complexity(llm_assisted_description) self.next_state.specification.complexity = complexity telemetry.set("initial_prompt", description) diff --git a/core/agents/wizard.py b/core/agents/wizard.py index 9fb404ab..bb5691f0 100644 --- a/core/agents/wizard.py +++ b/core/agents/wizard.py @@ -157,7 +157,6 @@ class Wizard(BaseAgent): session = inspect(self.next_state).async_session session.add(knowledge_base) self.next_state.knowledge_base = knowledge_base - self.next_state.knowledge_base.user_options = options self.next_state.epics = [ { diff --git a/core/cli/main.py b/core/cli/main.py index dea36883..b173a415 100644 --- a/core/cli/main.py +++ b/core/cli/main.py @@ -152,7 +152,7 @@ async def start_new_project(sm: StateManager, ui: UIBase) -> bool: {"language": stack.button}, ) - project_state = await sm.create_project(project_type=stack.button, create_dir=False) + project_state = await sm.create_project(project_type=stack.button) return project_state is not None diff --git a/core/db/models/project.py b/core/db/models/project.py index d89d3f88..0c5f32a3 100644 --- a/core/db/models/project.py +++ b/core/db/models/project.py @@ -70,9 +70,6 @@ class Project(Base): project.name = name project.folder_name = dir_name - # Commit changes - # await session.commit() - return project async def get_branch(self, name: Optional[str] = None) -> Optional["Branch"]: diff --git a/core/db/models/project_state.py b/core/db/models/project_state.py index e04ddb11..d2c8e8b8 100644 --- a/core/db/models/project_state.py +++ b/core/db/models/project_state.py @@ -124,6 +124,8 @@ class ProjectState(Base): :return: List of unfinished iterations. """ + if not self.iterations: + return [] return [ iteration for iteration in self.iterations if iteration.get("status") not in (None, IterationStatus.DONE) ] @@ -147,6 +149,8 @@ class ProjectState(Base): :return: List of unfinished tasks. """ + if not self.tasks: + return [] return [task for task in self.tasks if task.get("status") != TaskStatus.DONE] @property diff --git a/core/disk/vfs.py b/core/disk/vfs.py index 066ce6de..a35b0611 100644 --- a/core/disk/vfs.py +++ b/core/disk/vfs.py @@ -123,6 +123,8 @@ class LocalDiskVFS(VirtualFileSystem): if not os.path.isdir(root): if create: os.makedirs(root) + else: + raise ValueError(f"Root directory does not exist: {root}") else: if not allow_existing: raise FileExistsError(f"Root directory already exists: {root}") diff --git a/core/prompts/spec-writer/need_auth.prompt b/core/prompts/spec-writer/need_auth.prompt index 8b32cd7f..977c46dd 100644 --- a/core/prompts/spec-writer/need_auth.prompt +++ b/core/prompts/spec-writer/need_auth.prompt @@ -2,5 +2,5 @@ Decide if the user wants to use authentication (login and register) for the app ```text {{description}} ``` -Reply with "Yes" or "No" only, without any additional text or explanation. -If the description does not provide enough information to make a decision, reply with "Yes". \ No newline at end of file +Reply with Yes or No only (without quotation marks), and no additional text or explanation. +If the description does not provide enough information to make a decision, reply with Yes. \ No newline at end of file diff --git a/core/prompts/spec-writer/project_name.prompt b/core/prompts/spec-writer/project_name.prompt index e3b2df69..0e7e27b8 100644 --- a/core/prompts/spec-writer/project_name.prompt +++ b/core/prompts/spec-writer/project_name.prompt @@ -2,4 +2,4 @@ Generate a simple project name from the following description: ```text {{description}} ``` -Use a maximum of 2-3 words, no more than 15 characters, and avoid using special characters or spaces. \ No newline at end of file +Use a maximum of 2-3 words, no more than 15 characters, and avoid using special characters or spaces. Respond with only the project name, without any additional text or formatting. \ No newline at end of file diff --git a/core/state/state_manager.py b/core/state/state_manager.py index e5c8526e..1dba632b 100644 --- a/core/state/state_manager.py +++ b/core/state/state_manager.py @@ -1,6 +1,7 @@ import asyncio import os.path import re +import sys import traceback from contextlib import asynccontextmanager from typing import TYPE_CHECKING, Optional @@ -103,42 +104,6 @@ class StateManager: async def get_file_for_project(self, state_id: UUID, path: str): return await Project.get_file_for_project(self.current_session, state_id, path) - async def rename_project(self, id: UUID, project_name: str): - log.debug("Renaming project %s", project_name) - - new_dir_name = self.get_unique_folder_name(self.get_full_parent_project_root(), project_name) - - project = await Project.rename(self.current_session, id, project_name, new_dir_name) - - self.project = project - - def get_unique_folder_name(self, current_dir: str, folder_name: str) -> str: - """ - Generate a unique folder name based on the given name. - - If the folder already exists in the current directory, append a unique - identifier (7 characters from UUID) to ensure uniqueness. - - :param folder_name: Base folder name to check - :param current_dir: Current directory path (defaults to current working directory) - :return: A unique folder name that doesn't exist yet - """ - # Use current working directory if not specified - base_path = current_dir if current_dir else os.getcwd() - - # Full path to check - full_path = os.path.join(base_path, folder_name) - - # If the path doesn't exist, return the original folder name - if not os.path.exists(full_path): - return folder_name - - # Generate a unique name with UUID - unique_id = uuid4().hex[:7] - unique_folder_name = f"{folder_name}-{unique_id}" - - return unique_folder_name - async def get_project_state_by_id(self, state_id: UUID) -> Optional[ProjectState]: """ Get a project state by its ID. @@ -152,8 +117,7 @@ class StateManager: self, name: Optional[str] = "temp-project", project_type: Optional[str] = "node", - folder_name: Optional[str] = None, - create_dir: Optional[bool] = True, + folder_name: Optional[str] = "temp-project" if "pytest" not in sys.modules else None, ) -> Project: """ Create a new project and set it as the current one. @@ -171,7 +135,9 @@ class StateManager: # even for a new project, eg. offline changes check and stats updating await state.awaitable_attrs.files - await session.commit() + is_test = "pytest" in sys.modules + if is_test: + await session.commit() log.info( f'Created new project "{name}" (id={project.id}) ' @@ -185,7 +151,10 @@ class StateManager: self.next_state = state self.project = project self.branch = branch - self.file_system = await self.init_file_system(load_existing=False, create_dir=create_dir) + + if is_test: + self.file_system = await self.init_file_system(load_existing=False) + return project async def delete_project(self, project_id: UUID) -> bool: @@ -504,7 +473,7 @@ class StateManager: delta_lines = len(content.splitlines()) - len(original_content.splitlines()) telemetry.inc("created_lines", delta_lines) - async def init_file_system(self, load_existing: bool, create_dir: bool = True) -> VirtualFileSystem: + async def init_file_system(self, load_existing: bool) -> VirtualFileSystem: """ Initialize file system interface for the new or loaded project. @@ -517,7 +486,6 @@ class StateManager: ignored as configured. :param load_existing: Whether to load existing files from the file system. - :param create_dir: Whether to create the project directory if it doesn't exist. :return: The file system interface. """ config = get_config() @@ -537,10 +505,9 @@ class StateManager: ) try: - return LocalDiskVFS( - root, allow_existing=load_existing, ignore_matcher=ignore_matcher, create=create_dir - ) - except FileExistsError: + return LocalDiskVFS(root, allow_existing=load_existing, ignore_matcher=ignore_matcher) + except FileExistsError as e: + log.debug(e) self.project.folder_name = self.project.folder_name + "-" + uuid4().hex[:7] log.warning(f"Directory {root} already exists, changing project folder to {self.project.folder_name}") await self.current_session.commit() @@ -553,8 +520,8 @@ class StateManager: """ config = get_config() - if self.project is None: - raise ValueError("No project loaded") + if self.project is None or self.project.folder_name is None: + return os.path.join(config.fs.workspace_root, "") return os.path.join(config.fs.workspace_root, self.project.folder_name) def get_full_parent_project_root(self) -> str: @@ -655,6 +622,8 @@ class StateManager: :return: List of dictionaries containing paths, old content, and new content for new or modified files. """ + if not self.file_system: + return [] modified_files = [] files_in_workspace = self.file_system.list() @@ -695,6 +664,8 @@ class StateManager: """ Returns whether the workspace has any files in them or is empty. """ + if not self.file_system: + return False return not bool(self.file_system.list()) def get_implemented_pages(self) -> list[str]: