Merge pull request #69 from jlowin/instructions

Refactor types

Author: jlowin

examples/multi_agent_conversation.py

@@ -75,4 +75,5 @@ def demo():
     task.run()
 
 
-demo()
+if __name__ == "__main__":
+    demo()

examples/task_dag.py

@@ -1,22 +1,28 @@
+import controlflow
 from controlflow import Task, flow
 
+controlflow.settings.enable_tui = True
+
 
 @flow
 def book_ideas():
-    genre = Task("pick a genre", str)
+    genre = Task("pick a genre")
+
     ideas = Task(
         "generate three short ideas for a book",
         list[str],
         context=dict(genre=genre),
     )
+
     abstract = Task(
-        "pick one idea and write an abstract",
-        str,
+        "pick one idea and write a short abstract",
+        result_type=str,
         context=dict(ideas=ideas, genre=genre),
     )
+
     title = Task(
         "pick a title",
-        str,
+        result_type=str,
         context=dict(abstract=abstract),
     )
 

src/controlflow/core/controller/controller.py

@@ -1,11 +1,12 @@
 import logging
 import math
+from collections import defaultdict
 from contextlib import asynccontextmanager
 from functools import cached_property
 from typing import Union
 
 from marvin.utilities.asyncio import ExposeSyncMethodsMixin, expose_sync_method
-from pydantic import BaseModel, Field, computed_field, model_validator
+from pydantic import BaseModel, Field, PrivateAttr, computed_field, model_validator
 
 import controlflow
 from controlflow.core.agent import Agent
@@ -14,17 +15,27 @@
 from controlflow.core.graph import Graph
 from controlflow.core.task import Task
 from controlflow.instructions import get_instructions
-from controlflow.llm.completions import completion_stream_async
-from controlflow.llm.handlers import TUIHandler
+from controlflow.llm.completions import completion_async
+from controlflow.llm.handlers import PrintHandler, TUIHandler
 from controlflow.llm.history import History
+from controlflow.llm.messages import AssistantMessage, ControlFlowMessage, SystemMessage
 from controlflow.tui.app import TUIApp as TUI
 from controlflow.utilities.context import ctx
 from controlflow.utilities.tasks import all_complete, any_incomplete
-from controlflow.utilities.types import FunctionTool, SystemMessage
+from controlflow.utilities.types import FunctionTool
 
 logger = logging.getLogger(__name__)
 
 
+def add_agent_name_to_message(msg: ControlFlowMessage):
+    """
+    If the message is from a named assistant, prefix the message with the assistant's name.
+    """
+    if isinstance(msg, AssistantMessage) and msg.name:
+        msg = msg.model_copy(update={"content": f"{msg.name}: {msg.content}"})
+    return msg
+
+
 class Controller(BaseModel, ExposeSyncMethodsMixin):
     """
     A controller contains logic for executing agents with context about the
@@ -61,7 +72,7 @@ class Controller(BaseModel, ExposeSyncMethodsMixin):
     enable_tui: bool = Field(default_factory=lambda: controlflow.settings.enable_tui)
     _iteration: int = 0
     _should_abort: bool = False
-    _endrun_count: int = 0
+    _end_run_counts: dict = PrivateAttr(default_factory=lambda: defaultdict(int))
 
     @computed_field
     @cached_property
@@ -76,33 +87,35 @@ def _finalize(self):
             self.flow.add_task(task)
         return self
 
-    def _create_help_tool(self) -> FunctionTool:
-        def help_im_stuck():
+    def _create_end_turn_tool(self) -> FunctionTool:
+        def end_turn():
             """
