Completions working with new message types

src/controlflow/core/controller/controller.py

@@ -20,7 +20,7 @@
 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, Message
+from controlflow.utilities.types import FunctionTool, SystemMessage
 
 logger = logging.getLogger(__name__)
 
@@ -119,25 +119,23 @@ async def _run_agent(self, agent: Agent, tasks: list[Task] = None):
         instructions = instructions_template.render()
 
         # prepare messages
-        system_message = Message(content=instructions, role="system")
+        system_message = SystemMessage(content=instructions)
         messages = self.history.load_messages(thread_id=self.flow.thread_id)
 
         # call llm
-        r = []
-        async for _ in completion_stream_async(
+        response_messages = []
+        async for msg in completion_stream_async(
             messages=[system_message] + messages,
             model=agent.model,
             tools=tools,
             handlers=[TUIHandler()] if controlflow.settings.enable_tui else None,
             max_iterations=1,
-            response_callback=r.append,
         ):
-            pass
-        response = r[0]
+            response_messages.append(msg)
 
         # save history
         self.history.save_messages(
-            thread_id=self.flow.thread_id, messages=response.messages
+            thread_id=self.flow.thread_id, messages=response_messages
         )
 
         # create_json_artifact(

src/controlflow/core/flow.py

@@ -6,7 +6,7 @@
 
 from controlflow.utilities.context import ctx
 from controlflow.utilities.logging import get_logger
-from controlflow.utilities.types import ControlFlowModel, Message
+from controlflow.utilities.types import ControlFlowModel, MessageType
 
 if TYPE_CHECKING:
     from controlflow.core.agent import Agent
@@ -68,7 +68,7 @@ def get_flow() -> Optional[Flow]:
     return flow
 
 
-def get_flow_messages(limit: int = None) -> list[Message]:
+def get_flow_messages(limit: int = None) -> list[MessageType]:
     """
     Loads messages from the flow's thread.
 

src/controlflow/core/task.py

@@ -36,7 +36,6 @@
 from controlflow.utilities.types import (
     NOTSET,
     ControlFlowModel,
-    Message,
     PandasDataFrame,
     PandasSeries,
     ToolType,
@@ -58,54 +57,6 @@ class TaskStatus(Enum):
     SKIPPED = "SKIPPED"
 
 
-class LoadMessage(ControlFlowModel):
-    """
-    This special object can be used to indicate that a task result should be
-    loaded from a recent message posted to the flow's thread.
-    """
-
-    type: Literal["LoadMessage"] = Field(
-        'You must provide this value as "LoadMessage".'
-    )
-
-    num_messages_ago: int = Field(
-        1,
-        description="The number of messages ago to retrieve. Default is 1, or the most recent message.",
-    )
-
-    strip_prefix: str = Field(
-        None,
-        description="These characters will be removed from the start "
-        "of the message. For example, remove text like your name prefix.",
-    )
-
-    strip_suffix: Optional[str] = Field(
-        None,
-        description="These characters will be removed from the end of "
-        "the message. For example, remove comments like 'I'll mark the task complete now.'",
-    )
-
-    def trim_message(self, message: Message) -> str:
-        content = message.content[0].text.value
-        if self.strip_prefix:
-            if content.startswith(self.strip_prefix):
-                content = content[len(self.strip_prefix) :]
-            else:
-                raise ValueError(
-                    f'Invalid strip prefix "{self.strip_prefix}"; messages '
-                    f'starts with "{content[:len(self.strip_prefix) + 10]}"'
-                )
-        if self.strip_suffix:
-            if content.endswith(self.strip_suffix):
-                content = content[: -len(self.strip_suffix)]
-            else:
-                raise ValueError(
-                    f'Invalid strip suffix "{self.strip_suffix}"; messages '
-                    f'ends with "{content[-len(self.strip_suffix) - 10:]}"'
-                )
-        return content.strip()
-
-
 class Task(ControlFlowModel):
     id: str = Field(default_factory=lambda: str(uuid.uuid4().hex[:5]))
     objective: str = Field(