grasp_agents 0.5.3__tar.gz → 0.5.5__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 (64) hide show
  1. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/PKG-INFO +1 -1
  2. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/pyproject.toml +1 -1
  3. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/__init__.py +4 -6
  4. grasp_agents-0.5.5/src/grasp_agents/errors.py +156 -0
  5. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/llm_agent.py +106 -146
  6. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/llm_agent_memory.py +1 -1
  7. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/llm_policy_executor.py +17 -15
  8. grasp_agents-0.5.5/src/grasp_agents/packet.py +46 -0
  9. grasp_agents-0.5.5/src/grasp_agents/packet_pool.py +159 -0
  10. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/printer.py +9 -5
  11. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/processor.py +217 -166
  12. grasp_agents-0.5.5/src/grasp_agents/prompt_builder.py +155 -0
  13. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/run_context.py +3 -16
  14. grasp_agents-0.5.5/src/grasp_agents/runner.py +131 -0
  15. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/typing/events.py +8 -4
  16. grasp_agents-0.5.5/src/grasp_agents/typing/io.py +9 -0
  17. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/workflow/looped_workflow.py +13 -19
  18. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/workflow/sequential_workflow.py +6 -10
  19. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/workflow/workflow_processor.py +23 -16
  20. grasp_agents-0.5.3/src/grasp_agents/comm_processor.py +0 -214
  21. grasp_agents-0.5.3/src/grasp_agents/errors.py +0 -94
  22. grasp_agents-0.5.3/src/grasp_agents/packet.py +0 -27
  23. grasp_agents-0.5.3/src/grasp_agents/packet_pool.py +0 -92
  24. grasp_agents-0.5.3/src/grasp_agents/prompt_builder.py +0 -218
  25. grasp_agents-0.5.3/src/grasp_agents/runner.py +0 -42
  26. grasp_agents-0.5.3/src/grasp_agents/typing/io.py +0 -16
  27. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/.gitignore +0 -0
  28. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/LICENSE.md +0 -0
  29. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/README.md +0 -0
  30. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/cloud_llm.py +0 -0
  31. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/costs_dict.yaml +0 -0
  32. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/generics_utils.py +0 -0
  33. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/grasp_logging.py +0 -0
  34. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/http_client.py +0 -0
  35. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/litellm/__init__.py +0 -0
  36. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/litellm/completion_chunk_converters.py +0 -0
  37. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/litellm/completion_converters.py +0 -0
  38. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/litellm/converters.py +0 -0
  39. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/litellm/lite_llm.py +0 -0
  40. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/litellm/message_converters.py +0 -0
  41. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/llm.py +0 -0
  42. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/memory.py +0 -0
  43. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/openai/__init__.py +0 -0
  44. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/openai/completion_chunk_converters.py +0 -0
  45. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/openai/completion_converters.py +0 -0
  46. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/openai/content_converters.py +0 -0
  47. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/openai/converters.py +0 -0
  48. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/openai/message_converters.py +0 -0
  49. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/openai/openai_llm.py +0 -0
  50. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/openai/tool_converters.py +0 -0
  51. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/rate_limiting/__init__.py +0 -0
  52. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/rate_limiting/rate_limiter_chunked.py +0 -0
  53. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/rate_limiting/types.py +0 -0
  54. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/rate_limiting/utils.py +0 -0
  55. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/typing/__init__.py +0 -0
  56. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/typing/completion.py +0 -0
  57. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/typing/completion_chunk.py +0 -0
  58. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/typing/content.py +0 -0
  59. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/typing/converters.py +0 -0
  60. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/typing/message.py +0 -0
  61. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/typing/tool.py +0 -0
  62. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/usage_tracker.py +0 -0
  63. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/utils.py +0 -0
  64. {grasp_agents-0.5.3 → grasp_agents-0.5.5}/src/grasp_agents/workflow/__init__.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: grasp_agents
3
- Version: 0.5.3
3
+ Version: 0.5.5
4
4
  Summary: Grasp Agents Library
