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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (28) hide show
  1. aap_core-0.3.0/PKG-INFO +200 -0
  2. aap_core-0.3.0/README.md +169 -0
  3. {aap_core-0.1.3.dev1 → aap_core-0.3.0}/pyproject.toml +14 -2
  4. aap_core-0.3.0/src/aap_core/__init__.py +23 -0
  5. {aap_core-0.1.3.dev1 → aap_core-0.3.0}/src/aap_core/agent.py +4 -2
  6. {aap_core-0.1.3.dev1 → aap_core-0.3.0}/src/aap_core/chain.py +36 -39
  7. aap_core-0.3.0/src/aap_core/dedup.py +816 -0
  8. aap_core-0.3.0/src/aap_core/default_prompts/see_default_crossover.md +13 -0
  9. aap_core-0.3.0/src/aap_core/default_prompts/see_default_eda.md +7 -0
  10. aap_core-0.3.0/src/aap_core/default_prompts/see_default_examiner.md +10 -0
  11. aap_core-0.3.0/src/aap_core/default_prompts/see_default_improver.md +7 -0
  12. aap_core-0.3.0/src/aap_core/default_prompts/see_default_lamarckian.md +6 -0
  13. {aap_core-0.1.3.dev1 → aap_core-0.3.0}/src/aap_core/orchestration.py +60 -10
  14. aap_core-0.3.0/src/aap_core/policy.py +1018 -0
  15. aap_core-0.3.0/src/aap_core/policy_trainer.py +3114 -0
  16. aap_core-0.3.0/src/aap_core/prompt_augmenter.py +1640 -0
  17. {aap_core-0.1.3.dev1 → aap_core-0.3.0}/src/aap_core/retriever.py +12 -29
  18. {aap_core-0.1.3.dev1 → aap_core-0.3.0}/src/aap_core/thought.py +2 -2
  19. aap_core-0.3.0/src/aap_core/types.py +330 -0
  20. aap_core-0.3.0/src/aap_core/utils.py +839 -0
  21. aap_core-0.1.3.dev1/PKG-INFO +0 -25
  22. aap_core-0.1.3.dev1/README.md +0 -1
  23. aap_core-0.1.3.dev1/src/aap_core/__init__.py +0 -0
  24. aap_core-0.1.3.dev1/src/aap_core/prompt_augmenter.py +0 -935
  25. aap_core-0.1.3.dev1/src/aap_core/types.py +0 -106
  26. aap_core-0.1.3.dev1/src/aap_core/utils.py +0 -16
  27. {aap_core-0.1.3.dev1 → aap_core-0.3.0}/src/aap_core/guardrail.py +0 -0
  28. {aap_core-0.1.3.dev1 → aap_core-0.3.0}/src/aap_core/py.typed +0 -0
