Orchestrator Chat and OAI Assistant update (#31)

Author: ekzhu

examples/patterns.py

@@ -3,13 +3,12 @@
 import logging
 
 import openai
-from agnext.agent_components.model_client import OpenAI
 from agnext.application_components import (
     SingleThreadedAgentRuntime,
 )
 from agnext.chat.agents.oai_assistant import OpenAIAssistantAgent
 from agnext.chat.patterns.group_chat import GroupChat, GroupChatOutput
-from agnext.chat.patterns.orchestrator import Orchestrator
+from agnext.chat.patterns.orchestrator_chat import OrchestratorChat
 from agnext.chat.types import TextMessage
 from agnext.core._agent import Agent
 from agnext.core.intervention import DefaultInterventionHandler, DropMessage
@@ -38,20 +37,28 @@ def reset(self) -> None:
 
 
 class LoggingHandler(DefaultInterventionHandler):
+    send_color = "\033[31m"
+    response_color = "\033[34m"
+    reset_color = "\033[0m"
+
     @override
     async def on_send(self, message: Any, *, sender: Agent | None, recipient: Agent) -> Any | type[DropMessage]:
         if sender is None:
-            print(f"Sending message to {recipient.name}: {message}")
+            print(f"{self.send_color}Sending message to {recipient.name}:{self.reset_color} {message}")
         else:
-            print(f"Sending message from {sender.name} to {recipient.name}: {message}")
+            print(
+                f"{self.send_color}Sending message from {sender.name} to {recipient.name}:{self.reset_color} {message}"
+            )
         return message
 
     @override
     async def on_response(self, message: Any, *, sender: Agent, recipient: Agent | None) -> Any | type[DropMessage]:
         if recipient is None:
-            print(f"Received response from {sender.name}: {message}")
+            print(f"{self.response_color}Received response from {sender.name}:{self.reset_color} {message}")
         else:
-            print(f"Received response from {sender.name} to {recipient.name}: {message}")
+            print(
+                f"{self.response_color}Received response from {sender.name} to {recipient.name}:{self.reset_color} {message}"
+            )
         return message
 
 
@@ -131,19 +138,47 @@ async def orchestrator(message: str) -> None:
         thread_id=product_manager_oai_thread.id,
     )
 
-    chat = Orchestrator(
-        "Manager",
-        "A software development team manager.",
-        runtime,
-        [developer, product_manager],
-        model_client=OpenAI(model="gpt-3.5-turbo"),
+    planner_oai_assistant = openai.beta.assistants.create(
+        model="gpt-4-turbo",
+        name="Planner",
+        instructions="You are a planner of complex tasks.",
+    )
+    planner_oai_thread = openai.beta.threads.create()
+    planner = OpenAIAssistantAgent(
+        name="Planner",
+        description="A planner that organizes and schedules tasks.",
+        runtime=runtime,
+        client=openai.AsyncClient(),
+        assistant_id=planner_oai_assistant.id,
+        thread_id=planner_oai_thread.id,
+    )
+
+    orchestrator_oai_assistant = openai.beta.assistants.create(
+        model="gpt-4-turbo",
+        name="Orchestrator",
+        instructions="You are an orchestrator that coordinates the team to complete a complex task.",
+    )
+    orchestrator_oai_thread = openai.beta.threads.create()
+    orchestrator = OpenAIAssistantAgent(
+        name="Orchestrator",
+        description="An orchestrator that coordinates the team.",
+        runtime=runtime,
+        client=openai.AsyncClient(),
+        assistant_id=orchestrator_oai_assistant.id,
+        thread_id=orchestrator_oai_thread.id,
     )
 
-    response = runtime.send_message(
-        TextMessage(content=message, source="customer"),
-        chat,
+    chat = OrchestratorChat(
+        "Orchestrator Chat",
+        "A software development team.",
+        runtime,
+        orchestrator=orchestrator,
+        planner=planner,
+        specialists=[developer, product_manager],
     )
 
+    response = runtime.send_message(TextMessage(content=message, source="Customer"), chat)
+
     while not response.done():
         await runtime.process_next()
 

src/agnext/chat/agents/oai_assistant.py

@@ -1,4 +1,4 @@
-from typing import Callable, Dict
+from typing import Callable, Dict, List
 
 import openai
 
@@ -23,8 +23,6 @@ def __init__(
         self._client = client
         self._assistant_id = assistant_id
         self._thread_id = thread_id
-        # TODO: investigate why this is 1, as setting this to 0 causes the earlest message in the window to be ignored.
-        self._current_session_window_length = 1
         self._tools = tools or {}
 
     @message_handler(TextMessage)
@@ -36,32 +34,40 @@ async def on_text_message(self, message: TextMessage, cancellation_token: Cancel
             role="user",
             metadata={"sender": message.source},
         )
-        self._current_session_window_length += 1
 
     @message_handler(Reset)
     async def on_reset(self, message: Reset, cancellation_token: CancellationToken) -> None:
-        # Reset the current session window.
-        self._current_session_window_length = 1
+        # Get all messages in this thread.
+        all_msgs: List[str] = []
+        while True:
+            if not all_msgs:
+                msgs = await self._client.beta.threads.messages.list(self._thread_id)
+            else:
+                msgs = await self._client.beta.threads.messages.list(self._thread_id, after=all_msgs[-1])
+            for msg in msgs.data:
+                all_msgs.append(msg.id)
+            if not msgs.has_next_page():
+                break
+        # Delete all the messages.
+        for msg_id in all_msgs:
+            status = await self._client.beta.threads.messages.delete(message_id=msg_id, thread_id=self._thread_id)
+            assert status.deleted is True
 
     @message_handler(RespondNow)
     async def on_respond_now(self, message: RespondNow, cancellation_token: CancellationToken) -> TextMessage:
+        # Handle response format.
+
         # Create a run and wait until it finishes.
         run = await self._client.beta.threads.runs.create_and_poll(
             thread_id=self._thread_id,
             assistant_id=self._assistant_id,
-            truncation_strategy={
-                "type": "last_messages",
-                "last_messages": self._current_session_window_length,
-            },
+            response_format=message.response_format,
         )
 
         if run.status != "completed":
             # TODO: handle other statuses.
             raise ValueError(f"Run did not complete successfully: {run}")
 
-        # Increment the current session window length.
-        self._current_session_window_length += 1
-
         # Get the last message from the run.
         response = await self._client.beta.threads.messages.list(self._thread_id, run_id=run.id, order="desc", limit=1)
         last_message_content = response.data[0].content