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

letta-ai · Nov 17, 2024 · e5a7fc4 · e5a7fc4
1 parent 82058d2
commit e5a7fc4
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/letta/agent.py b/letta/agent.py
@@ -1582,6 +1582,10 @@ def get_context_window(self) -> ContextWindowOverview:
             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"""

diff --git a/letta/main.py b/letta/main.py
@@ -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 "):