From e5a7fc406d06f179a2ba47bf517a759055252c4d Mon Sep 17 00:00:00 2001
From: Charles Packer <packercharles@gmail.com>
Date: Sun, 17 Nov 2024 14:02:21 -0800
Subject: [PATCH] feat: added token counter command to CLI (#2047)

---
 letta/agent.py | 4 ++++
 letta/main.py  | 9 +++++++++
 2 files changed, 13 insertions(+)

diff --git a/letta/agent.py b/letta/agent.py
index f7bd6f0fb4..455d9eedc1 100644
--- 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
index 70595d6c80..abfd36ae9c 100644
--- 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 "):