feat: added token counter command to CLI (#2047)

This commit is contained in:
Charles Packer 2024-11-17 14:02:21 -08:00 committed by GitHub
parent 82058d2077
commit e5a7fc406d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 13 additions and 0 deletions

View File

@ -1582,6 +1582,10 @@ class Agent(BaseAgent):
functions_definitions=available_functions_definitions,
)
def count_tokens(self) -> int:
"""Count the tokens in the current context window"""
return self.get_context_window().context_window_size_current
def save_agent(agent: Agent, ms: MetadataStore):
"""Save agent to metadata store"""

View File

@ -254,6 +254,15 @@ def run_agent_loop(
)
continue
elif user_input.lower() == "/tokens":
tokens = letta_agent.count_tokens()
typer.secho(
f"{tokens}/{letta_agent.agent_state.llm_config.context_window}",
fg=typer.colors.GREEN,
bold=True,
)
continue
elif user_input.lower().startswith("/add_function"):
try:
if len(user_input) < len("/add_function "):