Merge pull request #123 from PrefectHQ/planning

Add basic planning

src/controlflow/llm/history.py

@@ -35,9 +35,6 @@ def save_messages(self, thread_id: str, messages: list[MessageType]):
 class InMemoryHistory(History):
     _history: ClassVar[dict[str, list[MessageType]]] = {}
 
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
     def load_messages(
         self,
         thread_id: str,

src/controlflow/planning/__init__.py

@@ -0,0 +1 @@
+from . import plan

src/controlflow/tasks/auto_tasks.py → src/controlflow/planning/auto_tasks.py

@@ -1,66 +1,37 @@
 from enum import Enum
-from typing import Any, Callable, Generic, Literal, Optional, TypeVar, Union
+from typing import Any, Callable, Literal, Optional, TypeVar, Union
 
 from pydantic import Field
 
 from controlflow.core.agent import Agent
 from controlflow.core.task import Task
-from controlflow.utilities.types import AssistantTool, ControlFlowModel
+from controlflow.utilities.types import ControlFlowModel
 
 ToolLiteral = TypeVar("ToolLiteral", bound=str)
 
 
-class SimpleType(Enum):
-    NONE = "NONE"
-    BOOLEAN = "BOOLEAN"
-    INTEGER = "INTEGER"
-    FLOAT = "FLOAT"
+class ResultType(Enum):
     STRING = "STRING"
-
-    def to_type(self):
-        return {
-            SimpleType.BOOLEAN: bool,
-            SimpleType.INTEGER: int,
-            SimpleType.FLOAT: float,
-            SimpleType.STRING: str,
-        }[self]
-
-
-class ListType(ControlFlowModel):
-    list_type: Union[SimpleType, "ListType", "DictType", "UnionType"]
-
-    def to_type(self):
-        return list[self.list_type.to_type()]
-
-
-class DictType(ControlFlowModel):
-    key_type: Union[SimpleType, "UnionType"]
-    value_type: Union[SimpleType, ListType, "UnionType", "DictType", None]
-
-    def to_type(self):
-        return dict[
-            self.key_type.to_type(),
-            self.value_type.to_type() if self.value_type is not None else None,
-        ]
-
-
-class UnionType(ControlFlowModel):
-    union_types: list[Union[SimpleType, ListType, DictType, None]]
-
-    def to_type(self):
-        types = [t.to_type() if t is not None else None for t in self.union_types]
-        return Union[*types]  # type: ignore
+    NONE = "NONE"
 
 
 class TaskReference(ControlFlowModel):
+    """
+    A reference to a task by its ID. Used for indicating task depenencies.
+    """
+
     id: int
 
 
 class AgentReference(ControlFlowModel):
+    """
+    A reference to an agent by its name. Used for assigning agents to tasks.
+    """
+
     name: str
 
 
-class AgentTemplate(ControlFlowModel, Generic[ToolLiteral]):
+class AgentTemplate(ControlFlowModel):
     name: str
     description: Optional[str] = Field(
         None,
@@ -71,37 +42,45 @@ class AgentTemplate(ControlFlowModel, Generic[ToolLiteral]):
         description="Private instructions for the agent to follow when completing tasks.",
     )
     user_access: bool = Field(
-        False, description="Whether the agent can interact with a human user."
-    )
-    tools: list[ToolLiteral] = Field(
-        [], description="The tools that the agent has access to."
+        False, description="If True, the agent can interact with a human user."
     )
+    tools: list[str] = Field([], description="The tools that the agent has access to.")
 
 
-class TaskTemplate(ControlFlowModel, Generic[ToolLiteral]):
+class TaskTemplate(ControlFlowModel):
     id: int
     objective: str = Field(description="The task's objective.")
     instructions: Optional[str] = Field(
         None, description="Instructions for completing the task."
     )
-    result_type: Optional[Union[SimpleType, ListType, DictType, UnionType]] = None
+    result_type: Union[ResultType, list[str]] = Field(
+        ResultType.STRING,
+        description="The type of result expected from the task, defaults to a string output. "
+        "Can also be `NONE` if the task does not produce a result (but may have side effects) or "
+        "a list of choices if the task has a discrete set of possible outputs.",
+    )
     context: dict[str, Union[TaskReference, Any]] = Field(
         default_factory=dict,
-        description="The task's context, which can include TaskReferences to create dependencies.",
+        description="The task's context. Values may be constants, TaskReferences, or "
+        "collections of either. Any `TaskReferences` will create upstream dependencies, meaning "
+        "this task will receive the referenced task's output as input.",
     )
     depends_on: list[TaskReference] = Field(
         default_factory=list,
-        description="Tasks that must be completed before this task can be started.",
+        description="Tasks that must be completed before this task can be started, "
+        "though their outputs are not used.",
     )
     parent: Optional[TaskReference] = Field(
-        None, description="The parent task that this task is a subtask of."
+        None,
+        description="Indicate that this task is a subtask of a parent. Not required for top-level tasks.",
     )
     agents: list[AgentReference] = Field(
         default_factory=list,
         description="Any agents assigned to the task. If not specified, the default agent will be used.",
     )
-    tools: list[ToolLiteral] = Field(
-        [], description="The tools available for this task."
+    tools: list[str] = Field([], description="The tools available for this task.")
+    user_access: bool = Field(
+        False, description="If True, the task requires interaction with a human user."
     )
 
 
@@ -134,15 +113,19 @@ def create_tasks(
 
     # create tasks from templates
     for task_template in task_templates.values():
+        if task_template.result_type == ResultType.NONE:
+            result_type = None
+        elif result_type == ResultType.STRING:
+            result_type = str
+        else:
+            result_type = task_template.result_type
+
         tasks[task_template.id] = Task(
             objective=task_template.objective,
             instructions=task_template.instructions,
-            result_type=(
-                task_template.result_type.to_type()
-                if task_template.result_type
-                else None
-            ),
+            result_type=result_type,
             tools=[tools[tool] for tool in task_template.tools],
+            use_access=task_template.user_access,
         )
 
     # resolve references
@@ -164,15 +147,15 @@ def create_tasks(
     return list(tasks.values())
 
 
-class Templates(ControlFlowModel, Generic[ToolLiteral]):
-    task_templates: list[TaskTemplate[ToolLiteral]]
-    agent_templates: list[AgentTemplate[ToolLiteral]]
+class Templates(ControlFlowModel):
+    task_templates: list[TaskTemplate]
+    agent_templates: list[AgentTemplate]
 
 
 def auto_tasks(
     description: str,
     available_agents: list[Agent] = None,
-    available_tools: list[Union[AssistantTool, Callable]] = None,
+    available_tools: list[Callable] = None,
 ) -> list[Task]:
     tool_names = []
     for tool in available_tools or []:
@@ -183,6 +166,20 @@ def auto_tasks(
     else:
         literal_tool_names = None
 
+    class TaskTemplate_Tools(TaskTemplate):
+        tools: list[literal_tool_names] = Field(
+            [], description="The tools available for this task."
+        )
+
+    class AgentTemplate_Tools(AgentTemplate):
+        tools: list[literal_tool_names] = Field(
+            [], description="The tools that the agent has access to."
+        )
+
+    class Templates_Tools(Templates):
+        task_templates: list[TaskTemplate_Tools]
+        agent_templates: list[AgentTemplate_Tools]
+
     task = Task(
         objective="""
         Generate the minimal set of tasks required to complete the provided
@@ -201,7 +198,7 @@ def auto_tasks(
         agents for your tasks, the default agent will be used. Do not post messages, just return your
         result.
         """,
-        result_type=Templates[literal_tool_names],
+        result_type=Templates,
         context=dict(
             description=description,
             available_agents=available_agents,

src/controlflow/planning/plan.py

@@ -0,0 +1,77 @@
+from typing import Optional, TypeVar, Union
+
+from pydantic import Field
+
+from controlflow.core.agent import Agent
+from controlflow.core.task import Task
+from controlflow.llm.tools import Tool, as_tools
+from controlflow.utilities.types import ControlFlowModel
+
+ToolLiteral = TypeVar("ToolLiteral", bound=str)
+
+
+class PlanTask(ControlFlowModel):
+    id: int
+    objective: str
+    instructions: Optional[str] = Field(
+        None,
+        description="Any additional instructions for completing the task objective.",
+    )
+    depends_on: list[int] = Field(
+        [], description="Tasks that must be completed before this task can be started."
+    )
+
+
+class Plan(ControlFlowModel):
+    tasks: list[PlanTask]
+
+
+def plan(
+    objective: str,
+    instructions: str = None,
+    agent: Agent = None,
+    tools: list[Union[callable, Tool]] = None,
+) -> Task:
+    """
+    Given an objective and instructions for achieving it, generate a plan for
+    completing the objective. Each step of the plan will be turned into a task
+    objective.
+    """
+    agents = [agent] if agent else None
+    plan_task = Task(
+        objective="""
+            Create a plan to complete the provided objective. Each step of your plan
+            will be turned into a task objective, like this one. After generating
+            the plan, you will be tasked with executing it, using your tools and any additional ones provided.
+            """,
+        instructions="""
+        Indicate dependencies between tasks, including sequential dependencies.
+        """,
+        result_type=Plan,
+        context=dict(
+            plan_objective=objective,
+            plan_instructions=instructions,
+            plan_tools=[
+                t.dict(include={"name", "description"}) for t in as_tools(tools or [])
+            ],
+        ),
+        agents=agents,
+    )
+
+    plan_task.run()
+
+    parent_task = Task(objective=objective, agents=agents)
+    task_ids = {}
+
+    subtask: PlanTask
+    for subtask in plan_task.result.tasks:
+        task_ids[subtask.id] = Task(
+            objective=subtask.objective,
+            instructions=subtask.instructions,
+            parent=parent_task,
+            depends_on=[task_ids[task_id] for task_id in subtask.depends_on],
+            agents=agents,
+            tools=tools,
+        )
+
+    return parent_task