aap-core 0.1.3.dev1__tar.gz → 0.2.0__tar.gz

{aap_core-0.1.3.dev1 → aap_core-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: aap-core
-Version: 0.1.3.dev1
+Version: 0.2.0
 Summary: Agent design pattern core framework
 Keywords: agent,ai,pattern,llm
 Author: Ly Hon Quang
@@ -15,9 +15,13 @@ Classifier: Programming Language :: Python :: 3.15
 Classifier: Framework :: Pydantic :: 2
 Classifier: Operating System :: OS Independent
 Requires-Dist: a2a-sdk>=0.3.12
+Requires-Dist: gymnasium>=1.2.3
+Requires-Dist: nltk>=3.9.2
 Requires-Dist: pandas>=2.0.0
 Requires-Dist: rouge-score>=0.1.2
 Requires-Dist: sacrebleu>=2.5.1
+Requires-Dist: scipy>=1.15.3
+Requires-Dist: wandb>=0.27.0
 Requires-Python: >=3.10.12
 Project-URL: Repository, https://github.com/quanghona/agent_design_pattern.git
 Description-Content-Type: text/markdown

{aap_core-0.1.3.dev1 → aap_core-0.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "aap_core"
-version = "0.1.3.dev1"
+version = "0.2.0"
 description = "Agent design pattern core framework"
 readme = "README.md"
 license = { text = "MIT" }
@@ -11,9 +11,13 @@ authors = [
 requires-python = ">=3.10.12"
 dependencies = [
     "a2a-sdk>=0.3.12",
+    "gymnasium>=1.2.3",
+    "nltk>=3.9.2",
     "pandas>=2.0.0",
     "rouge-score>=0.1.2",
     "sacrebleu>=2.5.1",
+    "scipy>=1.15.3",
+    "wandb>=0.27.0",
 ]
 classifiers = [
     "Programming Language :: Python :: 3.10",
@@ -30,6 +34,10 @@ classifiers = [
 Repository = "https://github.com/quanghona/agent_design_pattern.git"
 
 [dependency-groups]
+dev = [
+    "pytest>=9.0.3",
+    "ruff>=0.14.0",
+]
 lint = [
     "ruff==0.14.0",
 ]

aap_core-0.2.0/src/aap_core/__init__.py

@@ -0,0 +1,23 @@
+from aap_core.policy_trainer import (
+    BaseExplorationModule,
+    BasePolicyTrainer,
+    DPOTrainer,
+    EpsilonGreedyExploration,
+    GRPOTrainer,
+    PPOTrainer,
+    RandomExploration,
+    ReinforcePPTrainer,
+    SimpleReplayBuffer,
+)
+
+__all__ = [
+    "BaseExplorationModule",
+    "BasePolicyTrainer",
+    "DPOTrainer",
+    "EpsilonGreedyExploration",
+    "GRPOTrainer",
+    "PPOTrainer",
+    "RandomExploration",
+    "ReinforcePPTrainer",
+    "SimpleReplayBuffer",
+]

{aap_core-0.1.3.dev1 → aap_core-0.2.0}/src/aap_core/agent.py

@@ -5,8 +5,7 @@ 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
+from .types import AgentMessage, BaseLLMChain
 
 
 class BaseAgent(abc.ABC, BaseModel):

{aap_core-0.1.3.dev1 → aap_core-0.2.0}/src/aap_core/chain.py

@@ -9,28 +9,7 @@ from .prompt_augmenter import (
     BasePromptAugmenter,
     IdentityPromptAugmenter,
 )
-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)
+from .types import AgentMessage, BaseLLMChain, ChainMessage, ChainResponse, TokenUsage
 
 
 class BaseCausalMultiTurnsChain(BaseLLMChain, Generic[ChainMessage, ChainResponse]):
@@ -46,6 +25,10 @@ class BaseCausalMultiTurnsChain(BaseLLMChain, Generic[ChainMessage, ChainRespons
         default=False, description="store intermediate steps output"
     )
 
+    max_turns: int = Field(
+        default=50, description="maximum number of turns before stopping"
+    )
+
     def final_response_as_context(self, key: str) -> None:
         self._last_response_as_context = key.replace("context_", "")
 
@@ -56,7 +39,7 @@ class BaseCausalMultiTurnsChain(BaseLLMChain, Generic[ChainMessage, ChainRespons
     @abc.abstractmethod
     def _generate_response(
         self, conversation: List[ChainMessage], **kwargs
-    ) -> Tuple[List[ChainMessage], ChainResponse, bool]:
+    ) -> Tuple[List[ChainMessage], ChainResponse, bool, TokenUsage]:
         pass
 
     @abc.abstractmethod
@@ -74,14 +57,38 @@ class BaseCausalMultiTurnsChain(BaseLLMChain, Generic[ChainMessage, ChainRespons
         pass
 
     def invoke(self, message: AgentMessage, **kwargs) -> AgentMessage:
-        # Template method
+        if message.token_usage is None:
+            message.token_usage = {
+                "steps": [],
+                "total": TokenUsage(input_tokens=0, output_tokens=0, total_tokens=0),
+            }
+        if "steps" not in message.token_usage:
+            message.token_usage["steps"] = []
+        if "total" not in message.token_usage:
+            message.token_usage["total"] = TokenUsage(
+                input_tokens=0, output_tokens=0, total_tokens=0
+            )
+
+        def add_token(usage: TokenUsage):
+            message.token_usage["steps"].append(usage)  # type: ignore
+            message.token_usage["total"]["input_tokens"] += usage["input_tokens"]  # type: ignore
+            message.token_usage["total"]["output_tokens"] += usage["output_tokens"]  # type: ignore
+            message.token_usage["total"]["total_tokens"] += usage["total_tokens"]  # type: ignore
+
+        # Template method — ReAct loop: Thought → Act → Thought → ... → Final
         conversation = self._prepare_conversation(message)
-        conversation, response, has_tool = self._generate_response(
-            conversation, **kwargs
-        )
-        if has_tool:
+        turn = 0
+        while turn < self.max_turns:
+            conversation, response, has_tool, usage = self._generate_response(
+                conversation, **kwargs
+            )
+            add_token(usage)
+
+            if not has_tool:
+                break
+
             conversation = self._process_tools(conversation, response)
-            conversation, _, _ = self._generate_response(conversation, **kwargs)
+            turn += 1
 
         message = self._append_responses(message, conversation)
 
@@ -132,13 +139,3 @@ class TypicalLLMChain(BaseLLMChain):
         message = self.generate(message, **kwargs)
         message = self.output_guardrail(message)
         return message
-
-
-class MetaPromptAugmenter(BasePromptAugmenter):
-    """A prompt augmenter that use an LLM chain to rewrite the prompt."""
-
-    chain: BaseLLMChain = Field(..., description="LLM chain that rewrite the prompt")
-
-    def augment(self, message: AgentMessage, **kwargs) -> AgentMessage:
-        message = self.chain.invoke(message, **kwargs)
-        return message

{aap_core-0.1.3.dev1 → aap_core-0.2.0}/src/aap_core/orchestration.py

@@ -10,8 +10,8 @@ from sacrebleu import sentence_bleu
 
 from aap_core.types import AgentResponse
 
-from .agent import AgentMessage, BaseAgent
-from .chain import BaseLLMChain
+from .types import AgentMessage, BaseLLMChain
+from .agent import BaseAgent
 
 
 class ReflectionAgent(BaseAgent):
@@ -35,6 +35,10 @@ class ReflectionAgent(BaseAgent):
         ...,
         description="LLM chain that perform reflection on the result of the main task",
     )
+    task_response_key: str = Field(
+        default="context_response",
+        description="The key to use for the task response in the message context.",
+    )
 
     @field_validator("task_response_key")
     @classmethod
@@ -245,7 +249,7 @@ class ParallelAgent(BaseAgent):
             origin=self.card.name,
             responses=[message.responses[-1] for message in responses],
             execution_result="success",
-        )  # type: ignore
+        )
         self.state = "idle"
         return result_message
 
@@ -309,7 +313,7 @@ class CoordinatorAgent(BaseAgent):
         description="""The prompt that will be used by the summary chain""",
     )
     summary_steps_key: str = Field(
-        "context_results",
+        default="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_'""",
     )
@@ -330,11 +334,57 @@ class CoordinatorAgent(BaseAgent):
             return message
 
         steps = self.parse_plan(message, self.workers)
+
+        # Validate parsed plan
+        if not steps:
+            self.state = "idle"
+            message.execution_result = "error"
+            message.error_message = (
+                "Plan parsing failed: no steps returned from parse_plan"
+            )
+            return message
+
+        if not isinstance(steps, Sequence):
+            self.state = "idle"
+            message.execution_result = "error"
+            message.error_message = (
+                f"Plan parsing failed: expected Sequence, got {type(steps)}"
+            )
+            return message
+
+        for i, step in enumerate(steps):
+            if not isinstance(step, tuple) or len(step) != 3:
+                self.state = "idle"
+                message.execution_result = "error"
+                message.error_message = f"Plan parsing failed: step {i} is not a valid tuple of (message, agent, dependencies)"
+                return message
+            step_msg, step_agent, dependencies = step
+            if not isinstance(step_msg, AgentMessage):
+                self.state = "idle"
+                message.execution_result = "error"
+                message.error_message = (
+                    f"Plan parsing failed: step {i} message is not an AgentMessage"
+                )
+                return message
+            if not isinstance(step_agent, BaseAgent):
+                self.state = "idle"
+                message.execution_result = "error"
+                message.error_message = (
+                    f"Plan parsing failed: step {i} agent is not a BaseAgent"
+                )
+                return message
+            if not isinstance(dependencies, list):
+                self.state = "idle"
+                message.execution_result = "error"
+                message.error_message = (
+                    f"Plan parsing failed: step {i} dependencies is not a list"
+                )
+                return message
+
         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):
-            # TODO: deal with the dependencies
             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))
@@ -347,7 +397,7 @@ class CoordinatorAgent(BaseAgent):
                     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
@@ -358,7 +408,7 @@ class CoordinatorAgent(BaseAgent):
                 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"
@@ -521,7 +571,7 @@ class VotingAgent(BaseAgent):
                 scorer = rouge_scorer.RougeScorer([self.scorer], use_stemmer=True)
 
                 def rouge_score(hyp: str, ref: str) -> float:
-                    return scorer.score(ref, hyp)[self.scorer]
+                    return list(scorer.score(ref, hyp).values())[0].fmeasure
 
                 score_func = rouge_score
             else:
@@ -530,7 +580,7 @@ class VotingAgent(BaseAgent):
             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
+                total_score = 0.0
                 for other_agent in self.agents:
                     if agent_name == other_agent.card.name:
                         continue
@@ -560,7 +610,7 @@ class VotingAgent(BaseAgent):
                     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: