Task Decomposer in Machine Learning

import json
from openai import OpenAI
from typing import TypedDict
from dataclasses import dataclass, field


@dataclass
class Subtask:
    id: str
    description: str
    agent_type: str
    dependencies: list[str] = field(default_factory=list)
    inputs: list[str] = field(default_factory=list)
    outputs: list[str] = field(default_factory=list)
    estimated_seconds: int = 30


@dataclass
class TaskPlan:
    original_task: str
    subtasks: list[Subtask]
    critical_path: list[str] = field(default_factory=list)


DECOMPOSITION_PROMPT = """
You are a task decomposition engine for a multi-agent system.

Available agent types:
- researcher: Searches the web, reads documents, gathers information
- coder: Writes, debugs, and reviews code
- analyst: Analyzes data, generates insights, creates charts
- writer: Produces written content, summaries, reports
- reviewer: Validates outputs, checks quality, suggests improvements

Given the following task, decompose it into subtasks.
For each subtask, specify:
- id: unique snake_case identifier
- description: clear, actionable description
- agent_type: which agent type should handle this
- dependencies: list of subtask IDs that must complete first
- inputs: what data this subtask needs
- outputs: what data this subtask produces
- estimated_seconds: rough time estimate

Rules:
1. Each subtask must be completable in a single agent call
2. Maximize parallelism: only add dependencies when truly needed
3. Ensure the subtask graph is a DAG (no cycles)
4. Every subtask output that is needed downstream must be listed
5. Aim for 3-15 subtasks depending on complexity

Return valid JSON with this schema:
{"subtasks": [{"id": "...", "description": "...", "agent_type": "...", "dependencies": [...], "inputs": [...], "outputs": [...], "estimated_seconds": N}]}

Task: {task}
"""


def decompose_task(task: str, model: str = "gpt-4o") -> TaskPlan:
    client = OpenAI()
    response = client.chat.completions.create(
        model=model,
        response_format={"type": "json_object"},
        messages=[
            {"role": "system", "content": "You are a task decomposition engine. Always respond with valid JSON."},
            {"role": "user", "content": DECOMPOSITION_PROMPT.format(task=task)},
        ],
        temperature=0.2,
    )
    raw = json.loads(response.choices[0].message.content)
    subtasks = [Subtask(**s) for s in raw["subtasks"]]
    plan = TaskPlan(original_task=task, subtasks=subtasks)
    plan.critical_path = compute_critical_path(plan)
    return plan


def compute_critical_path(plan: TaskPlan) -> list[str]:
    """Compute the critical path through the subtask DAG."""
    task_map = {s.id: s for s in plan.subtasks}
    # Longest path via dynamic programming
    memo: dict[str, tuple[int, list[str]]] = {}

    def longest_path(task_id: str) -> tuple[int, list[str]]:
        if task_id in memo:
            return memo[task_id]
        task = task_map[task_id]
        if not task.dependencies:
            result = (task.estimated_seconds, [task_id])
        else:
            best_time, best_path = 0, []
            for dep_id in task.dependencies:
                dep_time, dep_path = longest_path(dep_id)
                if dep_time > best_time:
                    best_time, best_path = dep_time, dep_path
            result = (best_time + task.estimated_seconds, best_path + [task_id])
        memo[task_id] = result
        return result

    # Find the task with the longest path (no task depends on it = sink node)
    all_deps = {d for s in plan.subtasks for d in s.dependencies}
    sinks = [s.id for s in plan.subtasks if s.id not in all_deps]
    if not sinks:
        sinks = [s.id for s in plan.subtasks]

    best_time, best_path = 0, []
    for sink in sinks:
        time, path = longest_path(sink)
        if time > best_time:
            best_time, best_path = time, path
    return best_path


# Example usage
if __name__ == "__main__":
    plan = decompose_task(
        "Research the top 5 vector databases, compare their performance, "
        "write a technical blog post with benchmarks, and create a summary slide deck."
    )
    print(f"Task: {plan.original_task}")
    print(f"Subtasks: {len(plan.subtasks)}")
    print(f"Critical path: {' -> '.join(plan.critical_path)}")
    for s in plan.subtasks:
        print(f"  [{s.id}] ({s.agent_type}) {s.description}")
        print(f"    deps: {s.dependencies}, est: {s.estimated_seconds}s")

import asyncio
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable, Awaitable

class SubtaskStatus(Enum):
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"

