zrb 1.9.17__py3-none-any.whl → 1.10.0__py3-none-any.whl

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.
zrb/__init__.py CHANGED
@@ -107,7 +107,7 @@ if TYPE_CHECKING:
107
107
  from zrb.task.base_trigger import BaseTrigger
108
108
  from zrb.task.cmd_task import CmdTask
109
109
  from zrb.task.http_check import HttpCheck
110
- from zrb.task.llm.history import ConversationHistoryData
110
+ from zrb.task.llm.conversation_history import ConversationHistoryData
111
111
  from zrb.task.llm_task import LLMTask
112
112
  from zrb.task.make_task import make_task
113
113
  from zrb.task.rsync_task import RsyncTask
@@ -4,7 +4,7 @@ from typing import Any
4
4
 
5
5
  from zrb.config.config import CFG
6
6
  from zrb.context.any_shared_context import AnySharedContext
7
- from zrb.task.llm.history import ConversationHistoryData
7
+ from zrb.task.llm.conversation_history_model import ConversationHistory
8
8
  from zrb.util.file import read_file, write_file
9
9
 
10
10
 
@@ -51,9 +51,7 @@ def read_chat_conversation(ctx: AnySharedContext) -> dict[str, Any] | list | Non
51
51
  return None
52
52
 
53
53
 
54
- def write_chat_conversation(
55
- ctx: AnySharedContext, history_data: ConversationHistoryData
56
- ):
54
+ def write_chat_conversation(ctx: AnySharedContext, history_data: ConversationHistory):
57
55
  """Writes the conversation history data (including context) to a session file."""
58
56
  os.makedirs(CFG.LLM_HISTORY_DIR, exist_ok=True)
59
57
  current_session_name = ctx.session.name
@@ -1,3 +1,4 @@
1
+ import json
1
2
  import subprocess
2
3
 
3
4
 
@@ -13,17 +14,20 @@ def run_shell_command(command: str) -> str:
13
14
  command (str): The exact shell command to execute.
14
15
 
15
16
  Returns:
16
- str: The combined standard output (stdout) and standard error (stderr) from the command. If the command fails, this will contain the error message.
17
- Raises:
18
- subprocess.CalledProcessError: If the command returns a non-zero exit code, indicating an error.
17
+ str: A JSON string containing return code, standard output (stdout),
18
+ and standard error (stderr) from the command.
19
+ Example: {"return_code": 0, "stdout": "ok", "stderr": ""}
19
20
  """
20
- try:
21
- output = subprocess.check_output(
22
- command, shell=True, stderr=subprocess.STDOUT, text=True
23
- )
24
- return output
25
- except subprocess.CalledProcessError as e:
26
- # Include the error output in the exception message
27
- raise subprocess.CalledProcessError(
28
- e.returncode, e.cmd, e.output, e.stderr
29
- ) from None
21
+ result = subprocess.run(
22
+ command,
23
+ shell=True,
24
+ capture_output=True,
25
+ text=True,
26
+ )
27
+ return json.dumps(
28
+ {
29
+ "return_code": result.returncode,
30
+ "stdout": result.stdout,
31
+ "stderr": result.stderr,
32
+ }
33
+ )
@@ -218,7 +218,7 @@ def read_from_file(
218
218
 
219
219
  Returns:
220
220
  str: A JSON object containing the file path, the requested content with line numbers, the start and end lines, and the total number of lines in the file.
221
- Example: '{"path": "src/main.py", "content": "1: import os\n2: \n3: print(\"Hello, World!\")", "start_line": 1, "end_line": 3, "total_lines": 3}'
221
+ Example: '{"path": "src/main.py", "content": "1| import os\n2| \n3| print(\"Hello, World!\")", "start_line": 1, "end_line": 3, "total_lines": 3}'
222
222
  Raises:
223
223
  FileNotFoundError: If the specified file does not exist.
224
224
  """
@@ -492,7 +492,7 @@ def read_many_files(paths: List[str]) -> str:
492
492
 
493
493
  Returns:
494
494
  str: A JSON object where keys are the file paths and values are their corresponding contents, prefixed with line numbers. If a file cannot be read, its value will be an error message.
