aap-core 0.1.1.dev1__tar.gz

aap_core-0.1.1.dev1/PKG-INFO

@@ -0,0 +1,14 @@
+Metadata-Version: 2.3
+Name: aap-core
+Version: 0.1.1.dev1
+Summary: Agent design pattern core framework
+Author: Ly Hon Quang
+Author-email: Ly Hon Quang <lyhonquang@gmail.com>
+License: MIT
+Requires-Dist: a2a-sdk>=0.3.12
+Requires-Dist: pandas>=2.0.0
+Requires-Dist: rouge-score>=0.1.2
+Requires-Dist: sacrebleu>=2.5.1
+Requires-Python: >=3.10.12
+Description-Content-Type: text/markdown
+

aap_core-0.1.1.dev1/pyproject.toml

@@ -0,0 +1,25 @@
+[project]
+name = "aap_core"
+version = "0.1.1.dev1"
+description = "Agent design pattern core framework"
+readme = "README.md"
+license = { text = "MIT" }
+authors = [
+    { name = "Ly Hon Quang", email = "lyhonquang@gmail.com" }
+]
+requires-python = ">=3.10.12"
+dependencies = [
+    "a2a-sdk>=0.3.12",
+    "pandas>=2.0.0",
+    "rouge-score>=0.1.2",
+    "sacrebleu>=2.5.1",
+]
+
+[dependency-groups]
+lint = [
+    "ruff==0.14.0",
+]
+
+[build-system]
+requires = ["uv_build>=0.8.14,<0.9.0"]
+build-backend = "uv_build"

aap_core-0.1.1.dev1/src/aap_core/agent.py

