multi_agent.py

import logging
from dataclasses import dataclass, field
from typing import Annotated, Optional

from dotenv import load_dotenv
from livekit.agents import (
    Agent,
    AgentSession,
    ChatContext,
    JobContext,
    RoomInputOptions,
    RoomOutputOptions,
    RunContext,
    WorkerOptions,
    cli,
)
from livekit.agents.llm import function_tool
from livekit.plugins import openai
from pydantic import Field

logger = logging.getLogger("Tal")


load_dotenv()


@dataclass
class UserData:
    agents: dict[str, Agent] = field(default_factory=dict)
    summary_of_conversation: str = ""


class RememberAgent(Agent):
    def __init__(self) -> None:
        llm = openai.realtime.RealtimeModel()
        super().__init__(
            instructions="Your goal is remember facts the user tells you. Once the user requests handoff - handoff to the factoid agent.",
            llm=llm,
        )

    async def on_enter(self):
        # This doesn't work:
        # if self.session.userdata.summary_of_conversation:
        #     content = f"this is a summary of the conversation with the factoid agent  : {self.session.userdata.summary_of_conversation}. Now continue voice chatting."
        #     self.session.history.add_message(role="system", content=content)

        self.session.generate_reply()

    @function_tool
    async def handoff_to_factoid_agent(
        self,
        summary_of_conversation: Annotated[str, Field(description="A terse summary of the conversation")],
        context: RunContext[UserData],
    ):
        """Called when the user requests a handoff to the factoid agent.

        Args:
            summary_of_conversation: A terse summary of the conversation
        """
        context.userdata.summary_of_conversation = summary_of_conversation

        logger.info(f"switching to the factoid agent {summary_of_conversation=}")
        return context.userdata.agents["factoid"]

    @function_tool
    async def get_summary_of_previous_agent(
        self,
        context: RunContext[UserData],
    ) -> Annotated[str, Field(description="A terse summary of the conversation")]:
        """
        When you want to know the summary of the conversation with the previous agent, call this function.
        """
        summary_of_conversation = context.userdata.summary_of_conversation
        logger.info(f"get_summary_of_previous_agent(): {summary_of_conversation=}")
        return summary_of_conversation


class FactoidAgent(Agent):
    def __init__(self, *, chat_ctx: Optional[ChatContext] = None) -> None:
        llm = openai.realtime.RealtimeModel()
        super().__init__(
            instructions="Your job is to say a random fact AI."
            "When the user requests a handoff, return to the remember agent.",
            llm=llm,
            chat_ctx=chat_ctx,
        )

    async def on_enter(self):
        # This doesn't work:
        # if self.session.userdata.summary_of_conversation:
        #     content = (
        #         f"this is a summary of the conversation with the factoid agent  : {self.session.userdata.summary_of_conversation}."
        #         " Now continue audio chatting (no text messages!)."
        #     )
        #     self.session.history.add_message(role="user", content=content)
        self.session.generate_reply()

    @function_tool
    async def handoff_to_remember_agent(
        self,
        summary_of_conversation: Annotated[str, Field(description="A terse summary of the conversation")],
        context: RunContext[UserData],
    ):
        """
        When you are fininshed telling the fact (and the user acknowledges it),
        summarize the conversation (specifying which facts you told the user) and call this function handoff to the factoid agent
        """
        context.userdata.summary_of_conversation = summary_of_conversation
        logger.info(f"switching to the remember agent {summary_of_conversation=}")
        return context.userdata.agents["remember"]

    @function_tool
    async def get_summary_of_previous_agent(
        self,
        context: RunContext[UserData],
    ) -> Annotated[str, Field(description="A terse summary of the conversation")]:
        """
        When you want to know the summary of the conversation with the previous agent, call this function.
        """
        summary_of_conversation = context.userdata.summary_of_conversation
        logger.info(f"get_summary_of_previous_agent(): {summary_of_conversation=}")
        return summary_of_conversation


async def entrypoint(ctx: JobContext):
    await ctx.connect()

    logging.getLogger("livekit.plugins.openai.realtime").setLevel(logging.DEBUG)
    logging.getLogger("livekit.plugins.openai").setLevel(logging.DEBUG)
    logging.getLogger("livekit.agents.llm.realtime").setLevel(logging.DEBUG)
    import os

    os.environ.setdefault("LOG_OAI_EVENTS", "1")

    userdata = UserData()
    userdata.agents["remember"] = RememberAgent()
    userdata.agents["factoid"] = FactoidAgent()

    session = AgentSession[UserData](
        llm=openai.realtime.RealtimeModel(),
        userdata=userdata,
    )

    await session.start(
        agent=userdata.agents["remember"],
        room=ctx.room,
        room_input_options=RoomInputOptions(),
        room_output_options=RoomOutputOptions(transcription_enabled=True),
    )


if __name__ == "__main__":
    cli.run_app(WorkerOptions(entrypoint_fnc=entrypoint))