KeyArgo
/
smolagents
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Change rule color and end of step logs

This commit is contained in:
Aymeric 2024-12-24 13:00:24 +01:00
parent 48b7f53a54
commit 762ae9cfae
2 changed files with 13 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
src/smolagents/agents.py
View File

@ -171,6 +171,7 @@ def format_prompt_with_managed_agents_descriptions(
else: else:
return prompt_template.replace(agent_descriptions_placeholder, "") return prompt_template.replace(agent_descriptions_placeholder, "")
YELLOW_HEX = "#ffdd00"
class MultiStepAgent: class MultiStepAgent:
""" """
@ -498,7 +499,7 @@ class MultiStepAgent:
else: else:
self.logs.append(system_prompt_step) self.logs.append(system_prompt_step)
console.print(Group(Rule("[bold]New task", characters="="), Text(self.task))) console.print(Group(Rule("[bold]New run", characters="═", style=YELLOW_HEX), Text(self.task)))
self.logs.append(TaskStep(task=self.task)) self.logs.append(TaskStep(task=self.task))
if single_step: if single_step:
@ -533,7 +534,7 @@ class MultiStepAgent:
self.planning_step( self.planning_step(
task, is_first_step=(iteration == 0), iteration=iteration task, is_first_step=(iteration == 0), iteration=iteration
) )
console.rule(f"[bold]Step {iteration}") console.print(Rule(f"[bold]Step {iteration}", characters="━", style=YELLOW_HEX))
# Run one step! # Run one step!
final_answer = self.step(step_log) final_answer = self.step(step_log)
@ -579,7 +580,7 @@ class MultiStepAgent:
self.planning_step( self.planning_step(
task, is_first_step=(iteration == 0), iteration=iteration task, is_first_step=(iteration == 0), iteration=iteration
) )
console.rule(f"[bold]Step {iteration}") console.print(Rule(f"[bold]Step {iteration}", characters="━", style=YELLOW_HEX))
# Run one step! # Run one step!
final_answer = self.step(step_log) final_answer = self.step(step_log)

18
src/smolagents/monitoring.py
View File

@ -30,13 +30,15 @@ class Monitor:
self.total_input_token_count = 0 self.total_input_token_count = 0
self.total_output_token_count = 0 self.total_output_token_count = 0
def reset(self):
self.step_durations = []
self.total_input_token_count = 0
self.total_output_token_count = 0
def update_metrics(self, step_log): def update_metrics(self, step_log):
step_duration = step_log.duration step_duration = step_log.duration
self.step_durations.append(step_duration) self.step_durations.append(step_duration)
console_outputs = [ console_outputs = f"[Step {len(self.step_durations)-1}: Duration {step_duration:.2f} seconds"
Text(f"Step {len(self.step_durations)}:", style="bold"),
Text(f"- Time taken: {step_duration:.2f} seconds"),
]
if getattr(self.tracked_llm_engine, "last_input_token_count", None) is not None: if getattr(self.tracked_llm_engine, "last_input_token_count", None) is not None:
self.total_input_token_count += ( self.total_input_token_count += (
@ -45,11 +47,9 @@ class Monitor:
self.total_output_token_count += ( self.total_output_token_count += (
self.tracked_llm_engine.last_output_token_count self.tracked_llm_engine.last_output_token_count
) )
console_outputs += [ console_outputs += f"| Input tokens: {self.total_input_token_count:,} | Output tokens: {self.total_output_token_count:,}"
Text(f"- Input tokens: {self.total_input_token_count:,}"), console_outputs += "]"
Text(f"- Output tokens: {self.total_output_token_count:,}"), console.print(Text(console_outputs, style="dim"))
]
console.print(Group(*console_outputs))
__all__ = ["Monitor"] __all__ = ["Monitor"]