495
- Example: '{"results": {"src/api.py": "1: import ...", "config.yaml": "1: key: value"}}'
495
+ Example: '{"results": {"src/api.py": "1| import ...", "config.yaml": "1| key: value"}}'
496
496
  """
497
497
  results = {}
498
498
  for path in paths:
@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, Coroutine
6
6
  from zrb.context.any_context import AnyContext
7
7
  from zrb.task.llm.agent import create_agent_instance, run_agent_iteration
8
8
  from zrb.task.llm.config import get_model, get_model_settings
9
- from zrb.task.llm.prompt import get_combined_system_prompt
9
+ from zrb.task.llm.prompt import get_system_and_user_prompt
10
10
 
11
11
  if TYPE_CHECKING:
12
12
  from pydantic_ai import Tool
@@ -69,14 +69,12 @@ def create_sub_agent_tool(
69
69
  )
70
70
 
71
71
  if system_prompt is None:
72
- resolved_system_prompt = get_combined_system_prompt(
72
+ resolved_system_prompt, query = get_system_and_user_prompt(
73
73
  ctx=ctx,
74
+ user_message=query,
74
75
  persona_attr=None,
75
- render_persona=False,
76
76
  system_prompt_attr=None,
77
- render_system_prompt=False,
78
77
  special_instruction_prompt_attr=None,
79
- render_special_instruction_prompt=False,
80
78
  )
81
79
  else:
82
80
  resolved_system_prompt = system_prompt
zrb/config/config.py CHANGED
@@ -336,10 +336,6 @@ class Config:
336
336
  """Number of seconds to sleep when throttling is required."""
337
337
  return float(os.getenv("ZRB_LLM_THROTTLE_SLEEP", "1.0"))
338
338
 
339
- @property
340
- def LLM_CONTEXT_ENRICHMENT_PROMPT(self) -> str | None:
341
- return os.getenv("ZRB_LLM_CONTEXT_ENRICHMENT_PROMPT", None)
342
-
343
339
  @property
344
340
  def LLM_SUMMARIZE_HISTORY(self) -> bool:
345
341
  return to_boolean(os.getenv("ZRB_LLM_SUMMARIZE_HISTORY", "true"))
@@ -348,14 +344,6 @@ class Config:
348
344
  def LLM_HISTORY_SUMMARIZATION_TOKEN_THRESHOLD(self) -> int:
349
345
  return int(os.getenv("ZRB_LLM_HISTORY_SUMMARIZATION_TOKEN_THRESHOLD", "20000"))
350
346
 
351
- @property
352
- def LLM_ENRICH_CONTEXT(self) -> bool:
353
- return to_boolean(os.getenv("ZRB_LLM_ENRICH_CONTEXT", "true"))
354
-
355
- @property
356
- def LLM_CONTEXT_ENRICHMENT_TOKEN_THRESHOLD(self) -> int:
357
- return int(os.getenv("ZRB_LLM_CONTEXT_ENRICHMENT_TOKEN_THRESHOLD", "20000"))
358
-
359
347
  @property
360
348
  def LLM_REPO_ANALYSIS_EXTRACTION_TOKEN_THRESHOLD(self) -> int:
361
349
  return int(os.getenv("ZRB_LLM_REPO_ANALYSIS_EXTRACTION_TOKEN_LIMIT", "35000"))
@@ -445,5 +433,13 @@ class Config:
445
433
  {"VERSION": self.VERSION},
446
434
  )
447
435
 
436
+ @property
437
+ def LLM_CONTEXTUAL_NOTE_FILE(self) -> str:
438
+ return os.getenv("LLM_CONTEXTUAL_NOTE_FILE", "ZRB_README.md")
439
+
440
+ @property
441
+ def LLM_LONG_TERM_NOTE_PATH(self) -> str:
442
+ return os.getenv("LLM_LONG_TERM_NOTE_PATH", "~/ZRB_GLOBAL_README.md")
443
+
448
444
 
449
445
  CFG = Config()
zrb/config/llm_config.py CHANGED
@@ -34,13 +34,22 @@ _DEFAULT_INTERACTIVE_SYSTEM_PROMPT = (
34
34
  " * **CRITICAL:** Immediately after execution, you MUST use a tool "
35
35
  "to verify the outcome (e.g., after `write_file`, use `read_file`; "
36
36
  "after `rm`, use `ls` to confirm absence).\n\n"
37
- "4. **Report Results:**\n"
38
- " * Provide a concise summary of the action taken.\n"
39
- " * **You MUST explicitly state how you verified the action** (e.g., "
40
- "'I have deleted the file and verified its removal by listing the "
41
- "directory.').\n"
42
- " * If an error occurs, report the error and the failed verification "
43
- "step."
37
+ "4. **Report Results and Handle Errors:**\n"
38
+ " * **On Success:** Provide a concise summary of the action taken "
39
+ "and explicitly state how you verified it.\n"
40
+ " * **On Failure (The Debugging Loop):** If a tool call fails, you "
41
+ "MUST NOT give up. Instead, you will enter a debugging loop:\n"
42
+ " 1. **Analyze:** Scrutinize the complete error message, "
43
+ "including any `stdout` and `stderr`.\n"
44
+ " 2. **Hypothesize:** State a clear, specific hypothesis about "
45
+ "the root cause of the error.\n"
46
+ " 3. **Act:** Propose a concrete, single next step to fix the "
47
+ "issue. This could be running a command with different parameters, "
48
+ "modifying a file, or using another tool to gather more context.\n\n"
49
+ "---\n"
50
+ "**FINAL REMINDER:** Your last step before responding MUST be to ensure "
51
+ "you have followed the Execute and Verify (E+V) loop. Do not "
52
+ "hallucinate verifications."
44
53
  ).strip()
45
54
 
46
55
  _DEFAULT_SYSTEM_PROMPT = (
@@ -50,32 +59,36 @@ _DEFAULT_SYSTEM_PROMPT = (
50
59
  "1. **Plan:** Internally devise a step-by-step plan. This plan MUST "
51
60
  "include verification steps for each action.\n\n"
52
61
  "2. **Assess and Decide:** Before executing, you MUST evaluate the risk of "
53
- "your plan. For any destructive actions (modifying or deleting data), "
54
- "consider the command's nature and target. Based on your assessment, "
55
- "decide the appropriate course of action:\n"
62
+ "your plan. For any destructive actions, consider the command's nature "
63
+ "and target. Based on your assessment, decide the appropriate course of "
64
+ "action:\n"
56
65
  " * **Low/Moderate Risk:** Proceed directly.\n"
57
66
  " * **High Risk:** Refuse to execute, state your plan, and explain "
58
67
  "the risk to the user.\n"
59
- " * **Extreme Risk (e.g., operating on critical system files):** "
60
- "Refuse and explain the danger.\n\n"
61
- "3. **Execute and Verify (The E+V Loop):**\n"
68
+ " * **Extreme Risk:** Refuse and explain the danger.\n\n"
69
+ "3. **Execute and Verify (The E+V Loop):\n"
62
70
  " * Execute each step of your plan.\n"
63
71
  " * **CRITICAL:** After each step, you MUST use a tool to verify "
64
72
  "the outcome (e.g., check exit codes, read back file contents).\n\n"
65
- "4. **Report Final Outcome:**\n"
66
- " * Provide a concise summary of the result.\n"
67
- " * **You MUST explicitly state how you verified the final state**.\n"
68
- " * If an error occurred, report the error and the failed "
69
- "verification step.\n\n"
73
+ "4. **Report Final Outcome:\n"
74
+ " * **On Success:** Provide a concise summary of the result and "
75
+ "explicitly state how you verified the final state.\n"
76
+ " * **On Failure:** Report the complete error, including `stdout` "
77
+ "and `stderr`. Analyze the error and provide a corrected command or a "
78
+ "clear explanation of the root cause.\n\n"
70
79
  "---\n"
71
80
  "**FINAL REMINDER:** Your last step before responding MUST be to ensure "
72
- "you have followed the Execute and Verify (E+V) loop. If you are about "
73
- "to claim an action was taken, you MUST have already completed the "
74
- "corresponding verification tool call. Do not hallucinate verifications."
81
+ "you have followed the Execute and Verify (E+V) loop. Do not "
82
+ "hallucinate verifications."
75
83
  ).strip()
76
84
 
77
85
  _DEFAULT_SPECIAL_INSTRUCTION_PROMPT = (
78
86
  "## Guiding Principles\n"
87
+ "- **Clarify and Scope First:** Before undertaking any complex task (like "
88
+ "writing a new feature or a large test suite), you MUST ensure the request "
89
+ "is not ambiguous. If it is, ask clarifying questions. Propose a concise "
90
+ "plan or scope and ask for user approval before proceeding. Do not start a "
91
+ "multi-step task on a vague request.\n"
79
92
  "- **Safety First:** Never run commands that are destructive or could "
80
93
  "compromise the system without explicit user confirmation. When in "
81
94
  "doubt, ask.\n"
@@ -84,6 +97,17 @@ _DEFAULT_SPECIAL_INSTRUCTION_PROMPT = (
84
97
  "conventions.\n"
85
98
  "- **Efficiency:** Use your tools to get the job done with the minimum "
86
99
  "number of steps. Combine commands where possible.\n\n"
100
+ "## Critical Prohibitions\n"
101
+ "- **NEVER Assume Dependencies:** Do not use a library or framework unless "
102
+ "you have first verified it is an existing project dependency (e.g., in "
103
+ "`package.json`, `requirements.txt`).\n"
104
+ "- **NEVER Invent Conventions:** You MUST follow the existing conventions "
105
+ "discovered during your context-gathering phase. Do not introduce a new "
106
+ "style or pattern without a very good reason and, ideally, user "
107
+ "confirmation.\n"
108
+ "- **NEVER Commit Without Verification:** Do not use `git commit` until you "
109
+ "have staged the changes and run the project's own verification steps "
110
+ "(tests, linter, build).\n\n"
87
111
  "## Common Task Workflows\n\n"
88
112
  "**File System Operations:**\n"
89
113
  "1. **Analyze:** Before modifying, read the file or list the "
@@ -92,10 +116,24 @@ _DEFAULT_SPECIAL_INSTRUCTION_PROMPT = (
92
116
  "3. **Verify:** Check that the file/directory now exists (or doesn't) in "
93
117
  "its expected state.\n\n"
94
118
  "**Code & Software Development:**\n"
95
- "1. **CRITICAL: Gather Context First:** Before writing or modifying any code, "
96
- "you MUST gather context to ensure your changes are idiomatic and correct. "
97
- "Do not make assumptions. Your primary goal is to fit into the existing "
98
- "project seamlessly.\n"
119
+ "1. **CRITICAL: Gather Context First:** Before writing or modifying any "
120
+ "code, you MUST gather context to ensure your changes are idiomatic and "
121
+ "correct. Do not make assumptions. Your primary goal is to fit into the "
122
+ "existing project seamlessly.\n"
123
+ " * **Project Structure & Dependencies:** Check for `README.md`, "
124
+ "`CONTRIBUTING.md`, `package.json`, `pyproject.toml`, `build.gradle`, "
125
+ "etc., to understand the project's stated goals, dependencies, and "
126
+ "scripts (for linting, testing, building).\n"
127
+ " * **Code Style & Conventions:** Look for configuration files like "
128
+ "`.eslintrc`, `.prettierrc`, `.flake8`, or `ruff.toml`. Analyze "
129
+ "surrounding source files to determine:\n"
130
+ " * **Naming Conventions:** (e.g., `camelCase` vs. `snake_case`).\n"
131
+ " * **Typing Style:** (e.g., `List` from `typing` vs. built-in "
132
+ "`list`).\n"
133
+ " * **Error Handling:** (e.g., custom exceptions, `try/except` "
134
+ "blocks, returning error codes).\n"
135
+ " * **Architectural Patterns:** (e.g., is there a service layer? "
136
+ "Are components organized by feature or by type?).\n"
99
137
  " * **When writing a new test:** You MUST first read the full source "
100
138
  "code of the module(s) you are testing. This will inform you about the "
101
139
  "actual implementation, such as its logging methods, error handling, and "
@@ -108,65 +146,83 @@ _DEFAULT_SPECIAL_INSTRUCTION_PROMPT = (
108
146
  "context you gathered.\n"
109
147
  "3. **Implement:** Make the changes, strictly adhering to the patterns and "
110
148
  "conventions discovered in step 1.\n"
111
- "4. **Verify:** Run all relevant tests, linters, and build commands. If a "
112
- "test fails, analyze the error, read the relevant code again, and attempt "
113
- "to fix the underlying issue before trying again.\n\n"
149
+ "4. **Verify & Debug:** Run all relevant tests, linters, and build "
150
+ "commands. If a command fails, your immediate next action MUST be to "
151
+ "enter the **Debugging Loop**: analyze the complete error output (`stdout` "
152
+ "and `stderr`), hypothesize the root cause. Your next immediate action "
153
+ "MUST be to execute a single, concrete tool call that attempts to fix "
154
+ "the issue based on your hypothesis. Do not stop to ask the user for "
155
+ "confirmation. The goal is to resolve the error autonomously.\n\n"
114
156
  "**Research & Analysis:**\n"
115
157
  "1. **Clarify:** Understand the core question and the desired output "
116
158
  "format.\n"
117
159
  "2. **Search:** Use web search tools to gather information from multiple "
118
160
  "reputable sources.\n"
119
161
  "3. **Synthesize & Cite:** Present the information clearly. For factual "
120
- "claims, cite the source URL."
162
+ "claims, cite the source URL.\n\n"
163
+ "## Communicating with the User\n"
164
+ "- **Be Concise:** When reporting results, be brief. Focus on the outcome "
165
+ "and the verification step.\n"
166
+ "- **Explain 'Why,' Not Just 'What':** For complex changes or bug fixes, "
167
+ "briefly explain *why* the change was necessary (e.g., 'The previous code "
168
+ "was failing because it didn't handle null inputs. I've added a check to "
169
+ "prevent this.').\n"
170
+ "- **Structure Your Plans:** When you present a plan for approval, use a "
171
+ "numbered or bulleted list for clarity."
121
172
  ).strip()
122
173
 
123
174
 
124
175
  _DEFAULT_SUMMARIZATION_PROMPT = (
125
- "You are a Conversation Historian. Your task is to create a dense, "
126
- "structured snapshot of the conversation for the main assistant.\n\n"
127
- "You will receive a `Previous Summary` and the `Recent Conversation "
128
- "History`. Your goal is to produce an updated, rolling summary. Your "
129
- "output MUST be a single block of text with two sections:\n\n"
130
- "1. `## Narrative Summary`\n"
131
- " - **Identify Key Information:** From the `Recent Conversation "
132
- "History`, extract critical facts, user decisions, and final outcomes "
133
- "of tasks.\n"
134
- " - **Integrate and Condense:** Integrate these key facts into the "
135
- "`Previous Summary`. Discard conversational filler and intermediate "
136
- "steps of completed tasks (e.g., 'User asked to see file X, I showed "
137
- "them' can be discarded if the file was just a step towards a larger "
138
- "goal). The summary should reflect the current state of the project or "
139
- "conversation, not a log of every single turn.\n\n"
140
- "2. `## Transcript`\n"
141
- " - The Transcript is the assistant's working memory. It MUST "
142
- "contain the last few turns of the conversation in full detail.\n"
143
- " - **CRITICAL REQUIREMENT:** The assistant's and user's last response "
144
- "MUST be COPIED VERBATIM into the Transcript. NEVER alter or truncate them for any reason"
145
- ).strip()
146
-
147
-
148
- _DEFAULT_CONTEXT_ENRICHMENT_PROMPT = (
149
- "You are a ruthless Memory Curator. Your only goal is to produce a "
150
- "dense, concise, and up-to-date Markdown block of long-term context. "
151
- "You MUST be aggressive in your curation to keep the context small.\n\n"
152
- "You will be given the `Previous Long-Term Context` and the `Recent "
153
- "Conversation History`. Your job is to return a NEW, UPDATED version of "
154
- "the `Long-Term Context` by following these rules:\n\n"
155
- "**Curation Rules (You MUST follow these):**\n"
156
- "1. **Integrate, Don't Append:** You MUST merge new facts from the "
157
- "`Recent Conversation History` into the `Previous Long-Term Context`. "
158
- "Rewrite existing facts to incorporate new details. DO NOT simply "
159
- "append new information at the end.\n"
160
- "2. **Discard Ephemeral Details:** You MUST delete temporary states, "
161
- "resolved errors, one-off requests, and conversational filler (e.g., 'Okay, "
162
- "I will do that now'). The context should be a snapshot of the current "
163
- "project state and user preferences, not a log.\n"
164
- "3. **Retain Stable Facts:** You MUST retain durable information such as "
165
- "user preferences (e.g., 'I prefer tabs over spaces'), project-level "
166
- "decisions, or architectural choices.\n"
167
- "4. **Mark Dynamic Info:** For temporary information that is critical to "
168
- "retain (e.g., CWD, the name of a file being actively edited), you MUST "
169
- "add a note: `(short-term, must be re-verified)`."
176
+ "You are a meticulous Conversation Historian agent. Your purpose is to "
177
+ "process the conversation history and update the assistant's memory "
178
+ "using your available tools. You will be given the previous summary, "
179
+ "previous notes, and the latest conversation turns in JSON format.\n\n"
180
+ "Follow these steps:\n\n"
181
+ "1. **Analyze the Recent Conversation:** Review the `Recent Conversation "
182
+ "(JSON)` to understand what just happened. Identify key facts, user "
183
+ "intentions, decisions made, and the final outcomes of any tasks.\n\n"
184
+ "2. **Update Long-Term Note:**\n"
185
+ " - Read the existing `Long Term` note to understand what is already "
186
+ "known.\n"
187
+ " - Identify any new, stable, and globally relevant information from "
188
+ "the recent conversation. This includes user preferences, high-level "
189
+ "goals, or facts that will be true regardless of the current working "
190
+ "directory.\n"
191
+ " - If you find such information, use the `write_long_term_note` tool "
192
+ "to save a concise, updated version of the note. Keep it brief and "
193
+ "factual.\n\n"
194
+ "3. **Update Contextual Note:**\n"
195
+ " - Read the existing `Contextual` note.\n"
196
+ " - Identify new information relevant *only* to the current project "
197
+ "or directory. This could be the file the user is working on, the "
198
+ "specific bug they are fixing, or the feature they are building.\n"
199
+ " - Use the `write_contextual_note` tool to save a concise, updated "
200
+ "note about the current working context. This note should be focused on "
201
+ "the immediate task at hand.\n\n"
202
+ "4. **Update Narrative Summary:**\n"
203
+ " - Review the `Past Conversation` summary.\n"
204
+ " - Create a new, condensed narrative that integrates the key "
205
+ "outcomes and decisions from the recent conversation. Discard "
206
+ "conversational filler. The summary should be a brief story of the "
207
+ "project's progress.\n"
208
+ " - Use the `write_past_conversation_summary` tool to save this new "
209
+ "summary.\n\n"
210
+ "5. **Update Transcript:**\n"
211
+ " - **CRITICAL:** Your final and most important task is to create a "
212
+ "transcript of the last few turns (around 4 turns).\n"
213
+ " - From the `Recent Conversation (JSON)`, extract the messages with "
214
+ "the role `user` and `assistant`. Ignore roles `system` and `tool`.\n"
215
+ " - Format the extracted messages into a readable dialog. For example:\n"
216
+ " User: <content of user message>\n"
217
+ " Assistant: <content of assistant message>\n"
218
+ " - If an assistant message contains `tool_calls`, note it like this:\n"
219
+ " Assistant (calling tool <tool_name>): <content of assistant message>\n"
220
+ " - The content of the user and assistant messages MUST be copied "
221
+ "verbatim. DO NOT alter, shorten, or summarize them in any way.\n"
222
+ " - Use the `write_past_conversation_transcript` tool to save this "
223
+ "formatted dialog string.\n\n"
224
+ "Your primary goal is to use your tools to persist these four distinct "
225
+ "pieces of information accurately and concisely."
170
226
  ).strip()
171
227
 
172
228
 
@@ -293,14 +349,6 @@ class LLMConfig:
293
349
  return CFG.LLM_SUMMARIZATION_PROMPT
294
350
  return _DEFAULT_SUMMARIZATION_PROMPT
295
351
 
296
- @property
297
- def default_context_enrichment_prompt(self) -> str:
298
- if self._default_context_enrichment_prompt is not None:
299
- return self._default_context_enrichment_prompt
300
- if CFG.LLM_CONTEXT_ENRICHMENT_PROMPT is not None:
301
- return CFG.LLM_CONTEXT_ENRICHMENT_PROMPT
302
- return _DEFAULT_CONTEXT_ENRICHMENT_PROMPT
303
-
304
352
  @property
305
353
  def default_model(self) -> "Model | str | None":
306
354
  if self._default_model is not None:
@@ -327,18 +375,6 @@ class LLMConfig:
327
375
  return self._default_history_summarization_token_threshold
328
376
  return CFG.LLM_HISTORY_SUMMARIZATION_TOKEN_THRESHOLD
329
377
 
330
- @property
331
- def default_enrich_context(self) -> bool:
332
- if self._default_enrich_context is not None:
333
- return self._default_enrich_context
334
- return CFG.LLM_ENRICH_CONTEXT
335
-
336
- @property
337
- def default_context_enrichment_token_threshold(self) -> int:
338
- if self._default_context_enrichment_token_threshold is not None:
339
- return self._default_context_enrichment_token_threshold
340
- return CFG.LLM_CONTEXT_ENRICHMENT_TOKEN_THRESHOLD
341
-
342
378
  def set_default_persona(self, persona: str):
343
379
  self._default_persona = persona
344
380
 
@@ -354,9 +390,6 @@ class LLMConfig:
354
390
  def set_default_summarization_prompt(self, summarization_prompt: str):
355
391
  self._default_summarization_prompt = summarization_prompt
356
392
 
357
- def set_default_context_enrichment_prompt(self, context_enrichment_prompt: str):
358
- self._default_context_enrichment_prompt = context_enrichment_prompt
359
-
360
393
  def set_default_model_name(self, model_name: str):
361
394
  self._default_model_name = model_name
362
395
 
@@ -382,16 +415,6 @@ class LLMConfig:
382
415
  history_summarization_token_threshold
383
416
  )
384
417
 
385
- def set_default_enrich_context(self, enrich_context: bool):
386
- self._default_enrich_context = enrich_context
387
-
388
- def set_default_context_enrichment_token_threshold(
389
- self, context_enrichment_token_threshold: int
390
- ):
391
- self._default_context_enrichment_token_threshold = (
392
- context_enrichment_token_threshold
393
- )
394
-
395
418
  def set_default_model_settings(self, model_settings: "ModelSettings"):
396
419
  self._default_model_settings = model_settings
397
420
 
@@ -9,8 +9,19 @@ from zrb.config.config import CFG
9
9
 
10
10
 
11
11
  def _estimate_token(text: str) -> int:
12
- enc = tiktoken.encoding_for_model("gpt-4o")
13
- return len(enc.encode(text))
12
+ """
13
+ Estimates the number of tokens in a given text.
14
+ Tries to use the 'gpt-4o' model's tokenizer for an accurate count.
15
+ If the tokenizer is unavailable (e.g., due to network issues),
16
+ it falls back to a heuristic of 4 characters per token.
17
+ """
18
+ try:
19
+ # Primary method: Use tiktoken for an accurate count
20
+ enc = tiktoken.encoding_for_model("gpt-4o")
21
+ return len(enc.encode(text))
22
+ except Exception:
23
+ # Fallback method: Heuristic (4 characters per token)
24
+ return len(text) // 4
14
25
 
15
26
 
16
27
  class LLMRateLimiter:
@@ -0,0 +1,128 @@
1
+ import json
2
+ from collections.abc import Callable
3
+ from copy import deepcopy
4
+ from typing import Any
5
+
6
+ from zrb.attr.type import StrAttr
7
+ from zrb.context.any_context import AnyContext
8
+ from zrb.context.any_shared_context import AnySharedContext
9
+ from zrb.task.llm.conversation_history_model import ConversationHistory
10
+ from zrb.task.llm.typing import ListOfDict
11
+ from zrb.util.attr import get_str_attr
12
+ from zrb.util.file import write_file
13
+ from zrb.util.run import run_async
14
+
15
+
16
+ def get_history_file(
17
+ ctx: AnyContext,
18
+ conversation_history_file_attr: StrAttr | None,
19
+ render_history_file: bool,
20
+ ) -> str:
21
+ """Gets the path to the conversation history file, rendering if configured."""
22
+ return get_str_attr(
23
+ ctx,
24
+ conversation_history_file_attr,
25
+ "",
26
+ auto_render=render_history_file,
27
+ )
28
+
29
+
30
+ async def read_conversation_history(
31
+ ctx: AnyContext,
32
+ conversation_history_reader: (
33
+ Callable[[AnySharedContext], ConversationHistory | dict | list | None] | None
34
+ ),
35
+ conversation_history_file_attr: StrAttr | None,
36
+ render_history_file: bool,
37
+ conversation_history_attr: (
38
+ ConversationHistory
39
+ | Callable[[AnySharedContext], ConversationHistory | dict | list]
40
+ | dict
41
+ | list
42
+ ),
43
+ ) -> ConversationHistory:
44
+ """Reads conversation history from reader, file, or attribute, with validation."""
45
+ history_file = get_history_file(
46
+ ctx, conversation_history_file_attr, render_history_file
47
+ )
48
+ # Use the class method defined above
49
+ history_data = await ConversationHistory.read_from_source(
50
+ ctx=ctx,
51
+ reader=conversation_history_reader,
52
+ file_path=history_file,
53
+ )
54
+ if history_data:
55
+ return history_data
56
+ # Priority 3: Callable or direct conversation_history attribute
57
+ raw_data_attr: Any = None
58
+ if callable(conversation_history_attr):
59
+ try:
60
+ raw_data_attr = await run_async(conversation_history_attr(ctx))
61
+ except Exception as e:
62
+ ctx.log_warning(
63
+ f"Error executing callable conversation_history attribute: {e}. "
64
+ "Ignoring."
65
+ )
66
+ if raw_data_attr is None:
67
+ raw_data_attr = conversation_history_attr
68
+ if raw_data_attr:
69
+ # Use the class method defined above
70
+ history_data = ConversationHistory.parse_and_validate(
71
+ ctx, raw_data_attr, "attribute"
72
+ )
73
+ if history_data:
74
+ return history_data
75
+ # Fallback: Return default value
76
+ return ConversationHistory()
77
+
78
+
79
+ async def write_conversation_history(
80
+ ctx: AnyContext,
81
+ history_data: ConversationHistory,
82
+ conversation_history_writer: (
83
+ Callable[[AnySharedContext, ConversationHistory], None] | None
84
+ ),
85
+ conversation_history_file_attr: StrAttr | None,
86
+ render_history_file: bool,
87
+ ):
88
+ """Writes conversation history using the writer or to a file."""
89
+ if conversation_history_writer is not None:
90
+ await run_async(conversation_history_writer(ctx, history_data))
91
+ history_file = get_history_file(
92
+ ctx, conversation_history_file_attr, render_history_file
93
+ )
94
+ if history_file != "":
95
+ write_file(history_file, json.dumps(history_data.to_dict(), indent=2))
96
+
97
+
98
+ def replace_system_prompt_in_history(
99
+ history_list: ListOfDict, replacement: str = "<main LLM system prompt>"
100
+ ) -> ListOfDict:
101
+ """
102
+ Returns a new history list where any part with part_kind 'system-prompt'
103
+ has its 'content' replaced with the given replacement string.
104
+ Args:
105
+ history: List of history items (each item is a dict with a 'parts' list).
106
+ replacement: The string to use in place of system-prompt content.
107
+
108
+ Returns:
109
+ A deep-copied list of history items with system-prompt content replaced.
110
+ """
111
+ new_history = deepcopy(history_list)
112
+ for item in new_history:
113
+ parts = item.get("parts", [])
114
+ for part in parts:
115
+ if part.get("part_kind") == "system-prompt":
116
+ part["content"] = replacement
117
+ return new_history
118
+
119
+
120
+ def count_part_in_history_list(history_list: ListOfDict) -> int:
121
+ """Calculates the total number of 'parts' in a history list."""
122
+ history_part_len = 0
123
+ for history in history_list:
124
+ if "parts" in history:
125
+ history_part_len += len(history["parts"])
126
+ else:
127
+ history_part_len += 1
128
+ return history_part_len