5
5
  License-File: LICENSE.md
6
6
  Requires-Python: <4,>=3.11.4
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "grasp_agents"
3
- version = "0.5.3"
3
+ version = "0.5.5"
4
4
  description = "Grasp Agents Library"
5
5
  readme = "README.md"
6
6
  requires-python = ">=3.11.4,<4"
@@ -1,17 +1,16 @@
1
1
  # pyright: reportUnusedImport=false
2
2
 
3
3
 
4
- from .comm_processor import CommProcessor
5
4
  from .llm import LLM, LLMSettings
6
5
  from .llm_agent import LLMAgent
7
6
  from .llm_agent_memory import LLMAgentMemory
8
7
  from .memory import Memory
9
8
  from .packet import Packet
10
9
  from .processor import Processor
11
- from .run_context import RunArgs, RunContext
10
+ from .run_context import RunContext
12
11
  from .typing.completion import Completion
13
12
  from .typing.content import Content, ImageData
14
- from .typing.io import LLMPrompt, LLMPromptArgs, ProcName
13
+ from .typing.io import LLMPrompt, ProcName
15
14
  from .typing.message import AssistantMessage, Messages, SystemMessage, UserMessage
16
15
  from .typing.tool import BaseTool
17
16
 
@@ -19,20 +18,19 @@ __all__ = [
19
18
  "LLM",
20
19
  "AssistantMessage",
21
20
  "BaseTool",
22
- "CommProcessor",
23
21
  "Completion",
24
22
  "Content",
25
23
  "ImageData",
26
24
  "LLMAgent",
25
+ "LLMAgentMemory",
27
26
  "LLMPrompt",
28
- "LLMPromptArgs",
29
27
  "LLMSettings",
28
+ "Memory",
30
29
  "Messages",
31
30
  "Packet",
32
31
  "Packet",
33
32
  "ProcName",
34
33
  "Processor",
35
- "RunArgs",
36
34
  "RunContext",
37
35
  "SystemMessage",
38
36
  "UserMessage",
@@ -0,0 +1,156 @@
1
+ # from openai import APIResponseValidationError
2
+
3
+
4
+ class ProcRunError(Exception):
5
+ def __init__(
6
+ self, proc_name: str, call_id: str, message: str | None = None
7
+ ) -> None:
8
+ super().__init__(
9
+ message
10
+ or f"Processor run failed [proc_name: {proc_name}; call_id: {call_id}]."
11
+ )
12
+ self.proc_name = proc_name
13
+ self.call_id = call_id
14
+
15
+
16
+ class ProcInputValidationError(ProcRunError):
17
+ pass
18
+
19
+
20
+ class ProcOutputValidationError(ProcRunError):
21
+ def __init__(
22
+ self, schema: object, proc_name: str, call_id: str, message: str | None = None
23
+ ):
24
+ super().__init__(
25
+ proc_name=proc_name,
26
+ call_id=call_id,
27
+ message=message
28
+ or (
29
+ "Processor output validation failed "
30
+ f"[proc_name: {proc_name}; call_id: {call_id}]. "
31
+ f"Expected type:\n{schema}"
32
+ ),
33
+ )
34
+
35
+
36
+ class AgentFinalAnswerError(ProcRunError):
37
+ def __init__(
38
+ self, proc_name: str, call_id: str, message: str | None = None
39
+ ) -> None:
40
+ super().__init__(
41
+ proc_name=proc_name,
42
+ call_id=call_id,
43
+ message=message
44
+ or "Final answer tool call did not return a final answer message "
45
+ f"[proc_name={proc_name}; call_id={call_id}]",
46
+ )
47
+ self.message = message
48
+
49
+
50
+ class WorkflowConstructionError(Exception):
51
+ pass
52
+
53
+
54
+ class PacketRoutingError(ProcRunError):
55
+ def __init__(
56
+ self,
57
+ proc_name: str,
58
+ call_id: str,
59
+ selected_recipient: str | None = None,
60
+ allowed_recipients: list[str] | None = None,
61
+ message: str | None = None,
62
+ ) -> None:
63
+ default_message = (
64
+ f"Selected recipient '{selected_recipient}' is not in the allowed "
65
+ f"recipients: {allowed_recipients} "
66
+ f"[proc_name={proc_name}; call_id={call_id}]"
67
+ )
68
+ super().__init__(
69
+ proc_name=proc_name, call_id=call_id, message=message or default_message
70
+ )
71
+ self.selected_recipient = selected_recipient
72
+ self.allowed_recipients = allowed_recipients
73
+
74
+
75
+ class RunnerError(Exception):
76
+ pass
77
+
78
+
79
+ class PromptBuilderError(Exception):
80
+ def __init__(self, proc_name: str, message: str | None = None) -> None:
81
+ super().__init__(message or f"Prompt builder failed [proc_name={proc_name}]")
82
+ self.proc_name = proc_name
83
+ self.message = message
84
+
85
+
86
+ class SystemPromptBuilderError(PromptBuilderError):
87
+ def __init__(self, proc_name: str, message: str | None = None) -> None:
88
+ super().__init__(
89
+ proc_name=proc_name,
90
+ message=message
91
+ or "System prompt builder failed to make system prompt "
92
+ f"[proc_name={proc_name}]",
93
+ )
94
+ self.message = message
95
+
96
+
97
+ class InputPromptBuilderError(PromptBuilderError):
98
+ def __init__(self, proc_name: str, message: str | None = None) -> None:
99
+ super().__init__(
100
+ proc_name=proc_name,
101
+ message=message
102
+ or "Input prompt builder failed to make input content "
103
+ f"[proc_name={proc_name}]",
104
+ )
105
+ self.message = message
106
+
107
+
108
+ class PyJSONStringParsingError(Exception):
109
+ def __init__(self, s: str, message: str | None = None) -> None:
110
+ super().__init__(
111
+ message
112
+ or "Both ast.literal_eval and json.loads failed to parse the following "
113
+ f"JSON/Python string:\n{s}"
114
+ )
115
+ self.s = s
116
+
117
+
118
+ class JSONSchemaValidationError(Exception):
119
+ def __init__(self, s: str, schema: object, message: str | None = None) -> None:
120
+ super().__init__(
121
+ message
122
+ or f"JSON schema validation failed for:\n{s}\nExpected type: {schema}"
123
+ )
124
+ self.s = s
125
+ self.schema = schema
126
+
127
+
128
+ class CompletionError(Exception):
129
+ pass
130
+
131
+
132
+ class CombineCompletionChunksError(Exception):
133
+ pass
134
+
135
+
136
+ class LLMToolCallValidationError(Exception):
137
+ def __init__(
138
+ self, tool_name: str, tool_args: str, message: str | None = None
139
+ ) -> None:
140
+ super().__init__(
141
+ message
142
+ or f"Failed to validate tool call '{tool_name}' with arguments:"
143
+ f"\n{tool_args}."
144
+ )
145
+ self.tool_name = tool_name
146
+ self.tool_args = tool_args
147
+
148
+
149
+ class LLMResponseValidationError(JSONSchemaValidationError):
150
+ def __init__(self, s: str, schema: object, message: str | None = None) -> None:
151
+ super().__init__(
152
+ s,
153
+ schema,
154
+ message
155
+ or f"Failed to validate LLM response:\n{s}\nExpected type: {schema}",
156
+ )
@@ -1,22 +1,21 @@
1
1
  from collections.abc import AsyncIterator, Sequence
2
2
  from pathlib import Path
3
- from typing import Any, ClassVar, Generic, Protocol, TypeVar, cast
3
+ from typing import Any, ClassVar, Generic, Protocol, TypeVar, cast, final
4
4
 
5
5
  from pydantic import BaseModel
6
6
 
7
- from .comm_processor import CommProcessor
8
7
  from .llm import LLM, LLMSettings
9
- from .llm_agent_memory import LLMAgentMemory, PrepareMemoryHandler
8
+ from .llm_agent_memory import LLMAgentMemory, MemoryPreparator
10
9
  from .llm_policy_executor import (
11
- ExitToolCallLoopHandler,
12
10
  LLMPolicyExecutor,
13
- ManageMemoryHandler,
11
+ MemoryManager,
12
+ ToolCallLoopTerminator,
14
13
  )
15
- from .packet_pool import PacketPool
14
+ from .processor import Processor
16
15
  from .prompt_builder import (
17
- MakeInputContentHandler,
18
- MakeSystemPromptHandler,
16
+ InputContentBuilder,
19
17
  PromptBuilder,
18
+ SystemPromptBuilder,
20
19
  )
21
20
  from .run_context import CtxT, RunContext
22
21
  from .typing.content import Content, ImageData
@@ -27,7 +26,7 @@ from .typing.events import (
27
26
  SystemMessageEvent,
28
27
  UserMessageEvent,
29
28
  )
30
- from .typing.io import InT, LLMPrompt, LLMPromptArgs, OutT_co, ProcName
29
+ from .typing.io import InT, LLMPrompt, OutT, ProcName
31
30
  from .typing.message import Message, Messages, SystemMessage, UserMessage
32
31
  from .typing.tool import BaseTool
33
32
  from .utils import get_prompt, validate_obj_from_json_or_py_string
@@ -36,7 +35,7 @@ _InT_contra = TypeVar("_InT_contra", contravariant=True)
36
35
  _OutT_co = TypeVar("_OutT_co", covariant=True)
37
36
 
38
37
 
39
- class ParseOutputHandler(Protocol[_InT_contra, _OutT_co, CtxT]):
38
+ class OutputParser(Protocol[_InT_contra, _OutT_co, CtxT]):
40
39
  def __call__(
41
40
  self,
42
41
  conversation: Messages,
@@ -47,8 +46,8 @@ class ParseOutputHandler(Protocol[_InT_contra, _OutT_co, CtxT]):
47
46
 
48
47
 
49
48
  class LLMAgent(
50
- CommProcessor[InT, OutT_co, LLMAgentMemory, CtxT],
51
- Generic[InT, OutT_co, CtxT],
49
+ Processor[InT, OutT, LLMAgentMemory, CtxT],
50
+ Generic[InT, OutT, CtxT],
52
51
  ):
53
52
  _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
54
53
  0: "_in_type",
@@ -69,10 +68,6 @@ class LLMAgent(
69
68
  # System prompt template
70
69
  sys_prompt: LLMPrompt | None = None,
71
70
  sys_prompt_path: str | Path | None = None,
72
- # System args (static args provided via RunContext)
73
- sys_args_schema: type[LLMPromptArgs] | None = None,
74
- # User args (static args provided via RunContext)
75
- usr_args_schema: type[LLMPromptArgs] | None = None,
76
71
  # Agent loop settings
77
72
  max_turns: int = 100,
78
73
  react_mode: bool = False,
@@ -82,25 +77,31 @@ class LLMAgent(
82
77
  # Retries
83
78
  max_retries: int = 0,
84
79
  # Multi-agent routing
85
- packet_pool: PacketPool[CtxT] | None = None,
86
80
  recipients: list[ProcName] | None = None,
87
81
  ) -> None:
88
- super().__init__(
89
- name=name,
90
- packet_pool=packet_pool,
91
- recipients=recipients,
92
- max_retries=max_retries,
93
- )
82
+ super().__init__(name=name, recipients=recipients, max_retries=max_retries)
94
83
 
95
84
  # Agent memory
96
85
 
97
86
  self._memory: LLMAgentMemory = LLMAgentMemory()
98
87
  self._reset_memory_on_run = reset_memory_on_run
99
88
 
89
+ self.memory_preparator: MemoryPreparator | None
90
+ if not hasattr(type(self), "memory_preparator"):
91
+ self.memory_preparator = None
92
+
93
+ self.output_parser: OutputParser[InT, OutT, CtxT] | None
94
+ if not hasattr(type(self), "output_parser"):
95
+ self.output_parser = None
96
+
100
97
  # LLM policy executor
101
98
 
102
99
  self._used_default_llm_response_schema: bool = False
103
- if llm.response_schema is None and tools is None:
100
+ if (
101
+ llm.response_schema is None
102
+ and tools is None
103
+ and not hasattr(type(self), "output_parser")
104
+ ):
104
105
  llm.response_schema = self.out_type
105
106
  self._used_default_llm_response_schema = True
106
107
 
@@ -128,18 +129,11 @@ class LLMAgent(
128
129
 
129
130
  sys_prompt = get_prompt(prompt_text=sys_prompt, prompt_path=sys_prompt_path)
130
131
  in_prompt = get_prompt(prompt_text=in_prompt, prompt_path=in_prompt_path)
131
- self._prompt_builder: PromptBuilder[InT, CtxT] = PromptBuilder[
132
- self.in_type, CtxT
133
- ](
134
- agent_name=self._name,
135
- sys_prompt_template=sys_prompt,
136
- in_prompt_template=in_prompt,
137
- sys_args_schema=sys_args_schema,
138
- usr_args_schema=usr_args_schema,
132
+
133
+ self._prompt_builder = PromptBuilder[self.in_type, CtxT](
134
+ agent_name=self._name, sys_prompt=sys_prompt, in_prompt=in_prompt
139
135
  )
140
136
 
141
- self._prepare_memory_impl: PrepareMemoryHandler | None = None
142
- self._parse_output_impl: ParseOutputHandler[InT, OutT_co, CtxT] | None = None
143
137
  self._register_overridden_handlers()
144
138
 
145
139
  @property
@@ -154,22 +148,15 @@ class LLMAgent(
154
148
  def max_turns(self) -> int:
155
149
  return self._policy_executor.max_turns
156
150
 
157
- @property
158
- def sys_args_schema(self) -> type[LLMPromptArgs] | None:
159
- return self._prompt_builder.sys_args_schema
160
-
161
- @property
162
- def usr_args_schema(self) -> type[LLMPromptArgs] | None:
163
- return self._prompt_builder.usr_args_schema
164
-
165
151
  @property
166
152
  def sys_prompt(self) -> LLMPrompt | None:
167
- return self._prompt_builder.sys_prompt_template
153
+ return self._prompt_builder.sys_prompt
168
154
 
169
155
  @property
170
156
  def in_prompt(self) -> LLMPrompt | None:
171
- return self._prompt_builder.in_prompt_template
157
+ return self._prompt_builder.in_prompt
172
158
 
159
+ @final
173
160
  def _prepare_memory(
174
161
  self,
175
162
  memory: LLMAgentMemory,
@@ -177,35 +164,19 @@ class LLMAgent(
177
164
  sys_prompt: LLMPrompt | None = None,
178
165
  ctx: RunContext[Any] | None = None,
179
166
  ) -> None:
180
- if self._prepare_memory_impl:
181
- return self._prepare_memory_impl(
167
+ if self.memory_preparator:
168
+ return self.memory_preparator(
182
169
  memory=memory, in_args=in_args, sys_prompt=sys_prompt, ctx=ctx
183
170
  )
184
171
 
185
172
  def _memorize_inputs(
186
173
  self,
174
+ memory: LLMAgentMemory,
187
175
  chat_inputs: LLMPrompt | Sequence[str | ImageData] | None = None,
188
- *,
189
176
  in_args: InT | None = None,
190
- memory: LLMAgentMemory,
191
177
  ctx: RunContext[CtxT] | None = None,
192
178
  ) -> tuple[SystemMessage | None, UserMessage | None]:
193
- # 1. Get run arguments
194
- sys_args: LLMPromptArgs | None = None
195
- usr_args: LLMPromptArgs | None = None
196
- if ctx is not None:
197
- run_args = ctx.run_args.get(self.name)
198
- if run_args is not None:
199
- sys_args = run_args.sys
200
- usr_args = run_args.usr
201
-
202
- # 2. Make system prompt (can be None)
203
-
204
- formatted_sys_prompt = self._prompt_builder.make_system_prompt(
205
- sys_args=sys_args, ctx=ctx
206
- )
207
-
208
- # 3. Set agent memory
179
+ formatted_sys_prompt = self._prompt_builder.build_system_prompt(ctx=ctx)
209
180
 
210
181
  system_message: SystemMessage | None = None
211
182
  if self._reset_memory_on_run or memory.is_empty:
@@ -214,16 +185,11 @@ class LLMAgent(
214
185
  system_message = cast("SystemMessage", memory.message_history[0])
215
186
  else:
216
187
  self._prepare_memory(
217
- memory=memory,
218
- in_args=in_args,
219
- sys_prompt=formatted_sys_prompt,
220
- ctx=ctx,
188
+ memory=memory, in_args=in_args, sys_prompt=formatted_sys_prompt, ctx=ctx
221
189
  )
222
190
 
223
- # 3. Make and add input messages
224
-
225
- input_message = self._prompt_builder.make_input_message(
226
- chat_inputs=chat_inputs, in_args=in_args, usr_args=usr_args, ctx=ctx
191
+ input_message = self._prompt_builder.build_input_message(
192
+ chat_inputs=chat_inputs, in_args=in_args, ctx=ctx
227
193
  )
228
194
  if input_message:
229
195
  memory.update([input_message])
@@ -236,9 +202,9 @@ class LLMAgent(
236
202
  *,
237
203
  in_args: InT | None = None,
238
204
  ctx: RunContext[CtxT] | None = None,
239
- ) -> OutT_co:
240
- if self._parse_output_impl:
241
- return self._parse_output_impl(
205
+ ) -> OutT:
206
+ if self.output_parser:
207
+ return self.output_parser(
242
208
  conversation=conversation, in_args=in_args, ctx=ctx
243
209
  )
244
210
 
@@ -257,9 +223,12 @@ class LLMAgent(
257
223
  memory: LLMAgentMemory,
258
224
  call_id: str,
259
225
  ctx: RunContext[CtxT] | None = None,
260
- ) -> Sequence[OutT_co]:
226
+ ) -> OutT:
261
227
  system_message, input_message = self._memorize_inputs(
262
- chat_inputs=chat_inputs, in_args=in_args, memory=memory, ctx=ctx
228
+ memory=memory,
229
+ chat_inputs=chat_inputs,
230
+ in_args=in_args,
231
+ ctx=ctx,
263
232
  )
264
233
  if system_message:
265
234
  self._print_messages([system_message], call_id=call_id, ctx=ctx)
@@ -268,11 +237,9 @@ class LLMAgent(
268
237
 
269
238
  await self._policy_executor.execute(memory, call_id=call_id, ctx=ctx)
270
239
 
271
- return [
272
- self._parse_output(
273
- conversation=memory.message_history, in_args=in_args, ctx=ctx
274
- )
275
- ]
240
+ return self._parse_output(
241
+ conversation=memory.message_history, in_args=in_args, ctx=ctx
242
+ )
276
243
 
277
244
  async def _process_stream(
278
245
  self,
@@ -284,7 +251,10 @@ class LLMAgent(
284
251
  ctx: RunContext[CtxT] | None = None,
285
252
  ) -> AsyncIterator[Event[Any]]:
286
253
  system_message, input_message = self._memorize_inputs(
287
- chat_inputs=chat_inputs, in_args=in_args, memory=memory, ctx=ctx
254
+ memory=memory,
255
+ chat_inputs=chat_inputs,
256
+ in_args=in_args,
257
+ ctx=ctx,
288
258
  )
289
259
  if system_message:
290
260
  self._print_messages([system_message], call_id=call_id, ctx=ctx)
@@ -324,108 +294,98 @@ class LLMAgent(
324
294
 
325
295
  # Prompt builder
326
296
 
327
- if cur_cls._make_system_prompt is not base_cls._make_system_prompt: # noqa: SLF001
328
- self._prompt_builder.make_system_prompt_impl = self._make_system_prompt
297
+ if cur_cls.system_prompt_builder is not base_cls.system_prompt_builder:
298
+ self._prompt_builder.system_prompt_builder = self.system_prompt_builder
329
299
 
330
- if cur_cls._make_input_content is not base_cls._make_input_content: # noqa: SLF001
331
- self._prompt_builder.make_input_content_impl = self._make_input_content
300
+ if cur_cls.input_content_builder is not base_cls.input_content_builder:
301
+ self._prompt_builder.input_content_builder = self.input_content_builder
332
302
 
333
303
  # Policy executor
334
304
 
335
- if (
336
- cur_cls._exit_tool_call_loop is not base_cls._exit_tool_call_loop # noqa: SLF001
337
- ):
338
- self._policy_executor.exit_tool_call_loop_impl = self._exit_tool_call_loop
339
-
340
- if cur_cls._manage_memory is not base_cls._manage_memory: # noqa: SLF001
341
- self._policy_executor.manage_memory_impl = self._manage_memory
305
+ if cur_cls.tool_call_loop_terminator is not base_cls.tool_call_loop_terminator:
306
+ self._policy_executor.tool_call_loop_terminator = (
307
+ self.tool_call_loop_terminator
308
+ )
342
309
 
343
- # Make sure default LLM response schema is not used when custom output
344
- # parsing is provided
345
- if (
346
- cur_cls._parse_output is not base_cls._parse_output # noqa: SLF001
347
- and self._used_default_llm_response_schema
348
- ):
349
- self._policy_executor.llm.response_schema = None
310
+ if cur_cls.memory_manager is not base_cls.memory_manager:
311
+ self._policy_executor.memory_manager = self.memory_manager
350
312
 
351
- def _make_system_prompt(
352
- self, sys_args: LLMPromptArgs | None, *, ctx: RunContext[CtxT] | None = None
353
- ) -> str | None:
354
- return self._prompt_builder.make_system_prompt(sys_args=sys_args, ctx=ctx)
313
+ def system_prompt_builder(self, ctx: RunContext[CtxT] | None = None) -> str | None:
314
+ if self._prompt_builder.system_prompt_builder is not None:
315
+ return self._prompt_builder.system_prompt_builder(ctx=ctx)
316
+ raise NotImplementedError("System prompt builder is not implemented.")
355
317
 
356
- def _make_input_content(
357
- self,
358
- *,
359
- in_args: InT | None = None,
360
- usr_args: LLMPromptArgs | None = None,
361
- ctx: RunContext[CtxT] | None = None,
318
+ def input_content_builder(
319
+ self, in_args: InT | None = None, *, ctx: RunContext[CtxT] | None = None
362
320
  ) -> Content:
363
- return self._prompt_builder.make_input_content(
364
- in_args=in_args, usr_args=usr_args, ctx=ctx
365
- )
321
+ if self._prompt_builder.input_content_builder is not None:
322
+ return self._prompt_builder.input_content_builder(in_args=in_args, ctx=ctx)
323
+ raise NotImplementedError("Input content builder is not implemented.")
366
324
 
367
- def _exit_tool_call_loop(
325
+ def tool_call_loop_terminator(
368
326
  self,
369
327
  conversation: Messages,
370
328
  *,
371
329
  ctx: RunContext[CtxT] | None = None,
372
330
  **kwargs: Any,
373
331
  ) -> bool:
374
- return self._policy_executor._exit_tool_call_loop( # type: ignore[return-value]
375
- conversation=conversation, ctx=ctx, **kwargs
376
- )
332
+ if self._policy_executor.tool_call_loop_terminator is not None:
333
+ return self._policy_executor.tool_call_loop_terminator(
334
+ conversation=conversation, ctx=ctx, **kwargs
335
+ )
336
+ raise NotImplementedError("Tool call loop terminator is not implemented.")
377
337
 
378
- def _manage_memory(
338
+ def memory_manager(
379
339
  self,
380
340
  memory: LLMAgentMemory,
381
341
  *,
382
342
  ctx: RunContext[CtxT] | None = None,
383
343
  **kwargs: Any,
384
344
  ) -> None:
385
- return self._policy_executor._manage_memory( # type: ignore[return-value]
386
- memory=memory, ctx=ctx, **kwargs
387
- )
345
+ if self._policy_executor.memory_manager is not None:
346
+ return self._policy_executor.memory_manager(
347
+ memory=memory, ctx=ctx, **kwargs
348
+ )
349
+ raise NotImplementedError("Memory manager is not implemented.")
388
350
 
389
351
  # Decorators for custom implementations as an alternative to overriding methods
390
352
 
391
- def make_system_prompt(
392
- self, func: MakeSystemPromptHandler[CtxT]
393
- ) -> MakeSystemPromptHandler[CtxT]:
394
- self._prompt_builder.make_system_prompt_impl = func
353
+ def add_system_prompt_builder(
354
+ self, func: SystemPromptBuilder[CtxT]
355
+ ) -> SystemPromptBuilder[CtxT]:
356
+ self._prompt_builder.system_prompt_builder = func
395
357
 
396
358
  return func
397
359
 
398
- def make_input_content(
399
- self, func: MakeInputContentHandler[InT, CtxT]
400
- ) -> MakeInputContentHandler[InT, CtxT]:
401
- self._prompt_builder.make_input_content_impl = func
360
+ def add_input_content_builder(
361
+ self, func: InputContentBuilder[InT, CtxT]
362
+ ) -> InputContentBuilder[InT, CtxT]:
363
+ self._prompt_builder.input_content_builder = func
402
364
 
403
365
  return func
404
366
 
405
- def parse_output(
406
- self, func: ParseOutputHandler[InT, OutT_co, CtxT]
407
- ) -> ParseOutputHandler[InT, OutT_co, CtxT]:
408
- if self._used_default_llm_response_schema:
409
- self._policy_executor.llm.response_schema = None
410
- self._parse_output_impl = func
367
+ def add_memory_manager(self, func: MemoryManager[CtxT]) -> MemoryManager[CtxT]:
368
+ self._policy_executor.memory_manager = func
411
369
 
412
370
  return func
413
371
 
414
- def prepare_memory(self, func: PrepareMemoryHandler) -> PrepareMemoryHandler:
415
- self._prepare_memory_impl = func
372
+ def add_tool_call_loop_terminator(
373
+ self, func: ToolCallLoopTerminator[CtxT]
374
+ ) -> ToolCallLoopTerminator[CtxT]:
375
+ self._policy_executor.tool_call_loop_terminator = func
416
376
 
417
377
  return func
418
378
 
419
- def manage_memory(
420
- self, func: ManageMemoryHandler[CtxT]
421
- ) -> ManageMemoryHandler[CtxT]:
422
- self._policy_executor.manage_memory_impl = func
379
+ def add_output_parser(
380
+ self, func: OutputParser[InT, OutT, CtxT]
381
+ ) -> OutputParser[InT, OutT, CtxT]:
382
+ if self._used_default_llm_response_schema:
383
+ self._policy_executor.llm.response_schema = None
384
+ self.output_parser = func
423
385
 
424
386
  return func
425
387
 
426
- def exit_tool_call_loop(
427
- self, func: ExitToolCallLoopHandler[CtxT]
428
- ) -> ExitToolCallLoopHandler[CtxT]:
429
- self._policy_executor.exit_tool_call_loop_impl = func
388
+ def add_memory_preparator(self, func: MemoryPreparator) -> MemoryPreparator:
389
+ self.memory_preparator = func
430
390
 
431
391
  return func
@@ -9,7 +9,7 @@ from .typing.io import LLMPrompt
9
9
  from .typing.message import Message, Messages, SystemMessage
10
10
 
11
11
 
12
- class PrepareMemoryHandler(Protocol):
12
+ class MemoryPreparator(Protocol):
13
13
  def __call__(
14
14
  self,
15
15
  memory: "LLMAgentMemory",