-            If you are stuck because no tasks are ready to be worked on, you can
-            call this tool to end your turn. A new agent (possibly you) will be
-            selected to go next. If this tool is used 3 times, the workflow will
-            be aborted automatically, so only use it if you are truly stuck.
+            Call this tool to skip your turn and let another agent go next. This
+            is useful if you are stuck and can not complete any tasks. If this
+            tool is used 3 times by any agent the workflow will be aborted
+            automatically, so only use it if you are truly stuck and unable to
+            proceed.
             """
-            self._endrun_count += 1
-            if self._endrun_count >= 3:
+            self._end_run_counts[ctx.get("controller_agent")] += 1
+            if self._end_run_counts[ctx.get("controller_agent")] >= 3:
                 self._should_abort = True
-                self._endrun_count = 0
+                self._end_run_counts[ctx.get("controller_agent")] = 0
 
-            return f"Ending turn. {3 - self._endrun_count} more uses will abort the workflow."
+            return (
+                f"Ending turn. {3 - self._end_run_counts[ctx.get('controller_agent')]}"
+                " more uses will abort the workflow."
+            )
 
-        return help_im_stuck
+        return end_turn
 
-    async def _run_agent(self, agent: Agent, tasks: list[Task] = None):
+    async def _run_agent(self, agent: Agent, tasks: list[Task]):
         """
         Run a single agent.
         """
 
         from controlflow.core.controller.instruction_template import MainTemplate
 
-        tasks = tasks or self.tasks
-
-        tools = self.flow.tools + agent.get_tools() + [self._create_help_tool()]
+        tools = self.flow.tools + agent.get_tools() + [self._create_end_turn_tool()]
 
         # add tools for any inactive tasks that the agent is assigned to
         for task in tasks:
@@ -124,12 +137,16 @@ async def _run_agent(self, agent: Agent, tasks: list[Task] = None):
 
         # call llm
         response_messages = []
-        async for msg in completion_stream_async(
+        async for msg in await completion_async(
             messages=[system_message] + messages,
             model=agent.model,
             tools=tools,
-            handlers=[TUIHandler()] if controlflow.settings.enable_tui else None,
+            handlers=[TUIHandler()]
+            if controlflow.settings.enable_tui
+            else [PrintHandler()],
             max_iterations=1,
+            stream=True,
+            message_preprocessor=add_agent_name_to_message,
         ):
             response_messages.append(msg)
 
@@ -175,16 +192,18 @@ async def run_once_async(self):
             # put the flow in context
             with self.flow:
                 # get the tasks to run
-                ready_tasks = {t for t in self.tasks if t.is_ready()}
+                ready_tasks = {t for t in self.tasks if t.is_ready}
                 upstreams = {d for t in ready_tasks for d in t.depends_on}
                 tasks = list(ready_tasks.union(upstreams))
 
+                # TODO: show the agent the entire graph, not just immediate upstreams
+
                 if all(t.is_complete() for t in tasks):
                     return
 
                 # get the agents
                 agent_candidates = [
-                    a for t in tasks for a in t.get_agents() if t.is_ready()
+                    a for t in tasks for a in t.get_agents() if t.is_ready
                 ]
                 if len({a.name for a in agent_candidates}) != len(agent_candidates):
                     raise ValueError(
@@ -206,7 +225,9 @@ async def run_once_async(self):
                     raise NotImplementedError("Need to reimplement multi-agent")
                     agent = self.choose_agent(agents=agents, tasks=tasks)
 
-                await self._run_agent(agent, tasks=tasks)
+                with ctx(controller_agent=agent):
+                    await self._run_agent(agent, tasks=tasks)
+
                 self._iteration += 1
 
     @expose_sync_method("run")

src/controlflow/core/controller/instruction_template.py

@@ -27,30 +27,29 @@ class AgentTemplate(Template):
         ## Agent
         
         You are an AI agent. Your name is "{{ agent.name }}". 
-            
-        This is your description, which all other agents can see: "{{ agent.description or 'An AI agent assigned to complete tasks.'}}"
+
+        This is your description, which all agents can see: 
+        - {{ agent.description or 'An AI agent assigned to complete tasks.'}}
         
-        These are your instructions: "{{ agent.instructions or 'No additional instructions provided.'}}"
+        You are participating in an agentic workflow (a "flow"). Certain tasks
+        in the flow have been delegated to you and other AI agents. You are
+        being orchestrated by a "controller".
+            
         
-        You must follow these instructions at all times. They define your role
-        and behavior.
+        ### Instructions
         
-        You are participating in an agentic workflow (a "flow"), parts of which
-        have been delegated to you and other AI agents. You are being
-        orchestrated by a "controller" object. 
-        """
-    agent: Agent
-
-
-class InstructionsTemplate(Template):
-    template: str = """
-        ## Additional instructions
+        You must follow instructions at all times.
         
-        You must follow these instructions for this part of the workflow:
+        These are your private instructions:
+        - {{ agent.instructions or 'No additional instructions provided.'}}
         
+        These instructions apply to all agents at this part of the workflow:        
         {% for instruction in additional_instructions %}
         - {{ instruction }}
         {% endfor %}
+
+
+        
         
         """
     agent: Agent