@@ -0,0 +1,200 @@
1
+ Metadata-Version: 2.3
2
+ Name: aap-core
3
+ Version: 0.3.0
4
+ Summary: Agent design pattern core framework
5
+ Keywords: agent,ai,pattern,llm
6
+ Author: Ly Hon Quang
7
+ Author-email: Ly Hon Quang <lyhonquang@gmail.com>
8
+ License: MIT
9
+ Classifier: Programming Language :: Python :: 3.10
10
+ Classifier: Programming Language :: Python :: 3.11
11
+ Classifier: Programming Language :: Python :: 3.12
12
+ Classifier: Programming Language :: Python :: 3.13
13
+ Classifier: Programming Language :: Python :: 3.14
14
+ Classifier: Programming Language :: Python :: 3.15
15
+ Classifier: Framework :: Pydantic :: 2
16
+ Classifier: Operating System :: OS Independent
17
+ Requires-Dist: a2a-sdk>=0.3.12
18
+ Requires-Dist: gymnasium>=1.2.3
19
+ Requires-Dist: nltk>=3.9.2
20
+ Requires-Dist: pandas>=2.0.0
21
+ Requires-Dist: rouge-score>=0.1.2
22
+ Requires-Dist: sacrebleu>=2.5.1
23
+ Requires-Dist: scipy>=1.15.3
24
+ Requires-Dist: wandb>=0.27.0
25
+ Requires-Dist: torch
26
+ Requires-Dist: rensa>=0.4.0
27
+ Requires-Dist: rbloom>=1.5.4
28
+ Requires-Python: >=3.10.12
29
+ Project-URL: Repository, https://github.com/quanghona/aap_core.git
30
+ Description-Content-Type: text/markdown
31
+
32
+ # aap_core
33
+
34
+ Agent design pattern core framework for building, orchestrating, and training AI agents.
35
+
36
+ ![PyPI - License](https://img.shields.io/pypi/l/aap_core)
37
+ ![PyPI - Downloads](https://img.shields.io/pypi/dm/aap_core)
38
+ ![PyPI - Version](https://img.shields.io/pypi/v/aap_core?label=%20)
39
+ ![Python Version](https://img.shields.io/pypi/pyversions/aap_core)
40
+
41
+ ## What is this?
42
+
43
+ `aap_core` is the foundational framework within the **AI Agent Pattern** project. It provides a unified abstraction layer for building AI agents, orchestrating multi-agent workflows, and training policies via reinforcement learning — all built on Pydantic and compatible with the [A2A protocol](https://github.com/google/a2a-sdk).
44
+
45
+ The framework separates orchestration logic from LLM manipulation, delegating model interaction to integration packages (`aap_langchain`, `aap_llamaindex`, `aap_dspy`, `aap_transformers`). This lets you swap backends without changing agent architecture.
46
+
47
+ ## Quick Start
48
+
49
+ ```bash
50
+ pip install aap_core
51
+ # or
52
+ uv add aap_core
53
+ ```
54
+
55
+ ```python
56
+ from a2a.types import AgentCard
57
+ from aap_core.agent import BaseAgent
58
+ from aap_core.types import AgentMessage
59
+
60
+ # Define a minimal agent
61
+ class MyAgent(BaseAgent):
62
+ card: AgentCard = AgentCard(
63
+ name="my-agent",
64
+ description="A simple agent",
65
+ skills=[],
66
+ )
67
+
68
+ def execute(self, message: AgentMessage, **kwargs) -> AgentMessage:
69
+ # Your agent logic here
70
+ message.responses = [("my-agent", "Hello!")]
71
+ message.execution_result = "success"
72
+ return message
73
+ ```
74
+
75
+ ## Features
76
+
77
+ - **Agent Abstraction**: Clean `BaseAgent` interface with A2A `AgentCard` integration for self-describing agent metadata
78
+ - **Orchestration Patterns**: Built-in agents for reflection, looping, voting, coordination, and debate patterns
79
+ - **Chain System**: Generic multi-turn causal chains with history management, tool processing, and token tracking
80
+ - **Prompt Augmentation**: Data and structure-based prompt enhancement with configurable looping strategies
81
+ - **Retrieval System**: Pluggable retrievers with post-processing (e.g., rerankers, summarizers)
82
+ - **RL Policy Training**: Full RL loop with replay buffers, exploration modules, and trainers (PPO, DPO, GRPO, Reinforce++), powered by Gymnasium
83
+ - **Thought Chains**: Chain-of-Thought, Tree-of-Thought, Graph-of-Thought, and Self-Consistency patterns
84
+ - **Guardrails**: Extensible guardrail system for input/output filtering
85
+ - **Async Support**: All agent execution methods have async counterparts (`aexecute`)
86
+ - **Token Usage Tracking**: Automatic token accounting across multi-agent chains
87
+
88
+ ## Installation
89
+
90
+ ### Core
91
+
92
+ ```bash
93
+ pip install aap_core
94
+ ```
95
+
96
+ ### Integration Packages
97
+
98
+ ```bash
99
+ pip install aap_langchain # LangChain integration
100
+ pip install aap_llamaindex # LlamaIndex integration
101
+ pip install aap_dspy # DSPy integration
102
+ pip install aap_transformers # Hugging Face Transformers integration
103
+ ```
104
+
105
+ ## Usage
106
+
107
+ ### Building an Agent
108
+
109
+ ```python
110
+ from aap_core.agent import BaseAgent
111
+ from aap_core.types import AgentMessage
112
+ from a2a.types import AgentCard
113
+
114
+ class TaskAgent(BaseAgent):
115
+ def execute(self, message: AgentMessage, **kwargs) -> AgentMessage:
116
+ self.state = "running"
117
+ # Process the message...
118
+ message.responses = [(self.card.name, "Task completed")]
119
+ message.execution_result = "success"
120
+ self.state = "idle"
121
+ return message
122
+ ```
123
+
124
+ ### Using Orchestration Patterns
125
+
126
+ ```python
127
+ from aap_core.orchestration import ReflectionAgent, LoopAgent
128
+
129
+ # Self-reflection: execute then reflect on the result
130
+ reflector = ReflectionAgent(
131
+ chain_task=my_chain,
132
+ chain_reflection=reflection_chain,
133
+ )
134
+ result = reflector.execute(message)
135
+
136
+ # Loop until a condition is met
137
+ stopper = lambda msg: "done" in msg.query.lower()
138
+ loop = LoopAgent(agent=my_agent, stop_condition=stopper)
139
+ result = loop.execute(message)
140
+ ```
141
+
142
+ ### Prompt Augmentation
143
+
144
+ ```python
145
+ from aap_core.prompt_augmenter import SimplePromptAugmenter
146
+
147
+ # Augment prompts with external context data
148
+ augmenter = SimplePromptAugmenter(
149
+ format="Context:\n{context}\n\nQuery: {query}",
150
+ loop=3, # Apply 3 times
151
+ )
152
+ augmented = augmenter(message)
153
+ ```
154
+
155
+ ### RL-Based Prompt Optimization
156
+
157
+ ```python
158
+ from aap_core.policy_trainer import PPOTrainer, SimpleReplayBuffer
159
+ from aap_core.prompt_augmenter import PromptOptimizationEnv
160
+
161
+ # Create the environment with multiple augmenters
162
+ env = PromptOptimizationEnv(augmenters=[augmenter_a, augmenter_b])
163
+
164
+ # Create a replay buffer and trainer
165
+ buffer = SimpleReplayBuffer(capacity=10000)
166
+ trainer = PPOTrainer(policy=my_policy, buffer=buffer)
167
+ trainer.train(env, num_steps=10000)
168
+ ```
169
+
170
+ ## Architecture
171
+
172
+ `aap_core` is organized into the following modules:
173
+
174
+ | Module | Description |
175
+ |--------|-------------|
176
+ | `agent` | Base agent interface with A2A `AgentCard` support and state management |
177
+ | `orchestration` | Multi-agent patterns: reflection, looping, voting, coordination, debate |
178
+ | `chain` | Generic multi-turn causal chains with history, tool processing, and token tracking |
179
+ | `policy` | RL policy models (GPT-2 based) with RMSNorm, RoPE, and GQA |
180
+ | `policy_trainer` | Replay buffers, exploration modules, and trainers (PPO, DPO, GRPO, Reinforce++) |
181
+ | `prompt_augmenter` | Data and structure-based prompt enhancement with Gymnasium environments |
182
+ | `retriever` | Pluggable retrieval system with post-processing support |
183
+ | `thought` | Thought chain patterns: CoT, ToT, GoT, Self-Consistency |
184
+ | `guardrail` | Input/output filtering system |
185
+ | `types` | Core type definitions: `AgentMessage`, `TokenUsage`, `BaseLLMChain`, `BaseChain` |
186
+ | `utils` | Utility functions (e.g., thinking tag removal) |
187
+
188
+ ## Documentation
189
+
190
+ - Full project documentation: [aap_core docs](https://github.com/quanghona/aap_core)
191
+ - Example notebooks: [example](https://github.com/quanghona/aap_example.git)
192
+ - A2A Protocol: [google/a2a-sdk](https://github.com/google/a2a-sdk)
193
+
194
+ ## Contributing
195
+
196
+ See the [main project README](https://github.com/quanghona/aap_core) for contribution guidelines.
197
+
198
+ ## License
199
+
200
+ This project is licensed under the MIT License — see the [LICENSE](https://github.com/quanghona/aap_core/blob/master/LICENSE) file for details.
@@ -0,0 +1,169 @@
1
+ # aap_core
2
+
3
+ Agent design pattern core framework for building, orchestrating, and training AI agents.
4
+
5
+ ![PyPI - License](https://img.shields.io/pypi/l/aap_core)
6
+ ![PyPI - Downloads](https://img.shields.io/pypi/dm/aap_core)
7
+ ![PyPI - Version](https://img.shields.io/pypi/v/aap_core?label=%20)
8
+ ![Python Version](https://img.shields.io/pypi/pyversions/aap_core)
9
+
10
+ ## What is this?
11
+
12
+ `aap_core` is the foundational framework within the **AI Agent Pattern** project. It provides a unified abstraction layer for building AI agents, orchestrating multi-agent workflows, and training policies via reinforcement learning — all built on Pydantic and compatible with the [A2A protocol](https://github.com/google/a2a-sdk).
13
+
14
+ The framework separates orchestration logic from LLM manipulation, delegating model interaction to integration packages (`aap_langchain`, `aap_llamaindex`, `aap_dspy`, `aap_transformers`). This lets you swap backends without changing agent architecture.
15
+
16
+ ## Quick Start
17
+
18
+ ```bash
19
+ pip install aap_core
20
+ # or
21
+ uv add aap_core
22
+ ```
23
+
24
+ ```python
25
+ from a2a.types import AgentCard
26
+ from aap_core.agent import BaseAgent
27
+ from aap_core.types import AgentMessage
28
+
29
+ # Define a minimal agent
30
+ class MyAgent(BaseAgent):
31
+ card: AgentCard = AgentCard(
32
+ name="my-agent",
33
+ description="A simple agent",
34
+ skills=[],
35
+ )
36
+
37
+ def execute(self, message: AgentMessage, **kwargs) -> AgentMessage:
38
+ # Your agent logic here
39
+ message.responses = [("my-agent", "Hello!")]
40
+ message.execution_result = "success"
41
+ return message
42
+ ```
43
+
44
+ ## Features
45
+
46
+ - **Agent Abstraction**: Clean `BaseAgent` interface with A2A `AgentCard` integration for self-describing agent metadata
47
+ - **Orchestration Patterns**: Built-in agents for reflection, looping, voting, coordination, and debate patterns
48
+ - **Chain System**: Generic multi-turn causal chains with history management, tool processing, and token tracking
49
+ - **Prompt Augmentation**: Data and structure-based prompt enhancement with configurable looping strategies
50
+ - **Retrieval System**: Pluggable retrievers with post-processing (e.g., rerankers, summarizers)
51
+ - **RL Policy Training**: Full RL loop with replay buffers, exploration modules, and trainers (PPO, DPO, GRPO, Reinforce++), powered by Gymnasium
52
+ - **Thought Chains**: Chain-of-Thought, Tree-of-Thought, Graph-of-Thought, and Self-Consistency patterns
53
+ - **Guardrails**: Extensible guardrail system for input/output filtering
54
+ - **Async Support**: All agent execution methods have async counterparts (`aexecute`)
55
+ - **Token Usage Tracking**: Automatic token accounting across multi-agent chains
56
+
57
+ ## Installation
58
+
59
+ ### Core
60
+
61
+ ```bash
62
+ pip install aap_core
63
+ ```
64
+
65
+ ### Integration Packages
66
+
67
+ ```bash
68
+ pip install aap_langchain # LangChain integration
69
+ pip install aap_llamaindex # LlamaIndex integration
70
+ pip install aap_dspy # DSPy integration
71
+ pip install aap_transformers # Hugging Face Transformers integration
72
+ ```
73
+
74
+ ## Usage
75
+
76
+ ### Building an Agent
77
+
78
+ ```python
79
+ from aap_core.agent import BaseAgent
80
+ from aap_core.types import AgentMessage
81
+ from a2a.types import AgentCard
82
+
83
+ class TaskAgent(BaseAgent):
84
+ def execute(self, message: AgentMessage, **kwargs) -> AgentMessage:
85
+ self.state = "running"
86
+ # Process the message...
87
+ message.responses = [(self.card.name, "Task completed")]
88
+ message.execution_result = "success"
89
+ self.state = "idle"
90
+ return message
91
+ ```
92
+
93
+ ### Using Orchestration Patterns
94
+
95
+ ```python
96
+ from aap_core.orchestration import ReflectionAgent, LoopAgent
97
+
98
+ # Self-reflection: execute then reflect on the result
99
+ reflector = ReflectionAgent(
100
+ chain_task=my_chain,
101
+ chain_reflection=reflection_chain,
102
+ )
103
+ result = reflector.execute(message)
104
+
105
+ # Loop until a condition is met
106
+ stopper = lambda msg: "done" in msg.query.lower()
107
+ loop = LoopAgent(agent=my_agent, stop_condition=stopper)
108
+ result = loop.execute(message)
109
+ ```
110
+
111
+ ### Prompt Augmentation
112
+
113
+ ```python
114
+ from aap_core.prompt_augmenter import SimplePromptAugmenter
115
+
116
+ # Augment prompts with external context data
117
+ augmenter = SimplePromptAugmenter(
118
+ format="Context:\n{context}\n\nQuery: {query}",
119
+ loop=3, # Apply 3 times
120
+ )
121
+ augmented = augmenter(message)
122
+ ```
123
+
124
+ ### RL-Based Prompt Optimization
125
+
126
+ ```python
127
+ from aap_core.policy_trainer import PPOTrainer, SimpleReplayBuffer
128
+ from aap_core.prompt_augmenter import PromptOptimizationEnv
129
+
130
+ # Create the environment with multiple augmenters
131
+ env = PromptOptimizationEnv(augmenters=[augmenter_a, augmenter_b])
132
+
133
+ # Create a replay buffer and trainer
134
+ buffer = SimpleReplayBuffer(capacity=10000)
135
+ trainer = PPOTrainer(policy=my_policy, buffer=buffer)
136
+ trainer.train(env, num_steps=10000)
137
+ ```
138
+
139
+ ## Architecture
140
+
141
+ `aap_core` is organized into the following modules:
142
+
143
+ | Module | Description |
144
+ |--------|-------------|
145
+ | `agent` | Base agent interface with A2A `AgentCard` support and state management |
146
+ | `orchestration` | Multi-agent patterns: reflection, looping, voting, coordination, debate |
147
+ | `chain` | Generic multi-turn causal chains with history, tool processing, and token tracking |
148
+ | `policy` | RL policy models (GPT-2 based) with RMSNorm, RoPE, and GQA |
149
+ | `policy_trainer` | Replay buffers, exploration modules, and trainers (PPO, DPO, GRPO, Reinforce++) |
150
+ | `prompt_augmenter` | Data and structure-based prompt enhancement with Gymnasium environments |
151
+ | `retriever` | Pluggable retrieval system with post-processing support |
152
+ | `thought` | Thought chain patterns: CoT, ToT, GoT, Self-Consistency |
153
+ | `guardrail` | Input/output filtering system |
154
+ | `types` | Core type definitions: `AgentMessage`, `TokenUsage`, `BaseLLMChain`, `BaseChain` |
155
+ | `utils` | Utility functions (e.g., thinking tag removal) |
156
+
157
+ ## Documentation
158
+
159
+ - Full project documentation: [aap_core docs](https://github.com/quanghona/aap_core)
160
+ - Example notebooks: [example](https://github.com/quanghona/aap_example.git)
161
+ - A2A Protocol: [google/a2a-sdk](https://github.com/google/a2a-sdk)
162
+
163
+ ## Contributing
164
+
165
+ See the [main project README](https://github.com/quanghona/aap_core) for contribution guidelines.
166
+
167
+ ## License
168
+
169
+ This project is licensed under the MIT License — see the [LICENSE](https://github.com/quanghona/aap_core/blob/master/LICENSE) file for details.
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "aap_core"
3
- version = "0.1.3.dev1"
3
+ version = "0.3.0"
4
4
  description = "Agent design pattern core framework"
5
5
  readme = "README.md"
6
6
  license = { text = "MIT" }
@@ -11,9 +11,16 @@ authors = [
11
11
  requires-python = ">=3.10.12"
12
12
  dependencies = [
13
13
  "a2a-sdk>=0.3.12",
14
+ "gymnasium>=1.2.3",
15
+ "nltk>=3.9.2",
14
16
  "pandas>=2.0.0",
15
17
  "rouge-score>=0.1.2",
16
18
  "sacrebleu>=2.5.1",
19
+ "scipy>=1.15.3",
20
+ "wandb>=0.27.0",
21
+ "torch",
22
+ "rensa>=0.4.0",
23
+ "rbloom>=1.5.4",
17
24
  ]
18
25
  classifiers = [
19
26
  "Programming Language :: Python :: 3.10",
@@ -27,9 +34,14 @@ classifiers = [
27
34
  ]
28
35
 
29
36
  [project.urls]
30
- Repository = "https://github.com/quanghona/agent_design_pattern.git"
37
+ Repository = "https://github.com/quanghona/aap_core.git"
31
38
 
32
39
  [dependency-groups]
40
+ dev = [
41
+ "pytest>=9.0.3",
42
+ "ruff>=0.14.0",
43
+ "pytest-cov>=7.1.0",
44
+ ]
33
45
  lint = [
34
46
  "ruff==0.14.0",
35
47
  ]
@@ -0,0 +1,23 @@
1
+ from aap_core.policy_trainer import (
2
+ BaseExplorationModule,
3
+ BasePolicyTrainer,
4
+ DPOTrainer,
5
+ EpsilonGreedyExploration,
6
+ GRPOTrainer,
7
+ PPOTrainer,
8
+ RandomExploration,
9
+ ReinforcePPTrainer,
10
+ SimpleReplayBuffer,
11
+ )
12
+
13
+ __all__ = [
14
+ "BaseExplorationModule",
15
+ "BasePolicyTrainer",
16
+ "DPOTrainer",
17
+ "EpsilonGreedyExploration",
18
+ "GRPOTrainer",
19
+ "PPOTrainer",
20
+ "RandomExploration",
21
+ "ReinforcePPTrainer",
22
+ "SimpleReplayBuffer",
23
+ ]
@@ -5,13 +5,15 @@ from typing import Any, Dict, Literal
5
5
  from a2a.types import AgentCard
6
6
  from pydantic import BaseModel, Field, PrivateAttr
7
7
 
8
- from .chain import BaseLLMChain
9
- from .types import AgentMessage
8
+ from .types import AgentMessage, BaseLLMChain
10
9
 
11
10
 
12
11
  class BaseAgent(abc.ABC, BaseModel):
13
12
  """An interface for agent implementations."""
14
13
 
14
+ # a2a-sdk >= 1.0 exposes AgentCard as a protobuf message, not a pydantic model.
15
+ model_config = {"arbitrary_types_allowed": True}
16
+
15
17
  _state: str = PrivateAttr("idle")
16
18
  _composed_state: str = PrivateAttr("idle")
17
19
 
@@ -9,28 +9,7 @@ from .prompt_augmenter import (
9
9
  BasePromptAugmenter,
10
10
  IdentityPromptAugmenter,
11
11
  )
12
- from .types import AgentMessage, BaseChain, ChainMessage, ChainResponse
13
-
14
-
15
- class BaseLLMChain(BaseChain):
16
- """
17
- Base class for LLM chains.
18
- """
19
-
20
- name: str = Field(
21
- "chain",
22
- description="The name of the chain. Should be same at agent who hold this chain for easy to operate.",
23
- )
24
-
25
- @abc.abstractmethod
26
- def invoke(self, message: AgentMessage, **kwargs) -> AgentMessage:
27
- pass
28
-
29
- async def ainvoke(self, message: AgentMessage, **kwargs) -> AgentMessage:
30
- return self.invoke(message, **kwargs)
31
-
32
- def __call__(self, message: AgentMessage, **kwargs) -> AgentMessage:
33
- return self.invoke(message, **kwargs)
12
+ from .types import AgentMessage, BaseLLMChain, ChainMessage, ChainResponse, TokenUsage
34
13
 
35
14
 
36
15
  class BaseCausalMultiTurnsChain(BaseLLMChain, Generic[ChainMessage, ChainResponse]):
@@ -46,6 +25,10 @@ class BaseCausalMultiTurnsChain(BaseLLMChain, Generic[ChainMessage, ChainRespons
46
25
  default=False, description="store intermediate steps output"
47
26
  )
48
27
 
28
+ max_turns: int = Field(
29
+ default=50, description="maximum number of turns before stopping"
30
+ )
31
+
49
32
  def final_response_as_context(self, key: str) -> None:
50
33
  self._last_response_as_context = key.replace("context_", "")
51
34
 
@@ -56,7 +39,7 @@ class BaseCausalMultiTurnsChain(BaseLLMChain, Generic[ChainMessage, ChainRespons
56
39
  @abc.abstractmethod
57
40
  def _generate_response(
58
41
  self, conversation: List[ChainMessage], **kwargs
59
- ) -> Tuple[List[ChainMessage], ChainResponse, bool]:
42
+ ) -> Tuple[List[ChainMessage], ChainResponse, bool, TokenUsage]:
60
43
  pass
61
44
 
62
45
  @abc.abstractmethod
@@ -74,14 +57,38 @@ class BaseCausalMultiTurnsChain(BaseLLMChain, Generic[ChainMessage, ChainRespons
74
57
  pass
75
58
 
76
59
  def invoke(self, message: AgentMessage, **kwargs) -> AgentMessage:
77
- # Template method
60
+ if message.token_usage is None:
61
+ message.token_usage = {
62
+ "steps": [],
63
+ "total": TokenUsage(input_tokens=0, output_tokens=0, total_tokens=0),
64
+ }
65
+ if "steps" not in message.token_usage:
66
+ message.token_usage["steps"] = []
67
+ if "total" not in message.token_usage:
68
+ message.token_usage["total"] = TokenUsage(
69
+ input_tokens=0, output_tokens=0, total_tokens=0
70
+ )
71
+
72
+ def add_token(usage: TokenUsage):
73
+ message.token_usage["steps"].append(usage) # type: ignore
74
+ message.token_usage["total"]["input_tokens"] += usage["input_tokens"] # type: ignore
75
+ message.token_usage["total"]["output_tokens"] += usage["output_tokens"] # type: ignore
76
+ message.token_usage["total"]["total_tokens"] += usage["total_tokens"] # type: ignore
77
+
78
+ # Template method — ReAct loop: Thought → Act → Thought → ... → Final
78
79
  conversation = self._prepare_conversation(message)
79
- conversation, response, has_tool = self._generate_response(
80
- conversation, **kwargs
81
- )
82
- if has_tool:
80
+ turn = 0
81
+ while turn < self.max_turns:
82
+ conversation, response, has_tool, usage = self._generate_response(
83
+ conversation, **kwargs
84
+ )
85
+ add_token(usage)
86
+
87
+ if not has_tool:
88
+ break
89
+
83
90
  conversation = self._process_tools(conversation, response)
84
- conversation, _, _ = self._generate_response(conversation, **kwargs)
91
+ turn += 1
85
92
 
86
93
  message = self._append_responses(message, conversation)
87
94
 
@@ -132,13 +139,3 @@ class TypicalLLMChain(BaseLLMChain):
132
139
  message = self.generate(message, **kwargs)
133
140
  message = self.output_guardrail(message)
134
141
  return message
135
-
136
-
137
- class MetaPromptAugmenter(BasePromptAugmenter):
138
- """A prompt augmenter that use an LLM chain to rewrite the prompt."""
139
-
140
- chain: BaseLLMChain = Field(..., description="LLM chain that rewrite the prompt")
141
-
142
- def augment(self, message: AgentMessage, **kwargs) -> AgentMessage:
143
- message = self.chain.invoke(message, **kwargs)
144
- return message