DB-GPT/dbgpt/agent/agents/base_team.py

import logging
from typing import Dict, List, Optional, Union

from pydantic import BaseModel, Field

from dbgpt.agent.actions.action import Action, ActionOutput, T
from dbgpt.agent.agents.agent_new import Agent
from dbgpt.agent.agents.base_agent_new import ConversableAgent
from dbgpt.agent.resource.resource_api import AgentResource, ResourceType
from dbgpt.vis.tags.vis_agent_plans import Vis, VisAgentPlans

logger = logging.getLogger(__name__)


def content_str(content: Union[str, List, None]) -> str:
    """Converts `content` into a string format.

    This function processes content that may be a string, a list of mixed text and image URLs, or None,
    and converts it into a string. Text is directly appended to the result string, while image URLs are
    represented by a placeholder image token. If the content is None, an empty string is returned.

    Args:
        - content (Union[str, List, None]): The content to be processed. Can be a string, a list of dictionaries
                                      representing text and image URLs, or None.

    Returns:
        str: A string representation of the input content. Image URLs are replaced with an image token.

    Note:
    - The function expects each dictionary in the list to have a "type" key that is either "text" or "image_url".
      For "text" type, the "text" key's value is appended to the result. For "image_url", an image token is appended.
    - This function is useful for handling content that may include both text and image references, especially
      in contexts where images need to be represented as placeholders.
    """
    if content is None:
        return ""
    if isinstance(content, str):
        return content
    if not isinstance(content, list):
        raise TypeError(f"content must be None, str, or list, but got {type(content)}")

    rst = ""
    for item in content:
        if not isinstance(item, dict):
            raise TypeError(
                "Wrong content format: every element should be dict if the content is a list."
            )
        assert (
            "type" in item
        ), "Wrong content format. Missing 'type' key in content's dict."
        if item["type"] == "text":
            rst += item["text"]
        elif item["type"] == "image_url":
            rst += "<image>"
        else:
            raise ValueError(
                f"Wrong content format: unknown type {item['type']} within the content"
            )
    return rst


class Team(BaseModel):
    agents: List[Agent] = Field(default_factory=list)
    messages: List[Dict] = Field(default_factory=list)
    max_round: Optional[int] = 100
    is_team: bool = True

    class Config:
        arbitrary_types_allowed = True

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def hire(self, agents: List[Agent]):
        """Hire roles to cooperate"""
        self.agents.extend(agents)

    @property
    def agent_names(self) -> List[str]:
        """Return the names of the agents in the group chat."""
        return [agent.profile for agent in self.agents]

    def agent_by_name(self, name: str) -> Agent:
        """Returns the agent with a given name."""
        return self.agents[self.agent_names.index(name)]

    async def a_select_speaker(self, last_speaker: Agent, selector: Agent):
        pass

    def reset(self):
        """Reset the group chat."""
        self.messages.clear()

    def append(self, message: Dict):
        """Append a message to the group chat.
        We cast the content to str here so that it can be managed by text-based
        model.
        """
        message["content"] = content_str(message["content"])
        self.messages.append(message)


class ManagerAgent(ConversableAgent, Team):
    profile: str = "TeamManager"
    goal: str = "manage all hired intelligent agents to complete mission objectives"
    constraints: List[str] = []
    desc: str = goal
    is_team: bool = True

    # The management agent does not need to retry the exception. The actual execution of the agent has already been retried.
    max_retry_count: int = 1

    def __init__(self, **kwargs):
        ConversableAgent.__init__(self, **kwargs)
        Team.__init__(self, **kwargs)

    async def a_thinking(
        self, messages: Optional[List[Dict]], prompt: Optional[str] = None
    ) -> Union[str, Dict, None]:
        # TeamManager, which is based on processes and plans by default, only needs to ensure execution and does not require additional thinking.
        if messages is None or len(messages) <= 0:
            return None, None
        else:
            message = messages[-1]
            self.messages.append(message)
            return message["content"], None

    async def a_act(
        self,
        message: Optional[str],
        sender: Optional[ConversableAgent] = None,
        reviewer: Optional[ConversableAgent] = None,
    ) -> Optional[ActionOutput]:
        pass