diff --git a/examples/chat.py b/examples/chat.py index 3bdf3b31..3f1de333 100644 --- a/examples/chat.py +++ b/examples/chat.py @@ -1,6 +1,6 @@ import sys -from tapeagents.collective import CollectiveAgent, CollectiveTape +from tapeagents.team import TeamAgent, TeamTape from tapeagents.develop import Develop from tapeagents.llms import LLAMA, LLM from tapeagents.rendering import PrettyRenderer @@ -8,20 +8,20 @@ def try_chat(llm: LLM, develop: bool): # equilavent of https://microsoft.github.io/autogen/docs/tutorial/introduction - comedy_duo = CollectiveAgent.create_chat_initiator( + comedy_duo = TeamAgent.create_chat_initiator( name="Joe", llm=llm, system_prompt="Your name is Joe and you are a part of a duo of comedians.", - collective_manager=CollectiveAgent.create( + teammate=TeamAgent.create( name="Cathy", llm=llm, system_prompt="Your name is Cathy and you are a part of a duo of comedians." ), max_turns=3, init_message="Hey Cathy, tell me a joke", ) if develop: - Develop(comedy_duo, CollectiveTape(context=None, steps=[]), PrettyRenderer()).launch() + Develop(comedy_duo, TeamTape(context=None, steps=[]), PrettyRenderer()).launch() else: - for event in comedy_duo.run(CollectiveTape(context=None, steps=[])): + for event in comedy_duo.run(TeamTape(context=None, steps=[])): print(event.model_dump_json(indent=2)) diff --git a/examples/code_chat.py b/examples/code_chat.py index 25b86a15..59ee1af1 100644 --- a/examples/code_chat.py +++ b/examples/code_chat.py @@ -3,7 +3,7 @@ import sys from tapeagents.autogen_prompts import AUTOGEN_ASSISTANT_SYSTEM_MESSAGE -from tapeagents.collective import CollectiveAgent, CollectiveTape +from tapeagents.team import TeamAgent, TeamTape from tapeagents.container_executor import ContainerExecutor from tapeagents.develop import Develop from tapeagents.environment import CodeExecutionEnvironment @@ -16,11 +16,11 @@ def try_chat(llm: LLM, develop: bool): # equilavent of https://microsoft.github.io/autogen/docs/tutorial/introduction - org = CollectiveAgent.create_chat_initiator( + org = TeamAgent.create_chat_initiator( name="UserProxy", llm=llm, system_prompt="", - collective_manager=CollectiveAgent.create( + teammate=TeamAgent.create( name="Assistant", system_prompt=AUTOGEN_ASSISTANT_SYSTEM_MESSAGE, llm=llm, @@ -29,7 +29,7 @@ def try_chat(llm: LLM, develop: bool): init_message="compute 5 fibonacci numbers", execute_code=True, ) - start_tape = CollectiveTape(context=None, steps=[]) + start_tape = TeamTape(context=None, steps=[]) now = f"{datetime.datetime.now():%Y%m%d%H%M%S}" env = CodeExecutionEnvironment(ContainerExecutor(work_dir=f"outputs/chat_code/{now}")) if develop: diff --git a/examples/data_science.py b/examples/data_science.py index 1ecbd249..64b1e3fe 100644 --- a/examples/data_science.py +++ b/examples/data_science.py @@ -4,7 +4,7 @@ from tapeagents.agent import Agent from tapeagents.autogen_prompts import AUTOGEN_ASSISTANT_SYSTEM_MESSAGE -from tapeagents.collective import CollectiveAgent, CollectiveTape +from tapeagents.team import TeamAgent, TeamTape from tapeagents.container_executor import ContainerExecutor from tapeagents.core import Action, FinalStep, Observation, Tape from tapeagents.environment import CodeExecutionEnvironment, Environment @@ -19,33 +19,33 @@ def make_world(llm: LLM | None = None, env: Environment | None = None) -> tuple[Agent, Tape, Environment]: llm = llm or LiteLLM(model_name="gpt-4o", parameters={"timeout": 15.0}) - coder = CollectiveAgent.create( + coder = TeamAgent.create( name="SoftwareEngineer", system_prompt=( AUTOGEN_ASSISTANT_SYSTEM_MESSAGE + "Only say TERMINATE when your code was successfully executed." ), llm=llm, ) - code_executor = CollectiveAgent.create( + code_executor = TeamAgent.create( name="CodeExecutor", llm=llm, execute_code=True, ) - team = CollectiveAgent.create_collective_manager( + team = TeamAgent.create_team_manager( name="GroupChatManager", subagents=[coder, code_executor], max_calls=15, llm=llm, ) - org = CollectiveAgent.create_chat_initiator( + org = TeamAgent.create_chat_initiator( name="UserProxy", init_message=( "Make a plot comparing the stocks of ServiceNow and Salesforce" " since beginning of 2024. Save it to a PNG file." ), - collective_manager=team, + teammate=team, ) - start_tape = CollectiveTape(context=None, steps=[]) + start_tape = TeamTape(context=None, steps=[]) now = f"{datetime.datetime.now():%Y%m%d%H%M%S}" env = env or CodeExecutionEnvironment(ContainerExecutor(work_dir=f"outputs/data_science/{now}")) return org, start_tape, env diff --git a/examples/multi_chat.py b/examples/multi_chat.py index e58a4408..783de7b0 100644 --- a/examples/multi_chat.py +++ b/examples/multi_chat.py @@ -3,7 +3,7 @@ import sys from tapeagents.autogen_prompts import AUTOGEN_ASSISTANT_SYSTEM_MESSAGE -from tapeagents.collective import CollectiveAgent, CollectiveTape +from tapeagents.team import TeamAgent, TeamTape from tapeagents.container_executor import ContainerExecutor from tapeagents.core import FinalStep from tapeagents.develop import Develop @@ -18,33 +18,33 @@ def try_chat(develop: bool): llm = LiteLLM(model_name="gpt-4o", parameters={"timeout": 15.0}, use_cache=True) - product_manager = CollectiveAgent.create( + product_manager = TeamAgent.create( name="ProductManager", system_prompt="Creative in software product ideas.", llm=llm, ) - coder = CollectiveAgent.create( + coder = TeamAgent.create( name="SoftwareEngineer", system_prompt=AUTOGEN_ASSISTANT_SYSTEM_MESSAGE, llm=llm, ) - code_executor = CollectiveAgent.create( + code_executor = TeamAgent.create( name="CodeExecutor", llm=llm, execute_code=True, ) - team = CollectiveAgent.create_collective_manager( + team = TeamAgent.create_team_manager( name="GroupChatManager", subagents=[product_manager, coder, code_executor], max_calls=15, llm=llm, ) - org = CollectiveAgent.create_chat_initiator( + org = TeamAgent.create_chat_initiator( name="UserProxy", init_message="Find a latest paper about gpt-4 on arxiv and find its potential applications in software.", - collective_manager=team, + teammate=team, ) - start_tape = CollectiveTape(context=None, steps=[]) + start_tape = TeamTape(context=None, steps=[]) now = f"{datetime.datetime.now():%Y%m%d%H%M%S}" env = CodeExecutionEnvironment(ContainerExecutor(work_dir=f"outputs/multi_chat_code/{now}")) if develop: diff --git a/examples/tape_improver.py b/examples/tape_improver.py index 2caddc35..7e193e19 100644 --- a/examples/tape_improver.py +++ b/examples/tape_improver.py @@ -5,7 +5,7 @@ from tapeagents.agent import Agent from tapeagents.chain import Chain -from tapeagents.collective import CollectiveTape +from tapeagents.team import TeamTape from tapeagents.core import ( Action, AgentStep, @@ -113,7 +113,7 @@ class StepParsingError(Action): error: str -CodeImproverTape = Tape[CollectiveTape, SelectAgent | SelectStep | RewriteStep | FinalStep | Call | Respond] +CodeImproverTape = Tape[TeamTape, SelectAgent | SelectStep | RewriteStep | FinalStep | Call | Respond] def improver_tape_view(tape: Tape) -> str: @@ -186,7 +186,7 @@ def generate_steps(self, tape: Any, llm_stream: LLMStream): def make_world(llm: LLM | None = None) -> tuple[Agent, Tape, Tape]: res_dir = f"{pathlib.Path(__file__).parent.resolve()}/res" with open(f"{res_dir}/bad_tape.json", "r") as f: - bad_tape = CollectiveTape.model_validate(json.load(f)) + bad_tape = TeamTape.model_validate(json.load(f)) improver_tape = CodeImproverTape(context=bad_tape, steps=[]) llm = llm or LiteLLM( diff --git a/tapeagents/collective.py b/tapeagents/collective.py deleted file mode 100644 index 61e47cac..00000000 --- a/tapeagents/collective.py +++ /dev/null @@ -1,283 +0,0 @@ -from __future__ import annotations - -import logging -from enum import Enum - -from pydantic import ConfigDict - -from tapeagents.agent import DEFAULT, Agent -from tapeagents.autogen_prompts import SELECT_SPEAKER_MESSAGE_AFTER_TEMPLATE, SELECT_SPEAKER_MESSAGE_BEFORE_TEMPLATE -from tapeagents.container_executor import extract_code_blocks -from tapeagents.core import FinalStep, Jump, Pass, Prompt, Tape -from tapeagents.environment import CodeExecutionResult, ExecuteCode -from tapeagents.llms import LLM, LLMStream -from tapeagents.view import Broadcast, Call, Respond, TapeViewStack - -logger = logging.getLogger(__name__) - - -CollectiveTape = Tape[None, Call | Respond | Broadcast | FinalStep | Jump | ExecuteCode | CodeExecutionResult | Pass] - - -class Task(str, Enum): - """ - List of tasks that the agent from the collective can perform - """ - - broadcast_last_message = "broadcast_last_message" - call = "call" - select_and_call = "select_and_call" - execute_code = "execute_code" - respond = "respond" - terminate_or_repeat = "terminate_or_repeat" - respond_or_repeat = "respond_or_repeat" - - -class ActiveCollectiveAgentView: - def __init__(self, agent: CollectiveAgent, tape: CollectiveTape): - """ - CollectiveTapeView contains the ephemeral state computed from the tape. This class extracts the data relevant to - the given agent and also computes some additional information from it, e.g. whether the agent - should call the LLM to generate a message or respond with an already available one. - - """ - view = TapeViewStack.compute(tape) - self.messages = view.messages_by_agent[agent.full_name] - self.last_non_empty_message = next((m for m in reversed(self.messages) if m.content), None) - self.task = agent.get_task(view) - self.steps = view.top.steps - self.steps_by_kind = view.top.steps_by_kind - self.exec_result = self.steps[-1] if self.steps and isinstance(self.steps[-1], CodeExecutionResult) else None - self.should_generate_message = ( - self.task in {Task.call, Task.respond} - and self.messages - and not self.exec_result - and "system" in agent.templates - ) - self.should_stop = ( - agent.max_calls and (agent.max_calls and len(self.steps_by_kind.get("call", [])) >= agent.max_calls) - ) or (self.messages and ("TERMINATE" in self.messages[-1].content)) - - -class CollectiveAgent(Agent[CollectiveTape]): - """ - Agent designed to work in the collective with similar other agents performing different kinds - """ - - max_calls: int | None = None - init_message: str | None = None - tasks: list[Task] - - model_config = ConfigDict(use_enum_values=True) - - def get_task(self, view: TapeViewStack) -> Task: - # TODO: use nodes - return Task(self.tasks[view.top.next_node]) - - def make_prompt(self, tape: CollectiveTape) -> Prompt: - view = ActiveCollectiveAgentView(self, tape) - llm_messages = [] - for step in view.messages: - match step: - # When we make the LLM messages, we use "kind" == "user" for messages - # originating from other agents, and "kind" == "assistant" for messages by this agent. - case Call() if step.by == self.full_name: - # I called someone - llm_messages.append({"role": "assistant", "content": step.content}) - case Call(): - # someone called me - # we exclude empty call messages from the prompt - if not step.content: - continue - llm_messages.append( - { - "role": "user", - "content": step.content, - "name": step.by.split("/")[-1], - } - ) - case Respond() if step.by == self.full_name: - # I responded to someone - llm_messages.append({"role": "assistant", "content": step.content}) - case Respond(): - # someone responded to me - who_returned = step.by.split("/")[-1] - llm_messages.append({"role": "user", "content": step.content, "name": who_returned}) - case Broadcast(): - llm_messages.append({"role": "user", "content": step.content, "name": step.from_}) - match view.task: - case Task.select_and_call: - subagents = ", ".join(self.get_subagent_names()) - select_before = [ - { - "role": "system", - "content": self.templates["select_before"].format(subagents=subagents), - } - ] - select_after = [ - { - "role": "system", - "content": self.templates["select_after"].format(subagents=subagents), - } - ] - return Prompt(messages=select_before + llm_messages + select_after) - case _ if view.should_generate_message: - system = [{"role": "system", "content": self.templates["system"]}] - return Prompt(messages=system + llm_messages) - case _: - return Prompt() - - @classmethod - def create( - cls, - name: str, - system_prompt: str | None = None, - llm: LLM | None = None, - execute_code: bool = False, - ): # type: ignore - """ - Create a simple agent that can execute code, think and respond to messages - """ - return cls( - name=name, - templates={"system": system_prompt} if system_prompt else {}, - llms={DEFAULT: llm} if llm else {}, - tasks=([Task.execute_code] if execute_code else []) + [Task.respond], - ) - - @classmethod - def create_collective_manager( - cls, - name: str, - subagents: list[Agent[CollectiveTape]], - llm: LLM, - max_calls: int = 1, - ): - """ - Create a collective manager that broadcasts the last message to all subagents, selects one of them to call, call it and - responds to the last message if the termination message is not received. - """ - return cls( - name=name, - subagents=subagents, - tasks=[ - Task.broadcast_last_message, - Task.select_and_call, - Task.respond_or_repeat, - ], - max_calls=max_calls, - templates={ - "select_before": SELECT_SPEAKER_MESSAGE_BEFORE_TEMPLATE, - "select_after": SELECT_SPEAKER_MESSAGE_AFTER_TEMPLATE, - }, - llms={DEFAULT: llm}, - ) - - @classmethod - def create_chat_initiator( - cls, - name: str, - collective_manager: Agent[CollectiveTape], - init_message: str, - system_prompt: str = "", - llm: LLM | None = None, - max_calls: int = 1, - execute_code: bool = False, - ): - """ - Create an agent that sets the collective's initial message and calls the collective manager - """ - return cls( - name=name, - templates={ - "system": system_prompt, - }, - llms={DEFAULT: llm} if llm else {}, - subagents=[collective_manager], - tasks=([Task.execute_code] if execute_code else []) + [Task.call, Task.terminate_or_repeat], - max_calls=max_calls, - init_message=init_message, - ) - - def generate_steps(self, tape: CollectiveTape, llm_stream: LLMStream): - view = ActiveCollectiveAgentView(self, tape) - - exec_result_message = "" - if view.exec_result: - if output := view.exec_result.result.output.strip(): - exec_result_message = f"I ran the code and got the following output:\n\n{output}" - else: - exec_result_message = f"I ran the code, the exit code was {view.exec_result.result.exit_code}." - - def _implementation(): - match view.task: - case Task.broadcast_last_message: - recipients = self.get_subagent_names() - last = view.messages[-1] - from_ = last.by.split("/")[-1] - match last: - case Call(): - yield Broadcast(content=last.content, from_=from_, to=list(recipients)) - case Respond(): - recipients = [name for name in recipients if name != last.by.split("/")[-1]] - yield Broadcast( - content=view.messages[-1].content, - from_=from_, - to=list(recipients), - ) - case Broadcast(): - pass - case _: - assert False - case Task.select_and_call: - callee_name = llm_stream.get_text() - # check if the callee is an existing subagent - _ = self.find_subagent(callee_name) - yield Call(agent_name=callee_name) - case Task.call: - # if last task - (other,) = self.subagents - if view.should_generate_message: - yield Call(agent_name=other.name, content=llm_stream.get_text()) - elif view.exec_result: - yield Call(agent_name=other.name, content=exec_result_message) - else: - assert self.init_message and not view.messages - yield Call(agent_name=other.name, content=self.init_message) - case Task.execute_code: - assert not llm_stream - if view.last_non_empty_message is None: - yield Pass() - elif code := extract_code_blocks(view.last_non_empty_message.content): - yield ExecuteCode(code=code) - else: - yield Pass() - case Task.respond: - if view.should_generate_message: - yield Respond(content=llm_stream.get_text()) - elif view.exec_result: - yield Respond(content=exec_result_message) - else: - logger.info( - f"Agent {self.full_name} had to respond with an empty message." - f" You might want to optimize your orchestration logic." - ) - yield Respond() - case Task.terminate_or_repeat: - assert not llm_stream - if view.should_stop: - yield FinalStep(reason="Termination message received") - else: - yield Jump(next_node=0) - case Task.respond_or_repeat: - if view.should_stop: - yield Respond() - else: - yield Jump(next_node=0) - case _: - raise ValueError() - - for step in _implementation(): - step.task = view.task.value - yield step - return diff --git a/tapeagents/environment.py b/tapeagents/environment.py index 38fe3189..d24ed6cf 100644 --- a/tapeagents/environment.py +++ b/tapeagents/environment.py @@ -154,7 +154,7 @@ class CodeExecutionResult(Observation): class CodeExecutionEnvironment(Environment): """ - Environment for the collective agents + Environment for the team agents The only action that the environment can perform is to execute the code blocks """ diff --git a/tapeagents/rendering.py b/tapeagents/rendering.py index 92d69af2..dc855d3c 100644 --- a/tapeagents/rendering.py +++ b/tapeagents/rendering.py @@ -6,7 +6,7 @@ import yaml from .agent import Agent -from .collective import CodeExecutionResult, ExecuteCode +from .team import CodeExecutionResult, ExecuteCode from .observe import LLMCall, retrieve_tape_llm_calls from .container_executor import CodeBlock from .view import Call, Respond diff --git a/tapeagents/team.py b/tapeagents/team.py new file mode 100644 index 00000000..c2ce63d1 --- /dev/null +++ b/tapeagents/team.py @@ -0,0 +1,324 @@ +from __future__ import annotations + +import logging +from typing import Generator + +from pydantic import ConfigDict + +from tapeagents.agent import DEFAULT, Agent, AgentStep, Node +from tapeagents.autogen_prompts import SELECT_SPEAKER_MESSAGE_AFTER_TEMPLATE, SELECT_SPEAKER_MESSAGE_BEFORE_TEMPLATE +from tapeagents.container_executor import extract_code_blocks +from tapeagents.core import FinalStep, Jump, Pass, Prompt, Tape +from tapeagents.environment import CodeExecutionResult, ExecuteCode +from tapeagents.llms import LLM, LLMStream +from tapeagents.view import Broadcast, Call, Respond, TapeViewStack + +logger = logging.getLogger(__name__) + + +TeamTape = Tape[None, Call | Respond | Broadcast | FinalStep | Jump | ExecuteCode | CodeExecutionResult | Pass] + + +class ActiveTeamAgentView: + def __init__(self, agent: TeamAgent, tape: TeamTape): + """ + ActiveTeamAgentView contains the ephemeral state computed from the tape. This class extracts the data relevant to + the given agent and also computes some additional information from it, e.g. whether the agent + should call the LLM to generate a message or respond with an already available one. + """ + view = TapeViewStack.compute(tape) + self.messages = view.messages_by_agent[agent.full_name] + self.last_non_empty_message = next((m for m in reversed(self.messages) if m.content), None) + self.node = agent.get_node(view) + self.steps = view.top.steps + self.steps_by_kind = view.top.steps_by_kind + self.exec_result = self.steps[-1] if self.steps and isinstance(self.steps[-1], CodeExecutionResult) else None + self.should_generate_message = ( + isinstance(self.node, (CallNode, RespondNode)) + and self.messages + and not self.exec_result + and "system" in agent.templates + ) + self.should_stop = ( + agent.max_calls and (agent.max_calls and len(self.steps_by_kind.get("call", [])) >= agent.max_calls) + ) or (self.messages and ("TERMINATE" in self.messages[-1].content)) + + +class TeamAgent(Agent[TeamTape]): + """ + Agent designed to work in the team with similar other agents performing different kinds + """ + + max_calls: int | None = None + init_message: str | None = None + + model_config = ConfigDict(use_enum_values=True) + + def get_node(self, view: TapeViewStack) -> Node: + return self.flow[view.top.next_node] + + @classmethod + def create( + cls, + name: str, + system_prompt: str | None = None, + llm: LLM | None = None, + execute_code: bool = False, + ): # type: ignore + """ + Create a simple agent that can execute code, think and respond to messages + """ + return cls( + name=name, + templates={"system": system_prompt} if system_prompt else {}, + llms={DEFAULT: llm} if llm else {}, + flow=([ExecuteCodeNode()] if execute_code else []) + [RespondNode()], + ) + + @classmethod + def create_team_manager( + cls, + name: str, + subagents: list[Agent[TeamTape]], + llm: LLM, + max_calls: int = 1, + ): + """ + Create a team manager that broadcasts the last message to all subagents, selects one of them to call, call it and + responds to the last message if the termination message is not received. + """ + return cls( + name=name, + subagents=subagents, + flow=[ + BroadcastLastMessageNode(), + SelectAndCallNode(), + RespondOrRepeatNode(), + ], + max_calls=max_calls, + templates={ + "select_before": SELECT_SPEAKER_MESSAGE_BEFORE_TEMPLATE, + "select_after": SELECT_SPEAKER_MESSAGE_AFTER_TEMPLATE, + }, + llms={DEFAULT: llm}, + ) + + @classmethod + def create_chat_initiator( + cls, + name: str, + teammate: Agent[TeamTape], + init_message: str, + system_prompt: str = "", + llm: LLM | None = None, + max_calls: int = 1, + execute_code: bool = False, + ): + """ + Create an agent that sets the team's initial message and calls the team manager + """ + return cls( + name=name, + templates={ + "system": system_prompt, + }, + llms={DEFAULT: llm} if llm else {}, + subagents=[teammate], + flow=([ExecuteCodeNode()] if execute_code else []) + [CallNode(), TerminateOrRepeatNode()], + max_calls=max_calls, + init_message=init_message, + ) + + +class BroadcastLastMessageNode(Node): + name: str = "broadcast_last_message" + + def generate_steps( + self, agent: TeamAgent, tape: TeamTape, llm_stream: LLMStream + ) -> Generator[AgentStep, None, None]: + view = ActiveTeamAgentView(agent, tape) + recipients = agent.get_subagent_names() + last = view.messages[-1] + from_ = last.by.split("/")[-1] + match last: + case Call(): + yield Broadcast(task=self.name, content=last.content, from_=from_, to=list(recipients)) + case Respond(): + recipients = [name for name in recipients if name != last.by.split("/")[-1]] + yield Broadcast( + task=self.name, + content=view.messages[-1].content, + from_=from_, + to=list(recipients), + ) + case Broadcast(task=self.name): + pass + case _: + assert False + + +class CallNode(Node): + name: str = "call" + + def make_prompt(self, agent: TeamAgent, tape: TeamTape) -> Prompt: + p = Prompt() + view = ActiveTeamAgentView(agent, tape) + if view.should_generate_message: + system = [{"role": "system", "content": agent.templates["system"]}] + p = Prompt(messages=system + _llm_messages_from_tape(agent, tape)) + return p + + def generate_steps( + self, agent: TeamAgent, tape: TeamTape, llm_stream: LLMStream + ) -> Generator[AgentStep, None, None]: + view = ActiveTeamAgentView(agent, tape) + # if last node + (other,) = agent.subagents + if view.should_generate_message: + yield Call(task=self.name, agent_name=other.name, content=llm_stream.get_text()) + elif view.exec_result: + yield Call(task=self.name, agent_name=other.name, content=_exec_result_message(agent, tape)) + else: + assert agent.init_message and not view.messages + yield Call(task=self.name, agent_name=other.name, content=agent.init_message) + + +class SelectAndCallNode(Node): + name: str = "select_and_call" + + def make_prompt(self, agent: TeamAgent, tape: TeamTape) -> Prompt: + subagents = ", ".join(agent.get_subagent_names()) + select_before = [ + { + "role": "system", + "content": agent.templates["select_before"].format(subagents=subagents), + } + ] + select_after = [ + { + "role": "system", + "content": agent.templates["select_after"].format(subagents=subagents), + } + ] + return Prompt(messages=select_before + _llm_messages_from_tape(agent, tape) + select_after) + + def generate_steps( + self, agent: TeamAgent, tape: TeamTape, llm_stream: LLMStream + ) -> Generator[AgentStep, None, None]: + callee_name = llm_stream.get_text() + # check if the callee is an existing subagent + _ = agent.find_subagent(callee_name) + yield Call(task=self.name, agent_name=callee_name) + + +class ExecuteCodeNode(Node): + name: str = "execute_code" + + def generate_steps(self, agent: logging.Any, tape: Tape, llm_stream: LLMStream) -> Generator[AgentStep, None, None]: + assert not llm_stream + view = ActiveTeamAgentView(agent, tape) + if view.last_non_empty_message is None: + yield Pass(task=self.name) + elif code := extract_code_blocks(view.last_non_empty_message.content): + yield ExecuteCode(task=self.name, code=code) + else: + yield Pass(task=self.name) + + +class RespondNode(Node): + name: str = "respond" + + def make_prompt(self, agent: TeamAgent, tape: TeamTape) -> Prompt: + p = Prompt() + view = ActiveTeamAgentView(agent, tape) + if view.should_generate_message: + system = [{"role": "system", "content": agent.templates["system"]}] + p = Prompt(messages=system + _llm_messages_from_tape(agent, tape)) + return p + + def generate_steps( + self, agent: TeamAgent, tape: TeamTape, llm_stream: LLMStream + ) -> Generator[AgentStep, None, None]: + view = ActiveTeamAgentView(agent, tape) + if view.should_generate_message: + yield Respond(task=self.name, content=llm_stream.get_text()) + elif view.exec_result: + yield Respond(task=self.name, content=_exec_result_message(agent, tape)) + else: + logger.info( + f"Agent {agent.full_name} had to respond with an empty message." + f" You might want to optimize your orchestration logic." + ) + yield Respond(task=self.name) + + +class TerminateOrRepeatNode(Node): + name: str = "terminate_or_repeat" + + def generate_steps( + self, agent: TeamAgent, tape: TeamTape, llm_stream: LLMStream + ) -> Generator[AgentStep, None, None]: + assert not llm_stream + view = ActiveTeamAgentView(agent, tape) + if view.should_stop: + yield FinalStep(task=self.name, reason="Termination message received") + else: + yield Jump(task=self.name, next_node=0) + + +class RespondOrRepeatNode(Node): + name: str = "respond_or_repeat" + + def generate_steps( + self, agent: TeamAgent, tape: TeamTape, llm_stream: LLMStream + ) -> Generator[AgentStep, None, None]: + view = ActiveTeamAgentView(agent, tape) + if view.should_stop: + yield Respond(task=self.name) + else: + yield Jump(task=self.name, next_node=0) + + +def _exec_result_message(agent: TeamAgent, tape: TeamTape) -> str: + view = ActiveTeamAgentView(agent, tape) + exec_result_message = "" + if view.exec_result: + if output := view.exec_result.result.output.strip(): + exec_result_message = f"I ran the code and got the following output:\n\n{output}" + else: + exec_result_message = f"I ran the code, the exit code was {view.exec_result.result.exit_code}." + return exec_result_message + + +def _llm_messages_from_tape(agent: TeamAgent, tape: TeamTape) -> list[dict[str, str]]: + view = ActiveTeamAgentView(agent, tape) + llm_messages = [] + for step in view.messages: + match step: + # When we make the LLM messages, we use "kind" == "user" for messages + # originating from other agents, and "kind" == "assistant" for messages by this agent. + case Call() if step.by == agent.full_name: + # I called someone + llm_messages.append({"role": "assistant", "content": step.content}) + case Call(): + # someone called me + # we exclude empty call messages from the prompt + if not step.content: + continue + llm_messages.append( + { + "role": "user", + "content": step.content, + "name": step.by.split("/")[-1], + } + ) + case Respond() if step.by == agent.full_name: + # I responded to someone + llm_messages.append({"role": "assistant", "content": step.content}) + case Respond(): + # someone responded to me + who_returned = step.by.split("/")[-1] + llm_messages.append({"role": "user", "content": step.content, "name": who_returned}) + case Broadcast(): + llm_messages.append({"role": "user", "content": step.content, "name": step.from_}) + return llm_messages diff --git a/tests/test_examples.py b/tests/test_examples.py index 36ec5b86..c073267b 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -17,7 +17,7 @@ from examples.gaia_agent.eval import load_results from examples.gaia_agent.tape import GaiaTape from examples.llama_agent import LLAMAChatBot -from tapeagents.collective import CollectiveTape +from tapeagents.team import TeamTape from tapeagents.config import DB_DEFAULT_FILENAME from tapeagents.core import AgentStep, TrainingText from tapeagents.dialog import Dialog @@ -165,7 +165,7 @@ def test_data_science(): run_dir = f"{res_path}/data_science" llm = ReplayLLM.from_llm(llama(), run_dir) agent, start_tape, env = data_science.make_world(llm, EmptyEnvironment()) - final_tape = CollectiveTape.model_validate(load_tape_dict(run_dir, "final_tape.json")) + final_tape = TeamTape.model_validate(load_tape_dict(run_dir, "final_tape.json")) assert replay_tape(agent, final_tape, start_tape=start_tape, env=env, reuse_observations=True)