diff --git a/cecli/commands/tokens.py b/cecli/commands/tokens.py index 1ef30e6e825..349a89e00b2 100644 --- a/cecli/commands/tokens.py +++ b/cecli/commands/tokens.py @@ -1,4 +1,4 @@ -from typing import List +from typing import Any, Dict, List, Optional from cecli.commands.utils.base_command import BaseCommand from cecli.commands.utils.helpers import format_command_result @@ -9,6 +9,48 @@ class TokensCommand(BaseCommand): NORM_NAME = "tokens" DESCRIPTION = "Report on the number of tokens used by the current chat context" + @classmethod + def _extract_file_name(cls, msg: Dict[str, Any]) -> Optional[str]: + """Extract file name from a message dictionary.""" + if not isinstance(msg, dict): + return None + + # Check explicit image_file metadata first + fname = msg.get("image_file") + if fname: + return fname + + content = msg.get("content") + if isinstance(content, str): + if content.startswith(("Original File Contents For", "Current File Contents For")): + lines = content.split("\n", 3) + if len(lines) > 1: + return lines[1].strip() + elif content.startswith("Image file: "): + return content[len("Image file: "):].strip() + elif isinstance(content, list): + for part in content: + if isinstance(part, dict): + if part.get("image_file"): + return part.get("image_file") + text = part.get("text") + if isinstance(text, str): + if text.startswith(("Original File Contents For", "Current File Contents For")): + lines = text.split("\n", 3) + if len(lines) > 1: + return lines[1].strip() + elif text.startswith("Image file: "): + return text[len("Image file: "):].strip() + elif isinstance(part, str): + if part.startswith(("Original File Contents For", "Current File Contents For")): + lines = part.split("\n", 3) + if len(lines) > 1: + return lines[1].strip() + elif part.startswith("Image file: "): + return part[len("Image file: "):].strip() + + return None + @classmethod async def execute(cls, io, coder, args, **kwargs): res = [] @@ -123,27 +165,10 @@ async def execute(cls, io, coder, args, **kwargs): # Group messages by file (each file has user and assistant messages) file_tokens = {} for msg in readonly_msgs: - # Extract file name from message content - content = msg.get("content", "") - if content.startswith("Original File Contents For"): - # Extract file path from "File Contents {path}:" - lines = content.split("\n", 3) - if lines: - file_line = lines[1] - fname = file_line.strip() - # Calculate tokens for this message - tokens = coder.main_model.token_count([msg]) - if fname not in file_tokens: - file_tokens[fname] = 0 - file_tokens[fname] += tokens - elif "image_file" in msg: - # Handle image files - fname = msg.get("image_file") - if fname: - tokens = coder.main_model.token_count([msg]) - if fname not in file_tokens: - file_tokens[fname] = 0 - file_tokens[fname] += tokens + fname = cls._extract_file_name(msg) + if fname: + tokens = coder.main_model.token_count([msg]) + file_tokens[fname] = file_tokens.get(fname, 0) + tokens # Add to results for fname, tokens in file_tokens.items(): @@ -158,27 +183,10 @@ async def execute(cls, io, coder, args, **kwargs): msgs = ConversationService.get_manager(coder).get_messages_dict(tag=tag) if msgs: for msg in msgs: - # Extract file name from message content - content = msg.get("content", "") - if content.startswith("Original File Contents For"): - # Extract file path from "File Contents {path}:" - lines = content.split("\n", 3) - if lines: - file_line = lines[1] - fname = file_line.strip() - # Calculate tokens for this message - tokens = coder.main_model.token_count([msg]) - if fname not in editable_file_tokens: - editable_file_tokens[fname] = 0 - editable_file_tokens[fname] += tokens - elif "image_file" in msg: - # Handle image files - fname = msg.get("image_file") - if fname: - tokens = coder.main_model.token_count([msg]) - if fname not in editable_file_tokens: - editable_file_tokens[fname] = 0 - editable_file_tokens[fname] += tokens + fname = cls._extract_file_name(msg) + if fname: + tokens = coder.main_model.token_count([msg]) + editable_file_tokens[fname] = editable_file_tokens.get(fname, 0) + tokens # Add editable files to results for fname, tokens in editable_file_tokens.items(): diff --git a/tests/basic/test_commands.py b/tests/basic/test_commands.py index c367cafefd1..b20d723036f 100644 --- a/tests/basic/test_commands.py +++ b/tests/basic/test_commands.py @@ -420,6 +420,123 @@ async def test_cmd_tokens(self): self.assertIn("foo.txt", console_output) self.assertIn("bar.txt", console_output) + async def test_cmd_tokens_with_image_list_content(self): + from cecli.commands.tokens import TokensCommand + from cecli.helpers.conversation import ConversationService, MessageTag + + # Direct unit tests for _extract_file_name + self.assertIsNone(TokensCommand._extract_file_name(None)) + self.assertIsNone(TokensCommand._extract_file_name("not a dict")) + self.assertIsNone(TokensCommand._extract_file_name({})) + self.assertIsNone(TokensCommand._extract_file_name({"content": None})) + self.assertIsNone(TokensCommand._extract_file_name({"content": "Just a normal message"})) + + # String content + self.assertEqual( + TokensCommand._extract_file_name({ + "content": "Original File Contents For:\n/path/to/file.py\n\ncode..." + }), + "/path/to/file.py", + ) + self.assertEqual( + TokensCommand._extract_file_name({ + "content": "Current File Contents For:\n/path/to/file2.py\n\ncode..." + }), + "/path/to/file2.py", + ) + self.assertEqual( + TokensCommand._extract_file_name({"content": "Image file: photo.png"}), + "photo.png", + ) + + # image_file key + self.assertEqual( + TokensCommand._extract_file_name({"image_file": "photo.jpg"}), + "photo.jpg", + ) + self.assertEqual( + TokensCommand._extract_file_name({ + "image_file": "photo.jpg", + "content": [{"type": "image_url", "image_url": {}}], + }), + "photo.jpg", + ) + + # List content + self.assertEqual( + TokensCommand._extract_file_name({ + "content": [ + {"type": "text", "text": "Image file: nested.png"}, + {"type": "image_url", "image_url": {}}, + ] + }), + "nested.png", + ) + self.assertEqual( + TokensCommand._extract_file_name({ + "content": [ + {"type": "text", "text": "Original File Contents For:\nmodule.py\n\ncode"} + ] + }), + "module.py", + ) + self.assertEqual( + TokensCommand._extract_file_name({ + "content": [{"image_file": "from_part.png"}] + }), + "from_part.png", + ) + self.assertEqual( + TokensCommand._extract_file_name({ + "content": ["Image file: string_part.png"] + }), + "string_part.png", + ) + + # Integration test with Coder / Commands execution + io = InputOutput(pretty=False, fancy_input=False, yes=True) + coder = await Coder.create(self.GPT35, None, io) + commands = Commands(io, coder) + + manager = ConversationService.get_manager(coder) + img_msg = { + "role": "user", + "content": [ + {"type": "text", "text": "Image file: test_image.png"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,123"}}, + ], + "image_file": "test_image.png", + } + manager.add_message( + message_dict=img_msg, + tag=MessageTag.READONLY_FILES, + hash_key=("image_user", "test_image.png"), + ) + + img_msg_no_attr = { + "role": "user", + "content": [ + {"type": "text", "text": "Image file: another_image.png"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,456"}}, + ], + } + manager.add_message( + message_dict=img_msg_no_attr, + tag=MessageTag.CHAT_FILES, + hash_key=("image_user", "another_image.png"), + ) + + stdout = StringIO() + sys.stdout = stdout + try: + commands.execute("tokens", "") + finally: + sys.stdout = sys.__stdout__ + + console_output = stdout.getvalue() + self.assertIn("test_image.png", console_output) + self.assertIn("another_image.png", console_output) + async def test_cmd_add_from_subdir(self): repo = git.Repo.init() repo.config_writer().set_value("user", "name", "Test User").release()