devflow-engine 1.0.0__py3-none-any.whl

devin/dag_two_arm.py

@@ -0,0 +1,138 @@
+"""Devin two-arm chat DAG — ideation arm or insight arm (everything else)."""
+from __future__ import annotations
+
+from pathlib import Path
+
+from devflow_engine.stores.execution_store import ExecutionStore
+from devflow_engine.vendor.datalumina_genai.core.nodes.base import Node
+from devflow_engine.vendor.datalumina_genai.core.nodes.router import BaseRouter, RouterNode
+from devflow_engine.vendor.datalumina_genai.core.schema import NodeConfig, WorkflowSchema
+from devflow_engine.vendor.datalumina_genai.core.task import TaskContext
+from devflow_engine.vendor.datalumina_genai.core.workflow import Workflow
+
+from .nodes.ideation.node import IdeationAgentNode
+from .nodes.insight.node import InsightAgentNode
+from .nodes.intake.node import DevinIntakeNode
+from .nodes.shared.helpers import (
+    DevinChatDagResult,
+    build_pipeline_key,
+    dfs_running,
+    parse_list_option,
+    pipeline_root,
+    set_runtime_store,
+)
+from .nodes.shared.models import DevinChatDagEvent
+from .nodes.shared.post import DevinResponsePostNode
+
+DAG_ID = "devin_chat_two_arm_dag"
+
+
+class _TwoArmRouter(RouterNode):
+    """Routes to ideation arm if classify_route returned 'ideation', else insight arm."""
+
+    def determine_next_node(self, task_context: TaskContext) -> Node | None:
+        arm = str(
+            task_context.metadata.get("route_arm")
+            or task_context.metadata.get("route", {}).get("route_arm")
+            or ""
+        )
+        if arm == "ideation":
+            return IdeationAgentNode(task_context=task_context)
+        return InsightAgentNode(task_context=task_context)  # everything else
+
+
+class DevinIntakeRouter(BaseRouter):
+    def __init__(self, task_context: TaskContext | None = None):
+        super().__init__(task_context=task_context)
+        self.routes = [_TwoArmRouter()]
+        self.fallback = InsightAgentNode()
+
+
+class DevinTwoArmWorkflow(Workflow):
+    workflow_schema = WorkflowSchema(
+        description="Devin two-arm chat DAG: intake → ideation | insight → post",
+        event_schema=DevinChatDagEvent,
+        start=DevinIntakeNode,
+        nodes=[
+            NodeConfig(node=DevinIntakeNode, connections=[DevinIntakeRouter]),
+            NodeConfig(
+                node=DevinIntakeRouter,
+                connections=[InsightAgentNode, IdeationAgentNode],
+                is_router=True,
+            ),
+            NodeConfig(node=InsightAgentNode, connections=[DevinResponsePostNode]),
+            NodeConfig(node=IdeationAgentNode, connections=[DevinResponsePostNode]),
+            NodeConfig(node=DevinResponsePostNode, connections=[]),
+        ],
+    )
+
+
+def run_devin_two_arm_dag(
+    *,
+    repo_root: Path,
+    store: ExecutionStore,
+    idea_id: str,
+    text: str | None,
+    source_path: Path | None,
+    max_stories: int,
+    planes: list[str],
+    source_session_id: str | None = None,
+) -> DevinChatDagResult:
+    normalized_planes = parse_list_option(planes)
+    pipeline_key = build_pipeline_key(
+        repo_root=repo_root,
+        idea_id=idea_id,
+        text=text,
+        source_path=source_path,
+        max_stories=max_stories,
+        planes=normalized_planes,
+    )
+    pipeline_dir = pipeline_root(
+        repo_root, idea_id=idea_id, pipeline_key=pipeline_key
+    )
+    pipeline_dir.mkdir(parents=True, exist_ok=True)
+    run_id = store.create_run(
+        dag_id=DAG_ID,
+        dag_version="v1",
+        root_correlation_id=f"corr_{pipeline_key}",
+        config={
+            "idea_id": idea_id,
+            "pipeline_key": pipeline_key,
+            "max_stories": max_stories,
+            "planes": normalized_planes,
+        },
+    )
+    store.mark_run_started(run_id=run_id)
+    set_runtime_store(store, run_id)
+    try:
+        dfs_running(
+            project_id=f"unregistered:{idea_id}",
+            run_id=run_id,
+            summary="Running Devin two-arm chat DAG",
+            idea_id=idea_id,
+        )
+        ctx = DevinTwoArmWorkflow().run(
+            {
+                "repo_root": str(repo_root),
+                "idea_id": idea_id,
+                "raw_text": text,
+                "source_path": str(source_path) if source_path else None,
+                "max_stories": max_stories,
+                "planes": normalized_planes,
+                "pipeline_key": pipeline_key,
+                "source_session_id": source_session_id,
+            }
+        )
+        exit_code = int(ctx.metadata.get("exit_code") or 0)
+        store.mark_run_finished(
+            run_id=run_id, status="succeeded" if exit_code == 0 else "failed"
+        )
+        return DevinChatDagResult(
+            exit_code=exit_code,
+            run_id=run_id,
+            pipeline_dir=pipeline_dir,
+            message=str(ctx.metadata.get("message") or ""),
+            outcome=dict(ctx.metadata.get("outcome") or {}),
+        )
+    finally:
+        set_runtime_store(None, None)