Merge pull request #82 from PrefectHQ/ai-name

Restore ai name to messages

Author: jlowin

src/controlflow/core/controller/controller.py

@@ -217,14 +217,22 @@ async def run_once_async(self):
                     response_handler = ResponseHandler()
                     payload["handlers"].append(response_handler)
 
+                    messages = []
+                    for msg in payload["messages"]:
+                        if isinstance(msg, AIMessage) and msg.name:
+                            msg = msg.copy()
+                            msg.content = (
+                                f"Message from agent: {msg.name}\n\n{msg.content}"
+                            )
+                        messages.append(msg)
+
                     response_gen = await completion_async(
-                        messages=payload["messages"],
+                        messages=messages,
                         model=agent.model,
                         tools=payload["tools"],
                         handlers=payload["handlers"],
                         max_iterations=1,
-                        # assistant_name=agent.name,
-                        # message_preprocessor=payload["message_preprocessor"],
+                        ai_name=agent.name,
                         stream=True,
                     )
                     async for _ in response_gen:
@@ -245,14 +253,20 @@ def run_once(self):
                 response_handler = ResponseHandler()
                 payload["handlers"].append(response_handler)
 
+                messages = []
+                for msg in payload["messages"]:
+                    if isinstance(msg, AIMessage) and msg.name:
+                        msg = msg.copy()
+                        msg.content = f"Message from agent: {msg.name}\n\n{msg.content}"
+                    messages.append(msg)
+
                 response_gen = completion(
-                    messages=payload["messages"],
+                    messages=messages,
                     model=agent.model,
                     tools=payload["tools"],
                     handlers=payload["handlers"],
                     max_iterations=1,
-                    # assistant_name=agent.name,
-                    # message_preprocessor=payload["message_preprocessor"],
+                    ai_name=agent.name,
                     stream=True,
                 )
                 for _ in response_gen:

src/controlflow/llm/completions.py

@@ -1,6 +1,6 @@
 import datetime
 import math
-from typing import AsyncGenerator, Callable, Generator, Optional, Union, cast
+from typing import AsyncGenerator, Callable, Generator, Optional, Union
 
 import langchain_core.language_models as lc_models
 
@@ -24,6 +24,7 @@ def _completion_generator(
     model: lc_models.BaseChatModel,
     tools: Optional[list[Callable]],
     max_iterations: int,
+    ai_name: Optional[str],
     stream: bool,
     **kwargs,
 ) -> Generator[CompletionEvent, None, None]:
@@ -47,7 +48,9 @@ def _completion_generator(
                     input=messages + response_messages,
                     **kwargs,
                 )
-                response_message = AIMessage.from_message(response_message)
+                response_message = AIMessage.from_message(
+                    response_message, name=ai_name
+                )
 
             else:
                 deltas: list[AIMessageChunk] = []
@@ -57,7 +60,7 @@ def _completion_generator(
                     input=messages + response_messages,
                     **kwargs,
                 ):
-                    delta = AIMessageChunk.from_chunk(delta)
+                    delta = AIMessageChunk.from_chunk(delta, name=ai_name)
                     deltas.append(delta)
 
                     if snapshot is None:
@@ -127,6 +130,7 @@ async def _completion_async_generator(
     model: lc_models.BaseChatModel,
     tools: Optional[list[Callable]],
     max_iterations: int,
+    ai_name: Optional[str],
     stream: bool,
     **kwargs,
 ) -> AsyncGenerator[CompletionEvent, None]:
@@ -151,7 +155,9 @@ async def _completion_async_generator(
                     tools=tools or None,
                     **kwargs,
                 )
-                response_message = AIMessage.from_message(response_message)
+                response_message = AIMessage.from_message(
+                    response_message, name=ai_name
+                )
 
             else:
                 deltas: list[AIMessageChunk] = []
@@ -162,7 +168,7 @@ async def _completion_async_generator(
                     tools=tools or None,
                     **kwargs,
                 ):
-                    delta = cast(AIMessageChunk, delta)
+                    delta = AIMessageChunk.from_chunk(delta, name=ai_name)
                     deltas.append(delta)
 
                     if snapshot is None:
@@ -251,6 +257,7 @@ def completion(
     tools: list[Callable] = None,
     max_iterations: int = None,
     handlers: list[CompletionHandler] = None,
+    ai_name: Optional[str] = None,
     stream: bool = False,
     **kwargs,
 ) -> Union[list[MessageType], Generator[MessageType, None, None]]:
@@ -266,6 +273,7 @@ def completion(
         model=model,
         tools=tools,
         max_iterations=max_iterations,
+        ai_name=ai_name,
         stream=stream,
         **kwargs,
     )
@@ -286,6 +294,7 @@ async def completion_async(
     tools: list[Callable] = None,
     max_iterations: int = None,
     handlers: list[CompletionHandler] = None,
+    ai_name: Optional[str] = None,
     stream: bool = False,
     **kwargs,
 ) -> Union[list[MessageType], Generator[MessageType, None, None]]:
@@ -301,6 +310,7 @@ async def completion_async(
         model=model,
         tools=tools,
         max_iterations=max_iterations,
+        ai_name=ai_name,
         stream=stream,
         **kwargs,
     )

src/controlflow/llm/messages.py

@@ -63,6 +63,7 @@ def to_message(self, **kwargs) -> AIMessage:
     def __add__(self, other: Any) -> "AIMessageChunk":  # type: ignore
         result = super().__add__(other)
         result.timestamp = self.timestamp
+        result.name = self.name
         return result
 
 

src/controlflow/llm/tools.py

@@ -40,10 +40,12 @@ def wrapper(*args, **kwargs):
 
 class Tool(langchain_core.tools.StructuredTool):
     """
-    A subclass of StructuredTool that is compatible with Pydantic v1 models
-    (which Langchain uses) and v2 models (which ControlFlow users).
+    A subclass of StructuredTool that is compatible with functions whose
+    signatures include either Pydantic v1 models (which Langchain uses) or v2
+    models (which ControlFlow users).
 
-    Note that THIS is a Pydantic v1 model because it subclasses the Langchain class.
+    Note that THIS class is a Pydantic v1 model because it subclasses the Langchain
+    class.
     """
 
     tags: dict[str, Any] = pydantic.v1.Field(default_factory=dict)