Merge pull request #13 from minorcell/dev

feat: optimized token counting, added memory capability

Author: minorcell

docs/config-storage.md

@@ -10,6 +10,7 @@
   sessions/
     25/12/01/UUID.jsonl # 按 YY/MM/DD 分桶的会话日志（JSONL 事件）
     25/12/02/UUID.jsonl
+  memory.md             # 长期记忆（由 memory 工具追加，随系统提示词注入）
 ```
 
 - 路径可通过环境变量 `MEMO_HOME` 覆盖，默认为 `~/.memo`。
@@ -22,6 +23,8 @@
 current_provider = "deepseek"
 # 每个 trun 最大的步骤数
 max_steps = 100
+# 是否启用流式输出（默认 true）
+stream_output = true
 
 [[providers]]
 name = "deepseek"
@@ -44,6 +47,8 @@ base_url = "https://api.openai.com/v1"
     - `env_api_key`: 读取的环境变量名称（不在文件中存密钥）。
     - `model`: 默认模型名。
     - `base_url`: 可选，兼容 OpenAI 生态；缺省则使用 SDK 默认。
+- `stream_output`: 是否默认使用 LLM 流式输出（stream=true）。
+- `memory.md`: 独立文件存储长期记忆，由 `memory` 工具追加，Session 启动时附加到系统提示词。
 
 ## 核心改造要点
 

docs/core.md

@@ -6,16 +6,12 @@
 
 - `config/`：配置与路径
     - `config.ts`：读取/写入 `~/.memo/config.toml`，provider 选择（name/env_api_key/model/base_url）、会话路径 `sessions/YY/MM/DD/<uuid>.jsonl`、sessionId 生成。
-    - `constants.ts`：兜底默认（如 MAX_STEPS）。
 - `runtime/`：运行时与日志
     - `prompt.md/prompt.ts`：系统提示词加载（内容为 JSON 协议说明）。
     - `history.ts`：JSONL sink 与事件构造。
     - `defaults.ts`：补全工具集、LLM、prompt、history sink、tokenizer、maxSteps（基于配置）。
     - `session.ts`：Session/Turn 状态机，执行 ReAct 循环、写事件、统计 token、触发回调。
-- `llm/`：模型与 tokenizer
-    - `openai.ts`：OpenAI SDK 适配，`createOpenAIClient(provider)` 基于配置生成客户端，兼容 DeepSeek/OpenAI。
-    - `tokenizer.ts`：tiktoken 封装。
-- `utils/`：解析工具（assistant 输出解析、消息包装）。
+- `utils/`：解析工具（assistant 输出解析、消息包装）与 tokenizer 封装（tiktoken）。
 - `types.ts`：公共类型。
 - `index.ts`：包入口，导出上述模块。
 
@@ -55,10 +51,10 @@ await session.close()
 - 默认写入 `~/.memo/sessions/YY/MM/DD/<uuid>.jsonl`；包含 provider、模型、tokenizer、token 用量等元数据。
 - 仅写 JSONL 事件，默认写入 `~/.memo/sessions/YY/MM/DD/<uuid>.jsonl`。
 
-## LLM 适配（llm/openai.ts）
+## LLM 适配（runtime/defaults.ts）
 
-- `createOpenAIClient(provider)` 基于 provider 配置（env_api_key/model/base_url）构建 OpenAI SDK 客户端。
-- 环境变量只存 API key，不写入配置；可用 `OPENAI_BASE_URL/OPENAI_MODEL` 临时覆盖。
+- 通过 `withDefaultDeps` 内置基于 OpenAI SDK 的调用（按配置选择 provider、model、base_url、env_api_key）。
+- 优先使用传入的 `callLLM`，否则读取环境变量（当前 provider 的 env_api_key/OPENAI_API_KEY/DEEPSEEK_API_KEY）。
 
 ## 工具协议与注册表
 

docs/memory.md

@@ -0,0 +1,35 @@
+# Memory（长期记忆）设计
+
+## 目标
+
+- 为 memo-cli 提供跨 Session 的长期记忆，存储在用户目录（`~/.memo/memory.md`）。
+- 在每次 Session 启动时，将 memory 内容随系统提示词注入，让模型可利用用户画像/偏好。
+- 提供 `memory` 工具，允许模型在给出最终回答前写入一条极短的摘要。
+
+## 存储
+
+- 路径：`MEMO_HOME` 下的 `memory.md`（默认 `~/.memo/memory.md`）。
+- 结构：追加式 Markdown 列表（`- <ISO 时间> <摘要>`）。
+- 创建：不存在时自动创建所在目录。
+
+## 注入时机
+
+- `createAgentSession` 启动时读取 `memory.md`，非空则在系统提示词末尾附加「长期记忆」段落。
+- memory 内容不可在回答中直接泄露，只作为参考。
+
+## memory 工具
+
+- 名称：`memory`
+- 输入：`{"note":"一句极短摘要，≤32 字符"}`，自动清理多余空白/换行并追加到 `memory.md`。
+- 使用时机：仅在准备返回 final 之前，如发现可复用的用户画像/习惯/偏好信息可调用一次，避免频繁调用。
+- 输出：确认写入的文件路径，失败时返回错误信息。
+
+## 使用建议
+
+- 摘要要短、泛化，避免敏感信息与原文（例：`用户常用中文回答，偏好简洁步骤`）。
+- 可后续迭代：重复检测/压缩、清理/导出命令、权限控制（600）。
+
+## 测试覆盖
+
+- memory 工具：输入校验、写入与内容清理。
+- 注入：存在 `memory.md` 时系统提示词包含记忆内容。

docs/token-counting.md

@@ -0,0 +1,39 @@
+# memo-cli 的 Token 计数策略
+
+本文说明 memo-cli 如何估算与记录 token，用于提示词预算、超限防护和 usage 对账。
+
+## 计数实现
+
+- **底层编码器**：使用 `@dqbd/tiktoken`，默认 encoding `cl100k_base`；可通过 `tokenizerModel` 覆盖。
+- **纯文本计数**：`countText(text)` 直接对字符串编码，获得长度。
+- **消息数组计数（ChatML 近似）**：`countMessages(messages)` 采用 OpenAI ChatML 常用估算：
+  - 每条消息固定开销 4 tokens（role/name 封装等）。
+  - `content` 本身按 tiktoken 编码计数。
+  - 如果后续支持 `name` 字段，则额外加 1 token。
+  - 末尾补充 2 tokens 用于 assistant priming。
+
+该策略比单纯串联文本更贴近实际 ChatML 开销，但仍是近似值。
+
+## 使用场景
+
+- **提示词预算**：`runTurn` 会在每步前用 `countMessages` 估算 prompt tokens，触发：
+  - `warnPromptTokens`：打印警告。
+  - `maxPromptTokens`：超过时立即返回提示，避免向 LLM 发送超限请求。
+- **usage 对账**：每步会结合本地计数与模型返回的 `usage`（若提供），记录到 tokenUsage 与 JSONL 历史事件。
+
+## 精度与局限
+
+- ChatML 不同模型的固定开销略有差异，当前采用「每消息 4 + 末尾 2」的通用估算，可能与特定模型存在 ±几十 token 的偏差。
+- 未显式处理工具调用/函数调用的额外结构开销，如需精确对账可在未来针对具体模型定制 constants。
+- 若使用自定义 `callLLM`，建议传入对应模型的 encoding 名称或自定义 `tokenCounter` 以匹配真实开销。
+
+## 如何覆盖
+
+- 创建 Session 时传入 `tokenizerModel` 或直接注入自定义 `tokenCounter`：
+```ts
+import { createTokenCounter, createAgentSession } from '@memo/core'
+
+const tokenCounter = createTokenCounter('gpt-4o-mini')
+const session = await createAgentSession({ tokenCounter }, { warnPromptTokens: 8_000 })
+```
+- 自定义 counter 只需实现 `countText`、`countMessages`、`dispose` 接口。

packages/core/src/config/config.ts

@@ -15,15 +15,18 @@ export type ProviderConfig = {
 export type MemoConfig = {
     current_provider: string
     max_steps?: number
+    stream_output?: boolean
     providers: ProviderConfig[]
 }
 
 const DEFAULT_MEMO_HOME = join(homedir(), '.memo')
 const DEFAULT_SESSIONS_DIR = 'sessions'
+const DEFAULT_MEMORY_FILE = 'memory.md'
 
 const DEFAULT_CONFIG: MemoConfig = {
     current_provider: 'deepseek',
     max_steps: 100,
+    stream_output: true,
     providers: [
         {
             name: 'deepseek',
@@ -55,6 +58,7 @@ ${p.base_url ? `base_url = "${p.base_url}"\n` : ''}`,
     return `# Memo config. Edit to change provider/model/base_url.
 current_provider = "${config.current_provider}"
 max_steps = ${config.max_steps ?? 100}
+stream_output = ${config.stream_output ?? true}
 
 ${providers}`.trim()
 }