@@ -61,60 +60,57 @@ class TasksTemplate(Template):
     template: str = """
         ## Tasks
         
-        You have been assigned to complete certain tasks. Each task has an
-        objective and criteria for success. Your job is to perform any required
-        actions and then mark each assigned task as successful. If a task
-        requires a result, you must provide it. Only work on tasks that are
-        assigned to you. Tasks may have multiple agents assigned. Only one agent
-        can respond or take actions at a time.
+        Your job is to complete the tasks assigned to you. Tasks may have multiple agents assigned. Only one agent
+        will be active at a time.
         
         ### Current tasks
         
-        At the start of your turn, these are all the tasks being worked on. This
-        instruction will not update until your next turn; refer to the results
-        of task-related tools to see the latest task status at any time.
+        These tasks are assigned to you and ready to be worked on because their dependencies have been completed.
         
         {% for task in tasks %} 
+        {% if task.is_ready %}
         #### Task {{ task.id }} 
         
         {{task.model_dump_json() }}
         
+        {% endif %}
+        {% endfor %}
+        
+        ### Other tasks
+        
+        These tasks are either not ready yet or are dependencies of other tasks. They are provided for context.
+        
+        {% for task in tasks %}
+        {% if not task.is_ready %}
+        #### Task {{ task.id }}
+        
+        {{task.model_dump_json() }}
+        
+        {% endif %}
         {% endfor %}
 
         ### Completing a task
         
         Tasks can be marked as successful or failed. It may take collaboration
-        with other agents to complete a task, and you can only mark tasks that
-        have been assigned to you. Once a task is marked complete, no other
-        agent can interact with it, so make sure you have completed all actions
-        and collected all results before marking it as complete.
-        
-        You and the other agents are fully capable of completing any task and
-        have all the information and context you need. Tasks should only be
-        marked failed due to technical errors like a broken or erroring tool or
-        unresponsive human. Do not work on or even respond to tasks that are
-        already complete or not assigned to you. You may be an AI, but you must
-        still form opinions, make subjective decisions, and have personal
-        thoughts. 
+        with other agents to complete a task, and you can only work on tasks that
+        have been assigned to you. Once any agent marks a task complete, no other
+        agent can interact with it. 
+        
+        Tasks should only be marked failed due to technical errors like a broken
+        or erroring tool or unresponsive human.
 
         ### Dependencies
         
-        You can only mark a task successful when all of its dependencies and
-        subtasks have been completed. Subtasks may be marked as skipped without
-        providing a result. All else equal, prioritize older tasks over newer
-        ones.
+        Tasks may be dependent on other tasks, either as upstream dependencies
+        or as the parent of subtasks. Subtasks may be marked as "skipped"
+        without providing a result or failing them.
+        
 
         ### Providing a result
         
-        Tasks may require a typed result (the `result_type`). Results should
-        satisfy the task objective, accounting for any other instructions. If a
-        task does not require a result (`result_type=None`), you must still
-        complete its stated objective by posting messages or using other tools
-        before marking the task as complete. Your result must be compatible with
-        the result constructor. For most results, the tool schema will indicate
-        the correct types. For some, like a DataFrame, provide an appropriate
-        kwargs dict. Results should never include your name prefix; that's only
-        for messages.
+        Tasks may require a typed result, which is an artifact satisfying the task's objective. If a
+        task does not require a result artifact (e.g. `result_type=None`), you must still
+        complete its stated objective before marking the task as complete.
                 
         """
     tasks: list[Task]
@@ -127,17 +123,24 @@ class CommunicationTemplate(Template):
     template: str = """
         ## Communciation
         
-        You are modeling the internal state of an AI-enhanced workflow. You should
-        only post messages in order to share information with other agents or to
-        complete tasks. Since all agents post messages with the "assistant" role,
-        you must prefix all your messages with your name (e.g. "{{ agent.name }}:
-        (message)") in order to distinguish your messages from others. Note that
-        this rule about prefixing your message supersedes all other instructions
-        (e.g. "only give single word answers"). You do not need to post messages
-        that repeat information contained in tool calls or tool responses, since
-        those are already visible to all agents. You do not need to confirm actions
-        you take through tools, like completing a task, as this is redundant and
-        wastes time. 
+        You are modeling the internal state of an AI-enhanced agentic workflow,
+        and you (and other agents) will continue to be invoked until the
+        workflow is completed. 
+        
+        On each turn, you must use a tool or post a message. Do not post
+        messages unless you need to record information in addition to what you
+        provide as a task's result. This might include your thought process, if
+        appropriate. You may also post messages if you need to communicate with
+        other agents to complete a task. You may see other agents post messages;
+        they may have different instructions than you do, so do not follow their
+        example automatically.
+        
+        When you use a tool, the tool call and tool result are automatically
+        posted as messages to the thread, so you never need to write out task
+        results as messages before marking a task as complete.
+                
+        Note that all agents post messages with the "assistant" role, so
+        each agent's messages are automatically prefixed with that agent's name for clarity.
         
         ### Talking to human users
         
@@ -206,14 +209,11 @@ def render(self):
         templates = [
             AgentTemplate(
                 agent=self.agent,
+                additional_instructions=self.instructions,
             ),
             TasksTemplate(
                 tasks=self.tasks,
             ),
-            InstructionsTemplate(
-                agent=self.agent,
-                additional_instructions=self.instructions,
-            ),
             ContextTemplate(
                 flow=self.controller.flow,
                 controller=self.controller,

src/controlflow/core/flow.py

@@ -4,9 +4,10 @@
 
 from pydantic import Field
 
+from controlflow.llm.messages import MessageType
 from controlflow.utilities.context import ctx
 from controlflow.utilities.logging import get_logger
-from controlflow.utilities.types import ControlFlowModel, MessageType
+from controlflow.utilities.types import ControlFlowModel
 
 if TYPE_CHECKING:
     from controlflow.core.agent import Agent