@@ -0,0 +1,150 @@
+import abc
+from collections.abc import Callable, Sequence
+from typing import Any, Dict, Literal
+
+from a2a.types import AgentCard
+from pydantic import BaseModel, Field, PrivateAttr
+
+from .chain import BaseLLMChain
+from .types import AgentMessage
+
+
+class BaseAgent(abc.ABC, BaseModel):
+    """An interface for agent implementations."""
+
+    _state: str = PrivateAttr("idle")
+    _composed_state: str = PrivateAttr("idle")
+
+    card: AgentCard = Field(
+        ...,
+        description="""
+        A self-describing manifest for an agent.
+        It provides essential metadata including the agent's identity, capabilities, skills, supported communication methods, and security requirements.""",
+    )
+    state_change_callback: Callable[[str], None] | None = Field(
+        default=None,
+        description="The callback function to update the state of the agent.",
+    )
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._set_composed_state()
+        for _, value in self:
+            self.__value_deco(value)
+
+    @abc.abstractmethod
+    def execute(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        """The execution logic of the agent.
+        Depends on the detail implementation by subclass, agent can calls tools or RAG or MCP during execution.
+        And agent may or may not utilize LLM, depends on the implementation. A lot of agents are pure orchestration without LLM in it.
+        During execution, the implementation should update the state of the agent
+        and/or push notification to target to update the progress.
+        There are 3 types of agent:
+        - Remote agent: agent that is running on other server.
+        Example: For Google's A2A, we will act as a client agent and communicate with agent using A2A protocol and getback result
+        - Local agent: self-defined agent that is running on a this machine.
+        - Orchestration agent: itself is a agent, which can perform what normal agent can do, and orchestrate other agents to complete the provided task.
+
+        Args:
+            message (AgentMessage): The message to execute the agent with.
+            **kwargs: Additional keyword arguments to pass to the LLM generation method.
+
+        Returns:
+            AgentMessage: The response from the agent.
+        """
+        return message
+
+    async def aexecute(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        """The asynchronous version of method execute."""
+        return self.execute(message, **kwargs)
+
+    @property
+    def state(self) -> str:
+        """Get the current state of the agent."""
+        return self._state
+
+    @property
+    def composed_state(self) -> str:
+        """Get the current whole state of the agent hierarchy."""
+        return self._composed_state
+
+    def _set_composed_state(self) -> None:
+        # A subchild class that have child agent in their implementation may need to override this
+        self._composed_state = BaseAgent.build_composed_state(self, [], "sequential")
+
+    def _sync_state(self, state: str = "") -> None:
+        self._set_composed_state()
+        if self.state_change_callback:
+            self.state_change_callback(self._composed_state)
+
+    @state.setter
+    def state(self, value: str) -> None:
+        self._state = value
+        self._sync_state()
+
+    def __value_deco(self, value: Any):
+        if isinstance(value, str):
+            pass
+        elif isinstance(value, BaseAgent):
+            value.state_change_callback = self._sync_state
+        elif isinstance(
+            value, BaseLLMChain
+        ):  # force name of the chain is the same as the agent
+            value.name = self.card.name
+        elif isinstance(value, Sequence):
+            for agent in value:
+                if isinstance(agent, BaseAgent):
+                    agent.state_change_callback = self._sync_state
+                elif isinstance(value, BaseLLMChain):
+                    value.name = self.card.name
+        elif isinstance(value, Dict):
+            for agent in value.values():
+                if isinstance(agent, BaseAgent):
+                    agent.state_change_callback = self._sync_state
+                elif isinstance(value, BaseLLMChain):
+                    value.name = self.card.name
+
+    def __setattr__(self, name: str, value: Any):
+        self.__value_deco(value)
+        super().__setattr__(name, value)
+
+    @classmethod
+    def build_composed_state(
+        cls,
+        parent_agent: "BaseAgent",
+        child_agents: Sequence["BaseAgent"],
+        connect_type: Literal["sequential", "parallel"],
+    ) -> str:
+        """
+        Build a composed state string for the given agent and its children.
+        The composition rule is following:
+        - The order of the agent hierarchy is topdown.
+        - A state of a agent is with format: <agent_name>:<agent_state>
+        - The level separator is "/". For example: <parent_agent>:<parent_state>/<child_agent1>:<child_state1>
+        - If a child is made up of multiple agents, there are 2 type of connections which currently support are "sequential" and "parallel".
+            + "sequential": the children are connected with "-" separator and encapsulate with bracket, for example: <parent_agent>:<parent_state>/((<child_agent1>:<child_state1>)-(<child_agent2>:<child_state2>))
+            + "parallel": the children are connected with "|" separator and encapsulate with bracket, for example: <parent_agent>:<parent_state>/((<child_agent1>:<child_state1>)|(<child_agent2>:<child_state2>))
+        - Lower level of hierarchy will be constructed the same way to form a complete tree of agent state.
+
+        Args:
+            parent_agent: The parent agent.
+            child_agents: A list of child agents. This list is a direct child only. The function will use the composed_state of the direct child to form the tree. The composed state of the child is assumed to be already built by the child before.
+            connect_type: The type of connection between the parent and children, either "sequential" or "parallel".
+
+        Returns:
+            str: The composed state string.
+        """
+        self_state = f"{parent_agent.card.name}:{parent_agent.state}"
+        child_states = [child_agent.composed_state for child_agent in child_agents]
+        if len(child_states) == 0:
+            return self_state
+        elif len(child_states) == 1:
+            return f"{self_state}/{child_states[0]}"
+        else:
+            childs = [f"({state})" for state in child_states]
+            if connect_type == "sequential":
+                return f"{self_state}/({'-'.join(childs)})"
+            elif connect_type == "parallel":
+                return f"{self_state}/({'|'.join(childs)})"
+            else:
+                raise ValueError(f"Invalid connect_type: {connect_type}")

aap_core-0.1.1.dev1/src/aap_core/chain.py

@@ -0,0 +1,148 @@
+import abc
+from collections.abc import Sequence
+from typing import Callable, Generic, List, Tuple
+
+from pydantic import Field, PrivateAttr
+
+from .guardrail import BaseGuardRail, PassGuardRail
+from .prompt_enhancer import (
+    BasePromptEnhancer,
+    IdentityPromptEnhancer,
+)
+from .types import AgentMessage, BaseChain, ChainMessage, ChainResponse
+
+
+class BaseLLMChain(BaseChain):
+    """
+    Base class for LLM chains.
+    """
+
+    name: str = Field(
+        "chain",
+        description="The name of the chain. Should be same at agent who hold this chain for easy to operate.",
+    )
+
+    @abc.abstractmethod
+    def invoke(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        pass
+
+    async def ainvoke(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        return self.invoke(message, **kwargs)
+
+    def __call__(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        return self.invoke(message, **kwargs)
+
+
+class BaseCausalMultiTurnsChain(BaseLLMChain, Generic[ChainMessage, ChainResponse]):
+    """
+    Base class for multi-turns chains.
+    """
+
+    _last_response_as_context: str | None = PrivateAttr(default=None)
+    include_history: int = Field(
+        default=0, description="number of history message turns to include"
+    )
+    store_immediate_steps: bool = Field(
+        default=False, description="store intermediate steps output"
+    )
+
+    def final_response_as_context(self, key: str) -> None:
+        self._last_response_as_context = key.replace("context_", "")
+
+    @abc.abstractmethod
+    def _prepare_conversation(self, message: AgentMessage) -> List[ChainMessage]:
+        pass
+
+    @abc.abstractmethod
+    def _generate_response(
+        self, conversation: List[ChainMessage], **kwargs
+    ) -> Tuple[List[ChainMessage], ChainResponse, bool]:
+        pass
+
+    @abc.abstractmethod
+    def _process_tools(
+        self,
+        conversation: List[ChainMessage],
+        response: ChainResponse,
+    ) -> List[ChainMessage]:
+        pass
+
+    @abc.abstractmethod
+    def _append_responses(
+        self, message: AgentMessage, conversation: List[ChainMessage]
+    ) -> AgentMessage:
+        pass
+
+    def invoke(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        # Template method
+        conversation = self._prepare_conversation(message)
+        conversation, response, has_tool = self._generate_response(
+            conversation, **kwargs
+        )
+        if has_tool:
+            conversation = self._process_tools(conversation, response)
+            conversation, _, _ = self._generate_response(conversation, **kwargs)
+
+        message = self._append_responses(message, conversation)
+
+        if self._last_response_as_context is not None:
+            if message.context is None:
+                message.context = {}
+            message.context[self._last_response_as_context] = message.responses.pop()[1]
+
+        message.origin = self.name
+        message.execution_result = "success"
+        return message
+
+
+class TypicalLLMChain(BaseLLMChain):
+    """
+    A LLM chain that consists of all components in a typical response generation.
+    The chain  have 4 steps:
+    - input guardrail: validate the input message
+    - prompt enhancement: add more context to the prompt or rewrite / refine the prompt
+    - generate: generate the response. This steps should also handle the tool calling procedure and finalize the response
+    - output guardrail: validate the output message"""
+
+    input_guardrail: BaseGuardRail = Field(
+        default=PassGuardRail(),
+        description="validate the input message",
+    )
+    prompt_enhancer: BasePromptEnhancer = Field(
+        default=IdentityPromptEnhancer(),
+        description="add more context to the prompt or rewrite / refine the prompt",
+    )
+    tools: Sequence[Callable] = Field(
+        default=[],
+        description="tools to call for the LLM",
+    )
+    output_guardrail: BaseGuardRail = Field(
+        default=PassGuardRail(),
+        description="validate the output message",
+    )
+
+    @abc.abstractmethod
+    def generate(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        # call tool if any and finalize response
+        pass
+
+    def invoke(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        message = self.input_guardrail(message)
+        message = self.prompt_enhancer(message)
+        message = self.generate(message, **kwargs)
+        message = self.output_guardrail(message)
+        return message
+
+
+class LLMPromptEnhancer(BasePromptEnhancer):
+    """A prompt enhancer that use an LLM chain to rewrite the prompt."""
+
+    chain: BaseLLMChain = Field(..., description="LLM chain that rewrite the prompt")
+
+    def __call__(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        message = self.chain.invoke(message, **kwargs)
+        if message.execution_result != "success":
+            return message
+        _, message.query = message.responses[-1]
+        message.responses = []
+        return message

aap_core-0.1.1.dev1/src/aap_core/guardrail.py

@@ -0,0 +1,14 @@
+import abc
+
+from .types import AgentMessage, BaseChain
+
+
+class BaseGuardRail(BaseChain):
+    @abc.abstractmethod
+    def __call__(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        pass
+
+
+class PassGuardRail(BaseGuardRail):
+    def __call__(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        return message

aap_core-0.1.1.dev1/src/aap_core/orchestration.py

@@ -0,0 +1,597 @@
+import concurrent.futures
+import random
+import re
+from collections.abc import Generator, Iterator, Sequence
+from typing import Callable, List, Literal, Tuple
+
+from pydantic import Field, field_validator
+from rouge_score import rouge_scorer
+from sacrebleu import sentence_bleu
+
+from aap_core.types import AgentResponse
+
+from .agent import AgentMessage, BaseAgent
+from .chain import BaseLLMChain
+
+
+class ReflectionAgent(BaseAgent):
+    """
+    An agent that is capable of self-reflection.
+
+    This agent will first execute the given message and then reflect on the
+    result. If the reflection is successful, the agent will return a new
+    message with the reflection result. If the reflection is not successful,
+    the agent will return the original message.
+
+    Attributes:
+        chain: The BaseLLMChain to use for executing the message.
+        chain_reflection: The prompt to use for self-reflection.
+    """
+
+    chain_task: BaseLLMChain = Field(
+        ..., description="LLM chain that perform the main task"
+    )
+    chain_reflection: BaseLLMChain = Field(
+        ...,
+        description="LLM chain that perform reflection on the result of the main task",
+    )
+    task_response_key: str = Field(
+        default="context_response",
+        description="""The key of which to store the context (e.g. intermediate result) during generation.
+        Note that the key need to start with 'context_'""",
+    )
+
+    @field_validator("task_response_key")
+    @classmethod
+    def check_starts_with_prefix(cls, v: str) -> str:
+        if not v.startswith("context_"):
+            raise ValueError("task_response_key must start with 'context_'")
+        return v
+
+    def execute(
+        self, message: AgentMessage, keep_original_response: bool = True, **kwargs
+    ) -> AgentMessage:
+        self.state = "running"
+        message = self.chain_task.invoke(message, **kwargs)
+        message.execution_result = message.execution_result or "error"
+        if message.execution_result != "success":
+            message.error_message = (
+                "Call chain not success. Received: " + message.execution_result
+            )
+            return message
+
+        self.state = "reflecting"
+        task_response_key = self.task_response_key.replace("context_", "")
+        reflect_message = message.model_copy(deep=True)
+        reflect_message.context = {}
+        _, reflect_message.context[task_response_key] = reflect_message.responses[-1]  # type: ignore
+        reflect_message.responses = []
+        result_message = self.chain_reflection.invoke(reflect_message, **kwargs)
+        self.state = "idle"
+        result_message.origin = self.card.name
+        if result_message.execution_result == "success" and keep_original_response:
+            result_message.responses.insert(-1, message.responses[-1])
+
+        return result_message
+
+
+class LoopAgent(BaseAgent):
+    """
+    An agent that will loop until a certain condition is met.
+
+    This agent takes two parameters: an agent to loop and a condition to stop the loop.
+    The condition to stop the loop should be a function that takes an AgentMessage as an argument
+    and returns a boolean indicating whether the loop should stop or not.
+
+    Example:
+        def is_stop(message: AgentMessage) -> bool:
+            # stop when the message query is "stop"
+            return message.query == "stop"
+
+        agent = LoopAgent(
+            agent=Agent1(),
+            is_stop=is_stop
+        )
+
+        message = AgentMessage(query="hello")
+        result = agent.execute(message)
+
+    Attributes:
+        agent: The agent to loop.
+        is_stop: The condition to stop the loop.
+    """
+
+    agent: BaseAgent = Field(..., description="The agent to loop.")
+    is_stop: Callable[[AgentMessage], bool] | Generator[bool] = Field(
+        ...,
+        description="The condition to stop the loop. It should be a function that takes an AgentMessage as an argument and returns a boolean.",
+    )
+
+    def execute(
+        self,
+        message: AgentMessage,
+        keep_result: int | Callable[[Sequence[AgentResponse]], List[AgentResponse]] = 1,
+        **kwargs,
+    ) -> AgentMessage:
+        self.state = "running"
+        origin_len = len(message.responses)
+        responses = message.model_copy(deep=True)
+
+        def update_responses():
+            if isinstance(keep_result, int) and keep_result > 0:
+                responses.responses = responses.responses[-keep_result - origin_len :]
+            elif isinstance(keep_result, Callable):
+                responses.responses = keep_result(responses.responses)
+
+        i = 0
+        if isinstance(self.is_stop, Callable):
+            while self.is_stop(message) is False:
+                self.state = "running#" + str(i)
+                i += 1
+                message = self.agent.execute(message, **kwargs)
+                if message.execution_result != "success":
+                    break
+                else:
+                    responses.responses.append(message.responses.pop())
+                update_responses()
+        elif isinstance(self.is_stop, Iterator):
+            try:
+                while next(self.is_stop) is False:
+                    self.state = "running#" + str(i)
+                    i += 1
+                    message = self.agent.execute(message, **kwargs)
+                    if message.execution_result != "success":
+                        break
+                    else:
+                        responses.responses.append(message.responses.pop())
+                    update_responses()
+            except StopIteration:
+                pass
+        else:
+            raise TypeError(
+                f"is_stop must be a callable or a Generator, received {type(self.is_stop)}"
+            )
+
+        self.state = "idle"
+        responses.origin = self.card.name
+        return responses
+
+    def _set_composed_state(self) -> None:
+        self._composed_state = BaseAgent.build_composed_state(
+            self, [self.agent], "sequential"
+        )
+
+
+class SequentialAgent(BaseAgent):
+    """
+    A sequential agent is an agent that will execute a list of agents in a particular order.
+
+    The order of the agents is the order of the execution. This means that the first agent will be executed first, then the second agent and so on.
+
+    Example:
+        agent1 = Agent1()
+        agent2 = Agent2()
+
+        sequential_agent = SequentialAgent([agent1, agent2])
+
+        message = AgentMessage(query="hello")
+        result = sequential_agent.execute(message)
+
+    Attributes:
+        agents: The list of agents to execute in sequence.
+    """
+
+    agents: Sequence[BaseAgent] = Field(
+        ...,
+        description="""The list of agents to execute in sequence.
+            The order of the sequence is the order of the execution.""",
+        min_length=1,
+    )
+
+    def execute(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        self.state = "running"
+        for agent in self.agents:
+            message = agent.execute(message, **kwargs)
+            if message.execution_result != "success":
+                break
+
+        self.state = "idle"
+        message.origin = self.card.name
+        return message
+
+    def _set_composed_state(self) -> None:
+        self._composed_state = BaseAgent.build_composed_state(
+            self, self.agents, "sequential"
+        )
+
+
+class ParallelAgent(BaseAgent):
+    """
+    A parallel agent is an agent that will execute a list of agents in parallel.
+    The agent will return a list of AgentMessage where each message is the result of the corresponding agent in the list.
+
+    Example:
+        agent1 = Agent1()
+        agent2 = Agent2()
+
+        parallel_agent = ParallelAgent([agent1, agent2])
+
+        messages = [AgentMessage(query="hello"), AgentMessage(query="world")]
+        results = parallel_agent.execute(messages)
+
+    Attributes:
+        agents: The list of agents to execute in parallel.
+    """
+
+    agents: Sequence[BaseAgent] = Field(
+        ..., description="""The list of agents to execute in parallel""", min_length=1
+    )
+
+    def execute(  # type: ignore
+        self, messages: AgentMessage | Sequence[AgentMessage], **kwargs
+    ) -> AgentMessage:
+        # If received a single message, all agents will process the same message
+        if isinstance(messages, AgentMessage):
+            messages = [messages.model_copy(deep=True)] * len(self.agents)
+        elif len(messages) != len(self.agents):
+            raise ValueError(
+                "messages must be a list of AgentMessage with the same length as the number of agents"
+            )
+
+        self.state = "running"
+        # if torch is used, we can use torch.multiprocessing
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            futures = [
+                executor.submit(agent.execute, message, **kwargs)
+                for agent, message in zip(self.agents, messages)
+            ]
+            responses = [
+                future.result() for future in concurrent.futures.as_completed(futures)
+            ]
+
+        result_message = AgentMessage(
+            query="",
+            origin=self.card.name,
+            responses=[message.responses[-1] for message in responses],
+            execution_result="success",
+        )  # type: ignore
+        self.state = "idle"
+        return result_message
+
+    def _set_composed_state(self) -> None:
+        self._composed_state = BaseAgent.build_composed_state(
+            self, self.agents, "parallel"
+        )
+
+
+class CoordinatorAgent(BaseAgent):
+    """
+    This coordinator agent compose of 3 stages:
+    - Planning: a planner agent will plan the steps to be executed.
+    The planner recieves input query and returns a list of steps to be executed.
+    For each step, it contains 3 fields:
+        + worker: the worker agent that will execute the step
+        + message: the instruction for the worker agent
+        + dependencies: the dependencies of the step, this is picked from the result of the previous steps. The planner is responsible for choosing which results are needed for current step
+    - Execution: For each step, an assigned worker will execute the step. In current version, steps are performed sequentially.
+    - Summary: a summary agent will summarize the results so far and generate final answer.
+    The summary agent recieves all results and write a final answer.
+    This stage is optional.
+
+    Example:
+        planner_agent = PlannerAgent()
+        worker_agent = WorkerAgent()
+        summary_agent = SummaryAgent()
+
+        coordinator_agent = CoordinatorAgent(planner_agent, worker_agent, summary_agent)
+
+        message = AgentMessage(query="hello")
+        result = coordinator_agent.execute(message)
+
+    Attributes:
+        planner_agent: The agent that will plan the steps to be executed.
+        workers: The list of agents that will execute the steps.
+        summary_agent: The agent that will summarize the results so far and generate final answer.
+    """
+
+    planner_agent: BaseAgent = Field(
+        ..., description="""The agent that will plan the steps to be executed"""
+    )
+    parse_plan: Callable[
+        [AgentMessage, Sequence[BaseAgent]],
+        Sequence[Tuple[AgentMessage, BaseAgent, List]],
+    ] = Field(
+        ...,
+        description="""The function that will parse the plan and return the list of steps to be executed.""",
+    )
+    workers: Sequence[BaseAgent] = Field(
+        ...,
+        description="""The list of agents that will execute the steps based on the plan""",
+        min_length=1,
+    )
+    summary_chain: BaseLLMChain | None = Field(
+        default=None,
+        description="""The summary chain that will summarize the results so far and generate final answer""",
+    )
+    summary_prompt: str | None = Field(
+        default=None,
+        description="""The prompt that will be used by the summary chain""",
+    )
+    summary_steps_key: str = Field(
+        "context_results",
+        description="""The key of which to store the context (e.g. intermediate result) during generation.
+        Note that the key need to start with 'context_'""",
+    )
+
+    @field_validator("summary_steps_key")
+    @classmethod
+    def check_starts_with_prefix(cls, v: str) -> str:
+        if not v.startswith("context_"):
+            raise ValueError("summary_steps_key must start with 'context_'")
+        return v
+
+    def execute(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        self.state = "planning"
+        # TODO: abstract structured output
+        message = self.planner_agent.execute(message, **kwargs)
+
+        if message.execution_result != "success":
+            return message
+
+        steps = self.parse_plan(message, self.workers)
+        sub_task_result = []
+        # TODO: check step dependency and parallelize steps to speed up
+        summary_steps_key = self.summary_steps_key.replace("context_", "")
+        for i, (step_msg, step_agent, dependencies) in enumerate(steps):
+            step_msg.context = {summary_steps_key: sub_task_result}
+            self.state = f"step {i}: worker {step_agent.card.name} running"
+            sub_task_result.append(step_agent.execute(step_msg, **kwargs))
+
+        if self.summary_chain is None:
+            if self.summary_prompt is not None:
+                # Use coordinator as summary agent
+                self.state = f"{self.planner_agent.card.name} finalizing"
+                summary_message = AgentMessage(
+                    query=self.summary_prompt,
+                    context={summary_steps_key: sub_task_result},
+                    responses=[],
+                )  # type: ignore
+                summary_message = self.planner_agent.execute(summary_message, **kwargs)
+            else:
+                return message
+            # if no summary prompt is provided, just return the result from workers
+        else:
+            self.state = "finalizing"
+            summary_message = AgentMessage(
+                query=message.query,
+                context={summary_steps_key: sub_task_result},
+                responses=[],
+            )  # type: ignore
+            summary_message = self.summary_chain.invoke(summary_message, **kwargs)
+
+        self.state = "idle"
+        summary_message.origin = self.card.name
+        return summary_message
+
+    def _set_composed_state(self) -> None:
+        self._composed_state = BaseAgent.build_composed_state(
+            self, [self.planner_agent, *self.workers], "sequential"
+        )
+
+
+class DebateAgent(BaseAgent):
+    """
+    Multiple agents participate in a discussion to solve a problem.
+    The DebateAgent is designed to facilitate a discussion between multiple agents.
+    It acts like a conversation manager who will decide the conversation flow
+
+    Note that for the expandability, this agent is reponsible only for expanding the discussing topic, not the summary one.
+    For the complete flow that performs topic rollout and finalize the solution, it is necessary to combine with other agent.
+    There are several ways to achieve this: either by define a new agent that coodinate the actions or chain multiple agents together.
+    """
+
+    agents: Sequence[BaseAgent] = Field(
+        ...,
+        description="""The list of agents participate in the conversation""",
+        min_length=1,
+    )
+    pick_strategy: (
+        Literal["round_robin", "random", "simultaneous"]
+        | Callable[[Sequence[BaseAgent]], BaseAgent]
+    ) = Field(..., description="The strategy to pick the agent to run the next turn")
+    random_seed: int | None = Field(
+        default=None, description="The random seed for the 'random' pick strategy"
+    )
+    max_turns: int = Field(
+        5,
+        description="""The maximum number of debate turns to run.
+        Note: for the simultaneous strategy, the actual number of turn will be factor of the number of agents""",
+        ge=1,
+    )
+    should_stop: Callable[[AgentMessage], bool] | None = Field(
+        default=None, description="The condition to stop the debate"
+    )
+
+    def execute(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        n = 0
+        if self.pick_strategy == "random" and self.random_seed is not None:
+            random.seed(self.random_seed)
+
+        # All response turns are stored in the message.responses
+        while n < self.max_turns:
+            if self.pick_strategy == "round_robin":
+                next_turn_agent = self.agents[n % len(self.agents)]
+            elif self.pick_strategy == "random":
+                next_turn_agent = random.choice(self.agents)
+            elif isinstance(self.pick_strategy, Callable):
+                next_turn_agent = self.pick_strategy(self.agents)
+            elif self.pick_strategy == "simultaneous":
+                self.state = f"turn {n}: all agents running"
+                # TODO: parallelize the execution
+                current_turn_messages = [
+                    agent.execute(message.model_copy(deep=True), **kwargs)
+                    for agent in self.agents
+                ]
+                message.responses.extend(
+                    [msg.responses[-1] for msg in current_turn_messages]
+                )
+                n += len(self.agents)
+                message.execution_result = "success"
+            else:
+                raise ValueError(
+                    "pick_strategy must be one of 'round_robin', 'random', 'simutaneous' or a function"
+                )
+
+            if self.pick_strategy != "simultaneous":
+                self.state = f"turn {n}: agent {next_turn_agent.card.name} running"  # type: ignore
+                message = next_turn_agent.execute(message, **kwargs)  # type: ignore
+                if len(message.responses) > 0:
+                    pass
+                n += 1
+
+            if message.execution_result != "success" or (
+                self.should_stop and self.should_stop(message)
+            ):
+                break
+
+        self.state = "idle"
+        message.origin = self.card.name
+        message.execution_result = "success"
+        return message
+
+    def _set_composed_state(self) -> None:
+        self._composed_state = BaseAgent.build_composed_state(
+            self,
+            self.agents,
+            "parallel" if self.pick_strategy == "simultaneous" else "sequential",
+        )
+
+
+class VotingAgent(BaseAgent):
+    """
+    A voting agent selects the best answer among a list of candidates responses.
+    """
+
+    agents: Sequence[BaseAgent] = Field(
+        ...,
+        description="""The list of agents participate in the voting""",
+        min_length=1,
+    )
+    voting_method: Literal["majority_vote", "llm_score"] = Field(
+        ...,
+        description="""The voting method to use.
+        1. majority_vote: It uses bleu or rouge score to calculate score for candidates
+        2. llm_score: use LLM to score the candidate messages.
+            For each message, it will score messages generated by other agents.
+            Then score for each message will be aggregated (average) to get the final score.
+
+        Note:
+        - The majority_vote method can also be known as the metric-based scorer.
+            i.e. Use a metric to score the candidate with reference message. This approach is simple and fast.
+        - The llm_score method can also be known as the model-based scorer.
+            i.e. Use a model to score the candidate message.
+            This approach is more complex and slower. But it can capture more scenarios like semantic similarity, context similarity, etc.
+        """,
+    )
+    scorer: str | Callable[[str], float] = Field(
+        ...,
+        description="""1. If voting_method is 'llm_score', the function to use to convert LLM response to score.
+        Note that LLM may generate reponse that could not convert to score.
+        In these cases, there are various ways to handle.
+        For example, we just ignore it and assign the score to 0
+
+        2. If voting_method is 'majority_vote'. scorer accept bleu or rouge score
+            - 'bleu' or 'agent_forest': Use bleu score to calculate voting score. https://arxiv.org/pdf/2402.05120.
+            - 'rougeL': use rouge-L (rouge longest common sequence)
+            - 'rougeN': N can be an integer, e.g. rouge1, rouge2,...""",
+    )
+    voting_prompt: str | None = Field(
+        default=None,
+        description="The prompt to use for the voting. This parameter is used when voting_method is 'llm_score'",
+    )
+
+    def execute(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        self.state = "running"
+        if self.voting_method == "majority_vote":
+            # sanity check
+            if not isinstance(self.scorer, str):
+                raise TypeError(
+                    "When voting_method is 'majority_vote'. scorer need to be one of bleu, agent_forest, rougeL or rouge[N]"
+                )
+            elif self.scorer == "agent_forest" or self.scorer == "bleu":
+                # Use the same library and code snippet as the author
+                # Source: https://github.com/MoreAgentsIsAllYouNeed/AgentForest/blob/master/src/utils.py#L220
+                def bleu_score(hyp: str, ref: str) -> float:
+                    return sentence_bleu(hyp, [ref], lowercase=True).score
+
+                score_func = bleu_score
+            elif re.match(r"^rouge(L|\d+)$", self.scorer):
+                scorer = rouge_scorer.RougeScorer([self.scorer], use_stemmer=True)
+
+                def rouge_score(hyp: str, ref: str) -> float:
+                    return scorer.score(ref, hyp)[self.scorer]
+
+                score_func = rouge_score
+            else:
+                raise ValueError("Scorer not supported")
+
+            scores = {agent.card.name: 0.0 for agent in self.agents}
+            message_map = {name: msg for name, msg in message.responses}
+            for agent_name, msg in message.responses:
+                total_score = 0
+                for other_agent in self.agents:
+                    if agent_name == other_agent.card.name:
+                        continue
+                    total_score += score_func(message_map[other_agent.card.name], msg)
+                scores[agent_name] = total_score
+            highest_score_agent = max(scores, key=scores.get)  # type: ignore
+            message.responses.append((self.card.name, message_map[highest_score_agent]))
+
+        elif self.voting_method == "llm_score":
+            if self.voting_prompt is None:
+                raise ValueError(
+                    "voting_prompt is required for llm_score voting method"
+                )
+            elif not isinstance(self.scorer, Callable):
+                raise TypeError(
+                    "scorer must be a Callable when voting_method is 'llm_score'"
+                )
+
+            scores = {agent.card.name: 0.0 for agent in self.agents}
+            message_map = {name: msg for name, msg in message.responses}
+            for response in message.responses:
+                total_score = 0
+                for agent in self.agents:
+                    if response[0] == agent.card.name:
+                        continue
+
+                    try:
+                        voting_message = AgentMessage(
+                            query=self.voting_prompt, responses=[]
+                        )  # type: ignore
+                        voting_message = agent.execute(voting_message, **kwargs)
+                        total_score += self.scorer(voting_message.responses[-1][1])
+                    except Exception:
+                        # unable to parse the score
+                        pass
+                scores[response[0]] = total_score
+                # The highest or average score will return the same response
+                # Additional note: because of the score, we even can rank all candidate responses and output a leaderboard.
+            highest_score_agent = max(scores, key=scores.get)  # type: ignore
+            message.responses.append((self.card.name, message_map[highest_score_agent]))
+
+        else:
+            raise ValueError(
+                "voting_method must be one of 'llm_score', 'majority_vote'"
+            )
+
+        self.state = "idle"
+        message.origin = self.card.name
+        message.execution_result = "success"
+        return message
+
+    def _set_composed_state(self) -> None:
+        # TODO: find way to parallel voting of all agents
+        self._composed_state = BaseAgent.build_composed_state(
+            self, self.agents, "sequential"
+        )

aap_core-0.1.1.dev1/src/aap_core/prompt_enhancer.py

@@ -0,0 +1,209 @@
+import abc
+import json
+import os
+from collections.abc import Sequence
+from typing import List, Tuple
+
+import pandas as pd
+import tabulate
+from pydantic import Field, PrivateAttr, field_validator
+
+# import toon_format
+from .types import AgentMessage, BaseChain, ContentType
+
+
+class BasePromptEnhancer(BaseChain):
+    """A base class to enhance / rewrite the prompt.
+
+    There are two types of prompt enhancement:
+    - Data enhancement: Give more context to the prompt by adding external data, either by using files (CSV, JSON, Markdown, etc.), database (SQL,...) or more advanced techniques like RAG, web search.
+    - Structure enhancement: Rewrite / refine the prompt partially or entirely.
+    """
+
+    @abc.abstractmethod
+    def __call__(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        return super().__call__(message, **kwargs)
+
+
+class IdentityPromptEnhancer(BasePromptEnhancer):
+    """A prompt enhancer that does nothing.
+    This serves as a default prompt enhancer."""
+
+    def __call__(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        return message
+
+
+class DataFramePromptEnhancer(BasePromptEnhancer):
+    """A prompt enhancer that uses a pandas dataframe to give more context to the prompt."""
+
+    format: str = Field(
+        ...,
+        description="""
+        The format of the prompt.
+        The format must contain at least {prompt} and {data}.
+        Other parameters can be used and parsed""",
+    )
+    _data: str | None = PrivateAttr(None)
+
+    @field_validator("format")
+    @classmethod
+    def check_starts_with_prompt_and_data(cls, v: str) -> str:
+        if "{prompt}" not in v or "{data}" not in v:
+            raise ValueError("The format must contain at least {prompt} and {data}")
+        return v
+
+    @classmethod
+    def from_pandas(
+        cls,
+        format: str,
+        data: pd.DataFrame,
+        prettier: str | tabulate.TableFormat | None = None,
+    ):
+        """
+        Parse a pandas dataframe into a string format.
+
+        Args:
+            data (pd.DataFrame): The pandas dataframe to be parsed.
+            prettier (str | tabulate.TableFormat | None, optional): The prettier to be used. Defaults to None.
+            The available table formats are from the [tabulate](https://pypi.org/project/tabulate/) library.
+            [toon](https://github.com/toon-format/spec) are also supported
+
+        Returns:
+            self: The instance of the class.
+
+        Notes:
+            - If prettier is 'toon', the dataframe is converted to a string following toon specification
+            - If prettier is None, the dataframe is converted to a string using the to_string() method.
+            - If prettier is not None, the dataframe is converted to a string using the tabulate module with the specified tablefmt.
+        """
+        enhancer = DataFramePromptEnhancer(format=format)
+        if prettier is None:
+            enhancer._data = data.to_string()
+        # elif prettier == "toon":      # toon_format currently is beta release, waiting 1.0.0 release
+        #     enhancer._data = toon_format.encode(data.to_dict())
+        else:
+            enhancer._data = tabulate.tabulate(
+                data.to_dict(), headers="keys", tablefmt=prettier
+            )
+
+        return enhancer
+
+    @classmethod
+    def from_string(cls, format: str, data: str):
+        """
+        Parse a string into a format that can be used by the prompt enhancer.
+
+        Args:
+            data (str): The string to be parsed.
+
+        Returns:
+            self: The instance of the class.
+        """
+        enhancer = DataFramePromptEnhancer(format=format)
+        enhancer._data = data
+        return enhancer
+
+    @classmethod
+    def from_dict(
+        cls, format: str, data: dict, prettier: str | tabulate.TableFormat | None = None
+    ):
+        """
+        Parse a dictionary into a format that can be used by the prompt enhancer.
+
+        Args:
+            data (dict): The dictionary to be parsed.
+            prettier (str | tabulate.TableFormat | None, optional): The prettier to be used. Defaults to None.
+            The available table formats are from the [tabulate](https://pypi.org/project/tabulate/) library.
+            [toon](https://github.com/toon-format/spec) are also supported
+
+        Returns:
+            self: The instance of the class.
+
+        Notes:
+            - If prettier is 'toon', the dataframe is converted to a string following toon specification
+            - If prettier is None, the dictionary is converted to a string using the str() method.
+            - If prettier is not None, the dictionary is converted to a string using the tabulate module with the specified tablefmt.
+        """
+        enhancer = DataFramePromptEnhancer(format=format)
+        if prettier is None:
+            enhancer._data = str(data)
+        # elif prettier == "toon":      # toon_format currently is beta release, waiting 1.0.0 release
+        #     enhancer._data = toon_format.encode(data.to_dict())
+        else:
+            enhancer._data = tabulate.tabulate(data, headers="keys", tablefmt=prettier)
+
+        return enhancer
+
+    @classmethod
+    def from_list(cls, format: str, data: List[str], bullet_char: str = "-"):
+        """
+        Parse a list of strings into a format that can be used by the prompt enhancer.
+
+        Args:
+            data (List[str]): The list of strings to be parsed.
+            bullet_char (str, optional): The character to be used as the bullet. Defaults to "-".
+
+        Returns:
+            self: The instance of the class.
+
+        Notes:
+            - The list of strings is converted to a string using the join() method with the bullet character as the separator.
+        """
+        enhancer = DataFramePromptEnhancer(format=format)
+        enhancer._data = "\n".join([f"{bullet_char} {d}" for d in data])
+        return enhancer
+
+    @classmethod
+    def from_jsonl(cls, format: str, path: str | os.PathLike[str]):
+        """
+        Parse a JSONL file into a format that can be used by the prompt enhancer.
+
+        Args:
+            path (str): The path to the JSONL file.
+
+        Returns:
+            self: The instance of the class.
+
+        Raises:
+            FileNotFoundError: If the file specified by path does not exist.
+
+        Notes:
+            - The JSONL file is parsed line by line using the json.loads() method.
+            - The parsed data is converted to a pandas DataFrame and passed to the parse_pandas() method.
+        """
+        if not os.path.exists(path):
+            raise FileNotFoundError(f"File {path} does not exist")
+
+        # Implementation provided by https://stackoverflow.com/a/74406107
+        with open(path, "r") as f:
+            lines = f.read().splitlines()
+
+        line_dicts = [json.loads(line) for line in lines]
+        return DataFramePromptEnhancer.from_pandas(format, pd.DataFrame(line_dicts))
+
+    def __call__(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        if self._data is None:
+            raise ValueError("Data not parsed yet")
+        message.query = self.format.format(
+            prompt=message.query, data=self._data, **kwargs
+        )
+        return message
+
+
+class BaseRAGPromptEnhancer(BasePromptEnhancer):
+    @abc.abstractmethod
+    def _search(self, query: str, **kwargs) -> Sequence[Tuple[ContentType, str]]:
+        """"""
+        pass
+
+    @abc.abstractmethod
+    def _format(
+        self, message: AgentMessage, data: Sequence[Tuple[ContentType, str]]
+    ) -> AgentMessage:
+        # For string type, the field we need to touch is `query`. For other types, the media `query_media` and `query_media_type` needed to modify
+        pass
+
+    def __call__(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        data = self._search(message.query, **kwargs)
+        prompt = self._format(message, data)
+        return prompt

aap_core-0.1.1.dev1/src/aap_core/types.py

@@ -0,0 +1,106 @@
+import abc
+import json
+from typing import Any, Dict, List, Literal, Tuple, TypeVar
+
+from pydantic import BaseModel, Field
+
+
+ChainMessage = TypeVar("ChainMessage")
+ChainResponse = TypeVar("ChainResponse")
+AgentResponse = Tuple[str, str]
+ContentType = Literal["text", "image", "audio", "video", "document"]
+
+
+class AgentMessage(BaseModel):
+    query: str = Field(
+        ...,
+        description="""The user query consumed by agent's LLM.
+        In this message class, we don't store the system prompt and user template.
+        Because each agent have their own system prompt and user template, which is not shared between agents.""",
+    )
+    query_media: List[Tuple[ContentType, str]] | None = Field(
+        default=None, description="The media content associated with the query."
+    )
+    origin: str | None = Field(
+        default=None, description="The agent that send this message"
+    )
+    responses: List[AgentResponse] = Field(
+        default=[],
+        description="""
+        If an agent generate multiple responses, either by same or different subagents, all of them will be stored here.
+        In each response tuple, the first one should be agent name or index, and second is the response""",
+    )
+    context: dict | None = Field(
+        default=None,
+        description="""
+        Agent additional material, which probably is the output of other agent or user entered.
+        There are various types of context produced by user and other agents.
+        The prompt that comsume this context need to explicitly know the format of this context.""",
+    )
+    execution_result: Literal["success", "error"] | None = Field(
+        default=None,
+        description="The execution result of the agent. Can be success or error",
+    )
+    error_message: str | None = Field(
+        default=None, description="The error message if the execution result is error."
+    )
+    media: List[Tuple[ContentType, str]] | None = Field(
+        default=None,
+        description="The additional media content. Can be image, video, audio, etc.",
+    )
+
+    def flatten_dict(self, d: dict, parent_key: str = "", sep="_") -> Dict[str, Any]:
+        """
+        Flattens a nested dictionary into a single-level dictionary.
+
+        Args:
+            d (dict): The dictionary to flatten.
+            parent_key (str): The prefix for keys in the flattened dictionary.
+            sep (str): The separator used to join parent and child keys.
+
+        Returns:
+            dict: The flattened dictionary.
+        """
+        items = []
+        for k, v in d.items():
+            new_key = parent_key + sep + k if parent_key else k
+            if isinstance(v, dict):
+                items.extend(self.flatten_dict(v, new_key, sep=sep).items())
+            else:
+                items.append((new_key, v))
+        return dict(items)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Converts the AgentMessage object into a dictionary format.
+
+        Returns a dictionary that contains all the information in the AgentMessage object.
+        If the context is not None, it will be flattened and added to the returned dictionary.
+        All the data in context will have prefix of 'context_' by default.
+        For example, if the context is {'a': 1, 'b': 2}, the returned dictionary will contains {'context_a': 1, 'context_b': 2}
+
+        Returns:
+            dict: A dictionary that contains all the information in the AgentMessage object.
+        """
+        msg_json = self.model_dump(exclude_none=True, exclude={"context"})
+        if self.context:
+            context_dict = self.flatten_dict(self.context, parent_key="context")
+            for k, v in context_dict.items():
+                msg_json[k] = v
+        return msg_json
+
+    def dump_json(self) -> str:
+        """
+        Converts the AgentMessage object into a JSON string.
+
+        Returns:
+            str: A JSON string that contains all the information in the AgentMessage object.
+        """
+        msg_dict = self.to_dict()
+        return json.dumps(msg_dict)
+
+
+class BaseChain(abc.ABC, BaseModel):
+    @abc.abstractmethod
+    def __call__(self, message: AgentMessage, **kwargs) -> AgentMessage:
+        pass

aap_core-0.1.1.dev1/src/aap_core/utils.py

@@ -0,0 +1,16 @@
+import re
+
+
+def remove_thinking(response: str) -> str:
+    """A snippet to remove thinking
+    Note on the model Qwen 3: the openning <think> tag is missing. So we need to manually add it to beginning of the response
+
+    Args:
+        response (str): raw response
+
+    Returns:
+        str: response without thinking
+    """
+    if "</think>" in response and "<think>" not in response:
+        response = "<think>" + response
+    return re.sub(r"<think>.*?</think>", "", response, flags=re.DOTALL)