@@ -84,6 +88,7 @@ export async function loadMemoConfig(): Promise<LoadedConfig> {
         const merged: MemoConfig = {
             current_provider: parsed.current_provider ?? DEFAULT_CONFIG.current_provider,
             max_steps: parsed.max_steps ?? DEFAULT_CONFIG.max_steps,
+            stream_output: parsed.stream_output ?? DEFAULT_CONFIG.stream_output,
             providers: parsed.providers ?? [],
         }
         const needsSetup = !merged.providers.length
@@ -105,6 +110,10 @@ export function getSessionsDir(loaded: LoadedConfig, options: AgentSessionOption
     return expandHome(base)
 }
 
+export function getMemoryPath(loaded: LoadedConfig) {
+    return join(loaded.home, DEFAULT_MEMORY_FILE)
+}
+
 export function buildSessionPath(baseDir: string, sessionId: string) {
     const now = new Date()
     const yy = String(now.getFullYear()).slice(2)

packages/core/src/runtime/defaults.ts

@@ -4,6 +4,7 @@ import { createTokenCounter } from '@memo/core/utils/tokenizer'
 import {
     buildSessionPath,
     getSessionsDir,
+    getMemoryPath,
     loadMemoConfig,
     selectProvider,
 } from '@memo/core/config/config'
@@ -37,12 +38,28 @@ export async function withDefaultDeps(
     const loaded = await loadMemoConfig()
     const config = loaded.config
     const tools = deps.tools ?? TOOLKIT
-    const loadPrompt = deps.loadPrompt ?? defaultLoadPrompt
+    const loadPrompt = async () => {
+        const basePrompt = await (deps.loadPrompt ?? defaultLoadPrompt)()
+        const memoryPath = getMemoryPath(loaded)
+        try {
+            const file = Bun.file(memoryPath)
+            if (await file.exists()) {
+                const memory = (await file.text()).trim()
+                if (memory) {
+                    return `${basePrompt}\n\n# 长期记忆\n${memory}`
+                }
+            }
+        } catch (err) {
+            console.warn(`读取 memory 失败: ${(err as Error).message}`)
+        }
+        return basePrompt
+    }
+    const streamOutput = options.stream ?? config.stream_output ?? true
     return {
         tools,
         callLLM:
             deps.callLLM ??
-            (async (messages) => {
+            (async (messages, onChunk) => {
                 const provider = selectProvider(config, options.providerName)
                 const apiKey =
                     process.env[provider.env_api_key] ??
@@ -57,22 +74,40 @@ export async function withDefaultDeps(
                     apiKey,
                     baseURL: provider.base_url,
                 })
-                const data = await client.chat.completions.create({
-                    model: provider.model,
-                    messages,
-                    temperature: 0.35,
-                })
-                const content = data.choices?.[0]?.message?.content
-                if (typeof content !== 'string') {
-                    throw new Error('OpenAI 兼容接口返回内容为空')
-                }
-                return {
-                    content,
-                    usage: {
-                        prompt: data.usage?.prompt_tokens ?? undefined,
-                        completion: data.usage?.completion_tokens ?? undefined,
-                        total: data.usage?.total_tokens ?? undefined,
-                    },
+                if (streamOutput) {
+                    const stream = await client.chat.completions.create({
+                        model: provider.model,
+                        messages,
+                        temperature: 0.35,
+                        stream: true,
+                    })
+                    let content = ''
+                    for await (const part of stream) {
+                        const delta = part.choices?.[0]?.delta?.content
+                        if (delta) {
+                            content += delta
+                            onChunk?.(delta)
+                        }
+                    }
+                    return { content, streamed: true }
+                } else {
+                    const data = await client.chat.completions.create({
+                        model: provider.model,
+                        messages,
+                        temperature: 0.35,
+                    })
+                    const content = data.choices?.[0]?.message?.content
+                    if (typeof content !== 'string') {
+                        throw new Error('OpenAI 兼容接口返回内容为空')
+                    }
+                    return {
+                        content,
+                        usage: {
+                            prompt: data.usage?.prompt_tokens ?? undefined,
+                            completion: data.usage?.completion_tokens ?? undefined,
+                            total: data.usage?.total_tokens ?? undefined,
+                        },
+                    }
                 }
             }),
         loadPrompt,

packages/core/src/runtime/memory.test.ts

@@ -0,0 +1,57 @@
+import assert from 'node:assert'
+import { join } from 'node:path'
+import { tmpdir } from 'node:os'
+import { describe, test, beforeAll, afterAll } from 'bun:test'
+import { $ } from 'bun'
+import { createAgentSession, createTokenCounter } from '@memo/core'
+
+let tempHome: string
+let prevMemoHome: string | undefined
+
+async function makeTempDir(prefix: string) {
+    const dir = join(tmpdir(), `${prefix}-${crypto.randomUUID()}`)
+    await $`mkdir -p ${dir}`
+    return dir
+}
+
+async function removeDir(dir: string) {
+    await $`rm -rf ${dir}`
+}
+
+beforeAll(async () => {
+    tempHome = await makeTempDir('memo-core-memory')
+    prevMemoHome = process.env.MEMO_HOME
+    process.env.MEMO_HOME = tempHome
+})
+
+afterAll(async () => {
+    if (prevMemoHome === undefined) {
+        delete process.env.MEMO_HOME
+    } else {
+        process.env.MEMO_HOME = prevMemoHome
+    }
+    await removeDir(tempHome)
+})
+
+describe('memory injection', () => {
+    test('loads memory into system prompt when file exists', async () => {
+        const memoryPath = join(tempHome, 'memory.md')
+        await Bun.write(memoryPath, '用户偏好：中文回答\n')
+
+        const session = await createAgentSession(
+            {
+                callLLM: async () => ({ content: JSON.stringify({ final: 'ok' }) }),
+                historySinks: [],
+                tokenCounter: createTokenCounter('cl100k_base'),
+            },
+            { mode: 'once' },
+        )
+        try {
+            const systemPrompt = session.history[0]?.content ?? ''
+            assert.ok(systemPrompt.includes('长期记忆'), 'system prompt should include memory section')
+            assert.ok(systemPrompt.includes('用户偏好：中文回答'), 'memory content should be injected')
+        } finally {
+            await session.close()
+        }
+    })
+})

packages/core/src/runtime/prompt.md

@@ -13,6 +13,7 @@
 - `glob`: `{"pattern":"**/*.ts","path":"/repo"}`
 - `grep`: `{"pattern":"TODO","path":"/repo","output_mode":"content|files_with_matches|count","glob":"src/**/*.ts","-i":false,"-A":1,"-B":1}`
 - `fetch`: `{"url":"https://..."}`
+- `memory`: `{"note":"一句极短的用户画像/偏好/身份总结，≤32字符"}`（仅在准备返回 final 前追加）
 
 # 允许的输出格式
 
@@ -35,3 +36,4 @@
 - 每次仅调用一个工具。
 - 收到 observation 后，如已足够回答，直接返回 final。
 - 不要输出 XML/Markdown/额外文字；输出必须是 JSON 对象字符串。
+- 如发现可复用的用户画像/习惯/偏好等信息，且准备给出 final，可先调用一次 `memory` 工具写入摘要，避免重复调用；记忆内容不得包含敏感/原文，仅保留极短摘要。