@dataclass
class SubtaskState:
    id: str
    description: str
    agent_type: str
    dependencies: list[str]
    status: SubtaskStatus = SubtaskStatus.PENDING
    result: Any = None
    error: str | None = None

class DAGExecutor:
    def __init__(self, agent_registry: dict[str, Callable[..., Awaitable[Any]]],
                 max_concurrency: int = 5, max_retries: int = 2):
        self.agents = agent_registry
        self.semaphore = asyncio.Semaphore(max_concurrency)
        self.max_retries = max_retries
        self.states: dict[str, SubtaskState] = {}
        self.events: dict[str, asyncio.Event] = {}

    async def execute(self, subtasks: list[dict]) -> dict[str, Any]:
        for st in subtasks:
            self.states[st["id"]] = SubtaskState(**st)
            self.events[st["id"]] = asyncio.Event()
        await asyncio.gather(*[self._run(st["id"]) for st in subtasks],
                              return_exceptions=True)
        return {sid: s.result for sid, s in self.states.items()
                if s.status == SubtaskStatus.COMPLETED}

    async def _run(self, sid: str) -> None:
        state = self.states[sid]
        # Wait for dependencies
        for dep in state.dependencies:
            await self.events[dep].wait()
            if self.states[dep].status == SubtaskStatus.FAILED:
                state.status = SubtaskStatus.FAILED
                state.error = f"Dependency {dep} failed"
                self.events[sid].set()
                return
        dep_outputs = {d: self.states[d].result for d in state.dependencies}
        agent_fn = self.agents.get(state.agent_type)
        if not agent_fn:
            state.status = SubtaskStatus.FAILED
            self.events[sid].set()
            return
        for attempt in range(self.max_retries + 1):
            try:
                async with self.semaphore:
                    state.status = SubtaskStatus.RUNNING
                    state.result = await agent_fn(task=state.description, inputs=dep_outputs)
                    state.status = SubtaskStatus.COMPLETED
                    break
            except Exception as e:
                if attempt == self.max_retries:
                    state.status = SubtaskStatus.FAILED
                    state.error = str(e)
                else:
                    await asyncio.sleep(2 ** attempt)
        self.events[sid].set()

# Usage
async def main():
    async def researcher(task, inputs): return f"Researched: {task}"
    async def writer(task, inputs): return f"Wrote based on: {inputs}"
    executor = DAGExecutor({"researcher": researcher, "writer": writer})
    results = await executor.execute([
        {"id": "r1", "description": "Research A", "agent_type": "researcher", "dependencies": []},
        {"id": "r2", "description": "Research B", "agent_type": "researcher", "dependencies": []},
        {"id": "w1", "description": "Write report", "agent_type": "writer", "dependencies": ["r1", "r2"]},
    ])
    print(results)  # r1 and r2 run in parallel, w1 waits for both

asyncio.run(main())

from typing import Annotated, TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage
import json


class PlannerState(TypedDict):
    task: str
    subtasks: list[dict]
    current_subtask_idx: int
    results: dict[str, str]
    final_output: str


llm = ChatOpenAI(model="gpt-4o", temperature=0.2)


def decompose_node(state: PlannerState) -> PlannerState:
    """LLM-based task decomposition node."""
    response = llm.invoke([
        SystemMessage(content=(
            "Decompose the following task into 3-8 sequential subtasks. "
            "Return JSON: {\"subtasks\": [{\"id\": \"step_1\", "
            "\"description\": \"...\", \"agent_type\": \"...\"}]}"
        )),
        HumanMessage(content=state["task"]),
    ])
    subtasks = json.loads(response.content)["subtasks"]
    return {
        **state,
        "subtasks": subtasks,
        "current_subtask_idx": 0,
        "results": {},
    }


def execute_subtask_node(state: PlannerState) -> PlannerState:
    """Execute the current subtask using an LLM agent."""
    idx = state["current_subtask_idx"]
    subtask = state["subtasks"][idx]
    context = json.dumps(state["results"], indent=2)

    response = llm.invoke([
        SystemMessage(content=(
            f"You are a {subtask['agent_type']} agent. "
            f"Complete this subtask using the context from previous steps.\n"
            f"Previous results: {context}"
        )),
        HumanMessage(content=subtask["description"]),
    ])

    new_results = {**state["results"], subtask["id"]: response.content}
    return {
        **state,
        "results": new_results,
        "current_subtask_idx": idx + 1,
    }


def should_continue(state: PlannerState) -> str:
    """Check if there are more subtasks to execute."""
    if state["current_subtask_idx"] < len(state["subtasks"]):
        return "execute"
    return "aggregate"


