Add task summary

Author: lgxz

FEATURES.md

@@ -1,5 +1,12 @@
 # Features
 
+# usage-limit
+- 完善使用情况统计：支持各家 API，支持细粒度
+- 时间统计
+- 对话次数统计
+
+可以配置三种数据的最大值，超过就停止任务。
+
 ## shiv
 支持 shiv 打包
 

aipython/aipy/agent.py

@@ -11,6 +11,7 @@
 from rich.live import Live
 from rich.panel import Panel
 from rich.align import Align
+from rich.table import Table
 from rich.syntax import Syntax
 from rich.console import Console
 from rich.markdown import Markdown
@@ -173,6 +174,32 @@ def box(self, title, content, align=None, lang=None):
             content = Align(content, align=align)
         self._console.print(Panel(content, title=title))
 
+    def print_summary(self):
+        history = self.llm.history
+        table = Table(title=T("Task Summary"), show_lines=True)
+
+        table.add_column(T("Round"), justify="center", style="bold cyan", no_wrap=True)
+        table.add_column(T("Time(s)"), justify="right")
+        table.add_column(T("In Tokens"), justify="right")
+        table.add_column(T("Out Tokens"), justify="right")
+        table.add_column(T("Total Tokens"), justify="right", style="bold magenta")
+
+        round = 1
+        for row in history.get_usage():
+            table.add_row(
+                str(round),
+                str(row["time"]),
+                str(row["input_tokens"]),
+                str(row["output_tokens"]),
+                str(row["total_tokens"]),
+            )
+            round += 1
+        self._console.print("\n")
+        self._console.print(table)
+        summary = history.get_summary()
+        summary = "{rounds} | {time}s | Tokens: {input_tokens}/{output_tokens}/{total_tokens}".format(**summary)
+        self._console.print(f"\n⏹ [cyan]{T('end_instruction')} | {summary}")
+    
     def __call__(self, instruction, llm=None):
         """
         执行自动处理循环，直到 LLM 不再返回代码消息
@@ -188,8 +215,7 @@ def __call__(self, instruction, llm=None):
             if msg['type'] != MsgType.CODE:
                 break
             response = self.process_code_reply(msg, llm)
-        total_token = self.llm.history.total_tokens
-        self._console.print(f"\n⏹ {T('end_instruction')} | Tokens: {total_token}")
+        self.print_summary()
         os.write(1, b'\a\a\a')
 
     def chat(self, prompt):

aipython/aipy/i18n.py

@@ -29,7 +29,7 @@
         'publish_disabled': "当前环境不支持发布",
         'auto_confirm': '自动确认',
         'packages_exist': '申请的第三方包已经安装',
-        'thinking': '正在努力思考中，请稍等6-60秒',
+        'thinking': '正在绞尽脑汁思考中，请稍等6-60秒',
         'no_available_llm': '没有可用的 LLM，请检查配置文件',
         'banner1_python': "请用 ai('任务') 输入需要 AI 处理的任务 (输入 ai.use(llm) 切换 下述 LLM：",
         'banner1': "请输入需要 AI 处理的任务 (输入 /use llm 切换 下述LLM)",
@@ -38,6 +38,12 @@
         'ai_mode_enter': '进入 AI 模式，开始处理任务，输入 Ctrl+d 或 /done 结束任务',
         'ai_mode_exit': "[退出 AI 模式]",
         'ai_mode_unknown_command': "[AI 模式] 未知命令",
+        'Task Summary': '任务总结',
+        'Round': '轮次',
+        'Time(s)': '时间(秒)',
+        'In Tokens': '输入令牌数',
+        'Out Tokens': '输出令牌数',
+        'Total Tokens': '总令牌数',
     },
     'en': {
         'start_instruction': 'Start processing instruction',
@@ -73,6 +79,12 @@
         'ai_mode_enter': 'Enter AI mode, start processing tasks, enter Ctrl+d or /done to end the task',
         'ai_mode_exit': "[Exit AI mode]",
         'ai_mode_unknown_command': "[AI mode] Unknown command",
+        'Task Summary': 'Task Summary',
+        'Round': 'Round',
+        'Time(s)': 'Time(s)',
+        'In Tokens': 'In Tokens',
+        'Out Tokens': 'Out Tokens',
+        'Total Tokens': 'Total Tokens',
     }
 }
 

aipython/aipy/llm.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
+import time
 from collections import Counter, defaultdict
 from dataclasses import dataclass, field
 from abc import ABC, abstractmethod
@@ -21,7 +22,7 @@ class ChatMessage:
 class ChatHistory:
     def __init__(self):
         self.messages = []
-        self._total_tokens = 0
+        self._total_tokens = Counter()
 
     def __len__(self):
         return len(self.messages)
@@ -31,11 +32,15 @@ def add(self, role, content):
 
     def add_message(self, message: ChatMessage):
         self.messages.append(message)
-        self._total_tokens += message.usage["total_tokens"]
+        self._total_tokens += message.usage
 
-    @property
-    def total_tokens(self):
-        return self._total_tokens
+    def get_usage(self):
+        return iter(row.usage for row in self.messages if row.role == "assistant")
+    
+    def get_summary(self):
+        summary = dict(self._total_tokens)
+        summary['rounds'] = sum(1 for row in self.messages if row.role == "assistant")
+        return summary
 
     def get_messages(self):
         return [{"role": msg.role, "content": msg.content} for msg in self.messages]
@@ -60,6 +65,10 @@ def get_completion(self, messages):
     def add_system_prompt(self, history, system_prompt):
         history.add("system", system_prompt)
 
+    @abstractmethod
+    def parse_usage(self, response):
+        pass
+
     @abstractmethod
     def parse_response(self, response):
         pass
@@ -70,16 +79,23 @@ def __call__(self, history, prompt, system_prompt=None):
             self.add_system_prompt(history, system_prompt)
         history.add("user", prompt)
 
+        start = time.time()
         response = self.get_completion(history.get_messages())
+        end = time.time()
         if response:
             msg = self.parse_response(response)
+            usage = self.parse_usage(response)
+            usage['time'] = round(end - start, 3)
+            msg.usage = usage
             history.add_message(msg)
             if msg.reason:
                 response = f"{T('think')}:\n---\n{msg.reason}\n---\n{msg.content}"
             else:
                 response = msg.content
         return response
-    
+
+# https://platform.openai.com/docs/api-reference/chat/create
+# https://api-docs.deepseek.com/api/create-chat-completion
 class OpenAIClient(BaseClient):
     def __init__(self, config):
         super().__init__(config)
@@ -88,15 +104,20 @@ def __init__(self, config):
     def add_system_prompt(self, history, system_prompt):
         history.add("system", system_prompt)
 
+    def parse_usage(self, response):
+        usage = response.usage
+        return Counter({'total_tokens': usage.total_tokens,
+                'input_tokens': usage.prompt_tokens,
+                'output_tokens': usage.completion_tokens})
+    
     def parse_response(self, response):
-        usage = response.usage.model_dump()
         message = response.choices[0].message
         reason = getattr(message, "reasoning_content", None)
         return ChatMessage(
             role=message.role,
             content=message.content,
-            reason=reason,
-            usage=Counter(usage))
+            reason=reason
+        )
 
     def get_completion(self, messages):
         try:
@@ -110,12 +131,18 @@ def get_completion(self, messages):
             self.console.print(f"❌ [bold red]{self.name} API {T('call_failed')}: [yellow]{str(e)}")
             response = None
         return response
-    
+
+# https://github.com/ollama/ollama/blob/main/docs/api.md
 class OllamaClient(BaseClient):
     def __init__(self, config):
         super().__init__(config)
         self._session = requests.Session()
 
+    def parse_usage(self, response):
+        ret = Counter({'input_tokens': response['prompt_eval_count'], 'output_tokens': response['eval_count']})
+        ret['total_tokens'] = ret['input_tokens'] + ret['output_tokens']
+        return ret
+    
     def parse_response(self, response):
         msg = response["message"]
         return ChatMessage(role=msg['role'], content=msg['content'])
@@ -139,19 +166,22 @@ def get_completion(self, messages):
             response = None
         return response
 
+# https://docs.anthropic.com/en/api/messages
 class ClaudeClient(BaseClient):
     def __init__(self, config):
         super().__init__(config)
         self._client = anthropic.Anthropic(api_key=self._api_key, timeout=self._timeout)
 
+    def parse_usage(self, response):
+        usage = response.usage
+        ret = Counter({'input_tokens': usage.input_tokens, 'output_tokens': usage.output_tokens})
+        ret['total_tokens'] = ret['input_tokens'] + ret['output_tokens']
+        return ret
+    
     def parse_response(self, response):
-        usage = Counter(response.usage)
         content = response.content[0].text
         role = response.role
-        return ChatMessage(
-            role=role,
-            content=content,
-            usage=usage)
+        return ChatMessage(role=role, content=content)
 
     def add_system_prompt(self, history, system_prompt):
         self._system_prompt = system_prompt

aipython/saas.py

@@ -42,10 +42,16 @@ def input_with_possible_multiline(self, prompt_text, is_ai=False):
                 break
         return "\n".join(lines)
 
+    def run_ai_task(self, task):
+        try:
+            self.ai(task)
+        except Exception as e:
+            self.console.print(f"[bold red]Error: {e}")
+
     def run_ai_mode(self, initial_text):
         ai = self.ai
         self.console.print(f"{T('ai_mode_enter')}", style="cyan")
-        ai(initial_text)
+        self.run_ai_task(initial_text)
         while True:
             try:
                 user_input = self.input_with_possible_multiline(">>> ", is_ai=True).strip()
@@ -63,7 +69,7 @@ def run_ai_mode(self, initial_text):
                 else:
                     self.console.print(f"{T('ai_mode_unknown_command')}", style="cyan")
             else:
-                ai(user_input)
+                self.run_ai_task(user_input)
         try:
             ai.publish(verbose=False)
         except Exception as e:

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "aipython"
-version = "0.1.6"
+version = "0.1.7"
 description = "AIPython: AI-Powered Python & Python-Powered AI"
 readme = "README.md"
 requires-python = ">=3.11"