agent/app/modules/agent/engine/orchestrator/template_registry.py

from __future__ import annotations

from app.modules.agent.engine.orchestrator.models import ArtifactSpec, ArtifactType, ExecutionPlan, PlanStep, QualityGateRef, Scenario, TaskSpec


class ScenarioTemplateRegistry:
    def build(self, task: TaskSpec) -> ExecutionPlan:
        builders = {
            Scenario.EXPLAIN_PART: self._explain,
            Scenario.ANALYTICS_REVIEW: self._review,
            Scenario.DOCS_FROM_ANALYTICS: self._docs,
            Scenario.TARGETED_EDIT: self._edit,
            Scenario.GHERKIN_MODEL: self._gherkin,
            Scenario.GENERAL_QA: self._general,
        }
        return builders.get(task.scenario, self._general)(task)

    def _general(self, task: TaskSpec) -> ExecutionPlan:
        steps = [
            self._step("collect_state", "Collect state", "collect_state", outputs=[self._out("agent_state", ArtifactType.STRUCTURED_JSON)]),
            self._step(
                "execute_route_graph",
                "Execute selected graph",
                "execute_route_graph",
                executor="graph",
                graph_id="route",
                depends_on=["collect_state"],
                outputs=[self._out("graph_result", ArtifactType.STRUCTURED_JSON)],
                gates=[self._gate("required_outputs")],
            ),
            self._step(
                "finalize_graph_output",
                "Finalize graph output",
                "finalize_graph_output",
                depends_on=["execute_route_graph"],
                outputs=[self._out("final_answer", ArtifactType.TEXT, required=False)],
                gates=[self._gate("non_empty_answer_or_changeset")],
            ),
        ]
        return self._plan(task, "general_qa_v1", steps, [self._gate("non_empty_answer_or_changeset")])

    def _explain(self, task: TaskSpec) -> ExecutionPlan:
        steps = [
            self._step("collect_sources", "Collect sources", "collect_sources", outputs=[self._out("sources", ArtifactType.STRUCTURED_JSON)]),
            self._step("extract_logic", "Extract logic", "extract_logic", depends_on=["collect_sources"], outputs=[self._out("logic_model", ArtifactType.STRUCTURED_JSON)]),
            self._step("summarize", "Summarize", "summarize", depends_on=["extract_logic"], outputs=[self._out("final_answer", ArtifactType.TEXT)]),
        ]
        return self._plan(task, "explain_part_v1", steps, [self._gate("evidence_required"), self._gate("non_empty_answer_or_changeset")])

    def _review(self, task: TaskSpec) -> ExecutionPlan:
        steps = [
            self._step("fetch_source_doc", "Fetch source doc", "fetch_source_doc", outputs=[self._out("source_doc_raw", ArtifactType.TEXT)], side_effect="external"),
            self._step("normalize_document", "Normalize document", "normalize_document", depends_on=["fetch_source_doc"], outputs=[self._out("source_doc_text", ArtifactType.TEXT)]),
            self._step("structural_check", "Structural check", "structural_check", depends_on=["normalize_document"], outputs=[self._out("structural_findings", ArtifactType.STRUCTURED_JSON)]),
            self._step("semantic_consistency_check", "Semantic check", "semantic_consistency_check", depends_on=["normalize_document"], outputs=[self._out("semantic_findings", ArtifactType.STRUCTURED_JSON)]),
            self._step("architecture_fit_check", "Architecture fit", "architecture_fit_check", depends_on=["normalize_document"], outputs=[self._out("architecture_findings", ArtifactType.STRUCTURED_JSON)]),
            self._step("optimization_check", "Optimization check", "optimization_check", depends_on=["normalize_document"], outputs=[self._out("optimization_findings", ArtifactType.STRUCTURED_JSON)]),
            self._step(
                "compose_review_report",
                "Compose review report",
                "compose_review_report",
                depends_on=["structural_check", "semantic_consistency_check", "architecture_fit_check", "optimization_check"],
                outputs=[self._out("review_report", ArtifactType.REVIEW_REPORT), self._out("final_answer", ArtifactType.TEXT)],
                gates=[self._gate("review_report_schema")],
            ),
        ]
        return self._plan(task, "analytics_review_v1", steps, [self._gate("evidence_required"), self._gate("non_empty_answer_or_changeset")])

    def _docs(self, task: TaskSpec) -> ExecutionPlan:
        steps = [
            self._step("fetch_source_doc", "Fetch source doc", "fetch_source_doc", outputs=[self._out("source_doc_raw", ArtifactType.TEXT)], side_effect="external"),
            self._step("normalize_document", "Normalize document", "normalize_document", depends_on=["fetch_source_doc"], outputs=[self._out("source_doc_text", ArtifactType.TEXT)]),
            self._step("extract_change_intents", "Extract intents", "extract_change_intents", depends_on=["normalize_document"], outputs=[self._out("change_intents", ArtifactType.STRUCTURED_JSON)]),
            self._step("map_to_doc_tree", "Map to doc tree", "map_to_doc_tree", depends_on=["extract_change_intents"], outputs=[self._out("doc_targets", ArtifactType.STRUCTURED_JSON)]),
            self._step("load_current_docs_context", "Load current docs", "load_current_docs_context", depends_on=["map_to_doc_tree"], outputs=[self._out("current_docs_context", ArtifactType.STRUCTURED_JSON)]),
            self._step("generate_doc_updates", "Generate doc updates", "generate_doc_updates", depends_on=["load_current_docs_context"], outputs=[self._out("generated_doc_bundle", ArtifactType.DOC_BUNDLE)], side_effect="write"),
            self._step("cross_file_validation", "Cross-file validation", "cross_file_validation", depends_on=["generate_doc_updates"], outputs=[self._out("consistency_report", ArtifactType.STRUCTURED_JSON)], gates=[self._gate("cross_file_consistency")]),
            self._step("build_changeset", "Build changeset", "build_changeset", depends_on=["cross_file_validation"], outputs=[self._out("final_changeset", ArtifactType.CHANGESET)], side_effect="write"),
            self._step("compose_summary", "Compose summary", "compose_summary", depends_on=["build_changeset"], outputs=[self._out("final_answer", ArtifactType.TEXT)]),
        ]
        return self._plan(task, "docs_from_analytics_v1", steps, [self._gate("changeset_required_for_write"), self._gate("changeset_schema")])

    def _edit(self, task: TaskSpec) -> ExecutionPlan:
        steps = [
            self._step("resolve_target", "Resolve target", "resolve_target", outputs=[self._out("resolved_target", ArtifactType.STRUCTURED_JSON)], gates=[self._gate("target_path_must_exist_or_be_allowed")]),
            self._step("load_target_context", "Load target context", "load_target_context", depends_on=["resolve_target"], outputs=[self._out("target_context", ArtifactType.STRUCTURED_JSON)]),
            self._step("plan_minimal_patch", "Plan minimal patch", "plan_minimal_patch", depends_on=["load_target_context"], outputs=[self._out("patch_plan", ArtifactType.STRUCTURED_JSON)]),
            self._step("generate_patch", "Generate patch", "generate_patch", depends_on=["plan_minimal_patch"], outputs=[self._out("raw_changeset", ArtifactType.CHANGESET)], side_effect="write"),
            self._step("validate_patch_safety", "Validate patch", "validate_patch_safety", depends_on=["generate_patch"], outputs=[self._out("patch_validation_report", ArtifactType.STRUCTURED_JSON)], gates=[self._gate("minimal_patch_policy")]),
            self._step("finalize_changeset", "Finalize changeset", "finalize_changeset", depends_on=["validate_patch_safety"], outputs=[self._out("final_changeset", ArtifactType.CHANGESET)], side_effect="write"),
            self._step("compose_edit_summary", "Compose summary", "compose_edit_summary", depends_on=["finalize_changeset"], outputs=[self._out("final_answer", ArtifactType.TEXT)]),
        ]
        return self._plan(task, "targeted_edit_v1", steps, [self._gate("changeset_required_for_write"), self._gate("changeset_schema")])

    def _gherkin(self, task: TaskSpec) -> ExecutionPlan:
        steps = [
            self._step("fetch_source_doc", "Fetch source doc", "fetch_source_doc", outputs=[self._out("source_doc_raw", ArtifactType.TEXT)], side_effect="external"),
            self._step("normalize_document", "Normalize document", "normalize_document", depends_on=["fetch_source_doc"], outputs=[self._out("source_doc_text", ArtifactType.TEXT)]),
            self._step("extract_increment_scope", "Extract increment scope", "extract_increment_scope", depends_on=["normalize_document"], outputs=[self._out("increment_scope", ArtifactType.STRUCTURED_JSON)]),
            self._step("partition_features", "Partition features", "partition_features", depends_on=["extract_increment_scope"], outputs=[self._out("feature_groups", ArtifactType.STRUCTURED_JSON)]),
            self._step("generate_gherkin_bundle", "Generate gherkin", "generate_gherkin_bundle", depends_on=["partition_features"], outputs=[self._out("gherkin_bundle", ArtifactType.GHERKIN_BUNDLE)], side_effect="write"),
            self._step("lint_gherkin", "Lint gherkin", "lint_gherkin", depends_on=["generate_gherkin_bundle"], outputs=[self._out("gherkin_lint_report", ArtifactType.STRUCTURED_JSON)], gates=[self._gate("gherkin_syntax_lint")]),
            self._step("validate_coverage", "Validate coverage", "validate_coverage", depends_on=["generate_gherkin_bundle"], outputs=[self._out("coverage_report", ArtifactType.STRUCTURED_JSON)], gates=[self._gate("coverage_of_change_intents")]),
            self._step("compose_test_model_summary", "Compose summary", "compose_test_model_summary", depends_on=["lint_gherkin", "validate_coverage"], outputs=[self._out("final_answer", ArtifactType.TEXT), self._out("final_changeset", ArtifactType.CHANGESET)], side_effect="write"),
        ]
        return self._plan(task, "gherkin_model_v1", steps, [self._gate("changeset_schema"), self._gate("non_empty_answer_or_changeset")])

    def _plan(self, task: TaskSpec, template_id: str, steps: list[PlanStep], gates: list[QualityGateRef]) -> ExecutionPlan:
        return ExecutionPlan(
            plan_id=f"{task.task_id}:{template_id}",
            task_id=task.task_id,
            scenario=task.scenario,
            template_id=template_id,
            template_version="1.0",
            steps=steps,
            global_gates=gates,
        )

    def _step(
        self,
        step_id: str,
        title: str,
        action_id: str,
        *,
        executor: str = "function",
        graph_id: str | None = None,
        depends_on: list[str] | None = None,
        outputs: list[ArtifactSpec] | None = None,
        gates: list[QualityGateRef] | None = None,
        side_effect: str = "read",
    ) -> PlanStep:
        return PlanStep(
            step_id=step_id,
            title=title,
            action_id=action_id,
            executor=executor,
            graph_id=graph_id,
            depends_on=depends_on or [],
            outputs=outputs or [],
            quality_gates=gates or [],
            side_effect=side_effect,
        )

    def _out(self, key: str, artifact_type: ArtifactType, *, required: bool = True) -> ArtifactSpec:
        return ArtifactSpec(key=key, type=artifact_type, required=required)

    def _gate(self, gate_id: str, *, blocking: bool = True) -> QualityGateRef:
        return QualityGateRef(gate_id=gate_id, blocking=blocking)