From 7d225bbfe203dc40c2e1b77a7ef5fc485b3e1243 Mon Sep 17 00:00:00 2001 From: Senko Rasic Date: Thu, 13 Jun 2024 10:48:37 +0200 Subject: [PATCH] refactor example project so each agent does its job in turn --- core/agents/architect.py | 82 ++++++++++++++++--- core/agents/external_docs.py | 12 ++- core/agents/mixins.py | 53 ------------ core/agents/spec_writer.py | 52 +++++------- core/agents/tech_lead.py | 22 ++++- core/db/migrations/README | 11 ++- ...f891d366761_add_example_project_to_spec.py | 34 ++++++++ core/db/models/specification.py | 2 + core/telemetry/__init__.py | 4 +- core/templates/example_project.py | 14 ++++ tests/agents/test_spec_writer.py | 10 +-- 11 files changed, 185 insertions(+), 111 deletions(-) create mode 100644 core/db/migrations/versions/ff891d366761_add_example_project_to_spec.py diff --git a/core/agents/architect.py b/core/agents/architect.py index 712c70995c..c1ce78d72f 100644 --- a/core/agents/architect.py +++ b/core/agents/architect.py @@ -4,10 +4,12 @@ from pydantic import BaseModel, Field from core.agents.base import BaseAgent from core.agents.convo import AgentConvo -from core.agents.mixins import SystemDependencyCheckerMixin from core.agents.response import AgentResponse +from core.db.models import Specification from core.llm.parser import JSONParser +from core.log import get_logger from core.telemetry import telemetry +from core.templates.example_project import EXAMPLE_PROJECTS from core.templates.registry import PROJECT_TEMPLATES, ProjectTemplateEnum from core.ui.base import ProjectStage @@ -16,6 +18,8 @@ WARN_SYSTEM_DEPS = ["docker", "kubernetes", "microservices"] WARN_FRAMEWORKS = ["next.js", "vue", "vue.js", "svelte", "angular"] WARN_FRAMEWORKS_URL = "https://github.com/Pythagora-io/gpt-pilot/wiki/Using-GPT-Pilot-with-frontend-frameworks" +log = get_logger(__name__) + # FIXME: all the reponse pydantic models should be strict (see config._StrictModel), also check if we # can disallow adding custom Python attributes to the model @@ -68,33 +72,41 @@ class Architecture(BaseModel): ) -class Architect(SystemDependencyCheckerMixin, BaseAgent): +class Architect(BaseAgent): agent_type = "architect" display_name = "Architect" async def run(self) -> AgentResponse: await self.ui.send_project_stage(ProjectStage.ARCHITECTURE) + spec = self.current_state.specification.clone() + + if spec.example_project: + self.prepare_example_project(spec) + else: + await self.plan_architecture(spec) + + await self.check_system_dependencies(spec) + + self.next_state.specification = spec + telemetry.set("template", spec.template) + self.next_state.action = ARCHITECTURE_STEP_NAME + return AgentResponse.done(self) + + async def plan_architecture(self, spec: Specification): + await self.send_message("Planning project architecture ...") + llm = self.get_llm() convo = AgentConvo(self).template("technologies", templates=PROJECT_TEMPLATES).require_schema(Architecture) - - await self.send_message("Planning project architecture ...") arch: Architecture = await llm(convo, parser=JSONParser(Architecture)) await self.check_compatibility(arch) - spec = self.current_state.specification.clone() spec.architecture = arch.architecture spec.system_dependencies = [d.model_dump() for d in arch.system_dependencies] spec.package_dependencies = [d.model_dump() for d in arch.package_dependencies] spec.template = arch.template.value if arch.template else None - self.next_state.specification = spec - await self.check_system_dependencies(spec) - telemetry.set("template", spec.template) - self.next_state.action = ARCHITECTURE_STEP_NAME - return AgentResponse.done(self) - async def check_compatibility(self, arch: Architecture) -> bool: warn_system_deps = [dep.name for dep in arch.system_dependencies if dep.name.lower() in WARN_SYSTEM_DEPS] warn_package_deps = [dep.name for dep in arch.package_dependencies if dep.name.lower() in WARN_FRAMEWORKS] @@ -122,3 +134,51 @@ class Architect(SystemDependencyCheckerMixin, BaseAgent): # return AgentResponse.revise_spec() # that SpecWriter should catch and allow the user to reword the initial spec. return True + + def prepare_example_project(self, spec: Specification): + log.debug(f"Setting architecture for example project: {spec.example_project}") + arch = EXAMPLE_PROJECTS[spec.example_project]["architecture"] + + spec.architecture = arch["architecture"] + spec.system_dependencies = arch["system_dependencies"] + spec.package_dependencies = arch["package_dependencies"] + spec.template = arch["template"] + telemetry.set("template", spec.template) + + async def check_system_dependencies(self, spec: Specification): + """ + Check whether the required system dependencies are installed. + + This also stores the app architecture telemetry data, including the + information about whether each system dependency is installed. + + :param spec: Project specification. + """ + deps = spec.system_dependencies + + for dep in deps: + status_code, _, _ = await self.process_manager.run_command(dep["test"]) + dep["installed"] = bool(status_code == 0) + if status_code != 0: + if dep["required_locally"]: + remedy = "Please install it before proceeding with your app." + else: + remedy = "If you would like to use it locally, please install it before proceeding." + await self.send_message(f"❌ {dep['name']} is not available. {remedy}") + await self.ask_question( + f"Once you have installed {dep['name']}, please press Continue.", + buttons={"continue": "Continue"}, + buttons_only=True, + default="continue", + ) + else: + await self.send_message(f"✅ {dep['name']} is available.") + + telemetry.set( + "architecture", + { + "description": spec.architecture, + "system_dependencies": deps, + "package_dependencies": spec.package_dependencies, + }, + ) diff --git a/core/agents/external_docs.py b/core/agents/external_docs.py index 6e818df942..0da9bc5aca 100644 --- a/core/agents/external_docs.py +++ b/core/agents/external_docs.py @@ -26,7 +26,7 @@ class SelectedDocsets(BaseModel): class ExternalDocumentation(BaseAgent): """Agent in charge of collecting and storing additional documentation. - Docs are per task and are stores in the `tasks` variable in the project state. + Docs are per task and are stores in the `docs` variable in the project state. This agent ensures documentation is collected only once per task. Agent does 2 LLM interactions: @@ -44,7 +44,12 @@ class ExternalDocumentation(BaseAgent): display_name = "Documentation" async def run(self) -> AgentResponse: - available_docsets = await self._get_available_docsets() + if self.current_state.specification.example_project: + log.debug("Example project detected, no documentation selected.") + available_docsets = [] + else: + available_docsets = await self._get_available_docsets() + selected_docsets = await self._select_docsets(available_docsets) await telemetry.trace_code_event("docsets_used", selected_docsets) @@ -149,6 +154,8 @@ class ExternalDocumentation(BaseAgent): Documentation snippets are stored as a list of dictionaries: {"key": docset-key, "desc": documentation-description, "snippets": list-of-snippets} + :param snippets: List of tuples: (docset_key, snippets) + :param available_docsets: List of available docsets from the API. """ docsets_dict = dict(available_docsets) @@ -157,4 +164,3 @@ class ExternalDocumentation(BaseAgent): docs.append({"key": docset_key, "desc": docsets_dict[docset_key], "snippets": snip}) self.next_state.docs = docs - self.next_state.flag_tasks_as_modified() diff --git a/core/agents/mixins.py b/core/agents/mixins.py index bcaebddbfe..5ea0aae781 100644 --- a/core/agents/mixins.py +++ b/core/agents/mixins.py @@ -1,8 +1,6 @@ from typing import Optional from core.agents.convo import AgentConvo -from core.db.models.specification import Specification -from core.telemetry import telemetry class IterationPromptMixin: @@ -37,54 +35,3 @@ class IterationPromptMixin: ) llm_solution: str = await llm(convo) return llm_solution - - -class SystemDependencyCheckerMixin: - """ - Provides a method to check whether the required system dependencies are installed. - - Used by Architect and SpecWriter agents. Assumes the agent has access to UI - and ProcessManager. - """ - - async def check_system_dependencies(self, spec: Specification): - """ - Check whether the required system dependencies are installed. - - This also stores the app architecture telemetry data, including the - information about whether each system dependency is installed. - - :param spec: Project specification. - """ - deps = spec.system_dependencies - checked = {} - - for dep in deps: - status_code, _, _ = await self.process_manager.run_command(dep["test"]) - dep["installed"] = bool(status_code == 0) - if status_code != 0: - if dep["required_locally"]: - remedy = "Please install it before proceeding with your app." - else: - remedy = "If you would like to use it locally, please install it before proceeding." - await self.send_message(f"❌ {dep['name']} is not available. {remedy}") - await self.ask_question( - f"Once you have installed {dep['name']}, please press Continue.", - buttons={"continue": "Continue"}, - buttons_only=True, - default="continue", - ) - checked[dep["name"]] = "missing" - else: - await self.send_message(f"✅ {dep['name']} is available.") - checked[dep["name"]] = "present" - - telemetry.set( - "architecture", - { - "description": spec.architecture, - "system_dependencies": deps, - "package_dependencies": spec.package_dependencies, - "checked_system_dependencies": checked, - }, - ) diff --git a/core/agents/spec_writer.py b/core/agents/spec_writer.py index 1e2a9a2492..9e2f737312 100644 --- a/core/agents/spec_writer.py +++ b/core/agents/spec_writer.py @@ -1,14 +1,13 @@ from core.agents.base import BaseAgent from core.agents.convo import AgentConvo -from core.agents.mixins import SystemDependencyCheckerMixin from core.agents.response import AgentResponse from core.db.models import Complexity from core.llm.parser import StringParser +from core.log import get_logger from core.telemetry import telemetry from core.templates.example_project import ( - EXAMPLE_PROJECT_ARCHITECTURE, - EXAMPLE_PROJECT_DESCRIPTION, - EXAMPLE_PROJECT_PLAN, + DEFAULT_EXAMPLE_PROJECT, + EXAMPLE_PROJECTS, ) # If the project description is less than this, perform an analysis using LLM @@ -19,8 +18,10 @@ INITIAL_PROJECT_HOWTO_URL = ( ) SPEC_STEP_NAME = "Create specification" +log = get_logger(__name__) -class SpecWriter(SystemDependencyCheckerMixin, BaseAgent): + +class SpecWriter(BaseAgent): agent_type = "spec-writer" display_name = "Spec Writer" @@ -42,16 +43,15 @@ class SpecWriter(SystemDependencyCheckerMixin, BaseAgent): return AgentResponse.import_project(self) if response.button == "example": - await self.send_message("Starting example project with description:") - await self.send_message(EXAMPLE_PROJECT_DESCRIPTION) - await self.prepare_example_project() + await self.prepare_example_project(DEFAULT_EXAMPLE_PROJECT) return AgentResponse.done(self) + elif response.button == "continue": # FIXME: Workaround for the fact that VSCode "continue" button does # nothing but repeat the question. We reproduce this bug for bug here. return AgentResponse.done(self) - spec = response.text + spec = response.text.strip() complexity = await self.check_prompt_complexity(spec) await telemetry.trace_code_event( @@ -82,31 +82,21 @@ class SpecWriter(SystemDependencyCheckerMixin, BaseAgent): llm_response: str = await llm(convo, temperature=0, parser=StringParser()) return llm_response.lower() - async def prepare_example_project(self): + async def prepare_example_project(self, example_name: str): + example_description = EXAMPLE_PROJECTS[example_name]["description"].strip() + + log.debug(f"Starting example project: {example_name}") + await self.send_message(f"Starting example project with description:\n\n{example_description}") + spec = self.current_state.specification.clone() - spec.description = EXAMPLE_PROJECT_DESCRIPTION - spec.architecture = EXAMPLE_PROJECT_ARCHITECTURE["architecture"] - spec.system_dependencies = EXAMPLE_PROJECT_ARCHITECTURE["system_dependencies"] - spec.package_dependencies = EXAMPLE_PROJECT_ARCHITECTURE["package_dependencies"] - spec.template = EXAMPLE_PROJECT_ARCHITECTURE["template"] - spec.complexity = Complexity.SIMPLE - telemetry.set("initial_prompt", spec.description.strip()) - telemetry.set("is_complex_app", False) - telemetry.set("template", spec.template) - - await self.check_system_dependencies(spec) - + spec.example_project = example_name + spec.description = example_description + spec.complexity = EXAMPLE_PROJECTS[example_name]["complexity"] self.next_state.specification = spec - self.next_state.epics = [ - { - "name": "Initial Project", - "description": EXAMPLE_PROJECT_DESCRIPTION, - "completed": False, - "complexity": Complexity.SIMPLE, - } - ] - self.next_state.tasks = EXAMPLE_PROJECT_PLAN + telemetry.set("initial_prompt", spec.description) + telemetry.set("example_project", example_name) + telemetry.set("is_complex_app", spec.complexity != Complexity.SIMPLE) async def analyze_spec(self, spec: str) -> str: msg = ( diff --git a/core/agents/tech_lead.py b/core/agents/tech_lead.py index 1dc36a6fa4..4c11d55802 100644 --- a/core/agents/tech_lead.py +++ b/core/agents/tech_lead.py @@ -11,6 +11,7 @@ from core.db.models.project_state import TaskStatus from core.llm.parser import JSONParser from core.log import get_logger from core.telemetry import telemetry +from core.templates.example_project import EXAMPLE_PROJECTS from core.templates.registry import apply_project_template, get_template_description, get_template_summary from core.ui.base import ProjectStage, success_source @@ -42,8 +43,10 @@ class TechLead(BaseAgent): return await self.update_epic() if len(self.current_state.epics) == 0: - self.create_initial_project_epic() - # Orchestrator will rerun us to break down the initial project epic + if self.current_state.specification.example_project: + self.plan_example_project() + else: + self.create_initial_project_epic() return AgentResponse.done(self) await self.ui.send_project_stage(ProjectStage.CODING) @@ -203,3 +206,18 @@ class TechLead(BaseAgent): ] log.debug(f"Updated development plan for {epic['name']}, {len(response.plan)} tasks remaining") return AgentResponse.done(self) + + def plan_example_project(self): + example_name = self.current_state.specification.example_project + log.debug(f"Planning example project: {example_name}") + + example = EXAMPLE_PROJECTS[example_name] + self.next_state.epics = [ + { + "name": "Initial Project", + "description": example["description"], + "completed": False, + "complexity": example["complexity"], + } + ] + self.next_state.tasks = example["plan"] diff --git a/core/db/migrations/README b/core/db/migrations/README index 98e4f9c44e..5cdc74083b 100644 --- a/core/db/migrations/README +++ b/core/db/migrations/README @@ -1 +1,10 @@ -Generic single-database configuration. \ No newline at end of file +Pythagora uses Alembic for database migrations. + +After changing any of the database models, create a new migration: + + alembic -c core/db/alembic.ini revision --autogenerate -m "description" + +Migrations are applied automatically when the application starts, but can also be +run manually with: + + alembic -c core/db/alembic.ini upgrade head diff --git a/core/db/migrations/versions/ff891d366761_add_example_project_to_spec.py b/core/db/migrations/versions/ff891d366761_add_example_project_to_spec.py new file mode 100644 index 0000000000..f788761e2d --- /dev/null +++ b/core/db/migrations/versions/ff891d366761_add_example_project_to_spec.py @@ -0,0 +1,34 @@ +"""add example project to spec + +Revision ID: ff891d366761 +Revises: b760f66138c0 +Create Date: 2024-06-13 09:38:33.329161 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "ff891d366761" +down_revision: Union[str, None] = "b760f66138c0" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("specifications", schema=None) as batch_op: + batch_op.add_column(sa.Column("example_project", sa.String(), nullable=True)) + + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("specifications", schema=None) as batch_op: + batch_op.drop_column("example_project") + + # ### end Alembic commands ### diff --git a/core/db/models/specification.py b/core/db/models/specification.py index 9e2eb7c96e..e711f38755 100644 --- a/core/db/models/specification.py +++ b/core/db/models/specification.py @@ -31,6 +31,7 @@ class Specification(Base): package_dependencies: Mapped[list[dict]] = mapped_column(default=list) template: Mapped[Optional[str]] = mapped_column() complexity: Mapped[str] = mapped_column(server_default=Complexity.HARD) + example_project: Mapped[Optional[str]] = mapped_column() # Relationships project_states: Mapped[list["ProjectState"]] = relationship(back_populates="specification", lazy="raise") @@ -46,6 +47,7 @@ class Specification(Base): package_dependencies=self.package_dependencies, template=self.template, complexity=self.complexity, + example_project=self.example_project, ) return clone diff --git a/core/telemetry/__init__.py b/core/telemetry/__init__.py index 2270a83ad0..37a38b3fcb 100644 --- a/core/telemetry/__init__.py +++ b/core/telemetry/__init__.py @@ -73,7 +73,7 @@ class Telemetry: "python_version": sys.version, # GPT Pilot version "pilot_version": get_version(), - # GPT Pilot Extension version + # Pythagora VSCode Extension version "extension_version": None, # Is extension used "is_extension": False, @@ -86,6 +86,8 @@ class Telemetry: "is_complex_app": None, # Optional template used for the project "template": None, + # Optional, example project selected by the user + "example_project": None, # Optional user contact email "user_contact": None, # Unique project ID (app_id) diff --git a/core/templates/example_project.py b/core/templates/example_project.py index 6703a906d3..7cb47a6e3a 100644 --- a/core/templates/example_project.py +++ b/core/templates/example_project.py @@ -1,3 +1,5 @@ +from core.db.models import Complexity + EXAMPLE_PROJECT_DESCRIPTION = """ The application is a simple ToDo app built using React. Its primary function is to allow users to manage a list of tasks (todos). Each task has a description and a state (open or completed, with the default state being open). The application is frontend-only, with no user sign-up or authentication process. The goal is to provide a straightforward and user-friendly interface for task management. @@ -64,3 +66,15 @@ EXAMPLE_PROJECT_PLAN = [ "status": "todo", } ] + + +EXAMPLE_PROJECTS = { + "example-project": { + "description": EXAMPLE_PROJECT_DESCRIPTION, + "architecture": EXAMPLE_PROJECT_ARCHITECTURE, + "complexity": Complexity.SIMPLE, + "plan": EXAMPLE_PROJECT_PLAN, + } +} + +DEFAULT_EXAMPLE_PROJECT = "example-project" diff --git a/tests/agents/test_spec_writer.py b/tests/agents/test_spec_writer.py index 88f9495977..2140dee2b3 100644 --- a/tests/agents/test_spec_writer.py +++ b/tests/agents/test_spec_writer.py @@ -23,16 +23,8 @@ async def test_start_example_project(agentcontext): assert response.type == ResponseType.DONE assert sm.current_state.specification.description != "" - assert sm.current_state.specification.architecture != "" - assert sm.current_state.specification.system_dependencies != [] - assert sm.current_state.specification.package_dependencies != [] assert sm.current_state.specification.complexity == Complexity.SIMPLE - assert sm.current_state.epics != [] - assert sm.current_state.tasks != [] - pm.run_command.assert_awaited_once_with("node --version") - - assert telemetry.data["initial_prompt"] == sm.current_state.specification.description.strip() - assert telemetry.data["architecture"]["system_dependencies"][0]["installed"] is True + assert telemetry.data["initial_prompt"] == sm.current_state.specification.description @pytest.mark.asyncio