def aggregate_node(state: PlannerState) -> PlannerState:
    """Combine all subtask results into final output."""
    response = llm.invoke([
        SystemMessage(content="Synthesize all subtask results into a cohesive final output."),
        HumanMessage(content=json.dumps(state["results"], indent=2)),
    ])
    return {**state, "final_output": response.content}


# Build the LangGraph
graph = StateGraph(PlannerState)
graph.add_node("decompose", decompose_node)
graph.add_node("execute_subtask", execute_subtask_node)
graph.add_node("aggregate", aggregate_node)

graph.add_edge(START, "decompose")
graph.add_edge("decompose", "execute_subtask")
graph.add_conditional_edges("execute_subtask", should_continue, {
    "execute": "execute_subtask",
    "aggregate": "aggregate",
})
graph.add_edge("aggregate", END)

app = graph.compile()

# Run
result = app.invoke({
    "task": "Analyze the Indian e-commerce market and write a strategy memo",
    "subtasks": [],
    "current_subtask_idx": 0,
    "results": {},
    "final_output": "",
})
print(result["final_output"])

from crewai import Agent, Task, Crew, Process

# Define specialized agents
researcher = Agent(
    role="Senior Market Researcher",
    goal="Find accurate, up-to-date market data and competitive intelligence",
    backstory="You are an expert market researcher with 15 years of experience "
              "analyzing technology markets in India and globally.",
    verbose=True,
    allow_delegation=False,
)

analyst = Agent(
    role="Data Analyst",
    goal="Transform raw research into actionable insights with quantitative backing",
    backstory="You are a quantitative analyst who excels at finding patterns "
              "in data and creating compelling visualizations.",
    verbose=True,
    allow_delegation=False,
)

writer = Agent(
    role="Technical Writer",
    goal="Produce clear, well-structured technical content",
    backstory="You are a technical writer who specializes in making complex "
              "topics accessible to engineering leadership.",
    verbose=True,
    allow_delegation=False,
)

# Define tasks with explicit dependencies
research_task = Task(
    description=(
        "Research the current state of AI agent frameworks in 2026. "
        "Cover LangGraph, CrewAI, AutoGen, and OpenAI Agents SDK. "
        "Include pricing, features, community size, and production readiness."
    ),
    expected_output="Structured research report with data on each framework",
    agent=researcher,
)

analysis_task = Task(
    description=(
        "Analyze the research data to identify the best framework for "
        "an Indian fintech startup building a customer support automation system. "
        "Consider cost (budget: INR 5 lakh/month), scalability, and Hindi language support."
    ),
    expected_output="Comparative analysis with recommendation and justification",
    agent=analyst,
    context=[research_task],  # dependency: waits for research_task
)

writing_task = Task(
    description=(
        "Write a technical decision document based on the analysis. "
        "Include an executive summary, comparison table, recommendation, "
        "migration plan, and cost projections in INR."
    ),
    expected_output="Complete technical decision document (1500-2000 words)",
    agent=writer,
    context=[research_task, analysis_task],  # depends on both
)

# Create crew with planning enabled
crew = Crew(
    agents=[researcher, analyst, writer],
    tasks=[research_task, analysis_task, writing_task],
    process=Process.sequential,
    planning=True,  # enables automatic plan generation before execution
    verbose=True,
)

result = crew.kickoff()
print(result)

# Task decomposer configuration (YAML)
decomposer:
  model: gpt-4o
  temperature: 0.2
  max_subtasks: 15
  min_subtasks: 3
  max_depth: 3  # recursive decomposition depth limit
  
  complexity_threshold: 500  # tokens; subtasks above this get re-decomposed
  
  available_agents:
    - type: researcher
      capabilities: [web_search, document_reading, data_extraction]
      cost_per_call_usd: 0.05
    - type: coder
      capabilities: [code_generation, debugging, code_review]
      cost_per_call_usd: 0.08
    - type: analyst
      capabilities: [data_analysis, visualization, statistical_testing]
      cost_per_call_usd: 0.06
    - type: writer
      capabilities: [content_writing, summarization, editing]
      cost_per_call_usd: 0.04

  execution:
    max_concurrency: 5
    subtask_timeout_seconds: 120
    max_retries: 2
    retry_backoff: exponential

  replanning:
    enabled: true
    trigger_on: [subtask_failure, timeout, unexpected_output]
    max_replans: 3

  budget:
    max_total_cost_usd: 1.00  # ~INR 84
    max_subtask_cost_usd: 0.20