Merge branch 'sipeed:main' into main
This commit is contained in:
Harmoon
GitHub
commit
a81dd22b2d
186 changed files with 20563 additions and 2195 deletions
62
.github/workflows/create_dmg.yml
vendored
Normal file
62
.github/workflows/create_dmg.yml
vendored
Normal file
|
|
@ -0,0 +1,62 @@
|
|||
name: Create macOS DMG
|
||||
on:
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
build:
|
||||
name: Build ${{ matrix.arch }}
|
||||
runs-on: macos-latest
|
||||
strategy:
|
||||
matrix:
|
||||
# This creates two parallel jobs
|
||||
arch: [arm64, amd64]
|
||||
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
ref: main
|
||||
|
||||
# 1. 安装指定版本的 Go (可选，但推荐)
|
||||
- name: Setup Go
|
||||
uses: actions/setup-go@v6
|
||||
with:
|
||||
go-version-file: go.mod
|
||||
|
||||
# 2. 安装 pnpm
|
||||
- name: Install pnpm
|
||||
run: brew install pnpm
|
||||
|
||||
# 3. 运行你的 Makefile 编译二进制文件
|
||||
- name: Build with Make
|
||||
run: make build ARCH=${{ matrix.arch }} && make build-macos-app ARCH=${{ matrix.arch }}
|
||||
|
||||
# 4. 签名
|
||||
- name: Ad-hoc Sign
|
||||
run: codesign --force --deep --sign - "build/PicoClaw Launcher.app"
|
||||
|
||||
# 5. 安装打包工具
|
||||
- name: Install create-dmg
|
||||
run: brew install create-dmg
|
||||
|
||||
# 6. 执行打包命令
|
||||
- name: Create DMG
|
||||
run: |
|
||||
mkdir -p dist
|
||||
create-dmg \
|
||||
--volname "PicoClaw Installer" \
|
||||
--window-pos 200 120 \
|
||||
--window-size 800 400 \
|
||||
--icon-size 100 \
|
||||
--icon "PicoClaw Launcher.app" 200 190 \
|
||||
--hide-extension "PicoClaw Launcher.app" \
|
||||
--app-drop-link 600 185 \
|
||||
"dist/picoclaw-${{ matrix.arch }}.dmg" \
|
||||
"build/PicoClaw Launcher.app"
|
||||
|
||||
# 7. 上传文件到 GitHub Artifacts (供你下载)
|
||||
- name: Upload DMG
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: macos-dmg-${{ matrix.arch }}
|
||||
path: dist/*.dmg
|
||||
2
.gitignore
vendored
2
.gitignore
vendored
|
|
@ -67,3 +67,5 @@ web/backend/dist/*
|
|||
.claude/
|
||||
|
||||
docker/data
|
||||
|
||||
.omc/
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ linters:
|
|||
- exhaustruct
|
||||
- funcorder
|
||||
- gochecknoglobals
|
||||
- gosmopolitan # Project legitimately uses CJK text in tests (FTS5, token counting)
|
||||
- godot
|
||||
- intrange
|
||||
- ireturn
|
||||
|
|
|
|||
51
Makefile
51
Makefile
|
|
@ -5,6 +5,7 @@ BINARY_NAME=picoclaw
|
|||
BUILD_DIR=build
|
||||
CMD_DIR=cmd/$(BINARY_NAME)
|
||||
MAIN_GO=$(CMD_DIR)/main.go
|
||||
EXT=
|
||||
|
||||
# Version
|
||||
VERSION?=$(shell git describe --tags --always --dirty 2>/dev/null || echo "dev")
|
||||
|
|
@ -69,9 +70,11 @@ WORKSPACE_DIR?=$(PICOCLAW_HOME)/workspace
|
|||
WORKSPACE_SKILLS_DIR=$(WORKSPACE_DIR)/skills
|
||||
BUILTIN_SKILLS_DIR=$(CURDIR)/skills
|
||||
|
||||
LNCMD=ln -sf
|
||||
|
||||
# OS detection
|
||||
UNAME_S:=$(shell uname -s)
|
||||
UNAME_M:=$(shell uname -m)
|
||||
UNAME_S?=$(shell uname -s)
|
||||
UNAME_M?=$(shell uname -m)
|
||||
|
||||
# Platform-specific settings
|
||||
ifeq ($(UNAME_S),Linux)
|
||||
|
|
@ -93,17 +96,30 @@ ifeq ($(UNAME_S),Linux)
|
|||
endif
|
||||
else ifeq ($(UNAME_S),Darwin)
|
||||
PLATFORM=darwin
|
||||
WEB_GO=CGO_ENABLED=1 go
|
||||
WEB_GO=CGO_LDFLAGS="-mmacosx-version-min=10.11" CGO_CFLAGS="-mmacosx-version-min=10.11" CGO_ENABLED=1 go
|
||||
ifeq ($(UNAME_M),x86_64)
|
||||
ARCH=amd64
|
||||
ARCH?=amd64
|
||||
else ifeq ($(UNAME_M),arm64)
|
||||
ARCH=arm64
|
||||
ARCH?=arm64
|
||||
else
|
||||
ARCH=$(UNAME_M)
|
||||
ARCH?=$(UNAME_M)
|
||||
endif
|
||||
else
|
||||
PLATFORM=$(UNAME_S)
|
||||
ARCH=$(UNAME_M)
|
||||
ifeq ($(UNAME_M),x86_64)
|
||||
ARCH?=amd64
|
||||
else
|
||||
ARCH?=$(UNAME_M)
|
||||
endif
|
||||
# Detect Windows (Git Bash / MSYS2)
|
||||
IS_WINDOWS:=$(if $(findstring MINGW,$(UNAME_S)),yes,$(if $(findstring MSYS,$(UNAME_S)),yes,$(if $(findstring CYGWIN,$(UNAME_S)),yes,no)))
|
||||
ifeq ($(IS_WINDOWS),yes)
|
||||
EXT=.exe
|
||||
LNCMD=cp
|
||||
else ifeq ($(UNAME_S),windows) # failsafe for force windows build in other OS using UNAME_S=windows
|
||||
EXT=.exe
|
||||
endif
|
||||
|
||||
endif
|
||||
|
||||
BINARY_PATH=$(BUILD_DIR)/$(BINARY_NAME)-$(PLATFORM)-$(ARCH)
|
||||
|
|
@ -120,23 +136,23 @@ generate:
|
|||
|
||||
## build: Build the picoclaw binary for current platform
|
||||
build: generate
|
||||
@echo "Building $(BINARY_NAME) for $(PLATFORM)/$(ARCH)..."
|
||||
@echo "Building $(BINARY_NAME)$(EXT) for $(PLATFORM)/$(ARCH)..."
|
||||
@mkdir -p $(BUILD_DIR)
|
||||
@$(GO) build $(GOFLAGS) -ldflags "$(LDFLAGS)" -o $(BINARY_PATH) ./$(CMD_DIR)
|
||||
@echo "Build complete: $(BINARY_PATH)"
|
||||
@ln -sf $(BINARY_NAME)-$(PLATFORM)-$(ARCH) $(BUILD_DIR)/$(BINARY_NAME)
|
||||
@GOARCH=${ARCH} $(GO) build $(GOFLAGS) -ldflags "$(LDFLAGS)" -o $(BINARY_PATH)$(EXT) ./$(CMD_DIR)
|
||||
@echo "Build complete: $(BINARY_PATH)$(EXT)"
|
||||
@$(LNCMD) $(BINARY_NAME)-$(PLATFORM)-$(ARCH)$(EXT) $(BUILD_DIR)/$(BINARY_NAME)$(EXT)
|
||||
|
||||
## build-launcher: Build the picoclaw-launcher (web console) binary
|
||||
build-launcher:
|
||||
@echo "Building picoclaw-launcher for $(PLATFORM)/$(ARCH)..."
|
||||
@mkdir -p $(BUILD_DIR)
|
||||
@$(MAKE) -C web build \
|
||||
OUTPUT="$(CURDIR)/$(BUILD_DIR)/picoclaw-launcher-$(PLATFORM)-$(ARCH)" \
|
||||
@GOARCH=${ARCH} $(MAKE) -C web build \
|
||||
OUTPUT="$(CURDIR)/$(BUILD_DIR)/picoclaw-launcher-$(PLATFORM)-$(ARCH)$(EXT)" \
|
||||
WEB_GO='$(WEB_GO)' \
|
||||
GO_BUILD_TAGS='$(GO_BUILD_TAGS)' \
|
||||
LDFLAGS='$(LDFLAGS)'
|
||||
@ln -sf picoclaw-launcher-$(PLATFORM)-$(ARCH) $(BUILD_DIR)/picoclaw-launcher
|
||||
@echo "Build complete: $(BUILD_DIR)/picoclaw-launcher"
|
||||
@$(LNCMD) picoclaw-launcher-$(PLATFORM)-$(ARCH)$(EXT) $(BUILD_DIR)/picoclaw-launcher$(EXT)
|
||||
@echo "Build complete: $(BUILD_DIR)/picoclaw-launcher$(EXT)"
|
||||
|
||||
build-launcher-frontend:
|
||||
@$(MAKE) -C web build-frontend
|
||||
|
|
@ -324,14 +340,13 @@ docker-clean:
|
|||
|
||||
|
||||
## build-macos-app: Build PicoClaw macOS .app bundle (no terminal window)
|
||||
build-macos-app:
|
||||
build-macos-app:build-launcher
|
||||
@echo "Building macOS .app bundle..."
|
||||
@if [ "$(UNAME_S)" != "Darwin" ]; then \
|
||||
echo "Error: This target is only available on macOS"; \
|
||||
exit 1; \
|
||||
fi
|
||||
@cd web && $(MAKE) build && cd ..
|
||||
@./scripts/build-macos-app.sh $(BINARY_NAME)-$(PLATFORM)-$(ARCH)
|
||||
@./scripts/build-macos-app.sh $(PLATFORM)-$(ARCH)
|
||||
@echo "macOS .app bundle created: $(BUILD_DIR)/PicoClaw.app"
|
||||
|
||||
## help: Show this help message
|
||||
|
|
|
|||
27
README.fr.md
27
README.fr.md
|
|
@ -306,7 +306,25 @@ Pour la documentation détaillée du TUI, voir [docs.picoclaw.io](https://docs.p
|
|||
|
||||
Donnez une seconde vie à votre téléphone vieux de dix ans ! Transformez-le en assistant IA intelligent avec PicoClaw.
|
||||
|
||||
**Option 1 : Termux (disponible maintenant)**
|
||||
**Option 1 : Installation APK**
|
||||
|
||||
Aperçu :
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<td><img src="assets/fui_main_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_web_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_log_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_setting_page.jpg" width="200"></td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
Téléchargez l'APK depuis [picoclaw.io](https://picoclaw.io/download/) et installez-le directement. Pas besoin de Termux !
|
||||
|
||||
**Option 2 : Termux**
|
||||
|
||||
<details>
|
||||
<summary><b>Terminal Launcher (pour les environnements à ressources limitées)</b></summary>
|
||||
|
||||
1. Installez [Termux](https://github.com/termux/termux-app) (téléchargez depuis [GitHub Releases](https://github.com/termux/termux-app/releases), ou cherchez dans F-Droid / Google Play)
|
||||
2. Exécutez les commandes suivantes :
|
||||
|
|
@ -323,13 +341,6 @@ Suivez ensuite la section Terminal Launcher ci-dessous pour terminer la configur
|
|||
|
||||
<img src="assets/termux.jpg" alt="PicoClaw on Termux" width="512">
|
||||
|
||||
**Option 2 : Installation APK**
|
||||
|
||||
Téléchargez l'APK depuis [picoclaw.io](https://picoclaw.io/download/) et installez-le directement. Pas besoin de Termux !
|
||||
|
||||
<details>
|
||||
<summary><b>Terminal Launcher (pour les environnements à ressources limitées)</b></summary>
|
||||
|
||||
Pour les environnements minimaux où seul le binaire principal `picoclaw` est disponible (sans Launcher UI), vous pouvez tout configurer via la ligne de commande et un fichier de configuration JSON.
|
||||
|
||||
**1. Initialiser**
|
||||
|
|
|
|||
27
README.id.md
27
README.id.md
|
|
@ -303,7 +303,25 @@ Untuk dokumentasi TUI lengkap, lihat [docs.picoclaw.io](https://docs.picoclaw.io
|
|||
|
||||
Berikan kehidupan kedua untuk ponsel lama Anda! Ubah menjadi Asisten AI pintar dengan PicoClaw.
|
||||
|
||||
**Opsi 1: Termux (tersedia sekarang)**
|
||||
**Opsi 1: Instal APK**
|
||||
|
||||
Pratinjau:
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<td><img src="assets/fui_main_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_web_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_log_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_setting_page.jpg" width="200"></td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
Unduh APK dari [picoclaw.io](https://picoclaw.io/download/) dan instal langsung. Tanpa Termux!
|
||||
|
||||
**Opsi 2: Termux**
|
||||
|
||||
<details>
|
||||
<summary><b>Terminal Launcher (untuk lingkungan dengan sumber daya terbatas)</b></summary>
|
||||
|
||||
1. Instal [Termux](https://github.com/termux/termux-app) (unduh dari [GitHub Releases](https://github.com/termux/termux-app/releases), atau cari di F-Droid / Google Play)
|
||||
2. Jalankan perintah berikut:
|
||||
|
|
@ -320,13 +338,6 @@ Kemudian ikuti bagian Terminal Launcher di bawah untuk menyelesaikan konfigurasi
|
|||
|
||||
<img src="assets/termux.jpg" alt="PicoClaw on Termux" width="512">
|
||||
|
||||
**Opsi 2: Instal APK**
|
||||
|
||||
Unduh APK dari [picoclaw.io](https://picoclaw.io/download/) dan instal langsung. Tanpa Termux!
|
||||
|
||||
<details>
|
||||
<summary><b>Terminal Launcher (untuk lingkungan dengan sumber daya terbatas)</b></summary>
|
||||
|
||||
Untuk lingkungan minimal di mana hanya binary inti `picoclaw` yang tersedia (tanpa Launcher UI), Anda dapat mengonfigurasi semuanya melalui command line dan file konfigurasi JSON.
|
||||
|
||||
**1. Inisialisasi**
|
||||
|
|
|
|||
27
README.it.md
27
README.it.md
|
|
@ -303,7 +303,25 @@ Per la documentazione dettagliata del TUI, vedi [docs.picoclaw.io](https://docs.
|
|||
|
||||
Dai una seconda vita al tuo telefono di dieci anni fa! Trasformalo in un assistente IA intelligente con PicoClaw.
|
||||
|
||||
**Opzione 1: Termux (disponibile ora)**
|
||||
**Opzione 1: Installazione APK**
|
||||
|
||||
Anteprima:
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<td><img src="assets/fui_main_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_web_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_log_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_setting_page.jpg" width="200"></td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
Scarica l'APK da [picoclaw.io](https://picoclaw.io/download/) e installa direttamente. Senza Termux!
|
||||
|
||||
**Opzione 2: Termux**
|
||||
|
||||
<details>
|
||||
<summary><b>Terminal Launcher (per ambienti con risorse limitate)</b></summary>
|
||||
|
||||
1. Installa [Termux](https://github.com/termux/termux-app) (scarica da [GitHub Releases](https://github.com/termux/termux-app/releases), o cerca su F-Droid / Google Play)
|
||||
2. Esegui i seguenti comandi:
|
||||
|
|
@ -320,13 +338,6 @@ Poi segui la sezione Terminal Launcher qui sotto per completare la configurazion
|
|||
|
||||
<img src="assets/termux.jpg" alt="PicoClaw on Termux" width="512">
|
||||
|
||||
**Opzione 2: Installazione APK**
|
||||
|
||||
Scarica l'APK da [picoclaw.io](https://picoclaw.io/download/) e installa direttamente. Senza Termux!
|
||||
|
||||
<details>
|
||||
<summary><b>Terminal Launcher (per ambienti con risorse limitate)</b></summary>
|
||||
|
||||
Per ambienti minimali dove è disponibile solo il binario core `picoclaw` (senza Launcher UI), puoi configurare tutto tramite riga di comando e un file di configurazione JSON.
|
||||
|
||||
**1. Inizializza**
|
||||
|
|
|
|||
27
README.ja.md
27
README.ja.md
|
|
@ -303,7 +303,25 @@ TUI の詳細なドキュメントは [docs.picoclaw.io](https://docs.picoclaw.i
|
|||
|
||||
10 年前のスマホに第二の人生を！PicoClaw でスマート AI アシスタントに変身させましょう。
|
||||
|
||||
**オプション 1: Termux（現在利用可能）**
|
||||
**オプション 1: APK インストール**
|
||||
|
||||
プレビュー：
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<td><img src="assets/fui_main_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_web_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_log_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_setting_page.jpg" width="200"></td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
[picoclaw.io](https://picoclaw.io/download/) から APK をダウンロードして直接インストール。Termux 不要！
|
||||
|
||||
**オプション 2: Termux**
|
||||
|
||||
<details>
|
||||
<summary><b>Terminal Launcher（リソース制約環境向け）</b></summary>
|
||||
|
||||
1. [Termux](https://github.com/termux/termux-app) をインストール（[GitHub Releases](https://github.com/termux/termux-app/releases) からダウンロード、または F-Droid / Google Play で検索）
|
||||
2. 以下のコマンドを実行：
|
||||
|
|
@ -320,13 +338,6 @@ termux-chroot ./picoclaw onboard # chroot で標準的な Linux ファイル
|
|||
|
||||
<img src="assets/termux.jpg" alt="PicoClaw on Termux" width="512">
|
||||
|
||||
**オプション 2: APK インストール**
|
||||
|
||||
[picoclaw.io](https://picoclaw.io/download/) から APK をダウンロードして直接インストール。Termux 不要！
|
||||
|
||||
<details>
|
||||
<summary><b>Terminal Launcher（リソース制約環境向け）</b></summary>
|
||||
|
||||
`picoclaw` コアバイナリのみが利用可能な最小環境（Launcher UI なし）では、コマンドラインと JSON 設定ファイルですべてを設定できます。
|
||||
|
||||
**1. 初期化**
|
||||
|
|
|
|||
30
README.md
30
README.md
|
|
@ -303,7 +303,25 @@ For detailed TUI documentation, see [docs.picoclaw.io](https://docs.picoclaw.io)
|
|||
|
||||
Give your decade-old phone a second life! Turn it into a smart AI Assistant with PicoClaw.
|
||||
|
||||
**Option 1: Termux (available now)**
|
||||
**Option 1: APK Install**
|
||||
|
||||
Preview:
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<td><img src="assets/fui_main_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_web_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_log_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_setting_page.jpg" width="200"></td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
Download the APK from [picoclaw.io](https://picoclaw.io/download/) and install directly. No Termux required!
|
||||
|
||||
**Option 2: Termux**
|
||||
|
||||
<details>
|
||||
<summary><b>Terminal Launcher (for resource-constrained environments)</b></summary>
|
||||
|
||||
1. Install [Termux](https://github.com/termux/termux-app) (download from [GitHub Releases](https://github.com/termux/termux-app/releases), or search in F-Droid / Google Play)
|
||||
2. Run the following commands:
|
||||
|
|
@ -320,13 +338,6 @@ Then follow the Terminal Launcher section below to complete configuration.
|
|||
|
||||
<img src="assets/termux.jpg" alt="PicoClaw on Termux" width="512">
|
||||
|
||||
**Option 2: APK Install**
|
||||
|
||||
Download the APK from [picoclaw.io](https://picoclaw.io/download/) and install directly. No Termux required!
|
||||
|
||||
<details>
|
||||
<summary><b>Terminal Launcher (for resource-constrained environments)</b></summary>
|
||||
|
||||
For minimal environments where only the `picoclaw` core binary is available (no Launcher UI), you can configure everything via the command line and a JSON config file.
|
||||
|
||||
**1. Initialize**
|
||||
|
|
@ -443,7 +454,7 @@ For full provider configuration details, see [Providers & Models](docs/providers
|
|||
|
||||
## 💬 Channels (Chat Apps)
|
||||
|
||||
Talk to your PicoClaw through 17+ messaging platforms:
|
||||
Talk to your PicoClaw through 18+ messaging platforms:
|
||||
|
||||
| Channel | Setup | Protocol | Docs |
|
||||
|---------|-------|----------|------|
|
||||
|
|
@ -458,6 +469,7 @@ Talk to your PicoClaw through 17+ messaging platforms:
|
|||
| **Feishu / Lark** | Medium (App ID + Secret) | WebSocket/SDK | [Guide](docs/channels/feishu/README.md) |
|
||||
| **LINE** | Medium (credentials + webhook) | Webhook | [Guide](docs/channels/line/README.md) |
|
||||
| **WeCom** | Easy (QR login or manual) | WebSocket | [Guide](docs/channels/wecom/README.md) |
|
||||
| **VK** | Easy (group token) | Long Poll | [Guide](docs/channels/vk/README.md) |
|
||||
| **IRC** | Medium (server + nick) | IRC protocol | [Guide](docs/chat-apps.md#irc) |
|
||||
| **OneBot** | Medium (WebSocket URL) | OneBot v11 | [Guide](docs/channels/onebot/README.md) |
|
||||
| **MaixCam** | Easy (enable) | TCP socket | [Guide](docs/channels/maixcam/README.md) |
|
||||
|
|
|
|||
27
README.my.md
27
README.my.md
|
|
@ -300,7 +300,25 @@ Untuk dokumentasi TUI terperinci, lihat [docs.picoclaw.io](https://docs.picoclaw
|
|||
|
||||
Berikan telefon lama anda kehidupan baru! Jadikannya Pembantu AI pintar dengan PicoClaw.
|
||||
|
||||
**Pilihan 1: Termux (tersedia sekarang)**
|
||||
**Pilihan 1: Pasang APK**
|
||||
|
||||
Pratonton:
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<td><img src="assets/fui_main_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_web_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_log_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_setting_page.jpg" width="200"></td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
Muat turun APK dari [picoclaw.io](https://picoclaw.io/download/) dan pasang secara langsung. Tiada Termux diperlukan!
|
||||
|
||||
**Pilihan 2: Termux**
|
||||
|
||||
<details>
|
||||
<summary><b>Pelancar Terminal (untuk persekitaran terhad sumber)</b></summary>
|
||||
|
||||
1. Pasang [Termux](https://github.com/termux/termux-app) (muat turun dari [GitHub Releases](https://github.com/termux/termux-app/releases), atau cari di F-Droid / Google Play)
|
||||
2. Jalankan arahan berikut:
|
||||
|
|
@ -317,13 +335,6 @@ Kemudian ikuti bahagian Pelancar Terminal di bawah untuk melengkapkan konfiguras
|
|||
|
||||
<img src="assets/termux.jpg" alt="PicoClaw pada Termux" width="512">
|
||||
|
||||
**Pilihan 2: Pasang APK**
|
||||
|
||||
Muat turun APK dari [picoclaw.io](https://picoclaw.io/download/) dan pasang secara langsung. Tiada Termux diperlukan!
|
||||
|
||||
<details>
|
||||
<summary><b>Pelancar Terminal (untuk persekitaran terhad sumber)</b></summary>
|
||||
|
||||
Untuk persekitaran minimal di mana hanya binari teras `picoclaw` tersedia (tiada UI Pelancar), anda boleh mengkonfigurasi semua melalui baris arahan dan fail konfigurasi JSON.
|
||||
|
||||
**1. Mulakan**
|
||||
|
|
|
|||
|
|
@ -303,7 +303,25 @@ Para documentação detalhada do TUI, veja [docs.picoclaw.io](https://docs.picoc
|
|||
|
||||
Dê uma segunda vida ao seu celular de uma década! Transforme-o em um Assistente de IA inteligente com o PicoClaw.
|
||||
|
||||
**Opção 1: Termux (disponível agora)**
|
||||
**Opção 1: Instalação via APK**
|
||||
|
||||
Pré-visualização:
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<td><img src="assets/fui_main_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_web_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_log_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_setting_page.jpg" width="200"></td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
Baixe o APK de [picoclaw.io](https://picoclaw.io/download/) e instale diretamente. Sem necessidade de Termux!
|
||||
|
||||
**Opção 2: Termux**
|
||||
|
||||
<details>
|
||||
<summary><b>Terminal Launcher (para ambientes com recursos limitados)</b></summary>
|
||||
|
||||
1. Instale o [Termux](https://github.com/termux/termux-app) (baixe nas [GitHub Releases](https://github.com/termux/termux-app/releases), ou pesquise no F-Droid / Google Play)
|
||||
2. Execute os seguintes comandos:
|
||||
|
|
@ -320,13 +338,6 @@ Em seguida, siga a seção Terminal Launcher abaixo para concluir a configuraç
|
|||
|
||||
<img src="assets/termux.jpg" alt="PicoClaw on Termux" width="512">
|
||||
|
||||
**Opção 2: Instalação via APK**
|
||||
|
||||
Baixe o APK de [picoclaw.io](https://picoclaw.io/download/) e instale diretamente. Sem necessidade de Termux!
|
||||
|
||||
<details>
|
||||
<summary><b>Terminal Launcher (para ambientes com recursos limitados)</b></summary>
|
||||
|
||||
Para ambientes mínimos onde apenas o binário principal `picoclaw` está disponível (sem Launcher UI), você pode configurar tudo via linha de comando e um arquivo de configuração JSON.
|
||||
|
||||
**1. Inicializar**
|
||||
|
|
|
|||
27
README.vi.md
27
README.vi.md
|
|
@ -303,7 +303,25 @@ Sử dụng menu TUI để: **1)** Cấu hình Provider -> **2)** Cấu hình Ch
|
|||
|
||||
Hãy cho chiếc điện thoại cũ của bạn một cuộc sống mới! Biến nó thành Trợ lý AI thông minh với PicoClaw.
|
||||
|
||||
**Tùy chọn 1: Termux (có sẵn ngay)**
|
||||
**Tùy chọn 1: Cài đặt APK**
|
||||
|
||||
Xem trước:
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<td><img src="assets/fui_main_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_web_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_log_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_setting_page.jpg" width="200"></td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
Tải APK từ [picoclaw.io](https://picoclaw.io/download/) và cài đặt trực tiếp. Không cần Termux!
|
||||
|
||||
**Tùy chọn 2: Termux**
|
||||
|
||||
<details>
|
||||
<summary><b>Terminal Launcher (cho môi trường hạn chế tài nguyên)</b></summary>
|
||||
|
||||
1. Cài đặt [Termux](https://github.com/termux/termux-app) (tải từ [GitHub Releases](https://github.com/termux/termux-app/releases), hoặc tìm kiếm trong F-Droid / Google Play)
|
||||
2. Chạy các lệnh sau:
|
||||
|
|
@ -320,13 +338,6 @@ Sau đó làm theo phần Terminal Launcher bên dưới để hoàn tất cấu
|
|||
|
||||
<img src="assets/termux.jpg" alt="PicoClaw on Termux" width="512">
|
||||
|
||||
**Tùy chọn 2: Cài đặt APK**
|
||||
|
||||
Tải APK từ [picoclaw.io](https://picoclaw.io/download/) và cài đặt trực tiếp. Không cần Termux!
|
||||
|
||||
<details>
|
||||
<summary><b>Terminal Launcher (cho môi trường hạn chế tài nguyên)</b></summary>
|
||||
|
||||
Đối với các môi trường tối giản chỉ có binary lõi `picoclaw` (không có Launcher UI), bạn có thể cấu hình mọi thứ qua dòng lệnh và tệp cấu hình JSON.
|
||||
|
||||
**1. Khởi tạo**
|
||||
|
|
|
|||
30
README.zh.md
30
README.zh.md
|
|
@ -303,7 +303,25 @@ picoclaw-launcher-tui
|
|||
|
||||
让你十年前的旧手机焕发新生！将它变成你的 AI 助手。
|
||||
|
||||
**方式一：Termux（现已可用）**
|
||||
**方式一：APK 安装**
|
||||
|
||||
预览：
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<td><img src="assets/fui_main_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_web_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_log_page.jpg" width="200"></td>
|
||||
<td><img src="assets/fui_setting_page.jpg" width="200"></td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
从 [picoclaw.io](https://picoclaw.io/download/) 下载 APK 并直接安装，无需 Termux！
|
||||
|
||||
**方式二：Termux**
|
||||
|
||||
<details>
|
||||
<summary><b>Terminal Launcher（适用于资源受限环境）</b></summary>
|
||||
|
||||
1. 安装 [Termux](https://github.com/termux/termux-app)（可从 [GitHub Releases](https://github.com/termux/termux-app/releases) 下载，或在 F-Droid / Google Play 中搜索）
|
||||
2. 执行以下命令：
|
||||
|
|
@ -320,13 +338,6 @@ termux-chroot ./picoclaw onboard # chroot 提供标准 Linux 文件系统布
|
|||
|
||||
<img src="assets/termux.jpg" alt="PicoClaw on Termux" width="512">
|
||||
|
||||
**方式二：APK 安装**
|
||||
|
||||
从 [picoclaw.io](https://picoclaw.io/download/) 下载 APK 并直接安装，无需 Termux！
|
||||
|
||||
<details>
|
||||
<summary><b>Terminal Launcher（适用于资源受限环境）</b></summary>
|
||||
|
||||
对于只有 `picoclaw` 核心二进制文件的极简环境（无 Launcher UI），可通过命令行和 JSON 配置文件完成所有配置。
|
||||
|
||||
**1. 初始化**
|
||||
|
|
@ -437,7 +448,7 @@ PicoClaw 通过 `model_list` 配置支持 30+ LLM Provider，使用 `协议/模
|
|||
|
||||
## 💬 Channels（聊天应用）
|
||||
|
||||
通过 17+ 消息平台与你的 PicoClaw 对话：
|
||||
通过 18+ 消息平台与你的 PicoClaw 对话：
|
||||
|
||||
| Channel | 配置难度 | 协议 | 文档 |
|
||||
|---------|----------|------|------|
|
||||
|
|
@ -452,6 +463,7 @@ PicoClaw 通过 `model_list` 配置支持 30+ LLM Provider，使用 `协议/模
|
|||
| **飞书 / Lark** | 中等（App ID + Secret） | WebSocket/SDK | [指南](docs/channels/feishu/README.zh.md) |
|
||||
| **LINE** | 中等（credentials + webhook） | Webhook | [指南](docs/channels/line/README.zh.md) |
|
||||
| **企业微信** | 简单（扫码登录或手动配置） | WebSocket | [指南](docs/channels/wecom/README.zh.md) |
|
||||
| **VK** | 简单（群组 token） | Long Poll | [指南](docs/channels/vk/README.md) |
|
||||
| **IRC** | 中等（server + nick） | IRC 协议 | [指南](docs/zh/chat-apps.md#irc) |
|
||||
| **OneBot** | 中等（WebSocket URL） | OneBot v11 | [指南](docs/channels/onebot/README.zh.md) |
|
||||
| **MaixCam** | 简单（启用即可） | TCP socket | [指南](docs/channels/maixcam/README.zh.md) |
|
||||
|
|
|
|||
BIN
assets/fui_log_page.jpg
Normal file
BIN
assets/fui_log_page.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 11 KiB |
BIN
assets/fui_main_page.jpg
Normal file
BIN
assets/fui_main_page.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 34 KiB |
BIN
assets/fui_setting_page.jpg
Normal file
BIN
assets/fui_setting_page.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 46 KiB |
BIN
assets/fui_web_page.jpg
Normal file
BIN
assets/fui_web_page.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 20 KiB |
|
|
@ -9,6 +9,7 @@ package main
|
|||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
|
|
@ -24,10 +25,11 @@ import (
|
|||
"github.com/sipeed/picoclaw/cmd/picoclaw/internal/status"
|
||||
"github.com/sipeed/picoclaw/cmd/picoclaw/internal/version"
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
"github.com/sipeed/picoclaw/pkg/updater"
|
||||
)
|
||||
|
||||
func NewPicoclawCommand() *cobra.Command {
|
||||
short := fmt.Sprintf("%s picoclaw - Personal AI Assistant v%s\n\n", internal.Logo, config.GetVersion())
|
||||
short := fmt.Sprintf("%s picoclaw - Personal AI Assistant %s\n\n", internal.Logo, config.GetVersion())
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "picoclaw",
|
||||
|
|
@ -45,6 +47,7 @@ func NewPicoclawCommand() *cobra.Command {
|
|||
migrate.NewMigrateCommand(),
|
||||
skills.NewSkillsCommand(),
|
||||
model.NewModelCommand(),
|
||||
updater.NewUpdateCommand("picoclaw"),
|
||||
version.NewVersionCommand(),
|
||||
)
|
||||
|
||||
|
|
@ -66,6 +69,21 @@ const (
|
|||
|
||||
func main() {
|
||||
fmt.Printf("%s", banner)
|
||||
|
||||
tz_env := os.Getenv("TZ")
|
||||
if tz_env != "" {
|
||||
fmt.Println("TZ environment:", tz_env)
|
||||
zoneinfo_env := os.Getenv("ZONEINFO")
|
||||
fmt.Println("ZONEINFO environment:", zoneinfo_env)
|
||||
loc, err := time.LoadLocation(tz_env)
|
||||
if err != nil {
|
||||
fmt.Println("Error loading time zone:", err)
|
||||
} else {
|
||||
fmt.Println("Time zone loaded successfully:", loc)
|
||||
time.Local = loc //nolint:gosmopolitan // We intentionally set local timezone from TZ env
|
||||
}
|
||||
}
|
||||
|
||||
cmd := NewPicoclawCommand()
|
||||
if err := cmd.Execute(); err != nil {
|
||||
os.Exit(1)
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ func TestNewPicoclawCommand(t *testing.T) {
|
|||
|
||||
require.NotNil(t, cmd)
|
||||
|
||||
short := fmt.Sprintf("%s picoclaw - Personal AI Assistant v%s\n\n", internal.Logo, config.GetVersion())
|
||||
short := fmt.Sprintf("%s picoclaw - Personal AI Assistant %s\n\n", internal.Logo, config.GetVersion())
|
||||
|
||||
assert.Equal(t, "picoclaw", cmd.Use)
|
||||
assert.Equal(t, short, cmd.Short)
|
||||
|
|
@ -43,6 +43,7 @@ func TestNewPicoclawCommand(t *testing.T) {
|
|||
"onboard",
|
||||
"skills",
|
||||
"status",
|
||||
"update",
|
||||
"version",
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -48,6 +48,11 @@
|
|||
"model": "deepseek/deepseek-chat",
|
||||
"api_key": "sk-your-deepseek-key"
|
||||
},
|
||||
{
|
||||
"model_name": "venice-uncensored",
|
||||
"model": "venice/venice-uncensored",
|
||||
"api_key": "your-venice-api-key"
|
||||
},
|
||||
{
|
||||
"model_name": "lmstudio-local",
|
||||
"model": "lmstudio/openai/gpt-oss-20b"
|
||||
|
|
@ -416,7 +421,8 @@
|
|||
"enabled": true
|
||||
},
|
||||
"read_file": {
|
||||
"enabled": true
|
||||
"enabled": true,
|
||||
"mode": "bytes"
|
||||
},
|
||||
"send_tts": {
|
||||
"enabled": false
|
||||
|
|
|
|||
194
docs/channels/vk/README.md
Normal file
194
docs/channels/vk/README.md
Normal file
|
|
@ -0,0 +1,194 @@
|
|||
# VK (VKontakte)
|
||||
|
||||
The VK channel uses Bots Long Poll API for bot-based communication with VK social network. It supports text messages, media attachments (photos, videos, audio, documents, stickers), and group chat interactions.
|
||||
|
||||
## Configuration
|
||||
|
||||
```json
|
||||
{
|
||||
"channels": {
|
||||
"vk": {
|
||||
"enabled": true,
|
||||
"token": "NOT_HERE",
|
||||
"group_id": 123456789,
|
||||
"allow_from": ["123456789"],
|
||||
"group_trigger": {
|
||||
"mention_only": false,
|
||||
"prefixes": ["/bot", "!bot"]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
| Field | Type | Required | Description |
|
||||
| ---------------- | ------ | -------- | ------------------------------------------------------------------ |
|
||||
| enabled | bool | Yes | Whether to enable the VK channel |
|
||||
| token | string | Yes | Set to `NOT_HERE` - token is stored securely (see Token Storage) |
|
||||
| group_id | int | Yes | VK Community ID (Group ID) |
|
||||
| allow_from | array | No | Allowlist of user IDs; empty means all users are allowed |
|
||||
| group_trigger | object | No | Configuration for group chat triggers |
|
||||
|
||||
### Token Storage
|
||||
|
||||
For security reasons, the VK access token should not be stored directly in the configuration file. Instead:
|
||||
|
||||
1. Set `token` to `"NOT_HERE"` in the configuration
|
||||
2. Store the actual token using one of these methods:
|
||||
- **Environment variable**: Set `PICOCLAW_CHANNELS_VK_TOKEN` environment variable
|
||||
- **Secure storage**: Use PicoClaw's secure token storage mechanism
|
||||
|
||||
Example using environment variable:
|
||||
```bash
|
||||
export PICOCLAW_CHANNELS_VK_TOKEN="vk1.a.abc123..."
|
||||
```
|
||||
|
||||
### Group Trigger Configuration
|
||||
|
||||
| Field | Type | Description |
|
||||
| ------------ | -------- | ------------------------------------------------------------------ |
|
||||
| mention_only | bool | Only respond when bot is mentioned in group chats |
|
||||
| prefixes | []string | List of prefixes that trigger bot response in group chats |
|
||||
|
||||
## Setup
|
||||
|
||||
### 1. Create a VK Community
|
||||
|
||||
1. Go to [VK](https://vk.com) and log in
|
||||
2. Create a new community or use an existing one
|
||||
3. Note your Community ID (found in the community URL, e.g., `public123456789`)
|
||||
|
||||
### 2. Enable Messages
|
||||
|
||||
1. Go to your community page
|
||||
2. Click "Manage" → "Messages" → "Community Messages"
|
||||
3. Enable community messages
|
||||
|
||||
### 3. Create Access Token
|
||||
|
||||
1. Go to "Manage" → "API usage" → "Access tokens"
|
||||
2. Click "Create token"
|
||||
3. Select the following permissions:
|
||||
- `messages` - Access to messages
|
||||
- `photos` - Access to photos (optional)
|
||||
- `docs` - Access to documents (optional)
|
||||
4. Copy the generated access token
|
||||
5. Store the token securely (see Token Storage section below)
|
||||
|
||||
### 4. Configure PicoClaw
|
||||
|
||||
1. Add the token to your PicoClaw configuration
|
||||
2. Set the `group_id` to your community ID (numeric value)
|
||||
3. (Optional) Configure `allow_from` to restrict which user IDs can interact
|
||||
|
||||
## Features
|
||||
|
||||
### Supported Message Types
|
||||
|
||||
- **Text messages**: Full support for text messages
|
||||
- **Photos**: Photos are displayed as `[photo]` placeholder
|
||||
- **Videos**: Videos are displayed as `[video]` placeholder
|
||||
- **Audio**: Audio files are displayed as `[audio]` placeholder
|
||||
- **Voice messages**: Voice messages are displayed as `[voice]` placeholder and support transcription
|
||||
- **Documents**: Documents are displayed as `[document: filename]`
|
||||
- **Stickers**: Stickers are displayed as `[sticker]` placeholder
|
||||
|
||||
### Voice Support
|
||||
|
||||
The VK channel supports both voice message reception and text-to-speech capabilities:
|
||||
|
||||
- **ASR (Automatic Speech Recognition)**: Voice messages can be transcribed to text using configured voice models
|
||||
- **TTS (Text-to-Speech)**: Text responses can be converted to voice messages
|
||||
|
||||
To enable voice transcription, configure a voice model in your providers setup. See [Voice Transcription](../../providers.md#voice-transcription) for details.
|
||||
|
||||
### Group Chat Support
|
||||
|
||||
The VK channel supports group chats with configurable triggers:
|
||||
|
||||
- **Mention-only mode**: Bot only responds when mentioned
|
||||
- **Prefix mode**: Bot responds to messages starting with specified prefixes
|
||||
- **Permissive mode**: Bot responds to all messages (default)
|
||||
|
||||
### Message Length
|
||||
|
||||
VK has a maximum message length of 4000 characters. PicoClaw automatically splits longer messages into multiple parts.
|
||||
|
||||
## Example Configuration
|
||||
|
||||
### Basic Configuration
|
||||
|
||||
```json
|
||||
{
|
||||
"channels": {
|
||||
"vk": {
|
||||
"enabled": true,
|
||||
"token": "NOT_HERE",
|
||||
"group_id": 123456789
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### With User Whitelist
|
||||
|
||||
```json
|
||||
{
|
||||
"channels": {
|
||||
"vk": {
|
||||
"enabled": true,
|
||||
"token": "NOT_HERE",
|
||||
"group_id": 123456789,
|
||||
"allow_from": ["123456789", "987654321"]
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### With Group Chat Triggers
|
||||
|
||||
```json
|
||||
{
|
||||
"channels": {
|
||||
"vk": {
|
||||
"enabled": true,
|
||||
"token": "NOT_HERE",
|
||||
"group_id": 123456789,
|
||||
"group_trigger": {
|
||||
"prefixes": ["/bot", "!bot"]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Bot Not Responding
|
||||
|
||||
1. Check that the access token is valid
|
||||
2. Verify that the `group_id` is correct
|
||||
3. Ensure the user ID is in `allow_from` if configured
|
||||
4. Check PicoClaw logs for error messages
|
||||
|
||||
### Permission Errors
|
||||
|
||||
Make sure the access token has the necessary permissions:
|
||||
- `messages` - Required for sending and receiving messages
|
||||
- `photos` - Optional, for handling photo attachments
|
||||
- `docs` - Optional, for handling document attachments
|
||||
|
||||
### Group Chat Issues
|
||||
|
||||
If the bot doesn't respond in group chats:
|
||||
1. Check `group_trigger` configuration
|
||||
2. Try using a prefix to trigger the bot
|
||||
3. Check if the bot has permission to read group messages
|
||||
|
||||
## API Reference
|
||||
|
||||
The VK channel uses the [VK SDK for Go](https://github.com/SevereCloud/vksdk) library, which supports VK API version 5.199.
|
||||
|
||||
For more information about VK API, see:
|
||||
- [VK API Documentation](https://dev.vk.com/en)
|
||||
- [VK Bots Long Poll API](https://dev.vk.com/en/api/bots-long-poll/getting-started)
|
||||
|
|
@ -301,6 +301,66 @@ Even with `restrict_to_workspace: false`, the `exec` tool blocks these dangerous
|
|||
| `tools.allow_read_paths` | string[] | `[]` | Additional paths allowed for reading outside workspace |
|
||||
| `tools.allow_write_paths` | string[] | `[]` | Additional paths allowed for writing outside workspace |
|
||||
|
||||
### Read File Mode
|
||||
|
||||
`read_file` has two mutually exclusive implementations selected by config. PicoClaw registers exactly one of them at startup:
|
||||
|
||||
| Config Key | Type | Default | Description |
|
||||
|------------|------|---------|-------------|
|
||||
| `tools.read_file.enabled` | bool | `true` | Enables the `read_file` tool |
|
||||
| `tools.read_file.mode` | string | `bytes` | Selects the `read_file` implementation: `bytes` or `lines` |
|
||||
| `tools.read_file.max_read_file_size` | int | `65536` | Maximum bytes returned by `read_file` |
|
||||
|
||||
#### Mode: `bytes`
|
||||
|
||||
Optimized for arbitrary files and binary-safe pagination.
|
||||
|
||||
Parameters:
|
||||
|
||||
* `path` (required): File path
|
||||
* `offset` (optional): Starting byte offset, default `0`
|
||||
* `length` (optional): Maximum number of bytes to read, default `max_read_file_size`
|
||||
|
||||
Use `bytes` when:
|
||||
|
||||
* You may read binary files
|
||||
* You want deterministic byte-range pagination
|
||||
|
||||
#### Mode: `lines`
|
||||
|
||||
Text-oriented behavior, optimized for source files, markdown, logs, and configs. The tool reads sequentially by line and stops when the configured byte budget is reached.
|
||||
|
||||
Parameters:
|
||||
|
||||
* `path` (required): File path
|
||||
* `start_line` (optional): Starting line number, 1-indexed and inclusive, default `1`
|
||||
* `max_lines` (optional): Maximum number of lines to read, default = all remaining lines until EOF or byte budget
|
||||
|
||||
Behavior notes:
|
||||
|
||||
* Binary-looking files are rejected with guidance to switch `read_file` to `mode = bytes`
|
||||
* Extremely long single lines are truncated rather than skipped
|
||||
|
||||
Use `mode = lines` when:
|
||||
|
||||
* The agent mostly reads text files
|
||||
* You want line-based pagination in prompts and tool calls
|
||||
* You want cleaner chunks for code review, logs, and documentation
|
||||
|
||||
#### Example
|
||||
|
||||
```json
|
||||
{
|
||||
"tools": {
|
||||
"read_file": {
|
||||
"enabled": true,
|
||||
"mode": "lines",
|
||||
"max_read_file_size": 65536
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Exec Security
|
||||
|
||||
| Config Key | Type | Default | Description |
|
||||
|
|
|
|||
|
|
@ -99,6 +99,24 @@ Cette conception permet également le **support multi-agents** avec une sélecti
|
|||
}
|
||||
```
|
||||
|
||||
#### Champs d'entrée `model_list`
|
||||
|
||||
| Champ | Type | Requis | Description |
|
||||
|-------|------|--------|-------------|
|
||||
| `model_name` | string | Oui | Nom unique pour référencer ce modèle dans la config agent |
|
||||
| `model` | string | Oui | Identifiant fournisseur/modèle (ex : `openai/gpt-5.4`, `azure/gpt-5.4`, `anthropic/claude-sonnet-4.6`) |
|
||||
| `api_keys` | string[] | Oui* | Clé(s) API pour l'authentification. Plusieurs clés permettent la rotation par requête. Non requis pour les fournisseurs locaux (Ollama, LM Studio, VLLM) |
|
||||
| `api_base` | string | Non | Remplace l'URL de base API par défaut |
|
||||
| `proxy` | string | Non | URL du proxy HTTP pour cette entrée de modèle |
|
||||
| `user_agent` | string | Non | En-tête `User-Agent` personnalisé pour les requêtes API (supporté par les providers OpenAI-compatible, Anthropic et Azure) |
|
||||
| `request_timeout` | int | Non | Délai d'expiration de la requête en secondes (la valeur par défaut varie selon le provider) |
|
||||
| `max_tokens_field` | string | Non | Remplace le nom du champ max tokens dans le corps de la requête (ex : `max_completion_tokens` pour les modèles o1) |
|
||||
| `thinking_level` | string | Non | Niveau de pensée étendue : `off`, `low`, `medium`, `high`, `xhigh` ou `adaptive` |
|
||||
| `extra_body` | object | Non | Champs supplémentaires à injecter dans chaque corps de requête |
|
||||
| `rpm` | int | Non | Limite de requêtes par minute |
|
||||
| `fallbacks` | string[] | Non | Noms des modèles de secours pour le basculement automatique |
|
||||
| `enabled` | bool | Non | Activer ou désactiver cette entrée de modèle (par défaut : `true`) |
|
||||
|
||||
#### Exemples par Vendor
|
||||
|
||||
**OpenAI**
|
||||
|
|
@ -190,6 +208,7 @@ Pour l'accès direct à l'API Anthropic ou les endpoints personnalisés qui ne p
|
|||
"model": "openai/custom-model",
|
||||
"api_base": "https://my-proxy.com/v1",
|
||||
"api_keys": ["sk-..."],
|
||||
"user_agent": "MyApp/1.0",
|
||||
"request_timeout": 300
|
||||
}
|
||||
```
|
||||
|
|
|
|||
|
|
@ -1,219 +0,0 @@
|
|||
# ⚙️ Guida alla Configurazione
|
||||
|
||||
> Torna al [README](../../README.md)
|
||||
|
||||
## ⚙️ Configurazione
|
||||
|
||||
File di configurazione: `~/.picoclaw/config.json`
|
||||
|
||||
### Variabili d'Ambiente
|
||||
|
||||
Puoi sovrascrivere i percorsi predefiniti usando variabili d'ambiente. Questo è utile per installazioni portatili, distribuzioni containerizzate, o per eseguire picoclaw come servizio di sistema. Queste variabili sono indipendenti e controllano percorsi diversi.
|
||||
|
||||
| Variabile | Descrizione | Percorso Predefinito |
|
||||
|-------------------|-----------------------------------------------------------------------------------------------------------------------------------------|---------------------------|
|
||||
| `PICOCLAW_CONFIG` | Sovrascrive il percorso al file di configurazione. Indica direttamente a picoclaw quale `config.json` caricare, ignorando tutte le altre posizioni. | `~/.picoclaw/config.json` |
|
||||
| `PICOCLAW_HOME` | Sovrascrive la directory radice per i dati di picoclaw. Modifica la posizione predefinita del `workspace` e delle altre directory dati. | `~/.picoclaw` |
|
||||
|
||||
**Esempi:**
|
||||
|
||||
```bash
|
||||
# Esegui picoclaw usando un file di configurazione specifico
|
||||
# Il percorso del workspace verrà letto da quel file di configurazione
|
||||
PICOCLAW_CONFIG=/etc/picoclaw/production.json picoclaw gateway
|
||||
|
||||
# Esegui picoclaw con tutti i dati salvati in /opt/picoclaw
|
||||
# La configurazione verrà caricata dal percorso predefinito ~/.picoclaw/config.json
|
||||
# Il workspace verrà creato in /opt/picoclaw/workspace
|
||||
PICOCLAW_HOME=/opt/picoclaw picoclaw agent
|
||||
|
||||
# Usa entrambi per un setup completamente personalizzato
|
||||
PICOCLAW_HOME=/srv/picoclaw PICOCLAW_CONFIG=/srv/picoclaw/main.json picoclaw gateway
|
||||
```
|
||||
|
||||
### Struttura del Workspace
|
||||
|
||||
PicoClaw salva i dati nel workspace configurato (predefinito: `~/.picoclaw/workspace`):
|
||||
|
||||
```
|
||||
~/.picoclaw/workspace/
|
||||
├── sessions/ # Sessioni di conversazione e cronologia
|
||||
├── memory/ # Memoria a lungo termine (MEMORY.md)
|
||||
├── state/ # Stato persistente (ultimo canale, ecc.)
|
||||
├── cron/ # Database dei job pianificati
|
||||
├── skills/ # Skill personalizzate
|
||||
├── AGENTS.md # Guida al comportamento dell'agent
|
||||
├── HEARTBEAT.md # Prompt per task periodici (controllato ogni 30 min)
|
||||
├── IDENTITY.md # Identità dell'agent
|
||||
├── SOUL.md # Anima dell'agent
|
||||
└── USER.md # Preferenze dell'utente
|
||||
```
|
||||
|
||||
> **Nota:** Le modifiche a `AGENTS.md`, `SOUL.md`, `USER.md`, `IDENTITY.md` e `memory/MEMORY.md` vengono rilevate automaticamente a runtime tramite il tracciamento della data di modifica (mtime). **Non è necessario riavviare il gateway** dopo aver modificato questi file — l'agent caricherà il nuovo contenuto alla prossima richiesta.
|
||||
|
||||
### Sorgenti delle Skill
|
||||
|
||||
Per impostazione predefinita, le skill vengono caricate da:
|
||||
|
||||
1. `~/.picoclaw/workspace/skills` (workspace)
|
||||
2. `~/.picoclaw/skills` (globale)
|
||||
3. `<current-working-directory>/skills` (builtin)
|
||||
|
||||
Per configurazioni avanzate/di test, puoi sovrascrivere la directory radice delle skill builtin con:
|
||||
|
||||
```bash
|
||||
export PICOCLAW_BUILTIN_SKILLS=/path/to/skills
|
||||
```
|
||||
|
||||
### Politica Unificata di Esecuzione dei Comandi
|
||||
|
||||
- I comandi slash generici vengono eseguiti tramite un unico percorso in `pkg/agent/loop.go` via `commands.Executor`.
|
||||
- Gli adattatori dei canali non consumano più localmente i comandi generici; inoltrano il testo in entrata al percorso bus/agent. Telegram registra ancora automaticamente i comandi supportati all'avvio.
|
||||
- Un comando slash sconosciuto (ad esempio `/foo`) viene passato all'elaborazione LLM come se fosse un messaggio dell'utente.
|
||||
- Un comando registrato ma non supportato sul canale corrente (ad esempio `/show` su WhatsApp) restituisce un errore esplicito all'utente e interrompe l'elaborazione.
|
||||
|
||||
### 🔒 Sandbox di Sicurezza
|
||||
|
||||
PicoClaw esegue in un ambiente sandboxed per impostazione predefinita. L'agent può accedere solo ai file ed eseguire comandi all'interno del workspace configurato.
|
||||
|
||||
#### Configurazione Predefinita
|
||||
|
||||
```json
|
||||
{
|
||||
"agents": {
|
||||
"defaults": {
|
||||
"workspace": "~/.picoclaw/workspace",
|
||||
"restrict_to_workspace": true
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
| Opzione | Predefinito | Descrizione |
|
||||
| ----------------------- | ----------------------- | ---------------------------------------------------- |
|
||||
| `workspace` | `~/.picoclaw/workspace` | Directory di lavoro dell'agent |
|
||||
| `restrict_to_workspace` | `true` | Limita l'accesso a file/comandi al workspace |
|
||||
|
||||
#### Strumenti Protetti
|
||||
|
||||
Quando `restrict_to_workspace: true`, i seguenti strumenti sono in sandbox:
|
||||
|
||||
| Strumento | Funzione | Restrizione |
|
||||
| ------------- | ------------------------- | ---------------------------------------------------- |
|
||||
| `read_file` | Legge file | Solo file all'interno del workspace |
|
||||
| `write_file` | Scrive file | Solo file all'interno del workspace |
|
||||
| `list_dir` | Elenca directory | Solo directory all'interno del workspace |
|
||||
| `edit_file` | Modifica file | Solo file all'interno del workspace |
|
||||
| `append_file` | Aggiunge ai file | Solo file all'interno del workspace |
|
||||
| `exec` | Esegue comandi | I percorsi dei comandi devono essere nel workspace |
|
||||
|
||||
#### Protezione Exec Aggiuntiva
|
||||
|
||||
Anche con `restrict_to_workspace: false`, lo strumento `exec` blocca questi comandi pericolosi:
|
||||
|
||||
* `rm -rf`, `del /f`, `rmdir /s` — Cancellazione di massa
|
||||
* `format`, `mkfs`, `diskpart` — Formattazione del disco
|
||||
* `dd if=` — Imaging del disco
|
||||
* Scrittura su `/dev/sd[a-z]` — Scritture dirette su disco
|
||||
* `shutdown`, `reboot`, `poweroff` — Spegnimento del sistema
|
||||
* Fork bomb `:(){ :|:& };:`
|
||||
|
||||
### Controllo Accesso ai File
|
||||
|
||||
| Chiave di configurazione | Tipo | Predefinito | Descrizione |
|
||||
|--------------------------|------|-------------|-------------|
|
||||
| `tools.allow_read_paths` | string[] | `[]` | Percorsi aggiuntivi consentiti per la lettura al di fuori del workspace |
|
||||
| `tools.allow_write_paths` | string[] | `[]` | Percorsi aggiuntivi consentiti per la scrittura al di fuori del workspace |
|
||||
|
||||
### Sicurezza Exec
|
||||
|
||||
| Chiave di configurazione | Tipo | Predefinito | Descrizione |
|
||||
|--------------------------|------|-------------|-------------|
|
||||
| `tools.exec.allow_remote` | bool | `false` | Consente lo strumento exec da canali remoti (Telegram/Discord ecc.) |
|
||||
| `tools.exec.enable_deny_patterns` | bool | `true` | Abilita l'intercettazione dei comandi pericolosi |
|
||||
| `tools.exec.custom_deny_patterns` | string[] | `[]` | Pattern regex personalizzati da bloccare |
|
||||
| `tools.exec.custom_allow_patterns` | string[] | `[]` | Pattern regex personalizzati da consentire |
|
||||
|
||||
> **Nota di sicurezza:** La protezione dei symlink è abilitata per impostazione predefinita — tutti i percorsi file vengono risolti tramite `filepath.EvalSymlinks` prima del confronto con la whitelist, prevenendo attacchi di escape tramite symlink.
|
||||
|
||||
#### Limitazione Nota: Processi Figlio degli Strumenti di Build
|
||||
|
||||
Il controllo di sicurezza exec ispeziona solo la riga di comando avviata direttamente da PicoClaw. Non ispeziona ricorsivamente i processi figlio generati da strumenti di sviluppo consentiti come `make`, `go run`, `cargo`, `npm run` o script di build personalizzati.
|
||||
|
||||
Ciò significa che un comando di primo livello può comunque compilare o avviare altri binari dopo aver superato il controllo iniziale. In pratica, tratta gli script di build, i Makefile, gli script di pacchetti e i binari generati come codice eseguibile che richiede lo stesso livello di revisione di un comando shell diretto.
|
||||
|
||||
Per ambienti ad alto rischio:
|
||||
|
||||
* Esamina gli script di build prima dell'esecuzione.
|
||||
* Preferisci l'approvazione/revisione manuale per i workflow di compilazione ed esecuzione.
|
||||
* Esegui PicoClaw in un container o VM se hai bisogno di un isolamento più forte di quello fornito dal controllo integrato.
|
||||
|
||||
#### Esempi di Errore
|
||||
|
||||
```
|
||||
[ERROR] tool: Tool execution failed
|
||||
{tool=exec, error=Command blocked by safety guard (path outside working dir)}
|
||||
```
|
||||
|
||||
```
|
||||
[ERROR] tool: Tool execution failed
|
||||
{tool=exec, error=Command blocked by safety guard (dangerous pattern detected)}
|
||||
```
|
||||
|
||||
#### Disabilitare le Restrizioni (Rischio di Sicurezza)
|
||||
|
||||
Se hai bisogno che l'agent acceda a percorsi al di fuori del workspace:
|
||||
|
||||
**Metodo 1: File di configurazione**
|
||||
|
||||
```json
|
||||
{
|
||||
"agents": {
|
||||
"defaults": {
|
||||
"restrict_to_workspace": false
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Metodo 2: Variabile d'ambiente**
|
||||
|
||||
```bash
|
||||
export PICOCLAW_AGENTS_DEFAULTS_RESTRICT_TO_WORKSPACE=false
|
||||
```
|
||||
|
||||
> ⚠️ **Attenzione**: Disabilitare questa restrizione consente all'agent di accedere a qualsiasi percorso sul tuo sistema. Usare con cautela solo in ambienti controllati.
|
||||
|
||||
#### Coerenza dei Confini di Sicurezza
|
||||
|
||||
L'impostazione `restrict_to_workspace` si applica in modo coerente a tutti i percorsi di esecuzione:
|
||||
|
||||
| Percorso di esecuzione | Confine di sicurezza |
|
||||
| ---------------------- | --------------------------------- |
|
||||
| Main Agent | `restrict_to_workspace` ✅ |
|
||||
| Subagent / Spawn | Eredita la stessa restrizione ✅ |
|
||||
| Heartbeat tasks | Eredita la stessa restrizione ✅ |
|
||||
|
||||
Tutti i percorsi condividono la stessa restrizione del workspace — non è possibile aggirare il confine di sicurezza tramite subagent o task pianificati.
|
||||
|
||||
### Heartbeat (Task Periodici)
|
||||
|
||||
PicoClaw può eseguire task periodici automaticamente. Crea un file `HEARTBEAT.md` nel tuo workspace:
|
||||
|
||||
```markdown
|
||||
# Periodic Tasks
|
||||
|
||||
- Check my email for important messages
|
||||
- Review my calendar for upcoming events
|
||||
- Check the weather forecast
|
||||
```
|
||||
|
||||
L'agent leggerà questo file ogni 30 minuti (configurabile) ed eseguirà tutti i task usando gli strumenti disponibili.
|
||||
|
||||
#### Task Asincroni con Spawn
|
||||
|
||||
Per task di lunga durata (ricerca web, chiamate API), usa lo strumento `spawn` per creare un **subagent**:
|
||||
|
||||
```markdown
|
||||
# Periodic Tasks
|
||||
```
|
||||
|
|
@ -99,6 +99,24 @@
|
|||
}
|
||||
```
|
||||
|
||||
#### `model_list` エントリフィールド
|
||||
|
||||
| フィールド | 型 | 必須 | 説明 |
|
||||
|-----------|------|------|------|
|
||||
| `model_name` | string | はい | agent 設定でこのモデルを参照するための一意の名前 |
|
||||
| `model` | string | はい | ベンダー/モデル識別子（例：`openai/gpt-5.4`、`azure/gpt-5.4`、`anthropic/claude-sonnet-4.6`） |
|
||||
| `api_keys` | string[] | はい* | 認証キー。複数キーでリクエストごとのローテーションが可能。ローカル provider（Ollama、LM Studio、VLLM）には不要 |
|
||||
| `api_base` | string | いいえ | デフォルトの API エンドポイント URL を上書き |
|
||||
| `proxy` | string | いいえ | このモデルエントリの HTTP プロキシ URL |
|
||||
| `user_agent` | string | いいえ | カスタム `User-Agent` リクエストヘッダー（OpenAI 互換、Anthropic、Azure provider で対応） |
|
||||
| `request_timeout` | int | いいえ | リクエストタイムアウト（秒）。デフォルト値は provider により異なる |
|
||||
| `max_tokens_field` | string | いいえ | リクエストボディの max tokens フィールド名を上書き（例：o1 モデルでは `max_completion_tokens`） |
|
||||
| `thinking_level` | string | いいえ | 拡張思考レベル：`off`、`low`、`medium`、`high`、`xhigh`、`adaptive` |
|
||||
| `extra_body` | object | いいえ | 各リクエストボディに注入する追加フィールド |
|
||||
| `rpm` | int | いいえ | 1 分あたりのリクエストレート制限 |
|
||||
| `fallbacks` | string[] | いいえ | 自動フェイルオーバーのフォールバックモデル名 |
|
||||
| `enabled` | bool | いいえ | このモデルエントリを有効にするかどうか（デフォルト：`true`） |
|
||||
|
||||
#### ベンダー別設定例
|
||||
|
||||
**OpenAI**
|
||||
|
|
@ -201,6 +219,7 @@ Anthropic API への直接アクセスや、Anthropic のネイティブメッ
|
|||
"model": "openai/custom-model",
|
||||
"api_base": "https://my-proxy.com/v1",
|
||||
"api_keys": ["sk-..."],
|
||||
"user_agent": "MyApp/1.0",
|
||||
"request_timeout": 300
|
||||
}
|
||||
```
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@
|
|||
| `openrouter` | LLM (recommended, access to all models) | [openrouter.ai](https://openrouter.ai) |
|
||||
| `anthropic` | LLM (Claude direct) | [console.anthropic.com](https://console.anthropic.com) |
|
||||
| `openai` | LLM (GPT direct) | [platform.openai.com](https://platform.openai.com) |
|
||||
| `venice` | LLM (Venice AI direct) | [venice.ai](https://venice.ai) |
|
||||
| `deepseek` | LLM (DeepSeek direct) | [platform.deepseek.com](https://platform.deepseek.com) |
|
||||
| `qwen` | LLM (Qwen direct) | [dashscope.console.aliyun.com](https://dashscope.console.aliyun.com) |
|
||||
| `groq` | LLM + **Voice transcription** (Whisper) | [console.groq.com](https://console.groq.com) |
|
||||
|
|
@ -46,6 +47,7 @@ This design also enables **multi-agent support** with flexible provider selectio
|
|||
| Vendor | `model` Prefix | Default API Base | Protocol | API Key |
|
||||
| ------------------- | ----------------- |-----------------------------------------------------| --------- | ---------------------------------------------------------------- |
|
||||
| **OpenAI** | `openai/` | `https://api.openai.com/v1` | OpenAI | [Get Key](https://platform.openai.com) |
|
||||
| **Venice AI** | `venice/` | `https://api.venice.ai/api/v1` | OpenAI | [Get Key](https://venice.ai) |
|
||||
| **Anthropic** | `anthropic/` | `https://api.anthropic.com/v1` | Anthropic | [Get Key](https://console.anthropic.com) |
|
||||
| **智谱 AI (GLM)** | `zhipu/` | `https://open.bigmodel.cn/api/paas/v4` | OpenAI | [Get Key](https://open.bigmodel.cn/usercenter/proj-mgmt/apikeys) |
|
||||
| **Z.AI Coding Plan** | `openai/` | `https://api.z.ai/api/coding/paas/v4` | OpenAI | [Get Key](https://z.ai/manage-apikey/apikey-list) |
|
||||
|
|
@ -106,6 +108,24 @@ This design also enables **multi-agent support** with flexible provider selectio
|
|||
}
|
||||
```
|
||||
|
||||
#### `model_list` Entry Fields
|
||||
|
||||
| Field | Type | Required | Description |
|
||||
|-------|------|----------|-------------|
|
||||
| `model_name` | string | Yes | Unique name used to reference this model in agent config |
|
||||
| `model` | string | Yes | Vendor/model identifier (e.g., `openai/gpt-5.4`, `azure/gpt-5.4`, `anthropic/claude-sonnet-4.6`) |
|
||||
| `api_keys` | string[] | Yes* | API key(s) for authentication. Multiple keys enable per-request rotation. Not required for local providers (Ollama, LM Studio, VLLM) |
|
||||
| `api_base` | string | No | Override the default API endpoint URL |
|
||||
| `proxy` | string | No | HTTP proxy URL for this model entry |
|
||||
| `user_agent` | string | No | Custom `User-Agent` header sent with API requests (supported by OpenAI-compatible, Anthropic, and Azure providers) |
|
||||
| `request_timeout` | int | No | Request timeout in seconds (default varies by provider) |
|
||||
| `max_tokens_field` | string | No | Override the max tokens field name in request body (e.g., `max_completion_tokens` for o1 models) |
|
||||
| `thinking_level` | string | No | Extended thinking level: `off`, `low`, `medium`, `high`, `xhigh`, or `adaptive` |
|
||||
| `extra_body` | object | No | Additional fields to inject into every request body |
|
||||
| `rpm` | int | No | Per-minute request rate limit |
|
||||
| `fallbacks` | string[] | No | Fallback model names for automatic failover |
|
||||
| `enabled` | bool | No | Whether this model entry is active (default: `true`) |
|
||||
|
||||
#### Voice Transcription
|
||||
|
||||
You can configure a dedicated model for audio transcription with `voice.model_name`. This lets you reuse existing multimodal providers that support audio input instead of relying only on Groq.
|
||||
|
|
@ -247,6 +267,7 @@ PicoClaw sends OpenAI-compatible requests to LM Studio, and strips the `lmstudio
|
|||
"model": "openai/custom-model",
|
||||
"api_base": "https://my-proxy.com/v1",
|
||||
"api_keys": ["sk-..."],
|
||||
"user_agent": "MyApp/1.0",
|
||||
"request_timeout": 300
|
||||
}
|
||||
```
|
||||
|
|
|
|||
|
|
@ -99,6 +99,24 @@ Este design também permite **suporte multi-agente** com seleção flexível de
|
|||
}
|
||||
```
|
||||
|
||||
#### Campos de entrada `model_list`
|
||||
|
||||
| Campo | Tipo | Obrigatório | Descrição |
|
||||
|-------|------|-------------|-----------|
|
||||
| `model_name` | string | Sim | Nome único para referenciar este modelo na config do agent |
|
||||
| `model` | string | Sim | Identificador fornecedor/modelo (ex: `openai/gpt-5.4`, `azure/gpt-5.4`, `anthropic/claude-sonnet-4.6`) |
|
||||
| `api_keys` | string[] | Sim* | Chave(s) API para autenticação. Múltiplas chaves permitem rotação por requisição. Não necessário para providers locais (Ollama, LM Studio, VLLM) |
|
||||
| `api_base` | string | Não | Substitui a URL base da API padrão |
|
||||
| `proxy` | string | Não | URL do proxy HTTP para esta entrada de modelo |
|
||||
| `user_agent` | string | Não | Cabeçalho `User-Agent` personalizado enviado com requisições API (suportado por providers OpenAI-compatible, Anthropic e Azure) |
|
||||
| `request_timeout` | int | Não | Timeout de requisição em segundos (o padrão varia por provider) |
|
||||
| `max_tokens_field` | string | Não | Substitui o nome do campo max tokens no corpo da requisição (ex: `max_completion_tokens` para modelos o1) |
|
||||
| `thinking_level` | string | Não | Nível de pensamento estendido: `off`, `low`, `medium`, `high`, `xhigh` ou `adaptive` |
|
||||
| `extra_body` | object | Não | Campos adicionais para injetar em cada corpo de requisição |
|
||||
| `rpm` | int | Não | Limite de requisições por minuto |
|
||||
| `fallbacks` | string[] | Não | Nomes dos modelos de fallback para failover automático |
|
||||
| `enabled` | bool | Não | Ativar ou desativar esta entrada de modelo (padrão: `true`) |
|
||||
|
||||
#### Exemplos por Vendor
|
||||
|
||||
**OpenAI**
|
||||
|
|
@ -190,6 +208,7 @@ Para acesso direto à API Anthropic ou endpoints personalizados que suportam ape
|
|||
"model": "openai/custom-model",
|
||||
"api_base": "https://my-proxy.com/v1",
|
||||
"api_keys": ["sk-..."],
|
||||
"user_agent": "MyApp/1.0",
|
||||
"request_timeout": 300
|
||||
}
|
||||
```
|
||||
|
|
|
|||
95
docs/rate-limiting.md
Normal file
95
docs/rate-limiting.md
Normal file
|
|
@ -0,0 +1,95 @@
|
|||
# Dynamic Rate Limiting
|
||||
|
||||
PicoClaw prevents 429 errors from LLM provider APIs by enforcing configurable per-model request-rate limits **before** sending each request. Unlike the reactive cooldown/fallback system (which activates *after* a 429 is received), rate limiting is **proactive**: it keeps outbound QPS within the provider's free-tier or plan limits.
|
||||
|
||||
## How it works
|
||||
|
||||
### Token-bucket algorithm
|
||||
|
||||
Each rate-limited model gets a token bucket:
|
||||
|
||||
- **Capacity** = `rpm` (burst size equals the per-minute limit)
|
||||
- **Refill rate** = `rpm / 60` tokens per second
|
||||
- Tokens are consumed one per LLM call; if the bucket is empty, the call blocks until a token refills or the request context is cancelled
|
||||
|
||||
### Call chain integration
|
||||
|
||||
```
|
||||
AgentLoop.callLLM()
|
||||
└─ FallbackChain.Execute() ← iterate candidates
|
||||
├─ CooldownTracker.IsAvailable() ← skip if post-429 cooldown active
|
||||
├─ RateLimiterRegistry.Wait() ← NEW: block until token available
|
||||
└─ provider.Chat() ← actual LLM HTTP call
|
||||
```
|
||||
|
||||
The rate limiter runs **after** the cooldown check and **before** the provider call, so:
|
||||
- Candidates already in cooldown are skipped entirely (no token consumed)
|
||||
- Candidates that are available get throttled to the configured RPM
|
||||
|
||||
The same check applies in `ExecuteImage`.
|
||||
|
||||
### Thread safety
|
||||
|
||||
`RateLimiterRegistry` is safe for concurrent use. The per-limiter token bucket uses a fine-grained mutex so concurrent goroutines each acquire their own token independently.
|
||||
|
||||
## Configuration
|
||||
|
||||
Set `rpm` on any model in `model_list`:
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: gpt-4o-free
|
||||
model: openai/gpt-4o
|
||||
api_base: https://api.openai.com/v1
|
||||
rpm: 3 # max 3 requests per minute
|
||||
api_keys:
|
||||
- sk-...
|
||||
|
||||
- model_name: claude-haiku
|
||||
model: anthropic/claude-haiku-4-5
|
||||
rpm: 60 # 60 rpm (Anthropic free tier)
|
||||
api_keys:
|
||||
- sk-ant-...
|
||||
|
||||
- model_name: local-llm
|
||||
model: openai/llama3
|
||||
api_base: http://localhost:11434/v1
|
||||
# no rpm → unrestricted
|
||||
```
|
||||
|
||||
| Field | Type | Default | Description |
|
||||
|---|---|---|---|
|
||||
| `rpm` | `int` | `0` | Requests per minute. `0` means no limit. |
|
||||
|
||||
### Interaction with fallbacks
|
||||
|
||||
When a model has fallbacks configured, each candidate is rate-limited **independently**:
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: gpt4-with-fallback
|
||||
model: openai/gpt-4o
|
||||
rpm: 5
|
||||
fallbacks:
|
||||
- gpt-4o-mini # must also be in model_list; its own rpm applies
|
||||
```
|
||||
|
||||
If the current candidate's bucket is empty and there are more candidates available, PicoClaw skips the locally saturated candidate and tries the next fallback immediately. Only the last remaining candidate waits for a token to refill. If the context deadline is hit while waiting on that last candidate, the wait error propagates.
|
||||
|
||||
For `model_list` aliases that resolve to the same underlying provider/model, rate limiting is keyed by the stable config identity (for example `model_name`) rather than the resolved runtime model string. This preserves distinct RPM settings for multi-key and alias-based configurations.
|
||||
|
||||
### Burst behaviour
|
||||
|
||||
The bucket starts **full** (burst = RPM). For `rpm: 3`, the first 3 requests fire instantly; subsequent requests are spaced ~20 s apart.
|
||||
|
||||
To reduce burstiness for strict APIs, set a lower `rpm` and rely on the steady-state refill.
|
||||
|
||||
## Files changed
|
||||
|
||||
| File | What |
|
||||
|---|---|
|
||||
| `pkg/providers/ratelimiter.go` | `RateLimiter` (token bucket) + `RateLimiterRegistry` |
|
||||
| `pkg/providers/ratelimiter_test.go` | Unit tests for limiter and registry |
|
||||
| `pkg/providers/fallback.go` | `FallbackCandidate.RPM` field; `FallbackChain.rl`; `Wait()` call in `Execute`/`ExecuteImage` |
|
||||
| `pkg/agent/model_resolution.go` | Resolves candidates from `model_list`, preserving stable config identity and propagating `RPM` into `FallbackCandidate` |
|
||||
| `pkg/agent/loop.go` | Build `RateLimiterRegistry`, register all agents' candidates, pass to `NewFallbackChain` |
|
||||
|
|
@ -528,6 +528,9 @@ For example:
|
|||
- `PICOCLAW_TOOLS_EXEC_ENABLE_DENY_PATTERNS=false`
|
||||
- `PICOCLAW_TOOLS_CRON_EXEC_TIMEOUT_MINUTES=10`
|
||||
- `PICOCLAW_TOOLS_MCP_ENABLED=true`
|
||||
- `PICOCLAW_TOOLS_MCP_MAX_INLINE_TEXT_CHARS=16384`
|
||||
|
||||
Note: Nested map-style config (for example `tools.mcp.servers.<name>.*`) is configured in `config.json` rather than
|
||||
environment variables.
|
||||
|
||||
For MCP tools, `tools.mcp.max_inline_text_chars` controls how much text result is kept inline in model context. The threshold is counted in Unicode characters (Go runes), not bytes. For example, `16384` means up to 16,384 characters inline, which may occupy more than 16 KB for multibyte text such as CJK. Above this threshold, PicoClaw saves the MCP text result as a local artifact in the agent workspace and gives the model a short note plus a structured `[file:...]` artifact path instead of injecting the full payload into context.
|
||||
|
|
|
|||
|
|
@ -99,6 +99,24 @@ Thiết kế này cũng cho phép **hỗ trợ đa agent** với lựa chọn pr
|
|||
}
|
||||
```
|
||||
|
||||
#### Các trường entry `model_list`
|
||||
|
||||
| Trường | Kiểu | Bắt buộc | Mô tả |
|
||||
|--------|------|----------|------|
|
||||
| `model_name` | string | Có | Tên duy nhất để tham chiếu model này trong cấu hình agent |
|
||||
| `model` | string | Có | Định danh nhà cung cấp/model (ví dụ: `openai/gpt-5.4`, `azure/gpt-5.4`, `anthropic/claude-sonnet-4.6`) |
|
||||
| `api_keys` | string[] | Có* | Khóa API xác thực. Nhiều khóa cho phép xoay vòng theo yêu cầu. Không cần thiết cho provider nội bộ (Ollama, LM Studio, VLLM) |
|
||||
| `api_base` | string | Không | Ghi đè URL endpoint API mặc định |
|
||||
| `proxy` | string | Không | URL proxy HTTP cho entry model này |
|
||||
| `user_agent` | string | Không | Header `User-Agent` tùy chỉnh gửi với yêu cầu API (được hỗ trợ bởi provider OpenAI-compatible, Anthropic và Azure) |
|
||||
| `request_timeout` | int | Không | Timeout yêu cầu tính bằng giây (mặc định khác nhau tùy provider) |
|
||||
| `max_tokens_field` | string | Không | Ghi đè tên trường max tokens trong request body (ví dụ: `max_completion_tokens` cho model o1) |
|
||||
| `thinking_level` | string | Không | Mức độ tư duy mở rộng: `off`, `low`, `medium`, `high`, `xhigh` hoặc `adaptive` |
|
||||
| `extra_body` | object | Không | Các trường bổ sung để chèn vào mỗi request body |
|
||||
| `rpm` | int | Không | Giới hạn tốc độ yêu cầu mỗi phút |
|
||||
| `fallbacks` | string[] | Không | Tên model dự phòng cho failover tự động |
|
||||
| `enabled` | bool | Không | Kích hoạt hay vô hiệu hóa entry model này (mặc định: `true`) |
|
||||
|
||||
#### Ví Dụ Theo Vendor
|
||||
|
||||
**OpenAI**
|
||||
|
|
@ -190,6 +208,7 @@ Thiết kế này cũng cho phép **hỗ trợ đa agent** với lựa chọn pr
|
|||
"model": "openai/custom-model",
|
||||
"api_base": "https://my-proxy.com/v1",
|
||||
"api_keys": ["sk-..."],
|
||||
"user_agent": "MyApp/1.0",
|
||||
"request_timeout": 300
|
||||
}
|
||||
```
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@
|
|||
| `openrouter` | LLM (推荐，可访问所有模型) | [openrouter.ai](https://openrouter.ai) |
|
||||
| `anthropic` | LLM (Claude 直连) | [console.anthropic.com](https://console.anthropic.com) |
|
||||
| `openai` | LLM (GPT 直连) | [platform.openai.com](https://platform.openai.com) |
|
||||
| `venice` | LLM (Venice AI 直连) | [venice.ai](https://venice.ai) |
|
||||
| `deepseek` | LLM (DeepSeek 直连) | [platform.deepseek.com](https://platform.deepseek.com) |
|
||||
| `qwen` | LLM (通义千问) | [dashscope.console.aliyun.com](https://dashscope.console.aliyun.com) |
|
||||
| `groq` | LLM + **语音转录** (Whisper) | [console.groq.com](https://console.groq.com) |
|
||||
|
|
@ -44,6 +45,7 @@
|
|||
| 厂商 | `model` 前缀 | 默认 API Base | 协议 | 获取 API Key |
|
||||
| ------------------- | ----------------- | --------------------------------------------------- | --------- | ----------------------------------------------------------------- |
|
||||
| **OpenAI** | `openai/` | `https://api.openai.com/v1` | OpenAI | [获取密钥](https://platform.openai.com) |
|
||||
| **Venice AI** | `venice/` | `https://api.venice.ai/api/v1` | OpenAI | [获取密钥](https://venice.ai) |
|
||||
| **Anthropic** | `anthropic/` | `https://api.anthropic.com/v1` | Anthropic | [获取密钥](https://console.anthropic.com) |
|
||||
| **智谱 AI (GLM)** | `zhipu/` | `https://open.bigmodel.cn/api/paas/v4` | OpenAI | [获取密钥](https://open.bigmodel.cn/usercenter/proj-mgmt/apikeys) |
|
||||
| **DeepSeek** | `deepseek/` | `https://api.deepseek.com/v1` | OpenAI | [获取密钥](https://platform.deepseek.com) |
|
||||
|
|
@ -102,6 +104,24 @@
|
|||
}
|
||||
```
|
||||
|
||||
#### `model_list` 条目字段
|
||||
|
||||
| 字段 | 类型 | 必填 | 说明 |
|
||||
|------|------|------|------|
|
||||
| `model_name` | string | 是 | 在 agent 配置中引用此模型的唯一名称 |
|
||||
| `model` | string | 是 | 厂商/模型标识符（如 `openai/gpt-5.4`、`azure/gpt-5.4`、`anthropic/claude-sonnet-4.6`） |
|
||||
| `api_keys` | string[] | 是* | 认证密钥。多个密钥可按请求轮换。本地 provider（Ollama、LM Studio、VLLM）不需要 |
|
||||
| `api_base` | string | 否 | 覆盖默认的 API 端点 URL |
|
||||
| `proxy` | string | 否 | 此模型条目的 HTTP 代理 URL |
|
||||
| `user_agent` | string | 否 | 自定义 `User-Agent` 请求头（支持 OpenAI 兼容、Anthropic 和 Azure provider） |
|
||||
| `request_timeout` | int | 否 | 请求超时时间（秒），默认值因 provider 而异 |
|
||||
| `max_tokens_field` | string | 否 | 覆盖请求体中 max tokens 的字段名（如 o1 模型使用 `max_completion_tokens`） |
|
||||
| `thinking_level` | string | 否 | 扩展思考级别：`off`、`low`、`medium`、`high`、`xhigh` 或 `adaptive` |
|
||||
| `extra_body` | object | 否 | 注入到每个请求体中的额外字段 |
|
||||
| `rpm` | int | 否 | 每分钟请求速率限制 |
|
||||
| `fallbacks` | string[] | 否 | 自动故障转移的备用模型名称 |
|
||||
| `enabled` | bool | 否 | 是否启用此模型条目（默认：`true`） |
|
||||
|
||||
#### 语音转录
|
||||
|
||||
你可以通过 `voice.model_name` 为语音转录指定一个专用模型。这样可以直接复用已经配置好的、支持音频输入的多模态 provider，而不必只依赖 Groq。
|
||||
|
|
@ -232,6 +252,7 @@ PicoClaw 向 LM Studio 的 OpenAI 兼容终结点发送请求，且将移除首
|
|||
"model": "openai/custom-model",
|
||||
"api_base": "https://my-proxy.com/v1",
|
||||
"api_keys": ["sk-..."],
|
||||
"user_agent": "MyApp/1.0",
|
||||
"request_timeout": 300
|
||||
}
|
||||
```
|
||||
|
|
|
|||
9
go.mod
9
go.mod
|
|
@ -5,6 +5,7 @@ go 1.25.8
|
|||
require (
|
||||
fyne.io/systray v1.12.0
|
||||
github.com/BurntSushi/toml v1.6.0
|
||||
github.com/SevereCloud/vksdk/v3 v3.3.1
|
||||
github.com/adhocore/gronx v1.19.6
|
||||
github.com/anthropics/anthropic-sdk-go v1.26.0
|
||||
github.com/atotto/clipboard v0.1.4
|
||||
|
|
@ -23,6 +24,7 @@ require (
|
|||
github.com/h2non/filetype v1.1.3
|
||||
github.com/larksuite/oapi-sdk-go/v3 v3.5.3
|
||||
github.com/mdp/qrterminal/v3 v3.2.1
|
||||
github.com/minio/selfupdate v0.6.0
|
||||
github.com/modelcontextprotocol/go-sdk v1.4.1
|
||||
github.com/mymmrac/telego v1.7.0
|
||||
github.com/open-dingtalk/dingtalk-stream-sdk-go v0.9.1
|
||||
|
|
@ -30,7 +32,7 @@ require (
|
|||
github.com/pion/rtp v1.8.7
|
||||
github.com/pion/webrtc/v3 v3.3.6
|
||||
github.com/rivo/tview v0.42.0
|
||||
github.com/rs/zerolog v1.34.0
|
||||
github.com/rs/zerolog v1.35.0
|
||||
github.com/slack-go/slack v0.17.3
|
||||
github.com/spf13/cobra v1.10.2
|
||||
github.com/stretchr/testify v1.11.1
|
||||
|
|
@ -48,6 +50,7 @@ require (
|
|||
)
|
||||
|
||||
require (
|
||||
aead.dev/minisign v0.2.0 // indirect
|
||||
filippo.io/edwards25519 v1.2.0 // indirect
|
||||
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.8 // indirect
|
||||
github.com/aws/aws-sdk-go-v2/credentials v1.19.12 // indirect
|
||||
|
|
@ -87,6 +90,8 @@ require (
|
|||
github.com/segmentio/encoding v0.5.4 // indirect
|
||||
github.com/spf13/pflag v1.0.10 // indirect
|
||||
github.com/vektah/gqlparser/v2 v2.5.27 // indirect
|
||||
github.com/vmihailenco/msgpack/v5 v5.4.1 // indirect
|
||||
github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect
|
||||
go.mau.fi/libsignal v0.2.1 // indirect
|
||||
go.opentelemetry.io/auto/sdk v1.1.0 // indirect
|
||||
go.opentelemetry.io/otel v1.35.0 // indirect
|
||||
|
|
@ -124,7 +129,7 @@ require (
|
|||
golang.org/x/arch v0.24.0 // indirect
|
||||
golang.org/x/crypto v0.49.0
|
||||
golang.org/x/net v0.52.0
|
||||
golang.org/x/sync v0.20.0 // indirect
|
||||
golang.org/x/sync v0.20.0
|
||||
golang.org/x/sys v0.42.0
|
||||
)
|
||||
|
||||
|
|
|
|||
29
go.sum
29
go.sum
|
|
@ -1,3 +1,5 @@
|
|||
aead.dev/minisign v0.2.0 h1:kAWrq/hBRu4AARY6AlciO83xhNnW9UaC8YipS2uhLPk=
|
||||
aead.dev/minisign v0.2.0/go.mod h1:zdq6LdSd9TbuSxchxwhpA9zEb9YXcVGoE8JakuiGaIQ=
|
||||
cloud.google.com/go/compute/metadata v0.3.0/go.mod h1:zFmK7XCadkQkj6TtorcaGlCW1hT1fIilQDwofLpJ20k=
|
||||
filippo.io/edwards25519 v1.2.0 h1:crnVqOiS4jqYleHd9vaKZ+HKtHfllngJIiOpNpoJsjo=
|
||||
filippo.io/edwards25519 v1.2.0/go.mod h1:xzAOLCNug/yB62zG1bQ8uziwrIqIuxhctzJT18Q77mc=
|
||||
|
|
@ -7,6 +9,8 @@ github.com/BurntSushi/toml v1.6.0 h1:dRaEfpa2VI55EwlIW72hMRHdWouJeRF7TPYhI+AUQjk
|
|||
github.com/BurntSushi/toml v1.6.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho=
|
||||
github.com/DATA-DOG/go-sqlmock v1.5.2 h1:OcvFkGmslmlZibjAjaHm3L//6LiuBgolP7OputlJIzU=
|
||||
github.com/DATA-DOG/go-sqlmock v1.5.2/go.mod h1:88MAG/4G7SMwSE3CeA0ZKzrT5CiOU3OJ+JlNzwDqpNU=
|
||||
github.com/SevereCloud/vksdk/v3 v3.3.1 h1:O86zsp5LQnHE+O5acvuXM/s6S1LyxzVTkF6+Lup0Jyg=
|
||||
github.com/SevereCloud/vksdk/v3 v3.3.1/go.mod h1:c6WaA5aocUYsXfkcUbg2qy45V9M1VDcqHHmHIN14NAw=
|
||||
github.com/adhocore/gronx v1.19.6 h1:5KNVcoR9ACgL9HhEqCm5QXsab/gI4QDIybTAWcXDKDc=
|
||||
github.com/adhocore/gronx v1.19.6/go.mod h1:7oUY1WAU8rEJWmAxXR2DN0JaO4gi9khSgKjiRypqteg=
|
||||
github.com/agnivade/levenshtein v1.2.1 h1:EHBY3UOn1gwdy/VbFwgo4cxecRznFk7fKWN1KOX7eoM=
|
||||
|
|
@ -69,7 +73,6 @@ github.com/cloudwego/base64x v0.1.6 h1:t11wG9AECkCDk5fMSoxmufanudBtJ+/HemLstXDLI
|
|||
github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gEHfghB2IPU=
|
||||
github.com/coder/websocket v1.8.14 h1:9L0p0iKiNOibykf283eHkKUHHrpG7f65OE3BhhO7v9g=
|
||||
github.com/coder/websocket v1.8.14/go.mod h1:NX3SzP+inril6yawo5CQXx8+fk145lPDC6pumgx0mVg=
|
||||
github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
|
||||
github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
|
||||
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
|
||||
github.com/creack/pty v1.1.24 h1:bJrF4RRfyJnbTJqzRLHzcGaZK1NeM5kTC9jGgovnR1s=
|
||||
|
|
@ -108,7 +111,6 @@ github.com/go-resty/resty/v2 v2.17.1/go.mod h1:kCKZ3wWmwJaNc7S29BRtUhJwy7iqmn+2m
|
|||
github.com/go-task/slim-sprig v0.0.0-20210107165309-348f09dbbbc0/go.mod h1:fyg7847qk6SyHyPtNmDHnmrv/HOrqktSC+C9fM+CJOE=
|
||||
github.com/go-test/deep v1.1.1 h1:0r/53hagsehfO4bzD2Pgr/+RgHqhmf+k1Bpse2cTu1U=
|
||||
github.com/go-test/deep v1.1.1/go.mod h1:5C2ZWiW0ErCdrYzpqxLbTX7MG14M9iiw8DgHncVwcsE=
|
||||
github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
|
||||
github.com/godbus/dbus/v5 v5.1.0 h1:4KLkAxT3aOY8Li4FRJe/KvhoNFFxo0m6fNuFUO8QJUk=
|
||||
github.com/godbus/dbus/v5 v5.1.0/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
|
||||
github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
|
||||
|
|
@ -173,17 +175,16 @@ github.com/larksuite/oapi-sdk-go/v3 v3.5.3 h1:xvf8Dv29kBXC5/DNDCLhHkAFW8l/0LlQJi
|
|||
github.com/larksuite/oapi-sdk-go/v3 v3.5.3/go.mod h1:ZEplY+kwuIrj/nqw5uSCINNATcH3KdxSN7y+UxYY5fI=
|
||||
github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag=
|
||||
github.com/lucasb-eyer/go-colorful v1.3.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
|
||||
github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
|
||||
github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE=
|
||||
github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8=
|
||||
github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
|
||||
github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
|
||||
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
|
||||
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
|
||||
github.com/mattn/go-sqlite3 v1.14.34 h1:3NtcvcUnFBPsuRcno8pUtupspG/GM+9nZ88zgJcp6Zk=
|
||||
github.com/mattn/go-sqlite3 v1.14.34/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y=
|
||||
github.com/mdp/qrterminal/v3 v3.2.1 h1:6+yQjiiOsSuXT5n9/m60E54vdgFsw0zhADHhHLrFet4=
|
||||
github.com/mdp/qrterminal/v3 v3.2.1/go.mod h1:jOTmXvnBsMy5xqLniO0R++Jmjs2sTm9dFSuQ5kpz/SU=
|
||||
github.com/minio/selfupdate v0.6.0 h1:i76PgT0K5xO9+hjzKcacQtO7+MjJ4JKA8Ak8XQ9DDwU=
|
||||
github.com/minio/selfupdate v0.6.0/go.mod h1:bO02GTIPCMQFTEvE5h4DjYB58bCoZ35XLeBf0buTDdM=
|
||||
github.com/modelcontextprotocol/go-sdk v1.4.1 h1:M4x9GyIPj+HoIlHNGpK2hq5o3BFhC+78PkEaldQRphc=
|
||||
github.com/modelcontextprotocol/go-sdk v1.4.1/go.mod h1:Bo/mS87hPQqHSRkMv4dQq1XCu6zv4INdXnFZabkNU6s=
|
||||
github.com/mymmrac/telego v1.7.0 h1:yRO/l00tFGG4nY66ufUKb4ARqv7qx9+LsjQv/b0NEyo=
|
||||
|
|
@ -211,7 +212,6 @@ github.com/pion/rtp v1.8.7/go.mod h1:pBGHaFt/yW7bf1jjWAoUjpSNoDnw98KTMg+jWWvziqU
|
|||
github.com/pion/webrtc/v3 v3.3.6 h1:7XAh4RPtlY1Vul6/GmZrv7z+NnxKA6If0KStXBI2ZLE=
|
||||
github.com/pion/webrtc/v3 v3.3.6/go.mod h1:zyN7th4mZpV27eXybfR/cnUf3J2DRy8zw/mdjD9JTNM=
|
||||
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
|
||||
github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
|
||||
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
|
||||
|
|
@ -224,9 +224,8 @@ github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTE
|
|||
github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
|
||||
github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII=
|
||||
github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o=
|
||||
github.com/rs/xid v1.6.0/go.mod h1:7XoLgs4eV+QndskICGsho+ADou8ySMSjJKDIan90Nz0=
|
||||
github.com/rs/zerolog v1.34.0 h1:k43nTLIwcTVQAncfCw4KZ2VY6ukYoZaBPNOE8txlOeY=
|
||||
github.com/rs/zerolog v1.34.0/go.mod h1:bJsvje4Z08ROH4Nhs5iH600c3IkWhwp44iRc54W6wYQ=
|
||||
github.com/rs/zerolog v1.35.0 h1:VD0ykx7HMiMJytqINBsKcbLS+BJ4WYjz+05us+LRTdI=
|
||||
github.com/rs/zerolog v1.35.0/go.mod h1:EjML9kdfa/RMA7h/6z6pYmq1ykOuA8/mjWaEvGI+jcw=
|
||||
github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
|
||||
github.com/segmentio/asm v1.1.3 h1:WM03sfUOENvvKexOLp+pCqgb/WDjsi7EK8gIsICtzhc=
|
||||
github.com/segmentio/asm v1.1.3/go.mod h1:Ld3L4ZXGNcSLRg4JBsZ3//1+f/TjYl0Mzen/DQy1EJg=
|
||||
|
|
@ -277,6 +276,10 @@ github.com/valyala/fastjson v1.6.10 h1:/yjJg8jaVQdYR3arGxPE2X5z89xrlhS0eGXdv+ADT
|
|||
github.com/valyala/fastjson v1.6.10/go.mod h1:e6FubmQouUNP73jtMLmcbxS6ydWIpOfhz34TSfO3JaE=
|
||||
github.com/vektah/gqlparser/v2 v2.5.27 h1:RHPD3JOplpk5mP5JGX8RKZkt2/Vwj/PZv0HxTdwFp0s=
|
||||
github.com/vektah/gqlparser/v2 v2.5.27/go.mod h1:D1/VCZtV3LPnQrcPBeR/q5jkSQIPti0uYCP/RI0gIeo=
|
||||
github.com/vmihailenco/msgpack/v5 v5.4.1 h1:cQriyiUvjTwOHg8QZaPihLWeRAAVoCpE00IUPn0Bjt8=
|
||||
github.com/vmihailenco/msgpack/v5 v5.4.1/go.mod h1:GaZTsDaehaPpQVyxrf5mtQlH+pc21PIudVV/E3rRQok=
|
||||
github.com/vmihailenco/tagparser/v2 v2.0.0 h1:y09buUbR+b5aycVFQs/g70pqKVZNBmxwAhO7/IwNM9g=
|
||||
github.com/vmihailenco/tagparser/v2 v2.0.0/go.mod h1:Wri+At7QHww0WTrCBeu4J6bNtoV6mEfg5OIWRZA9qds=
|
||||
github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU=
|
||||
github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E=
|
||||
github.com/yeongaori/discordgo-fork v0.0.0-20260319072544-e8e546f5d532 h1:gxFHYeUDGziRb0zXYEqBFohC+NJbIW9L0tddaXMWr2o=
|
||||
|
|
@ -308,7 +311,9 @@ golang.org/x/arch v0.24.0/go.mod h1:dNHoOeKiyja7GTvF9NJS1l3Z2yntpQNzgrjh1cU103A=
|
|||
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
|
||||
golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
|
||||
golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
|
||||
golang.org/x/crypto v0.0.0-20210220033148-5ea612d1eb83/go.mod h1:jdWPYTVW3xRLrWPugEBEK3UY2ZEsg3UU495nc5E+M+I=
|
||||
golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
|
||||
golang.org/x/crypto v0.0.0-20211209193657-4570a0811e8b/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
|
||||
golang.org/x/crypto v0.16.0/go.mod h1:gCAAfMLgwOJRpTjQ2zCCt2OcSfYMTeZVSRtQlPC7Nq4=
|
||||
golang.org/x/crypto v0.49.0 h1:+Ng2ULVvLHnJ/ZFEq4KdcDd/cfjrrjjNSXNzxg0Y4U4=
|
||||
golang.org/x/crypto v0.49.0/go.mod h1:ErX4dUh2UM+CFYiXZRTcMpEcN8b/1gxEuv3nODoYtCA=
|
||||
|
|
@ -329,6 +334,7 @@ golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwY
|
|||
golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
|
||||
golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM=
|
||||
golang.org/x/net v0.0.0-20210428140749-89ef3d95e781/go.mod h1:OJAsFXCWl8Ukc7SiCT/9KSuxbyM7479/AVlXFRxuMCk=
|
||||
golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
|
||||
golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
|
||||
golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
|
||||
golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg=
|
||||
|
|
@ -351,24 +357,25 @@ golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5h
|
|||
golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20190904154756-749cb33beabd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20191005200804-aed5e4c7ecf9/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20191120155948-bd437916bb0e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20210112080510-489259a85091/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20210228012217-479acdf4ea46/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20210330210617-4fbd30eecc44/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.15.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo=
|
||||
golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
||||
golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw=
|
||||
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
|
||||
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
|
||||
golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
|
||||
|
|
|
|||
|
|
@ -602,14 +602,16 @@ func (cb *ContextBuilder) BuildMessages(
|
|||
// Add conversation history
|
||||
messages = append(messages, history...)
|
||||
|
||||
// Add current user message
|
||||
if strings.TrimSpace(currentMessage) != "" {
|
||||
// Add current user message. Media-only turns must still be preserved so
|
||||
// multimodal providers receive the uploaded image even when the user sends
|
||||
// no accompanying text.
|
||||
if strings.TrimSpace(currentMessage) != "" || len(media) > 0 {
|
||||
msg := providers.Message{
|
||||
Role: "user",
|
||||
Content: currentMessage,
|
||||
}
|
||||
if len(media) > 0 {
|
||||
msg.Media = media
|
||||
msg.Media = append([]string(nil), media...)
|
||||
}
|
||||
messages = append(messages, msg)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,10 +6,8 @@
|
|||
package agent
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
"github.com/sipeed/picoclaw/pkg/tokenizer"
|
||||
)
|
||||
|
||||
// parseTurnBoundaries returns the starting index of each Turn in the history.
|
||||
|
|
@ -86,88 +84,16 @@ func findSafeBoundary(history []providers.Message, targetIndex int) int {
|
|||
return 0
|
||||
}
|
||||
|
||||
// estimateMessageTokens estimates the token count for a single message,
|
||||
// including Content, ReasoningContent, ToolCalls arguments, ToolCallID
|
||||
// metadata, and Media items. Uses a heuristic of 2.5 characters per token.
|
||||
func estimateMessageTokens(msg providers.Message) int {
|
||||
contentChars := utf8.RuneCountInString(msg.Content)
|
||||
|
||||
// SystemParts are structured system blocks used for cache-aware adapters.
|
||||
// They carry the same content as Content, but in multiple blocks.
|
||||
// We estimate them as an alternative representation, not additive.
|
||||
systemPartsChars := 0
|
||||
if len(msg.SystemParts) > 0 {
|
||||
for _, part := range msg.SystemParts {
|
||||
systemPartsChars += utf8.RuneCountInString(part.Text)
|
||||
}
|
||||
// Per-part overhead for JSON structure (type, text, cache_control).
|
||||
const perPartOverhead = 20
|
||||
systemPartsChars += len(msg.SystemParts) * perPartOverhead
|
||||
// EstimateMessageTokens estimates the token count for a single message.
|
||||
// Delegates to the shared tokenizer package for consistency across agent and seahorse.
|
||||
func EstimateMessageTokens(msg providers.Message) int {
|
||||
return tokenizer.EstimateMessageTokens(msg)
|
||||
}
|
||||
|
||||
// Use the larger of the two representations to stay conservative.
|
||||
chars := contentChars
|
||||
if systemPartsChars > chars {
|
||||
chars = systemPartsChars
|
||||
}
|
||||
|
||||
chars += utf8.RuneCountInString(msg.ReasoningContent)
|
||||
|
||||
for _, tc := range msg.ToolCalls {
|
||||
chars += len(tc.ID) + len(tc.Type)
|
||||
if tc.Function != nil {
|
||||
// Count function name + arguments (the wire format for most providers).
|
||||
// tc.Name mirrors tc.Function.Name — count only once to avoid double-counting.
|
||||
chars += len(tc.Function.Name) + len(tc.Function.Arguments)
|
||||
} else {
|
||||
// Fallback: some provider formats use top-level Name without Function.
|
||||
chars += len(tc.Name)
|
||||
}
|
||||
}
|
||||
|
||||
if msg.ToolCallID != "" {
|
||||
chars += len(msg.ToolCallID)
|
||||
}
|
||||
|
||||
// Per-message overhead for role label, JSON structure, separators.
|
||||
const messageOverhead = 12
|
||||
chars += messageOverhead
|
||||
|
||||
tokens := chars * 2 / 5
|
||||
|
||||
// Media items (images, files) are serialized by provider adapters into
|
||||
// multipart or image_url payloads. Add a fixed per-item token estimate
|
||||
// directly (not through the chars heuristic) since actual cost depends
|
||||
// on resolution and provider-specific image tokenization.
|
||||
const mediaTokensPerItem = 256
|
||||
tokens += len(msg.Media) * mediaTokensPerItem
|
||||
|
||||
return tokens
|
||||
}
|
||||
|
||||
// estimateToolDefsTokens estimates the total token cost of tool definitions
|
||||
// as they appear in the LLM request. Each tool's name, description, and
|
||||
// JSON schema parameters contribute to the context window budget.
|
||||
func estimateToolDefsTokens(defs []providers.ToolDefinition) int {
|
||||
if len(defs) == 0 {
|
||||
return 0
|
||||
}
|
||||
|
||||
totalChars := 0
|
||||
for _, d := range defs {
|
||||
totalChars += len(d.Function.Name) + len(d.Function.Description)
|
||||
|
||||
if d.Function.Parameters != nil {
|
||||
if paramJSON, err := json.Marshal(d.Function.Parameters); err == nil {
|
||||
totalChars += len(paramJSON)
|
||||
}
|
||||
}
|
||||
|
||||
// Per-tool overhead: type field, JSON structure, separators.
|
||||
totalChars += 20
|
||||
}
|
||||
|
||||
return totalChars * 2 / 5
|
||||
// EstimateToolDefsTokens estimates the total token cost of tool definitions
|
||||
// as they appear in the LLM request. Delegates to the shared tokenizer package.
|
||||
func EstimateToolDefsTokens(defs []providers.ToolDefinition) int {
|
||||
return tokenizer.EstimateToolDefsTokens(defs)
|
||||
}
|
||||
|
||||
// isOverContextBudget checks whether the assembled messages plus tool definitions
|
||||
|
|
@ -181,10 +107,10 @@ func isOverContextBudget(
|
|||
) bool {
|
||||
msgTokens := 0
|
||||
for _, m := range messages {
|
||||
msgTokens += estimateMessageTokens(m)
|
||||
msgTokens += EstimateMessageTokens(m)
|
||||
}
|
||||
|
||||
toolTokens := estimateToolDefsTokens(toolDefs)
|
||||
toolTokens := EstimateToolDefsTokens(toolDefs)
|
||||
total := msgTokens + toolTokens + maxTokens
|
||||
|
||||
return total > contextWindow
|
||||
|
|
|
|||
|
|
@ -417,9 +417,9 @@ func TestEstimateMessageTokens(t *testing.T) {
|
|||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := estimateMessageTokens(tt.msg)
|
||||
got := EstimateMessageTokens(tt.msg)
|
||||
if got < tt.want {
|
||||
t.Errorf("estimateMessageTokens() = %d, want >= %d", got, tt.want)
|
||||
t.Errorf("EstimateMessageTokens() = %d, want >= %d", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
|
@ -443,8 +443,8 @@ func TestEstimateMessageTokens_ToolCallsContribute(t *testing.T) {
|
|||
},
|
||||
}
|
||||
|
||||
plainTokens := estimateMessageTokens(plain)
|
||||
withTCTokens := estimateMessageTokens(withTC)
|
||||
plainTokens := EstimateMessageTokens(plain)
|
||||
withTCTokens := EstimateMessageTokens(withTC)
|
||||
|
||||
if withTCTokens <= plainTokens {
|
||||
t.Errorf("message with ToolCalls (%d tokens) should exceed plain message (%d tokens)",
|
||||
|
|
@ -457,7 +457,7 @@ func TestEstimateMessageTokens_MultibyteContent(t *testing.T) {
|
|||
// but may map to different token counts. The heuristic should still produce
|
||||
// reasonable estimates via RuneCountInString.
|
||||
msg := msgUser("caf\u00e9 na\u00efve r\u00e9sum\u00e9 \u00fcber stra\u00dfe")
|
||||
tokens := estimateMessageTokens(msg)
|
||||
tokens := EstimateMessageTokens(msg)
|
||||
if tokens <= 0 {
|
||||
t.Errorf("multibyte message should produce positive token count, got %d", tokens)
|
||||
}
|
||||
|
|
@ -481,7 +481,7 @@ func TestEstimateMessageTokens_LargeArguments(t *testing.T) {
|
|||
},
|
||||
}
|
||||
|
||||
tokens := estimateMessageTokens(msg)
|
||||
tokens := EstimateMessageTokens(msg)
|
||||
// 5000+ chars → at least 2000 tokens with the 2.5 char/token heuristic
|
||||
if tokens < 2000 {
|
||||
t.Errorf("large tool call arguments should produce significant token count, got %d", tokens)
|
||||
|
|
@ -496,8 +496,8 @@ func TestEstimateMessageTokens_ReasoningContent(t *testing.T) {
|
|||
ReasoningContent: strings.Repeat("thinking step ", 200),
|
||||
}
|
||||
|
||||
plainTokens := estimateMessageTokens(plain)
|
||||
reasoningTokens := estimateMessageTokens(withReasoning)
|
||||
plainTokens := EstimateMessageTokens(plain)
|
||||
reasoningTokens := EstimateMessageTokens(withReasoning)
|
||||
|
||||
if reasoningTokens <= plainTokens {
|
||||
t.Errorf("message with ReasoningContent (%d tokens) should exceed plain message (%d tokens)",
|
||||
|
|
@ -513,8 +513,8 @@ func TestEstimateMessageTokens_MediaItems(t *testing.T) {
|
|||
Media: []string{"media://img1.png", "media://img2.png"},
|
||||
}
|
||||
|
||||
plainTokens := estimateMessageTokens(plain)
|
||||
mediaTokens := estimateMessageTokens(withMedia)
|
||||
plainTokens := EstimateMessageTokens(plain)
|
||||
mediaTokens := EstimateMessageTokens(withMedia)
|
||||
|
||||
if mediaTokens <= plainTokens {
|
||||
t.Errorf("message with Media (%d tokens) should exceed plain message (%d tokens)",
|
||||
|
|
@ -540,8 +540,8 @@ func TestEstimateMessageTokens_SystemParts(t *testing.T) {
|
|||
},
|
||||
}
|
||||
|
||||
plainTokens := estimateMessageTokens(plain)
|
||||
partsTokens := estimateMessageTokens(withParts)
|
||||
plainTokens := EstimateMessageTokens(plain)
|
||||
partsTokens := EstimateMessageTokens(withParts)
|
||||
|
||||
if partsTokens <= plainTokens {
|
||||
t.Errorf("system message with SystemParts (%d) should exceed plain message (%d)",
|
||||
|
|
@ -549,7 +549,7 @@ func TestEstimateMessageTokens_SystemParts(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
// --- estimateToolDefsTokens tests ---
|
||||
// --- EstimateToolDefsTokens tests ---
|
||||
|
||||
func TestEstimateToolDefsTokens(t *testing.T) {
|
||||
tests := []struct {
|
||||
|
|
@ -599,9 +599,9 @@ func TestEstimateToolDefsTokens(t *testing.T) {
|
|||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := estimateToolDefsTokens(tt.defs)
|
||||
got := EstimateToolDefsTokens(tt.defs)
|
||||
if got < tt.want {
|
||||
t.Errorf("estimateToolDefsTokens() = %d, want >= %d", got, tt.want)
|
||||
t.Errorf("EstimateToolDefsTokens() = %d, want >= %d", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
|
@ -624,8 +624,8 @@ func TestEstimateToolDefsTokens_ScalesWithCount(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
one := estimateToolDefsTokens([]providers.ToolDefinition{makeTool("tool_a")})
|
||||
three := estimateToolDefsTokens([]providers.ToolDefinition{
|
||||
one := EstimateToolDefsTokens([]providers.ToolDefinition{makeTool("tool_a")})
|
||||
three := EstimateToolDefsTokens([]providers.ToolDefinition{
|
||||
makeTool("tool_a"), makeTool("tool_b"), makeTool("tool_c"),
|
||||
})
|
||||
|
||||
|
|
@ -770,7 +770,7 @@ func TestEstimateMessageTokens_WithReasoningAndMedia(t *testing.T) {
|
|||
},
|
||||
}
|
||||
|
||||
tokens := estimateMessageTokens(msg)
|
||||
tokens := EstimateMessageTokens(msg)
|
||||
|
||||
// ReasoningContent alone is ~1700 chars → ~680 tokens.
|
||||
// Content + TC + overhead adds more. Should be well above 500.
|
||||
|
|
@ -781,7 +781,7 @@ func TestEstimateMessageTokens_WithReasoningAndMedia(t *testing.T) {
|
|||
// Compare without reasoning to ensure it's counted.
|
||||
msgNoReasoning := msg
|
||||
msgNoReasoning.ReasoningContent = ""
|
||||
tokensNoReasoning := estimateMessageTokens(msgNoReasoning)
|
||||
tokensNoReasoning := EstimateMessageTokens(msgNoReasoning)
|
||||
|
||||
if tokens <= tokensNoReasoning {
|
||||
t.Errorf("reasoning content should add tokens: with=%d, without=%d", tokens, tokensNoReasoning)
|
||||
|
|
|
|||
|
|
@ -707,6 +707,38 @@ func TestEmptyWorkspaceBaselineDetectsNewFiles(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestBuildMessages_IncludesMediaOnlyCurrentMessage(t *testing.T) {
|
||||
tmpDir := setupWorkspace(t, nil)
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
cb := NewContextBuilder(tmpDir)
|
||||
msgs := cb.BuildMessages(
|
||||
nil,
|
||||
"",
|
||||
"",
|
||||
[]string{"data:image/png;base64,abc123"},
|
||||
"pico",
|
||||
"chat-1",
|
||||
"",
|
||||
"",
|
||||
)
|
||||
|
||||
if len(msgs) != 2 {
|
||||
t.Fatalf("len(msgs) = %d, want 2", len(msgs))
|
||||
}
|
||||
|
||||
userMsg := msgs[1]
|
||||
if userMsg.Role != "user" {
|
||||
t.Fatalf("userMsg.Role = %q, want %q", userMsg.Role, "user")
|
||||
}
|
||||
if userMsg.Content != "" {
|
||||
t.Fatalf("userMsg.Content = %q, want empty string", userMsg.Content)
|
||||
}
|
||||
if len(userMsg.Media) != 1 || userMsg.Media[0] != "data:image/png;base64,abc123" {
|
||||
t.Fatalf("userMsg.Media = %#v, want image payload", userMsg.Media)
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkBuildMessagesWithCache measures caching performance.
|
||||
func BenchmarkBuildMessagesWithCache(b *testing.B) {
|
||||
tmpDir, _ := os.MkdirTemp("", "picoclaw-bench-*")
|
||||
|
|
|
|||
379
pkg/agent/context_legacy.go
Normal file
379
pkg/agent/context_legacy.go
Normal file
|
|
@ -0,0 +1,379 @@
|
|||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
)
|
||||
|
||||
// legacyContextManager wraps the existing summarization/compression logic
|
||||
// as a ContextManager implementation. It is the default when no other
|
||||
// ContextManager is configured.
|
||||
type legacyContextManager struct {
|
||||
al *AgentLoop
|
||||
summarizing sync.Map // dedup for async Compact (post-turn)
|
||||
}
|
||||
|
||||
func (m *legacyContextManager) Assemble(_ context.Context, req *AssembleRequest) (*AssembleResponse, error) {
|
||||
// Legacy: read history from session, return as-is.
|
||||
// Budget enforcement happens in BuildMessages caller via
|
||||
// isOverContextBudget + forceCompression.
|
||||
agent := m.al.registry.GetDefaultAgent()
|
||||
if agent == nil {
|
||||
return &AssembleResponse{}, nil
|
||||
}
|
||||
history := agent.Sessions.GetHistory(req.SessionKey)
|
||||
summary := agent.Sessions.GetSummary(req.SessionKey)
|
||||
return &AssembleResponse{
|
||||
History: history,
|
||||
Summary: summary,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (m *legacyContextManager) Compact(_ context.Context, req *CompactRequest) error {
|
||||
switch req.Reason {
|
||||
case ContextCompressReasonProactive, ContextCompressReasonRetry:
|
||||
// Sync emergency compression — budget exceeded.
|
||||
if result, ok := m.forceCompression(req.SessionKey); ok {
|
||||
m.al.emitEvent(
|
||||
EventKindContextCompress,
|
||||
m.al.newTurnEventScope("", req.SessionKey).meta(0, "forceCompression", "turn.context.compress"),
|
||||
ContextCompressPayload{
|
||||
Reason: req.Reason,
|
||||
DroppedMessages: result.DroppedMessages,
|
||||
RemainingMessages: result.RemainingMessages,
|
||||
},
|
||||
)
|
||||
}
|
||||
case ContextCompressReasonSummarize:
|
||||
m.maybeSummarize(req.SessionKey)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *legacyContextManager) Ingest(_ context.Context, _ *IngestRequest) error {
|
||||
// Legacy: no-op. Messages are persisted by Sessions JSONL.
|
||||
return nil
|
||||
}
|
||||
|
||||
// maybeSummarize triggers summarization if the session history exceeds thresholds.
|
||||
// It runs asynchronously in a goroutine.
|
||||
func (m *legacyContextManager) maybeSummarize(sessionKey string) {
|
||||
agent := m.al.registry.GetDefaultAgent()
|
||||
if agent == nil {
|
||||
return
|
||||
}
|
||||
|
||||
newHistory := agent.Sessions.GetHistory(sessionKey)
|
||||
tokenEstimate := m.estimateTokens(newHistory)
|
||||
threshold := agent.ContextWindow * agent.SummarizeTokenPercent / 100
|
||||
|
||||
if len(newHistory) > agent.SummarizeMessageThreshold || tokenEstimate > threshold {
|
||||
summarizeKey := agent.ID + ":" + sessionKey
|
||||
if _, loading := m.summarizing.LoadOrStore(summarizeKey, true); !loading {
|
||||
go func() {
|
||||
defer m.summarizing.Delete(summarizeKey)
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
logger.WarnCF("agent", "Summarization panic recovered", map[string]any{
|
||||
"session_key": sessionKey,
|
||||
"panic": r,
|
||||
})
|
||||
}
|
||||
}()
|
||||
logger.Debug("Memory threshold reached. Optimizing conversation history...")
|
||||
m.summarizeSession(agent, sessionKey)
|
||||
}()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type compressionResult struct {
|
||||
DroppedMessages int
|
||||
RemainingMessages int
|
||||
}
|
||||
|
||||
// forceCompression aggressively reduces context when the limit is hit.
|
||||
// It drops the oldest ~50% of Turns (a Turn is a complete user→LLM→response
|
||||
// cycle, as defined in #1316), so tool-call sequences are never split.
|
||||
func (m *legacyContextManager) forceCompression(sessionKey string) (compressionResult, bool) {
|
||||
agent := m.al.registry.GetDefaultAgent()
|
||||
if agent == nil {
|
||||
return compressionResult{}, false
|
||||
}
|
||||
|
||||
history := agent.Sessions.GetHistory(sessionKey)
|
||||
if len(history) <= 2 {
|
||||
return compressionResult{}, false
|
||||
}
|
||||
|
||||
turns := parseTurnBoundaries(history)
|
||||
var mid int
|
||||
if len(turns) >= 2 {
|
||||
mid = turns[len(turns)/2]
|
||||
} else {
|
||||
mid = findSafeBoundary(history, len(history)/2)
|
||||
}
|
||||
var keptHistory []providers.Message
|
||||
if mid <= 0 {
|
||||
for i := len(history) - 1; i >= 0; i-- {
|
||||
if history[i].Role == "user" {
|
||||
keptHistory = []providers.Message{history[i]}
|
||||
break
|
||||
}
|
||||
}
|
||||
} else {
|
||||
keptHistory = history[mid:]
|
||||
}
|
||||
|
||||
droppedCount := len(history) - len(keptHistory)
|
||||
|
||||
existingSummary := agent.Sessions.GetSummary(sessionKey)
|
||||
compressionNote := fmt.Sprintf(
|
||||
"[Emergency compression dropped %d oldest messages due to context limit]",
|
||||
droppedCount,
|
||||
)
|
||||
if existingSummary != "" {
|
||||
compressionNote = existingSummary + "\n\n" + compressionNote
|
||||
}
|
||||
agent.Sessions.SetSummary(sessionKey, compressionNote)
|
||||
|
||||
agent.Sessions.SetHistory(sessionKey, keptHistory)
|
||||
agent.Sessions.Save(sessionKey)
|
||||
|
||||
logger.WarnCF("agent", "Forced compression executed", map[string]any{
|
||||
"session_key": sessionKey,
|
||||
"dropped_msgs": droppedCount,
|
||||
"new_count": len(keptHistory),
|
||||
})
|
||||
|
||||
return compressionResult{
|
||||
DroppedMessages: droppedCount,
|
||||
RemainingMessages: len(keptHistory),
|
||||
}, true
|
||||
}
|
||||
|
||||
func (m *legacyContextManager) summarizeSession(agent *AgentInstance, sessionKey string) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
|
||||
defer cancel()
|
||||
|
||||
history := agent.Sessions.GetHistory(sessionKey)
|
||||
summary := agent.Sessions.GetSummary(sessionKey)
|
||||
|
||||
if len(history) <= 4 {
|
||||
return
|
||||
}
|
||||
|
||||
safeCut := findSafeBoundary(history, len(history)-4)
|
||||
if safeCut <= 0 {
|
||||
return
|
||||
}
|
||||
keepCount := len(history) - safeCut
|
||||
toSummarize := history[:safeCut]
|
||||
|
||||
maxMessageTokens := agent.ContextWindow / 2
|
||||
validMessages := make([]providers.Message, 0)
|
||||
omitted := false
|
||||
|
||||
for _, msg := range toSummarize {
|
||||
if msg.Role != "user" && msg.Role != "assistant" {
|
||||
continue
|
||||
}
|
||||
msgTokens := len(msg.Content) / 2
|
||||
if msgTokens > maxMessageTokens {
|
||||
omitted = true
|
||||
continue
|
||||
}
|
||||
validMessages = append(validMessages, msg)
|
||||
}
|
||||
|
||||
if len(validMessages) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
const (
|
||||
maxSummarizationMessages = 10
|
||||
llmMaxRetries = 3
|
||||
)
|
||||
|
||||
var finalSummary string
|
||||
if len(validMessages) > maxSummarizationMessages {
|
||||
mid := len(validMessages) / 2
|
||||
mid = m.findNearestUserMessage(validMessages, mid)
|
||||
|
||||
part1 := validMessages[:mid]
|
||||
part2 := validMessages[mid:]
|
||||
|
||||
s1, _ := m.summarizeBatch(ctx, agent, part1, "")
|
||||
s2, _ := m.summarizeBatch(ctx, agent, part2, "")
|
||||
|
||||
mergePrompt := fmt.Sprintf(
|
||||
"Merge these two conversation summaries into one cohesive summary:\n\n1: %s\n\n2: %s",
|
||||
s1, s2,
|
||||
)
|
||||
|
||||
resp, err := m.retryLLMCall(ctx, agent, mergePrompt, llmMaxRetries)
|
||||
if err == nil && resp.Content != "" {
|
||||
finalSummary = resp.Content
|
||||
} else {
|
||||
finalSummary = s1 + " " + s2
|
||||
}
|
||||
} else {
|
||||
finalSummary, _ = m.summarizeBatch(ctx, agent, validMessages, summary)
|
||||
}
|
||||
|
||||
if omitted && finalSummary != "" {
|
||||
finalSummary += "\n[Note: Some oversized messages were omitted from this summary for efficiency.]"
|
||||
}
|
||||
|
||||
if finalSummary != "" {
|
||||
agent.Sessions.SetSummary(sessionKey, finalSummary)
|
||||
agent.Sessions.TruncateHistory(sessionKey, keepCount)
|
||||
agent.Sessions.Save(sessionKey)
|
||||
m.al.emitEvent(
|
||||
EventKindSessionSummarize,
|
||||
m.al.newTurnEventScope(agent.ID, sessionKey).meta(0, "summarizeSession", "turn.session.summarize"),
|
||||
SessionSummarizePayload{
|
||||
SummarizedMessages: len(validMessages),
|
||||
KeptMessages: keepCount,
|
||||
SummaryLen: len(finalSummary),
|
||||
OmittedOversized: omitted,
|
||||
},
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *legacyContextManager) findNearestUserMessage(messages []providers.Message, mid int) int {
|
||||
originalMid := mid
|
||||
|
||||
for mid > 0 && messages[mid].Role != "user" {
|
||||
mid--
|
||||
}
|
||||
|
||||
if messages[mid].Role == "user" {
|
||||
return mid
|
||||
}
|
||||
|
||||
mid = originalMid
|
||||
for mid < len(messages) && messages[mid].Role != "user" {
|
||||
mid++
|
||||
}
|
||||
|
||||
if mid < len(messages) {
|
||||
return mid
|
||||
}
|
||||
|
||||
return originalMid
|
||||
}
|
||||
|
||||
func (m *legacyContextManager) retryLLMCall(
|
||||
ctx context.Context,
|
||||
agent *AgentInstance,
|
||||
prompt string,
|
||||
maxRetries int,
|
||||
) (*providers.LLMResponse, error) {
|
||||
const llmTemperature = 0.3
|
||||
|
||||
var resp *providers.LLMResponse
|
||||
var err error
|
||||
|
||||
for attempt := 0; attempt < maxRetries; attempt++ {
|
||||
m.al.activeRequests.Add(1)
|
||||
resp, err = func() (*providers.LLMResponse, error) {
|
||||
defer m.al.activeRequests.Done()
|
||||
return agent.Provider.Chat(
|
||||
ctx,
|
||||
[]providers.Message{{Role: "user", Content: prompt}},
|
||||
nil,
|
||||
agent.Model,
|
||||
map[string]any{
|
||||
"max_tokens": agent.MaxTokens,
|
||||
"temperature": llmTemperature,
|
||||
"prompt_cache_key": agent.ID,
|
||||
},
|
||||
)
|
||||
}()
|
||||
|
||||
if err == nil && resp != nil && resp.Content != "" {
|
||||
return resp, nil
|
||||
}
|
||||
if attempt < maxRetries-1 {
|
||||
time.Sleep(time.Duration(attempt+1) * 100 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
return resp, err
|
||||
}
|
||||
|
||||
func (m *legacyContextManager) summarizeBatch(
|
||||
ctx context.Context,
|
||||
agent *AgentInstance,
|
||||
batch []providers.Message,
|
||||
existingSummary string,
|
||||
) (string, error) {
|
||||
const (
|
||||
llmMaxRetries = 3
|
||||
fallbackMinContentLength = 200
|
||||
fallbackMaxContentPercent = 10
|
||||
)
|
||||
|
||||
var sb strings.Builder
|
||||
sb.WriteString("Provide a concise summary of this conversation segment, preserving core context and key points.\n")
|
||||
if existingSummary != "" {
|
||||
sb.WriteString("Existing context: ")
|
||||
sb.WriteString(existingSummary)
|
||||
sb.WriteString("\n")
|
||||
}
|
||||
sb.WriteString("\nCONVERSATION:\n")
|
||||
for _, msg := range batch {
|
||||
fmt.Fprintf(&sb, "%s: %s\n", msg.Role, msg.Content)
|
||||
}
|
||||
prompt := sb.String()
|
||||
|
||||
response, err := m.retryLLMCall(ctx, agent, prompt, llmMaxRetries)
|
||||
if err == nil && response.Content != "" {
|
||||
return strings.TrimSpace(response.Content), nil
|
||||
}
|
||||
|
||||
var fallback strings.Builder
|
||||
fallback.WriteString("Conversation summary: ")
|
||||
for i, msg := range batch {
|
||||
if i > 0 {
|
||||
fallback.WriteString(" | ")
|
||||
}
|
||||
content := strings.TrimSpace(msg.Content)
|
||||
runes := []rune(content)
|
||||
if len(runes) == 0 {
|
||||
fallback.WriteString(fmt.Sprintf("%s: ", msg.Role))
|
||||
continue
|
||||
}
|
||||
|
||||
keepLength := len(runes) * fallbackMaxContentPercent / 100
|
||||
if keepLength < fallbackMinContentLength {
|
||||
keepLength = fallbackMinContentLength
|
||||
}
|
||||
if keepLength > len(runes) {
|
||||
keepLength = len(runes)
|
||||
}
|
||||
|
||||
content = string(runes[:keepLength])
|
||||
if keepLength < len(runes) {
|
||||
content += "..."
|
||||
}
|
||||
fallback.WriteString(fmt.Sprintf("%s: %s", msg.Role, content))
|
||||
}
|
||||
return fallback.String(), nil
|
||||
}
|
||||
|
||||
func (m *legacyContextManager) estimateTokens(messages []providers.Message) int {
|
||||
total := 0
|
||||
for _, msg := range messages {
|
||||
total += EstimateMessageTokens(msg)
|
||||
}
|
||||
return total
|
||||
}
|
||||
90
pkg/agent/context_manager.go
Normal file
90
pkg/agent/context_manager.go
Normal file
|
|
@ -0,0 +1,90 @@
|
|||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sync"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
)
|
||||
|
||||
// ContextManager manages conversation context via a pluggable strategy.
|
||||
// Exactly ONE ContextManager is active per AgentLoop, selected by config.
|
||||
// The default ("legacy") preserves current summarization behavior.
|
||||
type ContextManager interface {
|
||||
// Assemble builds budget-aware context from the ContextManager's own storage.
|
||||
// Called before BuildMessages. Returns assembled messages ready for LLM.
|
||||
Assemble(ctx context.Context, req *AssembleRequest) (*AssembleResponse, error)
|
||||
|
||||
// Compact compresses conversation history.
|
||||
// Called after turn completes (may be async internally) and on context overflow (sync).
|
||||
Compact(ctx context.Context, req *CompactRequest) error
|
||||
|
||||
// Ingest records a message into the ContextManager's own storage.
|
||||
// Called after each message is persisted to session JSONL.
|
||||
Ingest(ctx context.Context, req *IngestRequest) error
|
||||
}
|
||||
|
||||
// AssembleRequest is the input to Assemble.
|
||||
type AssembleRequest struct {
|
||||
SessionKey string // session identifier
|
||||
Budget int // context window in tokens
|
||||
MaxTokens int // max response tokens
|
||||
}
|
||||
|
||||
// AssembleResponse is the output of Assemble.
|
||||
type AssembleResponse struct {
|
||||
History []providers.Message // assembled conversation history for BuildMessages
|
||||
Summary string // conversation summary embedded into system prompt by BuildMessages
|
||||
}
|
||||
|
||||
// CompactRequest is the input to Compact.
|
||||
type CompactRequest struct {
|
||||
SessionKey string // session identifier
|
||||
Reason ContextCompressReason // proactive_budget | llm_retry | summarize
|
||||
Budget int // context window budget (used for retry aggressive compaction)
|
||||
}
|
||||
|
||||
// IngestRequest is the input to Ingest.
|
||||
type IngestRequest struct {
|
||||
SessionKey string // session identifier
|
||||
Message providers.Message // the message just persisted
|
||||
}
|
||||
|
||||
// ContextManagerFactory constructs a ContextManager from config.
|
||||
// al provides access to the AgentLoop's runtime resources (provider, model, workspace, etc.)
|
||||
// cfg is the raw JSON configuration from config.json (may be nil).
|
||||
type ContextManagerFactory func(cfg json.RawMessage, al *AgentLoop) (ContextManager, error)
|
||||
|
||||
var (
|
||||
cmRegistryMu sync.RWMutex
|
||||
cmRegistry = map[string]ContextManagerFactory{}
|
||||
)
|
||||
|
||||
// RegisterContextManager registers a named ContextManager factory.
|
||||
func RegisterContextManager(name string, factory ContextManagerFactory) error {
|
||||
if name == "" {
|
||||
return fmt.Errorf("context manager name is required")
|
||||
}
|
||||
if factory == nil {
|
||||
return fmt.Errorf("context manager %q factory is nil", name)
|
||||
}
|
||||
|
||||
cmRegistryMu.Lock()
|
||||
defer cmRegistryMu.Unlock()
|
||||
|
||||
if _, exists := cmRegistry[name]; exists {
|
||||
return fmt.Errorf("context manager %q is already registered", name)
|
||||
}
|
||||
cmRegistry[name] = factory
|
||||
return nil
|
||||
}
|
||||
|
||||
func lookupContextManager(name string) (ContextManagerFactory, bool) {
|
||||
cmRegistryMu.RLock()
|
||||
defer cmRegistryMu.RUnlock()
|
||||
|
||||
f, ok := cmRegistry[name]
|
||||
return f, ok
|
||||
}
|
||||
764
pkg/agent/context_manager_test.go
Normal file
764
pkg/agent/context_manager_test.go
Normal file
|
|
@ -0,0 +1,764 @@
|
|||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/bus"
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Factory registry tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestRegisterContextManager_Success(t *testing.T) {
|
||||
cleanup := resetCMRegistry()
|
||||
defer cleanup()
|
||||
|
||||
factory := func(cfg json.RawMessage, al *AgentLoop) (ContextManager, error) {
|
||||
return &noopContextManager{}, nil
|
||||
}
|
||||
if err := RegisterContextManager("test_cm", factory); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
f, ok := lookupContextManager("test_cm")
|
||||
if !ok {
|
||||
t.Fatal("expected factory to be registered")
|
||||
}
|
||||
if f == nil {
|
||||
t.Fatal("expected non-nil factory")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterContextManager_EmptyName(t *testing.T) {
|
||||
cleanup := resetCMRegistry()
|
||||
defer cleanup()
|
||||
|
||||
err := RegisterContextManager("", func(cfg json.RawMessage, al *AgentLoop) (ContextManager, error) {
|
||||
return &noopContextManager{}, nil
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected error for empty name")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "name is required") {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterContextManager_NilFactory(t *testing.T) {
|
||||
cleanup := resetCMRegistry()
|
||||
defer cleanup()
|
||||
|
||||
err := RegisterContextManager("nil_factory", nil)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for nil factory")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "factory is nil") {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterContextManager_Duplicate(t *testing.T) {
|
||||
cleanup := resetCMRegistry()
|
||||
defer cleanup()
|
||||
|
||||
factory := func(cfg json.RawMessage, al *AgentLoop) (ContextManager, error) {
|
||||
return &noopContextManager{}, nil
|
||||
}
|
||||
if err := RegisterContextManager("dup_cm", factory); err != nil {
|
||||
t.Fatalf("first registration failed: %v", err)
|
||||
}
|
||||
err := RegisterContextManager("dup_cm", factory)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for duplicate registration")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "already registered") {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLookupContextManager_Unknown(t *testing.T) {
|
||||
cleanup := resetCMRegistry()
|
||||
defer cleanup()
|
||||
|
||||
_, ok := lookupContextManager("nonexistent")
|
||||
if ok {
|
||||
t.Fatal("expected lookup to fail for unknown name")
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// resolveContextManager tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestResolveContextManager_Default(t *testing.T) {
|
||||
cleanup := resetCMRegistry()
|
||||
defer cleanup()
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: t.TempDir(),
|
||||
ModelName: "test-model",
|
||||
MaxTokens: 4096,
|
||||
MaxToolIterations: 10,
|
||||
ContextManager: "", // default → legacy
|
||||
},
|
||||
},
|
||||
}
|
||||
al := newCMTestAgentLoop(cfg)
|
||||
|
||||
cm := al.contextManager
|
||||
if cm == nil {
|
||||
t.Fatal("expected non-nil context manager")
|
||||
}
|
||||
if _, ok := cm.(*legacyContextManager); !ok {
|
||||
t.Fatalf("expected *legacyContextManager, got %T", cm)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveContextManager_ExplicitLegacy(t *testing.T) {
|
||||
cleanup := resetCMRegistry()
|
||||
defer cleanup()
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: t.TempDir(),
|
||||
ModelName: "test-model",
|
||||
MaxTokens: 4096,
|
||||
MaxToolIterations: 10,
|
||||
ContextManager: "legacy",
|
||||
},
|
||||
},
|
||||
}
|
||||
al := newCMTestAgentLoop(cfg)
|
||||
|
||||
if _, ok := al.contextManager.(*legacyContextManager); !ok {
|
||||
t.Fatalf("expected *legacyContextManager, got %T", al.contextManager)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveContextManager_UnknownFallsBackToLegacy(t *testing.T) {
|
||||
cleanup := resetCMRegistry()
|
||||
defer cleanup()
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: t.TempDir(),
|
||||
ModelName: "test-model",
|
||||
MaxTokens: 4096,
|
||||
MaxToolIterations: 10,
|
||||
ContextManager: "unknown_cm",
|
||||
},
|
||||
},
|
||||
}
|
||||
al := newCMTestAgentLoop(cfg)
|
||||
|
||||
if _, ok := al.contextManager.(*legacyContextManager); !ok {
|
||||
t.Fatalf("expected fallback to *legacyContextManager, got %T", al.contextManager)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveContextManager_RegisteredFactory(t *testing.T) {
|
||||
cleanup := resetCMRegistry()
|
||||
defer cleanup()
|
||||
|
||||
factory := func(cfg json.RawMessage, al *AgentLoop) (ContextManager, error) {
|
||||
return &noopContextManager{}, nil
|
||||
}
|
||||
if err := RegisterContextManager("custom_cm", factory); err != nil {
|
||||
t.Fatalf("register failed: %v", err)
|
||||
}
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: t.TempDir(),
|
||||
ModelName: "test-model",
|
||||
MaxTokens: 4096,
|
||||
MaxToolIterations: 10,
|
||||
ContextManager: "custom_cm",
|
||||
},
|
||||
},
|
||||
}
|
||||
al := newCMTestAgentLoop(cfg)
|
||||
|
||||
if _, ok := al.contextManager.(*noopContextManager); !ok {
|
||||
t.Fatalf("expected *noopContextManager, got %T", al.contextManager)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveContextManager_FactoryError(t *testing.T) {
|
||||
cleanup := resetCMRegistry()
|
||||
defer cleanup()
|
||||
|
||||
factory := func(cfg json.RawMessage, al *AgentLoop) (ContextManager, error) {
|
||||
return nil, os.ErrPermission
|
||||
}
|
||||
if err := RegisterContextManager("broken_cm", factory); err != nil {
|
||||
t.Fatalf("register failed: %v", err)
|
||||
}
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: t.TempDir(),
|
||||
ModelName: "test-model",
|
||||
MaxTokens: 4096,
|
||||
MaxToolIterations: 10,
|
||||
ContextManager: "broken_cm",
|
||||
},
|
||||
},
|
||||
}
|
||||
al := newCMTestAgentLoop(cfg)
|
||||
|
||||
// Should fall back to legacy when factory returns error
|
||||
if _, ok := al.contextManager.(*legacyContextManager); !ok {
|
||||
t.Fatalf("expected fallback to *legacyContextManager on factory error, got %T", al.contextManager)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Legacy Assemble tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestLegacyAssemble_Passthrough(t *testing.T) {
|
||||
cfg := testConfig(t)
|
||||
al := newCMTestAgentLoop(cfg)
|
||||
|
||||
agent := al.registry.GetDefaultAgent()
|
||||
if agent == nil {
|
||||
t.Fatal("expected default agent")
|
||||
}
|
||||
|
||||
history := []providers.Message{
|
||||
{Role: "user", Content: "hello"},
|
||||
{Role: "assistant", Content: "hi there"},
|
||||
}
|
||||
agent.Sessions.SetHistory("test-session", history)
|
||||
|
||||
resp, err := al.contextManager.Assemble(context.Background(), &AssembleRequest{
|
||||
SessionKey: "test-session",
|
||||
Budget: 8000,
|
||||
MaxTokens: 4096,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if len(resp.History) != len(history) {
|
||||
t.Fatalf("expected %d messages, got %d", len(history), len(resp.History))
|
||||
}
|
||||
for i, msg := range resp.History {
|
||||
if msg.Content != history[i].Content || msg.Role != history[i].Role {
|
||||
t.Fatalf("message %d mismatch: want %+v, got %+v", i, history[i], msg)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestLegacyAssemble_EmptyHistory(t *testing.T) {
|
||||
cfg := testConfig(t)
|
||||
al := newCMTestAgentLoop(cfg)
|
||||
|
||||
resp, err := al.contextManager.Assemble(context.Background(), &AssembleRequest{
|
||||
SessionKey: "test-session",
|
||||
Budget: 8000,
|
||||
MaxTokens: 4096,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if len(resp.History) != 0 {
|
||||
t.Fatalf("expected empty messages, got %d", len(resp.History))
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Legacy Compact overflow tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestLegacyCompact_Overflow(t *testing.T) {
|
||||
cfg := testConfig(t)
|
||||
al := newCMTestAgentLoop(cfg)
|
||||
|
||||
defaultAgent := al.registry.GetDefaultAgent()
|
||||
if defaultAgent == nil {
|
||||
t.Fatal("expected default agent")
|
||||
}
|
||||
|
||||
history := []providers.Message{
|
||||
{Role: "user", Content: "msg 1"},
|
||||
{Role: "assistant", Content: "resp 1"},
|
||||
{Role: "user", Content: "msg 2"},
|
||||
{Role: "assistant", Content: "resp 2"},
|
||||
{Role: "user", Content: "msg 3"},
|
||||
}
|
||||
defaultAgent.Sessions.SetHistory("session-overflow", history)
|
||||
|
||||
sub := al.SubscribeEvents(16)
|
||||
defer al.UnsubscribeEvents(sub.ID)
|
||||
|
||||
err := al.contextManager.Compact(context.Background(), &CompactRequest{
|
||||
SessionKey: "session-overflow",
|
||||
Reason: ContextCompressReasonRetry,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
// After overflow compression, history should be shorter
|
||||
newHistory := defaultAgent.Sessions.GetHistory("session-overflow")
|
||||
if len(newHistory) >= len(history) {
|
||||
t.Fatalf("expected compressed history, got %d messages (was %d)", len(newHistory), len(history))
|
||||
}
|
||||
|
||||
// Summary should contain compression note
|
||||
summary := defaultAgent.Sessions.GetSummary("session-overflow")
|
||||
if !strings.Contains(summary, "Emergency compression") {
|
||||
t.Fatalf("expected compression note in summary, got %q", summary)
|
||||
}
|
||||
|
||||
// Event should carry the proactive reason
|
||||
events := collectEventStream(sub.C)
|
||||
compressEvt, ok := findEvent(events, EventKindContextCompress)
|
||||
if !ok {
|
||||
t.Fatal("expected context compress event")
|
||||
}
|
||||
payload, ok := compressEvt.Payload.(ContextCompressPayload)
|
||||
if !ok {
|
||||
t.Fatalf("expected ContextCompressPayload, got %T", compressEvt.Payload)
|
||||
}
|
||||
if payload.Reason != ContextCompressReasonRetry {
|
||||
t.Fatalf("expected retry reason, got %q", payload.Reason)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLegacyCompact_Overflow_ProactiveReason(t *testing.T) {
|
||||
cfg := testConfig(t)
|
||||
al := newCMTestAgentLoop(cfg)
|
||||
|
||||
defaultAgent := al.registry.GetDefaultAgent()
|
||||
if defaultAgent == nil {
|
||||
t.Fatal("expected default agent")
|
||||
}
|
||||
|
||||
history := []providers.Message{
|
||||
{Role: "user", Content: "msg 1"},
|
||||
{Role: "assistant", Content: "resp 1"},
|
||||
{Role: "user", Content: "msg 2"},
|
||||
{Role: "assistant", Content: "resp 2"},
|
||||
{Role: "user", Content: "msg 3"},
|
||||
}
|
||||
defaultAgent.Sessions.SetHistory("session-proactive", history)
|
||||
|
||||
sub := al.SubscribeEvents(16)
|
||||
defer al.UnsubscribeEvents(sub.ID)
|
||||
|
||||
err := al.contextManager.Compact(context.Background(), &CompactRequest{
|
||||
SessionKey: "session-proactive",
|
||||
Reason: ContextCompressReasonProactive,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
events := collectEventStream(sub.C)
|
||||
compressEvt, ok := findEvent(events, EventKindContextCompress)
|
||||
if !ok {
|
||||
t.Fatal("expected context compress event")
|
||||
}
|
||||
payload, ok := compressEvt.Payload.(ContextCompressPayload)
|
||||
if !ok {
|
||||
t.Fatalf("expected ContextCompressPayload, got %T", compressEvt.Payload)
|
||||
}
|
||||
if payload.Reason != ContextCompressReasonProactive {
|
||||
t.Fatalf("expected proactive reason, got %q", payload.Reason)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLegacyCompact_Overflow_TooShortToCompress(t *testing.T) {
|
||||
cfg := testConfig(t)
|
||||
al := newCMTestAgentLoop(cfg)
|
||||
|
||||
defaultAgent := al.registry.GetDefaultAgent()
|
||||
if defaultAgent == nil {
|
||||
t.Fatal("expected default agent")
|
||||
}
|
||||
|
||||
history := []providers.Message{
|
||||
{Role: "user", Content: "only one"},
|
||||
}
|
||||
defaultAgent.Sessions.SetHistory("session-tiny", history)
|
||||
|
||||
err := al.contextManager.Compact(context.Background(), &CompactRequest{
|
||||
SessionKey: "session-tiny",
|
||||
Reason: ContextCompressReasonRetry,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
// History should be unchanged (too short to compress)
|
||||
newHistory := defaultAgent.Sessions.GetHistory("session-tiny")
|
||||
if len(newHistory) != len(history) {
|
||||
t.Fatalf("expected history unchanged, got %d messages (was %d)", len(newHistory), len(history))
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Legacy Compact post-turn tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestLegacyCompact_PostTurn_BelowThreshold(t *testing.T) {
|
||||
cfg := testConfig(t)
|
||||
al := newCMTestAgentLoop(cfg)
|
||||
|
||||
defaultAgent := al.registry.GetDefaultAgent()
|
||||
if defaultAgent == nil {
|
||||
t.Fatal("expected default agent")
|
||||
}
|
||||
|
||||
// Small history, below summarization thresholds
|
||||
history := []providers.Message{
|
||||
{Role: "user", Content: "hi"},
|
||||
{Role: "assistant", Content: "hello"},
|
||||
}
|
||||
defaultAgent.Sessions.SetHistory("session-small", history)
|
||||
|
||||
err := al.contextManager.Compact(context.Background(), &CompactRequest{
|
||||
SessionKey: "session-small",
|
||||
Reason: ContextCompressReasonSummarize,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
// History should remain unchanged
|
||||
newHistory := defaultAgent.Sessions.GetHistory("session-small")
|
||||
if len(newHistory) != len(history) {
|
||||
t.Fatalf("expected unchanged history, got %d messages (was %d)", len(newHistory), len(history))
|
||||
}
|
||||
}
|
||||
|
||||
func TestLegacyCompact_PostTurn_ExceedsMessageThreshold(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: t.TempDir(),
|
||||
ModelName: "test-model",
|
||||
MaxTokens: 4096,
|
||||
MaxToolIterations: 10,
|
||||
ContextWindow: 8000,
|
||||
SummarizeMessageThreshold: 2,
|
||||
SummarizeTokenPercent: 75,
|
||||
},
|
||||
},
|
||||
}
|
||||
msgBus := bus.NewMessageBus()
|
||||
al := NewAgentLoop(cfg, msgBus, &simpleMockProvider{response: "summary"})
|
||||
|
||||
defaultAgent := al.registry.GetDefaultAgent()
|
||||
if defaultAgent == nil {
|
||||
t.Fatal("expected default agent")
|
||||
}
|
||||
|
||||
// 6 messages > threshold of 2
|
||||
history := []providers.Message{
|
||||
{Role: "user", Content: "q1"},
|
||||
{Role: "assistant", Content: "a1"},
|
||||
{Role: "user", Content: "q2"},
|
||||
{Role: "assistant", Content: "a2"},
|
||||
{Role: "user", Content: "q3"},
|
||||
{Role: "assistant", Content: "a3"},
|
||||
}
|
||||
defaultAgent.Sessions.SetHistory("session-threshold", history)
|
||||
|
||||
err := al.contextManager.Compact(context.Background(), &CompactRequest{
|
||||
SessionKey: "session-threshold",
|
||||
Reason: ContextCompressReasonSummarize,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
// Wait for async summarization to complete via event
|
||||
sub := al.SubscribeEvents(16)
|
||||
defer al.UnsubscribeEvents(sub.ID)
|
||||
|
||||
waitForEvent(t, sub.C, 5*time.Second, func(evt Event) bool {
|
||||
return evt.Kind == EventKindSessionSummarize
|
||||
})
|
||||
|
||||
newHistory := defaultAgent.Sessions.GetHistory("session-threshold")
|
||||
if len(newHistory) >= len(history) {
|
||||
t.Fatalf("expected summarization to reduce history from %d messages, got %d", len(history), len(newHistory))
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Legacy Ingest tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestLegacyIngest_NoOp(t *testing.T) {
|
||||
cfg := testConfig(t)
|
||||
al := newCMTestAgentLoop(cfg)
|
||||
|
||||
err := al.contextManager.Ingest(context.Background(), &IngestRequest{
|
||||
SessionKey: "session-ingest",
|
||||
Message: providers.Message{Role: "user", Content: "test"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Mock ContextManager — verifies dispatch through AgentLoop
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestAgentLoop_UsesCustomContextManager(t *testing.T) {
|
||||
cleanup := resetCMRegistry()
|
||||
defer cleanup()
|
||||
|
||||
mock := &trackingContextManager{}
|
||||
factory := func(cfg json.RawMessage, al *AgentLoop) (ContextManager, error) {
|
||||
return mock, nil
|
||||
}
|
||||
if err := RegisterContextManager("tracking_cm", factory); err != nil {
|
||||
t.Fatalf("register failed: %v", err)
|
||||
}
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: t.TempDir(),
|
||||
ModelName: "test-model",
|
||||
MaxTokens: 4096,
|
||||
MaxToolIterations: 10,
|
||||
ContextManager: "tracking_cm",
|
||||
},
|
||||
},
|
||||
}
|
||||
al := newCMTestAgentLoop(cfg)
|
||||
|
||||
// Verify the mock was installed
|
||||
if al.contextManager != mock {
|
||||
t.Fatalf("expected mock context manager, got %T", al.contextManager)
|
||||
}
|
||||
|
||||
// Direct method calls
|
||||
_, err := mock.Assemble(context.Background(), &AssembleRequest{
|
||||
SessionKey: "s1",
|
||||
Budget: 8000,
|
||||
MaxTokens: 4096,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Assemble error: %v", err)
|
||||
}
|
||||
if mock.assembleCalls.Load() != 1 {
|
||||
t.Fatalf("expected 1 assemble call, got %d", mock.assembleCalls.Load())
|
||||
}
|
||||
|
||||
err = mock.Compact(context.Background(), &CompactRequest{
|
||||
SessionKey: "s1",
|
||||
Reason: ContextCompressReasonRetry,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Compact error: %v", err)
|
||||
}
|
||||
if mock.compactCalls.Load() != 1 {
|
||||
t.Fatalf("expected 1 compact call, got %d", mock.compactCalls.Load())
|
||||
}
|
||||
|
||||
err = mock.Ingest(context.Background(), &IngestRequest{
|
||||
SessionKey: "s1",
|
||||
Message: providers.Message{Role: "user", Content: "test"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Ingest error: %v", err)
|
||||
}
|
||||
if mock.ingestCalls.Load() != 1 {
|
||||
t.Fatalf("expected 1 ingest call, got %d", mock.ingestCalls.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestIngestCalledDuringTurn(t *testing.T) {
|
||||
cleanup := resetCMRegistry()
|
||||
defer cleanup()
|
||||
|
||||
mock := &trackingContextManager{}
|
||||
factory := func(cfg json.RawMessage, al *AgentLoop) (ContextManager, error) {
|
||||
return mock, nil
|
||||
}
|
||||
if err := RegisterContextManager("ingest_track_cm", factory); err != nil {
|
||||
t.Fatalf("register failed: %v", err)
|
||||
}
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: t.TempDir(),
|
||||
ModelName: "test-model",
|
||||
MaxTokens: 4096,
|
||||
MaxToolIterations: 10,
|
||||
ContextManager: "ingest_track_cm",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
msgBus := bus.NewMessageBus()
|
||||
al := NewAgentLoop(cfg, msgBus, &simpleMockProvider{response: "done"})
|
||||
defaultAgent := al.registry.GetDefaultAgent()
|
||||
if defaultAgent == nil {
|
||||
t.Fatal("expected default agent")
|
||||
}
|
||||
|
||||
// Run a turn — ingestMessage is called for user message and final assistant message
|
||||
_, err := al.runAgentLoop(context.Background(), defaultAgent, processOptions{
|
||||
SessionKey: "session-ingest-turn",
|
||||
Channel: "cli",
|
||||
ChatID: "direct",
|
||||
UserMessage: "test ingest",
|
||||
DefaultResponse: defaultResponse,
|
||||
EnableSummary: false,
|
||||
SendResponse: false,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("runAgentLoop failed: %v", err)
|
||||
}
|
||||
|
||||
// Should have at least 2 ingest calls: user message + final assistant message
|
||||
if mock.ingestCalls.Load() < 2 {
|
||||
t.Fatalf("expected >= 2 ingest calls during turn, got %d", mock.ingestCalls.Load())
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// forceCompression edge cases (via legacy Compact)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestLegacyCompact_Overflow_SingleTurnKeepsLastUserMessage(t *testing.T) {
|
||||
cfg := testConfig(t)
|
||||
al := newCMTestAgentLoop(cfg)
|
||||
|
||||
defaultAgent := al.registry.GetDefaultAgent()
|
||||
if defaultAgent == nil {
|
||||
t.Fatal("expected default agent")
|
||||
}
|
||||
|
||||
// History with only 2 messages — forceCompression should still handle it
|
||||
history := []providers.Message{
|
||||
{Role: "user", Content: "first question"},
|
||||
{Role: "assistant", Content: "first answer"},
|
||||
}
|
||||
defaultAgent.Sessions.SetHistory("session-2msg", history)
|
||||
|
||||
err := al.contextManager.Compact(context.Background(), &CompactRequest{
|
||||
SessionKey: "session-2msg",
|
||||
Reason: ContextCompressReasonRetry,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
newHistory := defaultAgent.Sessions.GetHistory("session-2msg")
|
||||
// With 2 messages, forceCompression returns false (len <= 2), so no compression
|
||||
if len(newHistory) != len(history) {
|
||||
t.Fatalf("expected no compression for 2-message history, got %d", len(newHistory))
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// noopContextManager is a minimal ContextManager that does nothing.
|
||||
type noopContextManager struct{}
|
||||
|
||||
func (m *noopContextManager) Assemble(_ context.Context, req *AssembleRequest) (*AssembleResponse, error) {
|
||||
return &AssembleResponse{}, nil
|
||||
}
|
||||
func (m *noopContextManager) Compact(_ context.Context, _ *CompactRequest) error { return nil }
|
||||
func (m *noopContextManager) Ingest(_ context.Context, _ *IngestRequest) error { return nil }
|
||||
|
||||
// trackingContextManager tracks call counts for each method.
|
||||
type trackingContextManager struct {
|
||||
assembleCalls atomic.Int64
|
||||
compactCalls atomic.Int64
|
||||
ingestCalls atomic.Int64
|
||||
mu sync.Mutex
|
||||
lastAssemble *AssembleRequest
|
||||
lastCompact *CompactRequest
|
||||
lastIngest *IngestRequest
|
||||
}
|
||||
|
||||
func (m *trackingContextManager) Assemble(_ context.Context, req *AssembleRequest) (*AssembleResponse, error) {
|
||||
m.assembleCalls.Add(1)
|
||||
m.mu.Lock()
|
||||
m.lastAssemble = req
|
||||
m.mu.Unlock()
|
||||
return &AssembleResponse{}, nil
|
||||
}
|
||||
|
||||
func (m *trackingContextManager) Compact(_ context.Context, req *CompactRequest) error {
|
||||
m.compactCalls.Add(1)
|
||||
m.mu.Lock()
|
||||
m.lastCompact = req
|
||||
m.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *trackingContextManager) Ingest(_ context.Context, req *IngestRequest) error {
|
||||
m.ingestCalls.Add(1)
|
||||
m.mu.Lock()
|
||||
m.lastIngest = req
|
||||
m.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
// resetCMRegistry clears the global factory registry and returns a cleanup
|
||||
// function that restores the original state after the test.
|
||||
func resetCMRegistry() func() {
|
||||
cmRegistryMu.Lock()
|
||||
original := make(map[string]ContextManagerFactory, len(cmRegistry))
|
||||
for k, v := range cmRegistry {
|
||||
original[k] = v
|
||||
}
|
||||
cmRegistry = make(map[string]ContextManagerFactory)
|
||||
cmRegistryMu.Unlock()
|
||||
|
||||
return func() {
|
||||
cmRegistryMu.Lock()
|
||||
cmRegistry = original
|
||||
cmRegistryMu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
func testConfig(t *testing.T) *config.Config {
|
||||
t.Helper()
|
||||
return &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: t.TempDir(),
|
||||
ModelName: "test-model",
|
||||
MaxTokens: 4096,
|
||||
MaxToolIterations: 10,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func newCMTestAgentLoop(cfg *config.Config) *AgentLoop {
|
||||
msgBus := bus.NewMessageBus()
|
||||
return NewAgentLoop(cfg, msgBus, &simpleMockProvider{response: "test"})
|
||||
}
|
||||
267
pkg/agent/context_seahorse.go
Normal file
267
pkg/agent/context_seahorse.go
Normal file
|
|
@ -0,0 +1,267 @@
|
|||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
"github.com/sipeed/picoclaw/pkg/providers/protocoltypes"
|
||||
"github.com/sipeed/picoclaw/pkg/seahorse"
|
||||
"github.com/sipeed/picoclaw/pkg/session"
|
||||
"github.com/sipeed/picoclaw/pkg/tokenizer"
|
||||
)
|
||||
|
||||
// seahorseContextManager adapts seahorse.Engine to agent.ContextManager.
|
||||
type seahorseContextManager struct {
|
||||
engine *seahorse.Engine
|
||||
sessions session.SessionStore // for startup bootstrap
|
||||
}
|
||||
|
||||
// newSeahorseContextManager creates a seahorse-backed ContextManager.
|
||||
func newSeahorseContextManager(_ json.RawMessage, al *AgentLoop) (ContextManager, error) {
|
||||
if al == nil {
|
||||
return nil, fmt.Errorf("seahorse: AgentLoop is required")
|
||||
}
|
||||
|
||||
// Resolve workspace for DB path
|
||||
// DB stores session data, so it goes in sessions/ directory
|
||||
agent := al.registry.GetDefaultAgent()
|
||||
dbPath := agent.Workspace + "/sessions/seahorse.db"
|
||||
|
||||
// Create CompleteFn from provider
|
||||
completeFn := providerToCompleteFn(agent.Provider, agent.Model)
|
||||
|
||||
// Create engine
|
||||
engine, err := seahorse.NewEngine(seahorse.Config{
|
||||
DBPath: dbPath,
|
||||
}, completeFn)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("seahorse: create engine: %w", err)
|
||||
}
|
||||
|
||||
mgr := &seahorseContextManager{
|
||||
engine: engine,
|
||||
sessions: agent.Sessions,
|
||||
}
|
||||
|
||||
// Register seahorse tools with the agent's tool registry
|
||||
retrieval := mgr.engine.GetRetrieval()
|
||||
al.RegisterTool(seahorse.NewGrepTool(retrieval))
|
||||
al.RegisterTool(seahorse.NewExpandTool(retrieval))
|
||||
|
||||
// Bootstrap all existing sessions at startup
|
||||
if agent.Sessions != nil {
|
||||
ctx := context.Background()
|
||||
for _, sessionKey := range agent.Sessions.ListSessions() {
|
||||
mgr.bootstrapSession(ctx, sessionKey)
|
||||
}
|
||||
}
|
||||
|
||||
return mgr, nil
|
||||
}
|
||||
|
||||
// providerToCompleteFn wraps providers.LLMProvider as a seahorse.CompleteFn.
|
||||
func providerToCompleteFn(provider providers.LLMProvider, model string) seahorse.CompleteFn {
|
||||
return func(ctx context.Context, prompt string, opts seahorse.CompleteOptions) (string, error) {
|
||||
resp, err := provider.Chat(
|
||||
ctx,
|
||||
[]providers.Message{{Role: "user", Content: prompt}},
|
||||
nil, // no tools for summarization
|
||||
model,
|
||||
map[string]any{
|
||||
"max_tokens": opts.MaxTokens,
|
||||
"temperature": opts.Temperature,
|
||||
"prompt_cache_key": "seahorse",
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return resp.Content, nil
|
||||
}
|
||||
}
|
||||
|
||||
// Assemble builds budget-aware context from seahorse SQLite.
|
||||
func (m *seahorseContextManager) Assemble(ctx context.Context, req *AssembleRequest) (*AssembleResponse, error) {
|
||||
if req == nil {
|
||||
return nil, fmt.Errorf("seahorse assemble: nil request")
|
||||
}
|
||||
|
||||
budget := req.Budget
|
||||
if budget <= 0 {
|
||||
budget = 100000
|
||||
}
|
||||
|
||||
// Reserve space for model response (spec lines 1400-1410)
|
||||
effectiveBudget := budget - req.MaxTokens
|
||||
if effectiveBudget <= 0 {
|
||||
// MaxTokens >= budget is a configuration problem
|
||||
// Use 50% as minimum to avoid guaranteed overflow
|
||||
logger.WarnCF("agent", "MaxTokens >= budget, using 50% fallback",
|
||||
map[string]any{"budget": budget, "max_tokens": req.MaxTokens})
|
||||
effectiveBudget = budget / 2
|
||||
}
|
||||
|
||||
result, err := m.engine.Assemble(ctx, req.SessionKey, seahorse.AssembleInput{
|
||||
Budget: effectiveBudget,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("seahorse assemble: %w", err)
|
||||
}
|
||||
|
||||
history := seahorseToProviderMessages(result)
|
||||
|
||||
// Summary is already formatted as XML with system prompt addition by assembler
|
||||
return &AssembleResponse{
|
||||
History: history,
|
||||
Summary: result.Summary,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Compact compresses conversation history via seahorse summarization.
|
||||
func (m *seahorseContextManager) Compact(ctx context.Context, req *CompactRequest) error {
|
||||
if req == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// For retry (LLM overflow), use aggressive CompactUntilUnder to guarantee
|
||||
// context shrinks below budget (spec lines ~1410).
|
||||
if req.Reason == ContextCompressReasonRetry && req.Budget > 0 {
|
||||
_, err := m.engine.CompactUntilUnder(ctx, req.SessionKey, req.Budget)
|
||||
return err
|
||||
}
|
||||
|
||||
_, err := m.engine.Compact(ctx, req.SessionKey, seahorse.CompactInput{
|
||||
Force: req.Reason == ContextCompressReasonRetry,
|
||||
Budget: &req.Budget,
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
// Ingest records a message into seahorse SQLite.
|
||||
// All existing sessions are bootstrapped at startup, so this only ingests new messages.
|
||||
func (m *seahorseContextManager) Ingest(ctx context.Context, req *IngestRequest) error {
|
||||
if req == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
msg := providerToSeahorseMessage(req.Message)
|
||||
_, err := m.engine.Ingest(ctx, req.SessionKey, []seahorse.Message{msg})
|
||||
return err
|
||||
}
|
||||
|
||||
// bootstrapSession reconciles JSONL session history into seahorse SQLite.
|
||||
func (m *seahorseContextManager) bootstrapSession(ctx context.Context, sessionKey string) {
|
||||
if m.sessions == nil {
|
||||
return
|
||||
}
|
||||
|
||||
history := m.sessions.GetHistory(sessionKey)
|
||||
if len(history) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Convert provider messages to seahorse messages
|
||||
msgs := make([]seahorse.Message, len(history))
|
||||
for i, h := range history {
|
||||
msgs[i] = providerToSeahorseMessage(h)
|
||||
}
|
||||
|
||||
if err := m.engine.Bootstrap(ctx, sessionKey, msgs); err != nil {
|
||||
logger.WarnCF("seahorse", "bootstrap", map[string]any{
|
||||
"session": sessionKey,
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// providerToSeahorseMessage converts a providers.Message to a seahorse.Message.
|
||||
func providerToSeahorseMessage(msg protocoltypes.Message) seahorse.Message {
|
||||
result := seahorse.Message{
|
||||
Role: msg.Role,
|
||||
Content: msg.Content,
|
||||
ReasoningContent: msg.ReasoningContent,
|
||||
TokenCount: tokenizer.EstimateMessageTokens(msg),
|
||||
}
|
||||
|
||||
// Convert ToolCalls → MessageParts
|
||||
for _, tc := range msg.ToolCalls {
|
||||
part := seahorse.MessagePart{
|
||||
Type: "tool_use",
|
||||
Name: tc.Function.Name,
|
||||
Arguments: tc.Function.Arguments,
|
||||
ToolCallID: tc.ID,
|
||||
}
|
||||
result.Parts = append(result.Parts, part)
|
||||
}
|
||||
|
||||
// Convert tool result
|
||||
if msg.ToolCallID != "" {
|
||||
part := seahorse.MessagePart{
|
||||
Type: "tool_result",
|
||||
ToolCallID: msg.ToolCallID,
|
||||
Text: msg.Content,
|
||||
}
|
||||
result.Parts = append(result.Parts, part)
|
||||
}
|
||||
|
||||
// Convert media attachments
|
||||
for _, mediaURI := range msg.Media {
|
||||
part := seahorse.MessagePart{
|
||||
Type: "media",
|
||||
MediaURI: mediaURI,
|
||||
}
|
||||
result.Parts = append(result.Parts, part)
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// seahorseToProviderMessages converts a seahorse.AssembleResult to []providers.Message.
|
||||
func seahorseToProviderMessages(result *seahorse.AssembleResult) []protocoltypes.Message {
|
||||
messages := make([]protocoltypes.Message, 0, len(result.Messages))
|
||||
|
||||
// Convert assembled messages (which already include summary XML messages)
|
||||
for _, msg := range result.Messages {
|
||||
pm := protocoltypes.Message{
|
||||
Role: msg.Role,
|
||||
Content: msg.Content,
|
||||
ReasoningContent: msg.ReasoningContent,
|
||||
}
|
||||
|
||||
// Reconstruct ToolCalls from parts
|
||||
for _, part := range msg.Parts {
|
||||
if part.Type == "tool_use" {
|
||||
pm.ToolCalls = append(pm.ToolCalls, protocoltypes.ToolCall{
|
||||
ID: part.ToolCallID,
|
||||
Type: "function", // Required by OpenAI-compatible APIs (GLM, etc.)
|
||||
Function: &protocoltypes.FunctionCall{
|
||||
Name: part.Name,
|
||||
Arguments: part.Arguments,
|
||||
},
|
||||
})
|
||||
}
|
||||
if part.Type == "tool_result" {
|
||||
pm.ToolCallID = part.ToolCallID
|
||||
if pm.Content == "" && part.Text != "" {
|
||||
pm.Content = part.Text
|
||||
}
|
||||
}
|
||||
if part.Type == "media" && part.MediaURI != "" {
|
||||
pm.Media = append(pm.Media, part.MediaURI)
|
||||
}
|
||||
}
|
||||
|
||||
messages = append(messages, pm)
|
||||
}
|
||||
|
||||
return messages
|
||||
}
|
||||
|
||||
func init() {
|
||||
if err := RegisterContextManager("seahorse", newSeahorseContextManager); err != nil {
|
||||
panic(fmt.Sprintf("register seahorse context manager: %v", err))
|
||||
}
|
||||
}
|
||||
1086
pkg/agent/context_seahorse_test.go
Normal file
1086
pkg/agent/context_seahorse_test.go
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -472,8 +472,9 @@ func TestAgentLoop_EmitsSessionSummarizeEvent(t *testing.T) {
|
|||
sub := al.SubscribeEvents(16)
|
||||
defer al.UnsubscribeEvents(sub.ID)
|
||||
|
||||
turnScope := al.newTurnEventScope(defaultAgent.ID, "session-1")
|
||||
al.summarizeSession(defaultAgent, "session-1", turnScope)
|
||||
// Use legacyContextManager's summarizeSession via contextManager interface
|
||||
lcm := &legacyContextManager{al: al}
|
||||
lcm.summarizeSession(defaultAgent, "session-1")
|
||||
|
||||
events := collectEventStream(sub.C)
|
||||
summaryEvt, ok := findEvent(events, EventKindSessionSummarize)
|
||||
|
|
|
|||
|
|
@ -167,6 +167,8 @@ const (
|
|||
ContextCompressReasonProactive ContextCompressReason = "proactive_budget"
|
||||
// ContextCompressReasonRetry indicates compression during context-error retry handling.
|
||||
ContextCompressReasonRetry ContextCompressReason = "llm_retry"
|
||||
// ContextCompressReasonSummarize indicates post-turn async summarization.
|
||||
ContextCompressReasonSummarize ContextCompressReason = "summarize"
|
||||
)
|
||||
|
||||
// ContextCompressPayload describes a forced history compression.
|
||||
|
|
|
|||
|
|
@ -77,7 +77,12 @@ func NewAgentInstance(
|
|||
|
||||
if cfg.Tools.IsToolEnabled("read_file") {
|
||||
maxReadFileSize := cfg.Tools.ReadFile.MaxReadFileSize
|
||||
toolsRegistry.Register(tools.NewReadFileTool(workspace, readRestrict, maxReadFileSize, allowReadPaths))
|
||||
switch cfg.Tools.ReadFile.EffectiveMode() {
|
||||
case config.ReadFileModeLines:
|
||||
toolsRegistry.Register(tools.NewReadFileLinesTool(workspace, readRestrict, maxReadFileSize, allowReadPaths))
|
||||
default:
|
||||
toolsRegistry.Register(tools.NewReadFileBytesTool(workspace, readRestrict, maxReadFileSize, allowReadPaths))
|
||||
}
|
||||
}
|
||||
if cfg.Tools.IsToolEnabled("write_file") {
|
||||
toolsRegistry.Register(tools.NewWriteFileTool(workspace, restrict, allowWritePaths))
|
||||
|
|
|
|||
|
|
@ -165,6 +165,58 @@ func TestNewAgentInstance_ResolveCandidatesFromModelListAlias(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestNewAgentInstance_PreservesDistinctLimiterIdentityForSharedResolvedModel(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: tmpDir,
|
||||
ModelName: "glm-4.7",
|
||||
ModelFallbacks: []string{"glm-4.7__key_1"},
|
||||
},
|
||||
},
|
||||
ModelList: []*config.ModelConfig{
|
||||
{
|
||||
ModelName: "glm-4.7",
|
||||
Model: "zhipu/glm-4.7",
|
||||
RPM: 1,
|
||||
},
|
||||
{
|
||||
ModelName: "glm-4.7__key_1",
|
||||
Model: "zhipu/glm-4.7",
|
||||
RPM: 3,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, &mockProvider{})
|
||||
if len(agent.Candidates) != 2 {
|
||||
t.Fatalf("len(Candidates) = %d, want 2", len(agent.Candidates))
|
||||
}
|
||||
|
||||
first := agent.Candidates[0]
|
||||
second := agent.Candidates[1]
|
||||
if first.Provider != "zhipu" || first.Model != "glm-4.7" {
|
||||
t.Fatalf("first candidate = %s/%s, want zhipu/glm-4.7", first.Provider, first.Model)
|
||||
}
|
||||
if second.Provider != "zhipu" || second.Model != "glm-4.7" {
|
||||
t.Fatalf("second candidate = %s/%s, want zhipu/glm-4.7", second.Provider, second.Model)
|
||||
}
|
||||
if first.IdentityKey != "model_name:glm-4.7" {
|
||||
t.Fatalf("first identity key = %q, want %q", first.IdentityKey, "model_name:glm-4.7")
|
||||
}
|
||||
if second.IdentityKey != "model_name:glm-4.7__key_1" {
|
||||
t.Fatalf("second identity key = %q, want %q", second.IdentityKey, "model_name:glm-4.7__key_1")
|
||||
}
|
||||
if first.RPM != 1 {
|
||||
t.Fatalf("first RPM = %d, want 1", first.RPM)
|
||||
}
|
||||
if second.RPM != 3 {
|
||||
t.Fatalf("second RPM = %d, want 3", second.RPM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewAgentInstance_AllowsMediaTempDirForReadListAndExec(t *testing.T) {
|
||||
workspace := t.TempDir()
|
||||
mediaDir := media.TempDir()
|
||||
|
|
@ -248,6 +300,47 @@ func TestNewAgentInstance_AllowsMediaTempDirForReadListAndExec(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestNewAgentInstance_ReadFileModeSelectsSchema(t *testing.T) {
|
||||
workspace := t.TempDir()
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: workspace,
|
||||
ModelName: "test-model",
|
||||
},
|
||||
},
|
||||
Tools: config.ToolsConfig{
|
||||
ReadFile: config.ReadFileToolConfig{
|
||||
Enabled: true,
|
||||
Mode: config.ReadFileModeLines,
|
||||
MaxReadFileSize: 4096,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, &mockProvider{})
|
||||
readTool, ok := agent.Tools.Get("read_file")
|
||||
if !ok {
|
||||
t.Fatal("read_file tool not registered")
|
||||
}
|
||||
|
||||
params := readTool.Parameters()
|
||||
props, _ := params["properties"].(map[string]any)
|
||||
if _, ok := props["start_line"]; !ok {
|
||||
t.Fatalf("expected line-mode schema to expose start_line, got %#v", props)
|
||||
}
|
||||
if _, ok := props["max_lines"]; !ok {
|
||||
t.Fatalf("expected line-mode schema to expose max_lines, got %#v", props)
|
||||
}
|
||||
if _, ok := props["offset"]; ok {
|
||||
t.Fatalf("did not expect line-mode schema to expose offset, got %#v", props)
|
||||
}
|
||||
if _, ok := props["length"]; ok {
|
||||
t.Fatalf("did not expect line-mode schema to expose length, got %#v", props)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewAgentInstance_InvalidExecConfigDoesNotExit(t *testing.T) {
|
||||
workspace := t.TempDir()
|
||||
|
||||
|
|
|
|||
|
|
@ -48,7 +48,7 @@ type AgentLoop struct {
|
|||
|
||||
// Runtime state
|
||||
running atomic.Bool
|
||||
summarizing sync.Map
|
||||
contextManager ContextManager
|
||||
fallback *providers.FallbackChain
|
||||
channelManager *channels.Manager
|
||||
mediaStore media.MediaStore
|
||||
|
|
@ -119,9 +119,18 @@ func NewAgentLoop(
|
|||
) *AgentLoop {
|
||||
registry := NewAgentRegistry(cfg, provider)
|
||||
|
||||
// Set up shared fallback chain
|
||||
// Set up shared fallback chain with rate limiting.
|
||||
cooldown := providers.NewCooldownTracker()
|
||||
fallbackChain := providers.NewFallbackChain(cooldown)
|
||||
rl := providers.NewRateLimiterRegistry()
|
||||
// Register rate limiters for all agents' candidates so that RPM limits
|
||||
// configured in ModelConfig are enforced before each LLM call.
|
||||
for _, agentID := range registry.ListAgentIDs() {
|
||||
if agent, ok := registry.GetAgent(agentID); ok {
|
||||
rl.RegisterCandidates(agent.Candidates)
|
||||
rl.RegisterCandidates(agent.LightCandidates)
|
||||
}
|
||||
}
|
||||
fallbackChain := providers.NewFallbackChain(cooldown, rl)
|
||||
|
||||
// Create state manager using default agent's workspace for channel recording
|
||||
defaultAgent := registry.GetDefaultAgent()
|
||||
|
|
@ -137,13 +146,13 @@ func NewAgentLoop(
|
|||
registry: registry,
|
||||
state: stateManager,
|
||||
eventBus: eventBus,
|
||||
summarizing: sync.Map{},
|
||||
fallback: fallbackChain,
|
||||
cmdRegistry: commands.NewRegistry(commands.BuiltinDefinitions()),
|
||||
steering: newSteeringQueue(parseSteeringMode(cfg.Agents.Defaults.SteeringMode)),
|
||||
}
|
||||
al.hooks = NewHookManager(eventBus)
|
||||
configureHookManagerFromConfig(al.hooks, cfg)
|
||||
al.contextManager = al.resolveContextManager()
|
||||
|
||||
// Register shared tools to all agents (now that al is created)
|
||||
registerSharedTools(al, cfg, msgBus, registry, provider)
|
||||
|
|
@ -281,6 +290,17 @@ func registerSharedTools(
|
|||
agent.Tools.Register(tools.NewSendTTSTool(ttsProvider, nil))
|
||||
}
|
||||
|
||||
if cfg.Tools.IsToolEnabled("load_image") {
|
||||
loadImageTool := tools.NewLoadImageTool(
|
||||
agent.Workspace,
|
||||
cfg.Agents.Defaults.RestrictToWorkspace,
|
||||
cfg.Agents.Defaults.GetMaxMediaSize(),
|
||||
nil,
|
||||
allowReadPaths,
|
||||
)
|
||||
agent.Tools.Register(loadImageTool)
|
||||
}
|
||||
|
||||
// Skill discovery and installation tools
|
||||
skills_enabled := cfg.Tools.IsToolEnabled("skills")
|
||||
find_skills_enable := cfg.Tools.IsToolEnabled("find_skills")
|
||||
|
|
@ -323,6 +343,14 @@ func registerSharedTools(
|
|||
subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Workspace)
|
||||
subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature)
|
||||
|
||||
// Inject a media resolver so the legacy RunToolLoop fallback path can
|
||||
// resolve media:// refs in the same way the main AgentLoop does.
|
||||
// This keeps subagent vision support working even when the optimized
|
||||
// sub-turn spawner path is unavailable.
|
||||
subagentManager.SetMediaResolver(func(msgs []providers.Message) []providers.Message {
|
||||
return resolveMediaRefs(msgs, al.mediaStore, cfg.Agents.Defaults.GetMaxMediaSize())
|
||||
})
|
||||
|
||||
// Set the spawner that links into AgentLoop's turnState
|
||||
subagentManager.SetSpawner(func(
|
||||
ctx context.Context,
|
||||
|
|
@ -972,6 +1000,7 @@ func (al *AgentLoop) ReloadProviderAndConfig(
|
|||
go func() {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
logger.RecoverPanicNoExit(r)
|
||||
panicErr = fmt.Errorf("panic during registry creation: %v", r)
|
||||
logger.ErrorCF("agent", "Panic during registry creation",
|
||||
map[string]any{"panic": r})
|
||||
|
|
@ -1012,8 +1041,15 @@ func (al *AgentLoop) ReloadProviderAndConfig(
|
|||
al.cfg = cfg
|
||||
al.registry = registry
|
||||
|
||||
// Also update fallback chain with new config
|
||||
al.fallback = providers.NewFallbackChain(providers.NewCooldownTracker())
|
||||
// Also update fallback chain with new config; rebuild rate limiter registry.
|
||||
newRL := providers.NewRateLimiterRegistry()
|
||||
for _, agentID := range registry.ListAgentIDs() {
|
||||
if agent, ok := registry.GetAgent(agentID); ok {
|
||||
newRL.RegisterCandidates(agent.Candidates)
|
||||
newRL.RegisterCandidates(agent.LightCandidates)
|
||||
}
|
||||
}
|
||||
al.fallback = providers.NewFallbackChain(providers.NewCooldownTracker(), newRL)
|
||||
|
||||
al.mu.Unlock()
|
||||
|
||||
|
|
@ -1670,8 +1706,15 @@ func (al *AgentLoop) runTurn(ctx context.Context, ts *turnState) (turnResult, er
|
|||
var history []providers.Message
|
||||
var summary string
|
||||
if !ts.opts.NoHistory {
|
||||
history = ts.agent.Sessions.GetHistory(ts.sessionKey)
|
||||
summary = ts.agent.Sessions.GetSummary(ts.sessionKey)
|
||||
// ContextManager assembles budget-aware history and summary.
|
||||
if resp, err := al.contextManager.Assemble(turnCtx, &AssembleRequest{
|
||||
SessionKey: ts.sessionKey,
|
||||
Budget: ts.agent.ContextWindow,
|
||||
MaxTokens: ts.agent.MaxTokens,
|
||||
}); err == nil && resp != nil {
|
||||
history = resp.History
|
||||
summary = resp.Summary
|
||||
}
|
||||
}
|
||||
ts.captureRestorePoint(history, summary)
|
||||
|
||||
|
|
@ -1696,22 +1739,28 @@ func (al *AgentLoop) runTurn(ctx context.Context, ts *turnState) (turnResult, er
|
|||
if isOverContextBudget(ts.agent.ContextWindow, messages, toolDefs, ts.agent.MaxTokens) {
|
||||
logger.WarnCF("agent", "Proactive compression: context budget exceeded before LLM call",
|
||||
map[string]any{"session_key": ts.sessionKey})
|
||||
if compression, ok := al.forceCompression(ts.agent, ts.sessionKey); ok {
|
||||
al.emitEvent(
|
||||
EventKindContextCompress,
|
||||
ts.eventMeta("runTurn", "turn.context.compress"),
|
||||
ContextCompressPayload{
|
||||
if err := al.contextManager.Compact(turnCtx, &CompactRequest{
|
||||
SessionKey: ts.sessionKey,
|
||||
Reason: ContextCompressReasonProactive,
|
||||
DroppedMessages: compression.DroppedMessages,
|
||||
RemainingMessages: compression.RemainingMessages,
|
||||
},
|
||||
)
|
||||
ts.refreshRestorePointFromSession(ts.agent)
|
||||
Budget: ts.agent.ContextWindow,
|
||||
}); err != nil {
|
||||
logger.WarnCF("agent", "Proactive compact failed", map[string]any{
|
||||
"session_key": ts.sessionKey,
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
ts.refreshRestorePointFromSession(ts.agent)
|
||||
// Re-assemble from CM after compact.
|
||||
if resp, err := al.contextManager.Assemble(turnCtx, &AssembleRequest{
|
||||
SessionKey: ts.sessionKey,
|
||||
Budget: ts.agent.ContextWindow,
|
||||
MaxTokens: ts.agent.MaxTokens,
|
||||
}); err == nil && resp != nil {
|
||||
history = resp.History
|
||||
summary = resp.Summary
|
||||
}
|
||||
newHistory := ts.agent.Sessions.GetHistory(ts.sessionKey)
|
||||
newSummary := ts.agent.Sessions.GetSummary(ts.sessionKey)
|
||||
messages = ts.agent.ContextBuilder.BuildMessages(
|
||||
newHistory, newSummary, ts.userMessage,
|
||||
history, summary, ts.userMessage,
|
||||
ts.media, ts.channel, ts.chatID,
|
||||
ts.opts.SenderID, ts.opts.SenderDisplayName,
|
||||
activeSkillNames(ts.agent, ts.opts)...,
|
||||
|
|
@ -1733,6 +1782,7 @@ func (al *AgentLoop) runTurn(ctx context.Context, ts *turnState) (turnResult, er
|
|||
ts.agent.Sessions.AddMessage(ts.sessionKey, rootMsg.Role, rootMsg.Content)
|
||||
}
|
||||
ts.recordPersistedMessage(rootMsg)
|
||||
ts.ingestMessage(turnCtx, al, rootMsg)
|
||||
}
|
||||
|
||||
activeCandidates, activeModel, usedLight := al.selectCandidates(ts.agent, ts.userMessage, messages)
|
||||
|
|
@ -1808,6 +1858,7 @@ turnLoop:
|
|||
if !ts.opts.NoHistory {
|
||||
ts.agent.Sessions.AddFullMessage(ts.sessionKey, pm)
|
||||
ts.recordPersistedMessage(pm)
|
||||
ts.ingestMessage(turnCtx, al, pm)
|
||||
}
|
||||
logger.InfoCF("agent", "Injected steering message into context",
|
||||
map[string]any{
|
||||
|
|
@ -1861,6 +1912,14 @@ turnLoop:
|
|||
providerToolDefs = filtered
|
||||
}
|
||||
|
||||
// Resolve media:// refs produced by tool results (e.g. load_image).
|
||||
// Skipped on iteration 1 because inbound user media is already resolved
|
||||
// before entering the loop; only subsequent iterations can contain new
|
||||
// tool-generated media refs that need base64 encoding.
|
||||
if iteration > 1 {
|
||||
messages = resolveMediaRefs(messages, al.mediaStore, maxMediaSize)
|
||||
}
|
||||
|
||||
callMessages := messages
|
||||
if gracefulTerminal {
|
||||
callMessages = append(append([]providers.Message(nil), messages...), ts.interruptHintMessage())
|
||||
|
|
@ -2068,23 +2127,28 @@ turnLoop:
|
|||
})
|
||||
}
|
||||
|
||||
if compression, ok := al.forceCompression(ts.agent, ts.sessionKey); ok {
|
||||
al.emitEvent(
|
||||
EventKindContextCompress,
|
||||
ts.eventMeta("runTurn", "turn.context.compress"),
|
||||
ContextCompressPayload{
|
||||
if compactErr := al.contextManager.Compact(turnCtx, &CompactRequest{
|
||||
SessionKey: ts.sessionKey,
|
||||
Reason: ContextCompressReasonRetry,
|
||||
DroppedMessages: compression.DroppedMessages,
|
||||
RemainingMessages: compression.RemainingMessages,
|
||||
},
|
||||
)
|
||||
ts.refreshRestorePointFromSession(ts.agent)
|
||||
Budget: ts.agent.ContextWindow,
|
||||
}); compactErr != nil {
|
||||
logger.WarnCF("agent", "Context overflow compact failed", map[string]any{
|
||||
"session_key": ts.sessionKey,
|
||||
"error": compactErr.Error(),
|
||||
})
|
||||
}
|
||||
ts.refreshRestorePointFromSession(ts.agent)
|
||||
// Re-assemble from CM after compact.
|
||||
if asmResp, asmErr := al.contextManager.Assemble(turnCtx, &AssembleRequest{
|
||||
SessionKey: ts.sessionKey,
|
||||
Budget: ts.agent.ContextWindow,
|
||||
MaxTokens: ts.agent.MaxTokens,
|
||||
}); asmErr == nil && asmResp != nil {
|
||||
history = asmResp.History
|
||||
summary = asmResp.Summary
|
||||
}
|
||||
|
||||
newHistory := ts.agent.Sessions.GetHistory(ts.sessionKey)
|
||||
newSummary := ts.agent.Sessions.GetSummary(ts.sessionKey)
|
||||
messages = ts.agent.ContextBuilder.BuildMessages(
|
||||
newHistory, newSummary, "",
|
||||
history, summary, "",
|
||||
nil, ts.channel, ts.chatID, ts.opts.SenderID, ts.opts.SenderDisplayName,
|
||||
activeSkillNames(ts.agent, ts.opts)...,
|
||||
)
|
||||
|
|
@ -2257,6 +2321,7 @@ turnLoop:
|
|||
if !ts.opts.NoHistory {
|
||||
ts.agent.Sessions.AddFullMessage(ts.sessionKey, assistantMsg)
|
||||
ts.recordPersistedMessage(assistantMsg)
|
||||
ts.ingestMessage(turnCtx, al, assistantMsg)
|
||||
}
|
||||
|
||||
ts.setPhase(TurnPhaseTools)
|
||||
|
|
@ -2636,6 +2701,13 @@ turnLoop:
|
|||
}
|
||||
|
||||
if len(toolResult.Media) > 0 && !toolResult.ResponseHandled {
|
||||
// For tools like load_image that produce media refs without sending them
|
||||
// to the user channel (ResponseHandled == false), both Media and ArtifactTags
|
||||
// coexist on the result:
|
||||
// - Media: carries media:// refs that resolveMediaRefs will base64-encode
|
||||
// into image_url parts in the next LLM iteration (enabling vision).
|
||||
// - ArtifactTags: exposes the local file path as a structured [file:…] tag
|
||||
// in the tool result text, so the LLM knows an artifact was produced.
|
||||
toolResult.ArtifactTags = buildArtifactTags(al.mediaStore, toolResult.Media)
|
||||
}
|
||||
|
||||
|
|
@ -2655,6 +2727,9 @@ turnLoop:
|
|||
Content: contentForLLM,
|
||||
ToolCallID: toolCallID,
|
||||
}
|
||||
if len(toolResult.Media) > 0 && !toolResult.ResponseHandled {
|
||||
toolResultMsg.Media = append(toolResultMsg.Media, toolResult.Media...)
|
||||
}
|
||||
al.emitEvent(
|
||||
EventKindToolExecEnd,
|
||||
ts.eventMeta("runTurn", "turn.tool.end"),
|
||||
|
|
@ -2671,6 +2746,7 @@ turnLoop:
|
|||
if !ts.opts.NoHistory {
|
||||
ts.agent.Sessions.AddFullMessage(ts.sessionKey, toolResultMsg)
|
||||
ts.recordPersistedMessage(toolResultMsg)
|
||||
ts.ingestMessage(turnCtx, al, toolResultMsg)
|
||||
}
|
||||
|
||||
if steerMsgs := al.dequeueSteeringMessagesForScope(ts.sessionKey); len(steerMsgs) > 0 {
|
||||
|
|
@ -2770,6 +2846,7 @@ turnLoop:
|
|||
if !ts.opts.NoHistory {
|
||||
ts.agent.Sessions.AddMessage(ts.sessionKey, summaryMsg.Role, summaryMsg.Content)
|
||||
ts.recordPersistedMessage(summaryMsg)
|
||||
ts.ingestMessage(turnCtx, al, summaryMsg)
|
||||
if err := ts.agent.Sessions.Save(ts.sessionKey); err != nil {
|
||||
turnStatus = TurnEndStatusError
|
||||
al.emitEvent(
|
||||
|
|
@ -2784,7 +2861,7 @@ turnLoop:
|
|||
}
|
||||
}
|
||||
if ts.opts.EnableSummary {
|
||||
al.maybeSummarize(ts.agent, ts.sessionKey, ts.scope)
|
||||
al.contextManager.Compact(turnCtx, &CompactRequest{SessionKey: ts.sessionKey, Reason: ContextCompressReasonSummarize, Budget: ts.agent.ContextWindow})
|
||||
}
|
||||
|
||||
ts.setPhase(TurnPhaseCompleted)
|
||||
|
|
@ -2839,6 +2916,7 @@ turnLoop:
|
|||
finalMsg := providers.Message{Role: "assistant", Content: finalContent}
|
||||
ts.agent.Sessions.AddMessage(ts.sessionKey, finalMsg.Role, finalMsg.Content)
|
||||
ts.recordPersistedMessage(finalMsg)
|
||||
ts.ingestMessage(turnCtx, al, finalMsg)
|
||||
if err := ts.agent.Sessions.Save(ts.sessionKey); err != nil {
|
||||
turnStatus = TurnEndStatusError
|
||||
al.emitEvent(
|
||||
|
|
@ -2854,7 +2932,14 @@ turnLoop:
|
|||
}
|
||||
|
||||
if ts.opts.EnableSummary {
|
||||
al.maybeSummarize(ts.agent, ts.sessionKey, ts.scope)
|
||||
al.contextManager.Compact(
|
||||
turnCtx,
|
||||
&CompactRequest{
|
||||
SessionKey: ts.sessionKey,
|
||||
Reason: ContextCompressReasonSummarize,
|
||||
Budget: ts.agent.ContextWindow,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
ts.setPhase(TurnPhaseCompleted)
|
||||
|
|
@ -2933,103 +3018,28 @@ func (al *AgentLoop) selectCandidates(
|
|||
return agent.LightCandidates, resolvedCandidateModel(agent.LightCandidates, agent.Router.LightModel()), true
|
||||
}
|
||||
|
||||
// maybeSummarize triggers summarization if the session history exceeds thresholds.
|
||||
func (al *AgentLoop) maybeSummarize(agent *AgentInstance, sessionKey string, turnScope turnEventScope) {
|
||||
newHistory := agent.Sessions.GetHistory(sessionKey)
|
||||
tokenEstimate := al.estimateTokens(newHistory)
|
||||
threshold := agent.ContextWindow * agent.SummarizeTokenPercent / 100
|
||||
|
||||
if len(newHistory) > agent.SummarizeMessageThreshold || tokenEstimate > threshold {
|
||||
summarizeKey := agent.ID + ":" + sessionKey
|
||||
if _, loading := al.summarizing.LoadOrStore(summarizeKey, true); !loading {
|
||||
go func() {
|
||||
defer al.summarizing.Delete(summarizeKey)
|
||||
logger.Debug("Memory threshold reached. Optimizing conversation history...")
|
||||
al.summarizeSession(agent, sessionKey, turnScope)
|
||||
}()
|
||||
// resolveContextManager selects the ContextManager implementation based on config.
|
||||
func (al *AgentLoop) resolveContextManager() ContextManager {
|
||||
name := al.cfg.Agents.Defaults.ContextManager
|
||||
if name == "" || name == "legacy" {
|
||||
return &legacyContextManager{al: al}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type compressionResult struct {
|
||||
DroppedMessages int
|
||||
RemainingMessages int
|
||||
}
|
||||
|
||||
// forceCompression aggressively reduces context when the limit is hit.
|
||||
// It drops the oldest ~50% of Turns (a Turn is a complete user→LLM→response
|
||||
// cycle, as defined in #1316), so tool-call sequences are never split.
|
||||
//
|
||||
// If the history is a single Turn with no safe split point, the function
|
||||
// falls back to keeping only the most recent user message. This breaks
|
||||
// Turn atomicity as a last resort to avoid a context-exceeded loop.
|
||||
//
|
||||
// Session history contains only user/assistant/tool messages — the system
|
||||
// prompt is built dynamically by BuildMessages and is NOT stored here.
|
||||
// The compression note is recorded in the session summary so that
|
||||
// BuildMessages can include it in the next system prompt.
|
||||
func (al *AgentLoop) forceCompression(agent *AgentInstance, sessionKey string) (compressionResult, bool) {
|
||||
history := agent.Sessions.GetHistory(sessionKey)
|
||||
if len(history) <= 2 {
|
||||
return compressionResult{}, false
|
||||
}
|
||||
|
||||
// Split at a Turn boundary so no tool-call sequence is torn apart.
|
||||
// parseTurnBoundaries gives us the start of each Turn; we drop the
|
||||
// oldest half of Turns and keep the most recent ones.
|
||||
turns := parseTurnBoundaries(history)
|
||||
var mid int
|
||||
if len(turns) >= 2 {
|
||||
mid = turns[len(turns)/2]
|
||||
} else {
|
||||
// Fewer than 2 Turns — fall back to message-level midpoint
|
||||
// aligned to the nearest Turn boundary.
|
||||
mid = findSafeBoundary(history, len(history)/2)
|
||||
}
|
||||
var keptHistory []providers.Message
|
||||
if mid <= 0 {
|
||||
// No safe Turn boundary — the entire history is a single Turn
|
||||
// (e.g. one user message followed by a massive tool response).
|
||||
// Keeping everything would leave the agent stuck in a context-
|
||||
// exceeded loop, so fall back to keeping only the most recent
|
||||
// user message. This breaks Turn atomicity as a last resort.
|
||||
for i := len(history) - 1; i >= 0; i-- {
|
||||
if history[i].Role == "user" {
|
||||
keptHistory = []providers.Message{history[i]}
|
||||
break
|
||||
}
|
||||
}
|
||||
} else {
|
||||
keptHistory = history[mid:]
|
||||
}
|
||||
|
||||
droppedCount := len(history) - len(keptHistory)
|
||||
|
||||
// Record compression in the session summary so BuildMessages includes it
|
||||
// in the system prompt. We do not modify history messages themselves.
|
||||
existingSummary := agent.Sessions.GetSummary(sessionKey)
|
||||
compressionNote := fmt.Sprintf(
|
||||
"[Emergency compression dropped %d oldest messages due to context limit]",
|
||||
droppedCount,
|
||||
)
|
||||
if existingSummary != "" {
|
||||
compressionNote = existingSummary + "\n\n" + compressionNote
|
||||
}
|
||||
agent.Sessions.SetSummary(sessionKey, compressionNote)
|
||||
|
||||
agent.Sessions.SetHistory(sessionKey, keptHistory)
|
||||
agent.Sessions.Save(sessionKey)
|
||||
|
||||
logger.WarnCF("agent", "Forced compression executed", map[string]any{
|
||||
"session_key": sessionKey,
|
||||
"dropped_msgs": droppedCount,
|
||||
"new_count": len(keptHistory),
|
||||
factory, ok := lookupContextManager(name)
|
||||
if !ok {
|
||||
logger.WarnCF("agent", "Unknown context manager, falling back to legacy", map[string]any{
|
||||
"name": name,
|
||||
})
|
||||
|
||||
return compressionResult{
|
||||
DroppedMessages: droppedCount,
|
||||
RemainingMessages: len(keptHistory),
|
||||
}, true
|
||||
return &legacyContextManager{al: al}
|
||||
}
|
||||
cm, err := factory(al.cfg.Agents.Defaults.ContextManagerConfig, al)
|
||||
if err != nil {
|
||||
logger.WarnCF("agent", "Failed to create context manager, falling back to legacy", map[string]any{
|
||||
"name": name,
|
||||
"error": err.Error(),
|
||||
})
|
||||
return &legacyContextManager{al: al}
|
||||
}
|
||||
return cm
|
||||
}
|
||||
|
||||
// GetStartupInfo returns information about loaded tools and skills for logging.
|
||||
|
|
@ -3121,247 +3131,13 @@ func formatToolsForLog(toolDefs []providers.ToolDefinition) string {
|
|||
}
|
||||
|
||||
// summarizeSession summarizes the conversation history for a session.
|
||||
func (al *AgentLoop) summarizeSession(agent *AgentInstance, sessionKey string, turnScope turnEventScope) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
|
||||
defer cancel()
|
||||
|
||||
history := agent.Sessions.GetHistory(sessionKey)
|
||||
summary := agent.Sessions.GetSummary(sessionKey)
|
||||
|
||||
// Keep the most recent Turns for continuity, aligned to a Turn boundary
|
||||
// so that no tool-call sequence is split.
|
||||
if len(history) <= 4 {
|
||||
return
|
||||
}
|
||||
|
||||
safeCut := findSafeBoundary(history, len(history)-4)
|
||||
if safeCut <= 0 {
|
||||
return
|
||||
}
|
||||
keepCount := len(history) - safeCut
|
||||
toSummarize := history[:safeCut]
|
||||
|
||||
// Oversized Message Guard
|
||||
maxMessageTokens := agent.ContextWindow / 2
|
||||
validMessages := make([]providers.Message, 0)
|
||||
omitted := false
|
||||
|
||||
for _, m := range toSummarize {
|
||||
if m.Role != "user" && m.Role != "assistant" {
|
||||
continue
|
||||
}
|
||||
msgTokens := len(m.Content) / 2
|
||||
if msgTokens > maxMessageTokens {
|
||||
omitted = true
|
||||
continue
|
||||
}
|
||||
validMessages = append(validMessages, m)
|
||||
}
|
||||
|
||||
if len(validMessages) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
const (
|
||||
maxSummarizationMessages = 10
|
||||
llmMaxRetries = 3
|
||||
llmTemperature = 0.3
|
||||
fallbackMaxContentLength = 200
|
||||
)
|
||||
|
||||
// Multi-Part Summarization
|
||||
var finalSummary string
|
||||
if len(validMessages) > maxSummarizationMessages {
|
||||
mid := len(validMessages) / 2
|
||||
|
||||
mid = al.findNearestUserMessage(validMessages, mid)
|
||||
|
||||
part1 := validMessages[:mid]
|
||||
part2 := validMessages[mid:]
|
||||
|
||||
s1, _ := al.summarizeBatch(ctx, agent, part1, "")
|
||||
s2, _ := al.summarizeBatch(ctx, agent, part2, "")
|
||||
|
||||
mergePrompt := fmt.Sprintf(
|
||||
"Merge these two conversation summaries into one cohesive summary:\n\n1: %s\n\n2: %s",
|
||||
s1,
|
||||
s2,
|
||||
)
|
||||
|
||||
resp, err := al.retryLLMCall(ctx, agent, mergePrompt, llmMaxRetries)
|
||||
if err == nil && resp.Content != "" {
|
||||
finalSummary = resp.Content
|
||||
} else {
|
||||
finalSummary = s1 + " " + s2
|
||||
}
|
||||
} else {
|
||||
finalSummary, _ = al.summarizeBatch(ctx, agent, validMessages, summary)
|
||||
}
|
||||
|
||||
if omitted && finalSummary != "" {
|
||||
finalSummary += "\n[Note: Some oversized messages were omitted from this summary for efficiency.]"
|
||||
}
|
||||
|
||||
if finalSummary != "" {
|
||||
agent.Sessions.SetSummary(sessionKey, finalSummary)
|
||||
agent.Sessions.TruncateHistory(sessionKey, keepCount)
|
||||
agent.Sessions.Save(sessionKey)
|
||||
al.emitEvent(
|
||||
EventKindSessionSummarize,
|
||||
turnScope.meta(0, "summarizeSession", "turn.session.summarize"),
|
||||
SessionSummarizePayload{
|
||||
SummarizedMessages: len(validMessages),
|
||||
KeptMessages: keepCount,
|
||||
SummaryLen: len(finalSummary),
|
||||
OmittedOversized: omitted,
|
||||
},
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// findNearestUserMessage finds the nearest user message to the given index.
|
||||
// It searches backward first, then forward if no user message is found.
|
||||
func (al *AgentLoop) findNearestUserMessage(messages []providers.Message, mid int) int {
|
||||
originalMid := mid
|
||||
|
||||
for mid > 0 && messages[mid].Role != "user" {
|
||||
mid--
|
||||
}
|
||||
|
||||
if messages[mid].Role == "user" {
|
||||
return mid
|
||||
}
|
||||
|
||||
mid = originalMid
|
||||
for mid < len(messages) && messages[mid].Role != "user" {
|
||||
mid++
|
||||
}
|
||||
|
||||
if mid < len(messages) {
|
||||
return mid
|
||||
}
|
||||
|
||||
return originalMid
|
||||
}
|
||||
|
||||
// retryLLMCall calls the LLM with retry logic.
|
||||
func (al *AgentLoop) retryLLMCall(
|
||||
ctx context.Context,
|
||||
agent *AgentInstance,
|
||||
prompt string,
|
||||
maxRetries int,
|
||||
) (*providers.LLMResponse, error) {
|
||||
const (
|
||||
llmTemperature = 0.3
|
||||
)
|
||||
|
||||
var resp *providers.LLMResponse
|
||||
var err error
|
||||
|
||||
for attempt := 0; attempt < maxRetries; attempt++ {
|
||||
al.activeRequests.Add(1)
|
||||
resp, err = func() (*providers.LLMResponse, error) {
|
||||
defer al.activeRequests.Done()
|
||||
return agent.Provider.Chat(
|
||||
ctx,
|
||||
[]providers.Message{{Role: "user", Content: prompt}},
|
||||
nil,
|
||||
agent.Model,
|
||||
map[string]any{
|
||||
"max_tokens": agent.MaxTokens,
|
||||
"temperature": llmTemperature,
|
||||
"prompt_cache_key": agent.ID,
|
||||
},
|
||||
)
|
||||
}()
|
||||
|
||||
if err == nil && resp != nil && resp.Content != "" {
|
||||
return resp, nil
|
||||
}
|
||||
if attempt < maxRetries-1 {
|
||||
time.Sleep(time.Duration(attempt+1) * 100 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
return resp, err
|
||||
}
|
||||
|
||||
// summarizeBatch summarizes a batch of messages.
|
||||
func (al *AgentLoop) summarizeBatch(
|
||||
ctx context.Context,
|
||||
agent *AgentInstance,
|
||||
batch []providers.Message,
|
||||
existingSummary string,
|
||||
) (string, error) {
|
||||
const (
|
||||
llmMaxRetries = 3
|
||||
llmTemperature = 0.3
|
||||
fallbackMinContentLength = 200
|
||||
fallbackMaxContentPercent = 10
|
||||
)
|
||||
|
||||
var sb strings.Builder
|
||||
sb.WriteString(
|
||||
"Provide a concise summary of this conversation segment, preserving core context and key points.\n",
|
||||
)
|
||||
if existingSummary != "" {
|
||||
sb.WriteString("Existing context: ")
|
||||
sb.WriteString(existingSummary)
|
||||
sb.WriteString("\n")
|
||||
}
|
||||
sb.WriteString("\nCONVERSATION:\n")
|
||||
for _, m := range batch {
|
||||
fmt.Fprintf(&sb, "%s: %s\n", m.Role, m.Content)
|
||||
}
|
||||
prompt := sb.String()
|
||||
|
||||
response, err := al.retryLLMCall(ctx, agent, prompt, llmMaxRetries)
|
||||
if err == nil && response.Content != "" {
|
||||
return strings.TrimSpace(response.Content), nil
|
||||
}
|
||||
|
||||
var fallback strings.Builder
|
||||
fallback.WriteString("Conversation summary: ")
|
||||
for i, m := range batch {
|
||||
if i > 0 {
|
||||
fallback.WriteString(" | ")
|
||||
}
|
||||
content := strings.TrimSpace(m.Content)
|
||||
runes := []rune(content)
|
||||
if len(runes) == 0 {
|
||||
fallback.WriteString(fmt.Sprintf("%s: ", m.Role))
|
||||
continue
|
||||
}
|
||||
|
||||
keepLength := len(runes) * fallbackMaxContentPercent / 100
|
||||
if keepLength < fallbackMinContentLength {
|
||||
keepLength = fallbackMinContentLength
|
||||
}
|
||||
|
||||
if keepLength > len(runes) {
|
||||
keepLength = len(runes)
|
||||
}
|
||||
|
||||
content = string(runes[:keepLength])
|
||||
if keepLength < len(runes) {
|
||||
content += "..."
|
||||
}
|
||||
fallback.WriteString(fmt.Sprintf("%s: %s", m.Role, content))
|
||||
}
|
||||
return fallback.String(), nil
|
||||
}
|
||||
|
||||
// estimateTokens estimates the number of tokens in a message list.
|
||||
// Counts Content, ToolCalls arguments, and ToolCallID metadata so that
|
||||
// tool-heavy conversations are not systematically undercounted.
|
||||
func (al *AgentLoop) estimateTokens(messages []providers.Message) int {
|
||||
total := 0
|
||||
for _, m := range messages {
|
||||
total += estimateMessageTokens(m)
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
func (al *AgentLoop) handleCommand(
|
||||
ctx context.Context,
|
||||
msg bus.InboundMessage,
|
||||
|
|
@ -3558,7 +3334,7 @@ func (al *AgentLoop) buildCommandsRuntime(agent *AgentInstance, opts *processOpt
|
|||
return "", fmt.Errorf("failed to initialize model %q: %w", value, err)
|
||||
}
|
||||
|
||||
nextCandidates := resolveModelCandidates(cfg, cfg.Agents.Defaults.Provider, modelCfg.Model, agent.Fallbacks)
|
||||
nextCandidates := resolveModelCandidates(cfg, cfg.Agents.Defaults.Provider, value, agent.Fallbacks)
|
||||
if len(nextCandidates) == 0 {
|
||||
return "", fmt.Errorf("model %q did not resolve to any provider candidates", value)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -126,6 +126,8 @@ func (al *AgentLoop) ensureMCPInitialized(ctx context.Context) error {
|
|||
}
|
||||
|
||||
mcpTool := tools.NewMCPTool(mcpManager, serverName, tool)
|
||||
mcpTool.SetWorkspace(agent.Workspace)
|
||||
mcpTool.SetMaxInlineTextRunes(al.cfg.Tools.MCP.GetMaxInlineTextChars())
|
||||
|
||||
if registerAsHidden {
|
||||
agent.Tools.RegisterHidden(mcpTool)
|
||||
|
|
|
|||
|
|
@ -8,8 +8,7 @@ import (
|
|||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
)
|
||||
|
||||
func buildModelListResolver(cfg *config.Config) func(raw string) (string, bool) {
|
||||
ensureProtocol := func(model string) string {
|
||||
func ensureProtocolModel(model string) string {
|
||||
model = strings.TrimSpace(model)
|
||||
if model == "" {
|
||||
return ""
|
||||
|
|
@ -20,32 +19,91 @@ func buildModelListResolver(cfg *config.Config) func(raw string) (string, bool)
|
|||
return "openai/" + model
|
||||
}
|
||||
|
||||
return func(raw string) (string, bool) {
|
||||
func modelConfigIdentityKey(mc *config.ModelConfig) string {
|
||||
if mc == nil {
|
||||
return ""
|
||||
}
|
||||
if name := strings.TrimSpace(mc.ModelName); name != "" {
|
||||
return "model_name:" + name
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func candidateFromModelConfig(
|
||||
defaultProvider string,
|
||||
mc *config.ModelConfig,
|
||||
) (providers.FallbackCandidate, bool) {
|
||||
if mc == nil {
|
||||
return providers.FallbackCandidate{}, false
|
||||
}
|
||||
|
||||
ref := providers.ParseModelRef(ensureProtocolModel(mc.Model), defaultProvider)
|
||||
if ref == nil {
|
||||
return providers.FallbackCandidate{}, false
|
||||
}
|
||||
|
||||
return providers.FallbackCandidate{
|
||||
Provider: ref.Provider,
|
||||
Model: ref.Model,
|
||||
RPM: mc.RPM,
|
||||
IdentityKey: modelConfigIdentityKey(mc),
|
||||
}, true
|
||||
}
|
||||
|
||||
func lookupModelConfigByRef(cfg *config.Config, raw string) *config.ModelConfig {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" || cfg == nil {
|
||||
return "", false
|
||||
return nil
|
||||
}
|
||||
|
||||
if mc, err := cfg.GetModelConfig(raw); err == nil && mc != nil && strings.TrimSpace(mc.Model) != "" {
|
||||
return ensureProtocol(mc.Model), true
|
||||
return mc
|
||||
}
|
||||
|
||||
for i := range cfg.ModelList {
|
||||
fullModel := strings.TrimSpace(cfg.ModelList[i].Model)
|
||||
mc := cfg.ModelList[i]
|
||||
if mc == nil {
|
||||
continue
|
||||
}
|
||||
fullModel := strings.TrimSpace(mc.Model)
|
||||
if fullModel == "" {
|
||||
continue
|
||||
}
|
||||
if fullModel == raw {
|
||||
return ensureProtocol(fullModel), true
|
||||
return mc
|
||||
}
|
||||
_, modelID := providers.ExtractProtocol(fullModel)
|
||||
if modelID == raw {
|
||||
return ensureProtocol(fullModel), true
|
||||
return mc
|
||||
}
|
||||
}
|
||||
|
||||
return "", false
|
||||
return nil
|
||||
}
|
||||
|
||||
func resolveModelCandidate(
|
||||
cfg *config.Config,
|
||||
defaultProvider string,
|
||||
raw string,
|
||||
) (providers.FallbackCandidate, bool) {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" {
|
||||
return providers.FallbackCandidate{}, false
|
||||
}
|
||||
|
||||
if mc := lookupModelConfigByRef(cfg, raw); mc != nil {
|
||||
return candidateFromModelConfig(defaultProvider, mc)
|
||||
}
|
||||
|
||||
ref := providers.ParseModelRef(raw, defaultProvider)
|
||||
if ref == nil {
|
||||
return providers.FallbackCandidate{}, false
|
||||
}
|
||||
|
||||
return providers.FallbackCandidate{
|
||||
Provider: ref.Provider,
|
||||
Model: ref.Model,
|
||||
}, true
|
||||
}
|
||||
|
||||
func resolveModelCandidates(
|
||||
|
|
@ -54,14 +112,29 @@ func resolveModelCandidates(
|
|||
primary string,
|
||||
fallbacks []string,
|
||||
) []providers.FallbackCandidate {
|
||||
return providers.ResolveCandidatesWithLookup(
|
||||
providers.ModelConfig{
|
||||
Primary: primary,
|
||||
Fallbacks: fallbacks,
|
||||
},
|
||||
defaultProvider,
|
||||
buildModelListResolver(cfg),
|
||||
)
|
||||
seen := make(map[string]bool)
|
||||
candidates := make([]providers.FallbackCandidate, 0, 1+len(fallbacks))
|
||||
|
||||
addCandidate := func(raw string) {
|
||||
candidate, ok := resolveModelCandidate(cfg, defaultProvider, raw)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
key := candidate.StableKey()
|
||||
if seen[key] {
|
||||
return
|
||||
}
|
||||
seen[key] = true
|
||||
candidates = append(candidates, candidate)
|
||||
}
|
||||
|
||||
addCandidate(primary)
|
||||
for _, fallback := range fallbacks {
|
||||
addCandidate(fallback)
|
||||
}
|
||||
|
||||
return candidates
|
||||
}
|
||||
|
||||
func resolvedCandidateModel(candidates []providers.FallbackCandidate, fallback string) string {
|
||||
|
|
|
|||
|
|
@ -427,6 +427,7 @@ func spawnSubTurn(
|
|||
// 7. Defer cleanup: deliver result (for async), emit End event, and recover from panics
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
logger.RecoverPanicNoExit(r)
|
||||
err = fmt.Errorf("subturn panicked: %v", r)
|
||||
result = nil
|
||||
logger.ErrorCF("subturn", "SubTurn panicked", map[string]any{
|
||||
|
|
@ -510,6 +511,7 @@ func deliverSubTurnResult(al *AgentLoop, parentTS *turnState, childID string, re
|
|||
// We use defer/recover to catch any unlikely channel panics if it were ever closed.
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
logger.RecoverPanicNoExit(r)
|
||||
logger.WarnCF("subturn", "recovered panic sending to pendingResults", map[string]any{
|
||||
"parent_id": parentTS.turnID,
|
||||
"child_id": childID,
|
||||
|
|
@ -602,6 +604,7 @@ type ephemeralSessionStoreIface interface {
|
|||
SetHistory(key string, history []providers.Message)
|
||||
TruncateHistory(key string, keepLast int)
|
||||
Save(key string) error
|
||||
ListSessions() []string
|
||||
Close() error
|
||||
}
|
||||
|
||||
|
|
@ -663,6 +666,7 @@ func (e *ephemeralSessionStore) TruncateHistory(_ string, keepLast int) {
|
|||
|
||||
func (e *ephemeralSessionStore) Save(_ string) error { return nil }
|
||||
func (e *ephemeralSessionStore) Close() error { return nil }
|
||||
func (e *ephemeralSessionStore) ListSessions() []string { return nil }
|
||||
|
||||
func (e *ephemeralSessionStore) truncateLocked() {
|
||||
if len(e.history) > maxEphemeralHistorySize {
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ import (
|
|||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/bus"
|
||||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
"github.com/sipeed/picoclaw/pkg/session"
|
||||
"github.com/sipeed/picoclaw/pkg/tools"
|
||||
|
|
@ -338,6 +339,23 @@ func (ts *turnState) refreshRestorePointFromSession(agent *AgentInstance) {
|
|||
ts.captureRestorePoint(history, summary)
|
||||
}
|
||||
|
||||
// ingestMessage calls the ContextManager's Ingest method for a persisted message.
|
||||
// Errors are logged but never block the turn.
|
||||
func (ts *turnState) ingestMessage(ctx context.Context, al *AgentLoop, msg providers.Message) {
|
||||
if al.contextManager == nil {
|
||||
return
|
||||
}
|
||||
if err := al.contextManager.Ingest(ctx, &IngestRequest{
|
||||
SessionKey: ts.sessionKey,
|
||||
Message: msg,
|
||||
}); err != nil {
|
||||
logger.WarnCF("agent", "Context manager ingest failed", map[string]any{
|
||||
"session_key": ts.sessionKey,
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (ts *turnState) restoreSession(agent *AgentInstance) error {
|
||||
ts.mu.RLock()
|
||||
history := append([]providers.Message(nil), ts.restorePointHistory...)
|
||||
|
|
|
|||
|
|
@ -120,6 +120,7 @@ func streamOggOpusToDiscord(ctx context.Context, vc *discordgo.VoiceConnection,
|
|||
defer func() {
|
||||
if rec := recover(); rec != nil {
|
||||
retErr = fmt.Errorf("voice connection closed during playback")
|
||||
logger.RecoverPanicNoExit(rec)
|
||||
}
|
||||
}()
|
||||
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import (
|
|||
"fmt"
|
||||
"math"
|
||||
"net/http"
|
||||
"sort"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
|
|
@ -425,6 +426,10 @@ func (m *Manager) initChannels(channels *config.ChannelsConfig) error {
|
|||
m.initChannel("irc", "IRC")
|
||||
}
|
||||
|
||||
if channels.VK.Enabled && channels.VK.Token.String() != "" && channels.VK.GroupID != 0 {
|
||||
m.initChannel("vk", "VK")
|
||||
}
|
||||
|
||||
logger.InfoCF("channels", "Channel initialization completed", map[string]any{
|
||||
"enabled_channels": len(m.channels),
|
||||
})
|
||||
|
|
@ -513,6 +518,8 @@ func (m *Manager) StartAll(ctx context.Context) error {
|
|||
|
||||
dispatchCtx, cancel := context.WithCancel(ctx)
|
||||
m.dispatchTask = &asyncTask{cancel: cancel}
|
||||
failedStarts := make([]error, 0, len(m.channels))
|
||||
failedNames := make([]string, 0, len(m.channels))
|
||||
|
||||
for name, channel := range m.channels {
|
||||
logger.InfoCF("channels", "Starting channel", map[string]any{
|
||||
|
|
@ -523,6 +530,8 @@ func (m *Manager) StartAll(ctx context.Context) error {
|
|||
"channel": name,
|
||||
"error": err.Error(),
|
||||
})
|
||||
failedStarts = append(failedStarts, fmt.Errorf("channel %s: %w", name, err))
|
||||
failedNames = append(failedNames, name)
|
||||
continue
|
||||
}
|
||||
// Lazily create worker only after channel starts successfully
|
||||
|
|
@ -532,6 +541,36 @@ func (m *Manager) StartAll(ctx context.Context) error {
|
|||
go m.runMediaWorker(dispatchCtx, name, w)
|
||||
}
|
||||
|
||||
if len(m.channels) > 0 && len(m.workers) == 0 {
|
||||
if m.dispatchTask != nil {
|
||||
m.dispatchTask.cancel()
|
||||
m.dispatchTask = nil
|
||||
}
|
||||
|
||||
sort.Strings(failedNames)
|
||||
if len(failedStarts) == 0 {
|
||||
return fmt.Errorf("failed to start any enabled channels")
|
||||
}
|
||||
|
||||
logger.ErrorCF("channels", "All enabled channels failed to start", map[string]any{
|
||||
"failed": len(failedNames),
|
||||
"total": len(m.channels),
|
||||
"failed_channels": failedNames,
|
||||
})
|
||||
|
||||
return fmt.Errorf("failed to start any enabled channels: %w", errors.Join(failedStarts...))
|
||||
}
|
||||
|
||||
if len(failedNames) > 0 {
|
||||
sort.Strings(failedNames)
|
||||
logger.WarnCF("channels", "Some channels failed to start", map[string]any{
|
||||
"failed": len(failedNames),
|
||||
"started": len(m.workers),
|
||||
"total": len(m.channels),
|
||||
"failed_channels": failedNames,
|
||||
})
|
||||
}
|
||||
|
||||
// Start the dispatcher that reads from the bus and routes to workers
|
||||
go m.dispatchOutbound(dispatchCtx)
|
||||
go m.dispatchOutboundMedia(dispatchCtx)
|
||||
|
|
@ -553,7 +592,11 @@ func (m *Manager) StartAll(ctx context.Context) error {
|
|||
}()
|
||||
}
|
||||
|
||||
logger.InfoC("channels", "All channels started")
|
||||
logger.InfoCF("channels", "Channel startup completed", map[string]any{
|
||||
"started": len(m.workers),
|
||||
"failed": len(failedNames),
|
||||
"total": len(m.channels),
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -19,6 +19,8 @@ import (
|
|||
type mockChannel struct {
|
||||
BaseChannel
|
||||
sendFn func(ctx context.Context, msg bus.OutboundMessage) error
|
||||
startFn func(ctx context.Context) error
|
||||
stopFn func(ctx context.Context) error
|
||||
sentMessages []bus.OutboundMessage
|
||||
placeholdersSent int
|
||||
editedMessages int
|
||||
|
|
@ -33,8 +35,19 @@ func (m *mockChannel) Send(ctx context.Context, msg bus.OutboundMessage) ([]stri
|
|||
return nil, m.sendFn(ctx, msg)
|
||||
}
|
||||
|
||||
func (m *mockChannel) Start(ctx context.Context) error { return nil }
|
||||
func (m *mockChannel) Stop(ctx context.Context) error { return nil }
|
||||
func (m *mockChannel) Start(ctx context.Context) error {
|
||||
if m.startFn != nil {
|
||||
return m.startFn(ctx)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mockChannel) Stop(ctx context.Context) error {
|
||||
if m.stopFn != nil {
|
||||
return m.stopFn(ctx)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mockChannel) SendPlaceholder(ctx context.Context, chatID string) (string, error) {
|
||||
m.placeholdersSent++
|
||||
|
|
@ -86,6 +99,101 @@ func newTestManager() *Manager {
|
|||
return &Manager{
|
||||
channels: make(map[string]Channel),
|
||||
workers: make(map[string]*channelWorker),
|
||||
bus: bus.NewMessageBus(),
|
||||
}
|
||||
}
|
||||
|
||||
func TestStartAll_AllChannelsFail_ReturnsJoinedError(t *testing.T) {
|
||||
m := newTestManager()
|
||||
errA := errors.New("channel-a start failed")
|
||||
errB := errors.New("channel-b start failed")
|
||||
|
||||
m.channels["a"] = &mockChannel{
|
||||
startFn: func(_ context.Context) error { return errA },
|
||||
}
|
||||
m.channels["b"] = &mockChannel{
|
||||
startFn: func(_ context.Context) error { return errB },
|
||||
}
|
||||
|
||||
err := m.StartAll(t.Context())
|
||||
if err == nil {
|
||||
t.Fatal("expected StartAll to fail when all channels fail")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "failed to start any enabled channels") {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if !errors.Is(err, errA) {
|
||||
t.Fatalf("expected error to wrap errA, got: %v", err)
|
||||
}
|
||||
if !errors.Is(err, errB) {
|
||||
t.Fatalf("expected error to wrap errB, got: %v", err)
|
||||
}
|
||||
if len(m.workers) != 0 {
|
||||
t.Fatalf("expected no workers on full startup failure, got %d", len(m.workers))
|
||||
}
|
||||
if m.dispatchTask != nil {
|
||||
t.Fatal("expected dispatch task to be cleared on full startup failure")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStartAll_PartialFailure_StartsSuccessfulWorkers(t *testing.T) {
|
||||
m := newTestManager()
|
||||
errBad := errors.New("bad channel start failed")
|
||||
processed := make(chan struct{}, 1)
|
||||
|
||||
m.channels["good"] = &mockChannel{
|
||||
sendFn: func(_ context.Context, msg bus.OutboundMessage) error {
|
||||
if msg.Channel == "good" {
|
||||
select {
|
||||
case processed <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
m.channels["bad"] = &mockChannel{
|
||||
startFn: func(_ context.Context) error { return errBad },
|
||||
}
|
||||
|
||||
err := m.StartAll(t.Context())
|
||||
if err != nil {
|
||||
t.Fatalf("expected StartAll to succeed with partial channel failures, got: %v", err)
|
||||
}
|
||||
if len(m.workers) != 1 {
|
||||
t.Fatalf("expected exactly 1 active worker, got %d", len(m.workers))
|
||||
}
|
||||
if _, ok := m.workers["good"]; !ok {
|
||||
t.Fatal("expected worker for successful channel 'good'")
|
||||
}
|
||||
if _, ok := m.workers["bad"]; ok {
|
||||
t.Fatal("did not expect worker for failed channel 'bad'")
|
||||
}
|
||||
if m.dispatchTask == nil {
|
||||
t.Fatal("expected dispatch task to run when at least one channel starts")
|
||||
}
|
||||
|
||||
pubCtx, pubCancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer pubCancel()
|
||||
if err := m.bus.PublishOutbound(pubCtx, bus.OutboundMessage{
|
||||
Channel: "good",
|
||||
ChatID: "chat-1",
|
||||
Content: "hello",
|
||||
}); err != nil {
|
||||
t.Fatalf("PublishOutbound() error = %v", err)
|
||||
}
|
||||
|
||||
select {
|
||||
case <-processed:
|
||||
// worker processed outbound message as expected
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("expected successful channel worker to process outbound message")
|
||||
}
|
||||
|
||||
stopCtx, stopCancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer stopCancel()
|
||||
if err := m.StopAll(stopCtx); err != nil {
|
||||
t.Fatalf("StopAll() error = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -262,3 +262,57 @@ func TestSend_ClosedConnection(t *testing.T) {
|
|||
|
||||
ch.Stop(ctx)
|
||||
}
|
||||
|
||||
func TestParseInlineImageMedia_Valid(t *testing.T) {
|
||||
media, err := parseInlineImageMedia(map[string]any{
|
||||
"media": []any{
|
||||
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+X2ioAAAAASUVORK5CYII=",
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("parseInlineImageMedia() error = %v", err)
|
||||
}
|
||||
if len(media) != 1 {
|
||||
t.Fatalf("len(media) = %d, want 1", len(media))
|
||||
}
|
||||
}
|
||||
|
||||
func TestPicoChannel_HandleMessageSend_AllowsMediaOnly(t *testing.T) {
|
||||
mb := bus.NewMessageBus()
|
||||
ch, err := NewPicoChannel(config.PicoConfig{
|
||||
Token: *config.NewSecureString("test-token"),
|
||||
}, mb)
|
||||
if err != nil {
|
||||
t.Fatalf("NewPicoChannel() error = %v", err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
if err := ch.Start(ctx); err != nil {
|
||||
t.Fatalf("Start() error = %v", err)
|
||||
}
|
||||
defer ch.Stop(ctx)
|
||||
|
||||
pc := &picoConn{id: "conn-1", sessionID: "sess-1"}
|
||||
ch.handleMessageSend(pc, PicoMessage{
|
||||
ID: "msg-1",
|
||||
Payload: map[string]any{
|
||||
"media": []any{
|
||||
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+X2ioAAAAASUVORK5CYII=",
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
select {
|
||||
case msg := <-mb.InboundChan():
|
||||
if msg.Content != "" {
|
||||
t.Fatalf("msg.Content = %q, want empty", msg.Content)
|
||||
}
|
||||
if len(msg.Media) != 1 || !strings.HasPrefix(msg.Media[0], "data:image/png;base64,") {
|
||||
t.Fatalf("msg.Media = %#v, want inline image payload", msg.Media)
|
||||
}
|
||||
case <-ctx.Done():
|
||||
t.Fatal("timed out waiting for inbound media message")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ package pico
|
|||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
|
@ -30,6 +31,14 @@ type picoConn struct {
|
|||
cancel context.CancelFunc // cancels per-connection goroutines (e.g. pingLoop)
|
||||
}
|
||||
|
||||
var allowedInlineImageMIMETypes = map[string]struct{}{
|
||||
"image/jpeg": {},
|
||||
"image/png": {},
|
||||
"image/gif": {},
|
||||
"image/webp": {},
|
||||
"image/bmp": {},
|
||||
}
|
||||
|
||||
// writeJSON sends a JSON message to the connection with write locking.
|
||||
func (pc *picoConn) writeJSON(v any) error {
|
||||
if pc.closed.Load() {
|
||||
|
|
@ -516,6 +525,9 @@ func (c *PicoChannel) handleMessage(pc *picoConn, msg PicoMessage) {
|
|||
case TypeMessageSend:
|
||||
c.handleMessageSend(pc, msg)
|
||||
|
||||
case TypeMediaSend:
|
||||
c.handleMessageSend(pc, msg)
|
||||
|
||||
default:
|
||||
errMsg := newError("unknown_type", fmt.Sprintf("unknown message type: %s", msg.Type))
|
||||
pc.writeJSON(errMsg)
|
||||
|
|
@ -525,8 +537,19 @@ func (c *PicoChannel) handleMessage(pc *picoConn, msg PicoMessage) {
|
|||
// handleMessageSend processes an inbound message.send from a client.
|
||||
func (c *PicoChannel) handleMessageSend(pc *picoConn, msg PicoMessage) {
|
||||
content, _ := msg.Payload["content"].(string)
|
||||
if strings.TrimSpace(content) == "" {
|
||||
errMsg := newError("empty_content", "message content is empty")
|
||||
media, err := parseInlineImageMedia(msg.Payload)
|
||||
if err != nil {
|
||||
errMsg := newErrorWithPayload("invalid_media", err.Error(), map[string]any{
|
||||
"request_id": msg.ID,
|
||||
})
|
||||
pc.writeJSON(errMsg)
|
||||
return
|
||||
}
|
||||
|
||||
if strings.TrimSpace(content) == "" && len(media) == 0 {
|
||||
errMsg := newErrorWithPayload("empty_content", "message content is empty", map[string]any{
|
||||
"request_id": msg.ID,
|
||||
})
|
||||
pc.writeJSON(errMsg)
|
||||
return
|
||||
}
|
||||
|
|
@ -550,6 +573,7 @@ func (c *PicoChannel) handleMessageSend(pc *picoConn, msg PicoMessage) {
|
|||
logger.DebugCF("pico", "Received message", map[string]any{
|
||||
"session_id": sessionID,
|
||||
"preview": truncate(content, 50),
|
||||
"media": len(media),
|
||||
})
|
||||
|
||||
sender := bus.SenderInfo{
|
||||
|
|
@ -562,7 +586,7 @@ func (c *PicoChannel) handleMessageSend(pc *picoConn, msg PicoMessage) {
|
|||
return
|
||||
}
|
||||
|
||||
c.HandleMessage(c.ctx, peer, msg.ID, senderID, chatID, content, nil, metadata, sender)
|
||||
c.HandleMessage(c.ctx, peer, msg.ID, senderID, chatID, content, media, metadata, sender)
|
||||
}
|
||||
|
||||
// truncate truncates a string to maxLen runes.
|
||||
|
|
@ -573,3 +597,99 @@ func truncate(s string, maxLen int) string {
|
|||
}
|
||||
return string(runes[:maxLen]) + "..."
|
||||
}
|
||||
|
||||
func parseInlineImageMedia(payload map[string]any) ([]string, error) {
|
||||
if len(payload) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
raw, ok := payload["media"]
|
||||
if !ok || raw == nil {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
switch values := raw.(type) {
|
||||
case []any:
|
||||
media := make([]string, 0, len(values))
|
||||
for i, item := range values {
|
||||
value, err := inlineImageValue(item)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("media[%d]: %w", i, err)
|
||||
}
|
||||
if err := validateInlineImageDataURL(value); err != nil {
|
||||
return nil, fmt.Errorf("media[%d]: %w", i, err)
|
||||
}
|
||||
media = append(media, value)
|
||||
}
|
||||
return media, nil
|
||||
case []string:
|
||||
media := make([]string, 0, len(values))
|
||||
for i, value := range values {
|
||||
value = strings.TrimSpace(value)
|
||||
if err := validateInlineImageDataURL(value); err != nil {
|
||||
return nil, fmt.Errorf("media[%d]: %w", i, err)
|
||||
}
|
||||
media = append(media, value)
|
||||
}
|
||||
return media, nil
|
||||
case string:
|
||||
value := strings.TrimSpace(values)
|
||||
if err := validateInlineImageDataURL(value); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return []string{value}, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("media must be a string or array of strings")
|
||||
}
|
||||
}
|
||||
|
||||
func inlineImageValue(item any) (string, error) {
|
||||
switch value := item.(type) {
|
||||
case string:
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
return "", fmt.Errorf("image payload is empty")
|
||||
}
|
||||
return value, nil
|
||||
case map[string]any:
|
||||
for _, key := range []string{"url", "data_url"} {
|
||||
if raw, ok := value[key].(string); ok && strings.TrimSpace(raw) != "" {
|
||||
return strings.TrimSpace(raw), nil
|
||||
}
|
||||
}
|
||||
return "", fmt.Errorf("image payload must include url or data_url")
|
||||
default:
|
||||
return "", fmt.Errorf("image payload must be a string or object")
|
||||
}
|
||||
}
|
||||
|
||||
func validateInlineImageDataURL(mediaURL string) error {
|
||||
if mediaURL == "" {
|
||||
return fmt.Errorf("image payload is empty")
|
||||
}
|
||||
if !strings.HasPrefix(mediaURL, "data:image/") {
|
||||
return fmt.Errorf("only inline image data URLs are supported")
|
||||
}
|
||||
|
||||
header, data, found := strings.Cut(mediaURL, ",")
|
||||
if !found || strings.TrimSpace(data) == "" {
|
||||
return fmt.Errorf("image data URL is malformed")
|
||||
}
|
||||
if !strings.Contains(header, ";base64") {
|
||||
return fmt.Errorf("image data URL must be base64 encoded")
|
||||
}
|
||||
mimeType, _, _ := strings.Cut(strings.TrimPrefix(header, "data:"), ";")
|
||||
if _, ok := allowedInlineImageMIMETypes[mimeType]; !ok {
|
||||
return fmt.Errorf("unsupported image format: %s", mimeType)
|
||||
}
|
||||
|
||||
data = strings.TrimSpace(data)
|
||||
if base64.StdEncoding.DecodedLen(len(data)) > config.DefaultMaxMediaSize {
|
||||
return fmt.Errorf("image exceeds %d byte limit", config.DefaultMaxMediaSize)
|
||||
}
|
||||
if _, err := base64.StdEncoding.DecodeString(data); err != nil {
|
||||
return fmt.Errorf("invalid base64 image data")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -39,10 +39,18 @@ func newMessage(msgType string, payload map[string]any) PicoMessage {
|
|||
}
|
||||
}
|
||||
|
||||
// newError creates an error PicoMessage.
|
||||
func newError(code, message string) PicoMessage {
|
||||
return newMessage(TypeError, map[string]any{
|
||||
func newErrorWithPayload(code, message string, extra map[string]any) PicoMessage {
|
||||
payload := map[string]any{
|
||||
"code": code,
|
||||
"message": message,
|
||||
})
|
||||
}
|
||||
for key, value := range extra {
|
||||
payload[key] = value
|
||||
}
|
||||
return newMessage(TypeError, payload)
|
||||
}
|
||||
|
||||
// newError creates an error PicoMessage.
|
||||
func newError(code, message string) PicoMessage {
|
||||
return newErrorWithPayload(code, message, nil)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import (
|
|||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
|
|
@ -377,9 +378,39 @@ func (c *TelegramChannel) EditMessage(ctx context.Context, chatID string, messag
|
|||
}
|
||||
_, err = c.bot.EditMessageText(ctx, editMsg)
|
||||
if err != nil {
|
||||
// If it failed because it was already modified (likely from a previous
|
||||
// attempt that timed out on our end but landed on Telegram), we treat
|
||||
// it as success to prevent the Manager from sending a duplicate message.
|
||||
if strings.Contains(err.Error(), "message is not modified") {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Only fallback to plain text if the error looks like a parsing failure (Bad Request).
|
||||
// Network errors or timeouts should NOT trigger a retry with different content.
|
||||
if strings.Contains(err.Error(), "Bad Request") {
|
||||
logParseFailed(err, useMarkdownV2)
|
||||
_, err = c.bot.EditMessageText(ctx, tu.EditMessageText(tu.ID(cid), mid, content))
|
||||
}
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
if strings.Contains(err.Error(), "message is not modified") {
|
||||
return nil
|
||||
}
|
||||
|
||||
if isPostConnectError(err) {
|
||||
logger.WarnCF(
|
||||
"telegram",
|
||||
"EditMessage likely landed but result is unknown; swallowing error to prevent duplicate",
|
||||
map[string]any{
|
||||
"chat_id": chatID,
|
||||
"mid": mid,
|
||||
"error": err.Error(),
|
||||
},
|
||||
)
|
||||
return nil // Swallow to prevent Manager fallback to a new SendMessage
|
||||
}
|
||||
}
|
||||
|
||||
return err
|
||||
}
|
||||
|
|
@ -1134,6 +1165,30 @@ func cryptoRandInt() int {
|
|||
return int(binary.BigEndian.Uint32(b[:])) | 1 // ensure non-zero
|
||||
}
|
||||
|
||||
// isPostConnectError identifies network errors that likely occurred after
|
||||
// the request was transmitted to Telegram (e.g. dropped connection while
|
||||
// waiting for response). Swallowing these for edits prevents duplicate
|
||||
// fallbacks, at the small risk of leaving a stale placeholder if the
|
||||
// edit never actually reached the server.
|
||||
func isPostConnectError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
|
||||
// Context errors (timeout/canceled) are too broad; they can be triggered
|
||||
// locally before any data is sent. Never swallow them.
|
||||
if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
|
||||
return false
|
||||
}
|
||||
|
||||
msg := strings.ToLower(err.Error())
|
||||
// Narrowly target connection dropouts where the request likely landed.
|
||||
return strings.Contains(msg, "connection reset by peer") ||
|
||||
strings.Contains(msg, "unexpected eof") ||
|
||||
strings.Contains(msg, "connection closed by foreign host") ||
|
||||
strings.Contains(msg, "broken pipe")
|
||||
}
|
||||
|
||||
// VoiceCapabilities returns the voice capabilities of the channel.
|
||||
func (c *TelegramChannel) VoiceCapabilities() channels.VoiceCapabilities {
|
||||
return channels.VoiceCapabilities{ASR: true, TTS: true}
|
||||
|
|
|
|||
13
pkg/channels/vk/init.go
Normal file
13
pkg/channels/vk/init.go
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
package vk
|
||||
|
||||
import (
|
||||
"github.com/sipeed/picoclaw/pkg/bus"
|
||||
"github.com/sipeed/picoclaw/pkg/channels"
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
)
|
||||
|
||||
func init() {
|
||||
channels.RegisterFactory("vk", func(cfg *config.Config, b *bus.MessageBus) (channels.Channel, error) {
|
||||
return NewVKChannel(cfg, b)
|
||||
})
|
||||
}
|
||||
286
pkg/channels/vk/vk.go
Normal file
286
pkg/channels/vk/vk.go
Normal file
|
|
@ -0,0 +1,286 @@
|
|||
package vk
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/SevereCloud/vksdk/v3/api"
|
||||
"github.com/SevereCloud/vksdk/v3/api/params"
|
||||
"github.com/SevereCloud/vksdk/v3/events"
|
||||
"github.com/SevereCloud/vksdk/v3/longpoll-bot"
|
||||
"github.com/SevereCloud/vksdk/v3/object"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/bus"
|
||||
"github.com/sipeed/picoclaw/pkg/channels"
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
"github.com/sipeed/picoclaw/pkg/identity"
|
||||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
)
|
||||
|
||||
type VKChannel struct {
|
||||
*channels.BaseChannel
|
||||
vk *api.VK
|
||||
lp *longpoll.LongPoll
|
||||
config *config.Config
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
}
|
||||
|
||||
func NewVKChannel(cfg *config.Config, bus *bus.MessageBus) (*VKChannel, error) {
|
||||
vkCfg := cfg.Channels.VK
|
||||
|
||||
vk := api.NewVK(vkCfg.Token.String())
|
||||
|
||||
base := channels.NewBaseChannel(
|
||||
"vk",
|
||||
vkCfg,
|
||||
bus,
|
||||
vkCfg.AllowFrom,
|
||||
channels.WithMaxMessageLength(4000),
|
||||
channels.WithGroupTrigger(vkCfg.GroupTrigger),
|
||||
channels.WithReasoningChannelID(vkCfg.ReasoningChannelID),
|
||||
)
|
||||
|
||||
return &VKChannel{
|
||||
BaseChannel: base,
|
||||
vk: vk,
|
||||
config: cfg,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c *VKChannel) Start(ctx context.Context) error {
|
||||
logger.InfoC("vk", "Starting VK bot (Long Poll mode)...")
|
||||
|
||||
c.ctx, c.cancel = context.WithCancel(ctx)
|
||||
|
||||
groupID := c.config.Channels.VK.GroupID
|
||||
if groupID == 0 {
|
||||
c.cancel()
|
||||
return fmt.Errorf("group_id is required for VK bot")
|
||||
}
|
||||
|
||||
lp, err := longpoll.NewLongPoll(c.vk, groupID)
|
||||
if err != nil {
|
||||
c.cancel()
|
||||
return fmt.Errorf("failed to create long poll: %w", err)
|
||||
}
|
||||
c.lp = lp
|
||||
|
||||
lp.MessageNew(func(_ context.Context, obj events.MessageNewObject) {
|
||||
c.handleMessage(obj.Message)
|
||||
})
|
||||
|
||||
c.SetRunning(true)
|
||||
|
||||
logger.InfoCF("vk", "VK bot connected", map[string]any{
|
||||
"group_id": groupID,
|
||||
})
|
||||
|
||||
go func() {
|
||||
if err := lp.Run(); err != nil {
|
||||
logger.ErrorCF("vk", "Long poll failed", map[string]any{
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
}()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *VKChannel) Stop(ctx context.Context) error {
|
||||
logger.InfoC("vk", "Stopping VK bot...")
|
||||
c.SetRunning(false)
|
||||
|
||||
if c.lp != nil {
|
||||
c.lp.Shutdown()
|
||||
}
|
||||
|
||||
if c.cancel != nil {
|
||||
c.cancel()
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *VKChannel) handleMessage(msg object.MessagesMessage) {
|
||||
if msg.Action.Type != "" {
|
||||
return
|
||||
}
|
||||
|
||||
if bool(msg.Out) {
|
||||
return
|
||||
}
|
||||
|
||||
peerID := msg.PeerID
|
||||
chatID := strconv.Itoa(peerID)
|
||||
|
||||
fromID := msg.FromID
|
||||
userID := strconv.Itoa(fromID)
|
||||
|
||||
platformID := userID
|
||||
sender := bus.SenderInfo{
|
||||
Platform: "vk",
|
||||
PlatformID: platformID,
|
||||
CanonicalID: identity.BuildCanonicalID("vk", platformID),
|
||||
DisplayName: c.getUserName(fromID),
|
||||
}
|
||||
|
||||
if !c.IsAllowedSender(sender) {
|
||||
logger.DebugCF("vk", "Message from unauthorized user", map[string]any{
|
||||
"peer_id": peerID,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
text := msg.Text
|
||||
if text == "" && len(msg.Attachments) > 0 {
|
||||
text = c.processAttachments(msg.Attachments)
|
||||
}
|
||||
|
||||
if text == "" {
|
||||
return
|
||||
}
|
||||
|
||||
groupTrigger := c.config.Channels.VK.GroupTrigger
|
||||
isGroupChat := peerID != fromID
|
||||
|
||||
if isGroupChat {
|
||||
isMentioned := c.isMentioned(msg)
|
||||
if isMentioned {
|
||||
text = c.stripBotMention(text)
|
||||
}
|
||||
respond, cleaned := c.ShouldRespondInGroup(isMentioned, text)
|
||||
if !respond {
|
||||
return
|
||||
}
|
||||
text = cleaned
|
||||
_ = groupTrigger
|
||||
}
|
||||
|
||||
peerKind := "direct"
|
||||
peerIDStr := userID
|
||||
if isGroupChat {
|
||||
peerKind = "group"
|
||||
peerIDStr = chatID
|
||||
}
|
||||
|
||||
peer := bus.Peer{Kind: peerKind, ID: peerIDStr}
|
||||
messageID := strconv.Itoa(msg.ConversationMessageID)
|
||||
|
||||
metadata := map[string]string{
|
||||
"user_id": userID,
|
||||
"is_group": fmt.Sprintf("%t", isGroupChat),
|
||||
}
|
||||
|
||||
c.HandleMessage(c.ctx,
|
||||
peer,
|
||||
messageID,
|
||||
userID,
|
||||
chatID,
|
||||
text,
|
||||
nil,
|
||||
metadata,
|
||||
sender,
|
||||
)
|
||||
}
|
||||
|
||||
func (c *VKChannel) Send(ctx context.Context, msg bus.OutboundMessage) ([]string, error) {
|
||||
if !c.IsRunning() {
|
||||
return nil, channels.ErrNotRunning
|
||||
}
|
||||
|
||||
peerID, err := strconv.Atoi(msg.ChatID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid chat ID %s: %w", msg.ChatID, channels.ErrSendFailed)
|
||||
}
|
||||
|
||||
if msg.Content == "" {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var messageIDs []string
|
||||
chunks := channels.SplitMessage(msg.Content, 4000)
|
||||
|
||||
for _, chunk := range chunks {
|
||||
if chunk == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
b := params.NewMessagesSendBuilder()
|
||||
b.Message(chunk)
|
||||
b.RandomID(0)
|
||||
b.PeerID(peerID)
|
||||
|
||||
if msg.ReplyToMessageID != "" {
|
||||
if replyID, err := strconv.Atoi(msg.ReplyToMessageID); err == nil {
|
||||
b.ReplyTo(replyID)
|
||||
}
|
||||
}
|
||||
|
||||
resp, err := c.vk.MessagesSend(b.Params)
|
||||
if err != nil {
|
||||
logger.ErrorCF("vk", "Failed to send message", map[string]any{
|
||||
"error": err.Error(),
|
||||
"peer_id": peerID,
|
||||
})
|
||||
return messageIDs, fmt.Errorf("failed to send message: %w", err)
|
||||
}
|
||||
|
||||
messageIDs = append(messageIDs, strconv.Itoa(resp))
|
||||
}
|
||||
|
||||
return messageIDs, nil
|
||||
}
|
||||
|
||||
func (c *VKChannel) isMentioned(msg object.MessagesMessage) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
func (c *VKChannel) stripBotMention(text string) string {
|
||||
return strings.TrimSpace(text)
|
||||
}
|
||||
|
||||
func (c *VKChannel) getUserName(userID int) string {
|
||||
users, err := c.vk.UsersGet(api.Params{
|
||||
"user_ids": userID,
|
||||
})
|
||||
if err != nil || len(users) == 0 {
|
||||
return strconv.Itoa(userID)
|
||||
}
|
||||
|
||||
user := users[0]
|
||||
return fmt.Sprintf("%s %s", user.FirstName, user.LastName)
|
||||
}
|
||||
|
||||
func (c *VKChannel) processAttachments(attachments []object.MessagesMessageAttachment) string {
|
||||
var parts []string
|
||||
|
||||
for _, att := range attachments {
|
||||
switch att.Type {
|
||||
case "photo":
|
||||
parts = append(parts, "[photo]")
|
||||
case "video":
|
||||
parts = append(parts, "[video]")
|
||||
case "audio":
|
||||
parts = append(parts, "[audio]")
|
||||
case "doc":
|
||||
if att.Doc.Title != "" {
|
||||
parts = append(parts, fmt.Sprintf("[document: %s]", att.Doc.Title))
|
||||
} else {
|
||||
parts = append(parts, "[document]")
|
||||
}
|
||||
case "audio_message":
|
||||
parts = append(parts, "[voice]")
|
||||
case "sticker":
|
||||
parts = append(parts, "[sticker]")
|
||||
}
|
||||
}
|
||||
|
||||
return strings.Join(parts, " ")
|
||||
}
|
||||
|
||||
func (c *VKChannel) VoiceCapabilities() channels.VoiceCapabilities {
|
||||
return channels.VoiceCapabilities{ASR: true, TTS: true}
|
||||
}
|
||||
260
pkg/channels/vk/vk_test.go
Normal file
260
pkg/channels/vk/vk_test.go
Normal file
|
|
@ -0,0 +1,260 @@
|
|||
package vk
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/bus"
|
||||
"github.com/sipeed/picoclaw/pkg/channels"
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
)
|
||||
|
||||
func TestNewVKChannel(t *testing.T) {
|
||||
msgBus := bus.NewMessageBus()
|
||||
|
||||
t.Run("missing group_id", func(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
Channels: config.ChannelsConfig{
|
||||
VK: config.VKConfig{
|
||||
Enabled: true,
|
||||
Token: *config.NewSecureString("test_token"),
|
||||
},
|
||||
},
|
||||
}
|
||||
ch, err := NewVKChannel(cfg, msgBus)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error during creation: %v", err)
|
||||
}
|
||||
if ch.Name() != "vk" {
|
||||
t.Errorf("Name() = %q, want %q", ch.Name(), "vk")
|
||||
}
|
||||
if ch.IsRunning() {
|
||||
t.Error("new channel should not be running")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("valid config with group_id", func(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
Channels: config.ChannelsConfig{
|
||||
VK: config.VKConfig{
|
||||
Enabled: true,
|
||||
Token: *config.NewSecureString("test_token"),
|
||||
GroupID: 123456789,
|
||||
},
|
||||
},
|
||||
}
|
||||
ch, err := NewVKChannel(cfg, msgBus)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if ch.Name() != "vk" {
|
||||
t.Errorf("Name() = %q, want %q", ch.Name(), "vk")
|
||||
}
|
||||
if ch.IsRunning() {
|
||||
t.Error("new channel should not be running")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("with allow_from", func(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
Channels: config.ChannelsConfig{
|
||||
VK: config.VKConfig{
|
||||
Enabled: true,
|
||||
Token: *config.NewSecureString("test_token"),
|
||||
GroupID: 123456789,
|
||||
AllowFrom: []string{"123456789"},
|
||||
},
|
||||
},
|
||||
}
|
||||
ch, err := NewVKChannel(cfg, msgBus)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if !ch.IsAllowedSender(bus.SenderInfo{PlatformID: "123456789"}) {
|
||||
t.Error("user 123456789 should be allowed")
|
||||
}
|
||||
if ch.IsAllowedSender(bus.SenderInfo{PlatformID: "999999999"}) {
|
||||
t.Error("user 999999999 should not be allowed")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("with group_trigger", func(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
Channels: config.ChannelsConfig{
|
||||
VK: config.VKConfig{
|
||||
Enabled: true,
|
||||
Token: *config.NewSecureString("test_token"),
|
||||
GroupID: 123456789,
|
||||
GroupTrigger: config.GroupTriggerConfig{
|
||||
MentionOnly: false,
|
||||
Prefixes: []string{"/bot", "!bot"},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
ch, err := NewVKChannel(cfg, msgBus)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if ch.Name() != "vk" {
|
||||
t.Errorf("Name() = %q, want %q", ch.Name(), "vk")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestVKChannel_MaxMessageLength(t *testing.T) {
|
||||
msgBus := bus.NewMessageBus()
|
||||
cfg := &config.Config{
|
||||
Channels: config.ChannelsConfig{
|
||||
VK: config.VKConfig{
|
||||
Enabled: true,
|
||||
Token: *config.NewSecureString("test_token"),
|
||||
GroupID: 123456789,
|
||||
},
|
||||
},
|
||||
}
|
||||
ch, err := NewVKChannel(cfg, msgBus)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
maxLen := ch.MaxMessageLength()
|
||||
if maxLen != 4000 {
|
||||
t.Errorf("MaxMessageLength() = %d, want 4000", maxLen)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVKChannel_SplitMessage(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
content string
|
||||
maxLen int
|
||||
want int
|
||||
}{
|
||||
{
|
||||
name: "short message",
|
||||
content: "hello",
|
||||
maxLen: 4000,
|
||||
want: 1,
|
||||
},
|
||||
{
|
||||
name: "exact length",
|
||||
content: string(make([]byte, 4000)),
|
||||
maxLen: 4000,
|
||||
want: 1,
|
||||
},
|
||||
{
|
||||
name: "needs split",
|
||||
content: string(make([]byte, 5000)),
|
||||
maxLen: 4000,
|
||||
want: 2,
|
||||
},
|
||||
{
|
||||
name: "empty message",
|
||||
content: "",
|
||||
maxLen: 4000,
|
||||
want: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := channels.SplitMessage(tt.content, tt.maxLen)
|
||||
if len(got) != tt.want {
|
||||
t.Errorf("SplitMessage() got %d parts, want %d parts", len(got), tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestVKChannel_ProcessAttachments(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
attachments []string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "empty attachments",
|
||||
attachments: []string{},
|
||||
want: "",
|
||||
},
|
||||
{
|
||||
name: "photo attachment",
|
||||
attachments: []string{"photo"},
|
||||
want: "[photo]",
|
||||
},
|
||||
{
|
||||
name: "video attachment",
|
||||
attachments: []string{"video"},
|
||||
want: "[video]",
|
||||
},
|
||||
{
|
||||
name: "audio attachment",
|
||||
attachments: []string{"audio"},
|
||||
want: "[audio]",
|
||||
},
|
||||
{
|
||||
name: "document attachment",
|
||||
attachments: []string{"doc"},
|
||||
want: "[doc]",
|
||||
},
|
||||
{
|
||||
name: "sticker attachment",
|
||||
attachments: []string{"sticker"},
|
||||
want: "[sticker]",
|
||||
},
|
||||
{
|
||||
name: "audio_message attachment",
|
||||
attachments: []string{"audio_message"},
|
||||
want: "[voice]",
|
||||
},
|
||||
{
|
||||
name: "multiple attachments",
|
||||
attachments: []string{"photo", "video", "audio"},
|
||||
want: "[photo] [video] [audio]",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
var result string
|
||||
for i, att := range tt.attachments {
|
||||
if i > 0 {
|
||||
result += " "
|
||||
}
|
||||
if att == "audio_message" {
|
||||
result += "[voice]"
|
||||
} else {
|
||||
result += "[" + att + "]"
|
||||
}
|
||||
}
|
||||
if result != tt.want {
|
||||
t.Errorf("processAttachments() = %q, want %q", result, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestVKChannel_VoiceCapabilities(t *testing.T) {
|
||||
msgBus := bus.NewMessageBus()
|
||||
cfg := &config.Config{
|
||||
Channels: config.ChannelsConfig{
|
||||
VK: config.VKConfig{
|
||||
Enabled: true,
|
||||
Token: *config.NewSecureString("test_token"),
|
||||
GroupID: 123456789,
|
||||
},
|
||||
},
|
||||
}
|
||||
ch, err := NewVKChannel(cfg, msgBus)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
caps := ch.VoiceCapabilities()
|
||||
if !caps.ASR {
|
||||
t.Error("VoiceCapabilities().ASR should be true")
|
||||
}
|
||||
if !caps.TTS {
|
||||
t.Error("VoiceCapabilities().TTS should be true")
|
||||
}
|
||||
}
|
||||
|
|
@ -7,6 +7,7 @@ import (
|
|||
"math/rand"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
|
|
@ -246,6 +247,8 @@ type AgentDefaults struct {
|
|||
SubTurn SubTurnConfig `json:"subturn" envPrefix:"PICOCLAW_AGENTS_DEFAULTS_SUBTURN_"`
|
||||
ToolFeedback ToolFeedbackConfig `json:"tool_feedback,omitempty"`
|
||||
SplitOnMarker bool `json:"split_on_marker" env:"PICOCLAW_AGENTS_DEFAULTS_SPLIT_ON_MARKER"` // split messages on <|[SPLIT]|> marker
|
||||
ContextManager string `json:"context_manager,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_CONTEXT_MANAGER"`
|
||||
ContextManagerConfig json.RawMessage `json:"context_manager_config,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_CONTEXT_MANAGER_CONFIG"`
|
||||
}
|
||||
|
||||
const DefaultMaxMediaSize = 20 * 1024 * 1024 // 20 MB
|
||||
|
|
@ -293,6 +296,7 @@ type ChannelsConfig struct {
|
|||
Pico PicoConfig `json:"pico" yaml:"pico,omitempty"`
|
||||
PicoClient PicoClientConfig `json:"pico_client" yaml:"pico_client,omitempty"`
|
||||
IRC IRCConfig `json:"irc" yaml:"irc,omitempty"`
|
||||
VK VKConfig `json:"vk" yaml:"vk,omitempty"`
|
||||
}
|
||||
|
||||
// GroupTriggerConfig controls when the bot responds in group chats.
|
||||
|
|
@ -547,6 +551,21 @@ type IRCConfig struct {
|
|||
ReasoningChannelID string `json:"reasoning_channel_id" yaml:"-"`
|
||||
}
|
||||
|
||||
type VKConfig struct {
|
||||
Enabled bool `json:"enabled" yaml:"-" env:"PICOCLAW_CHANNELS_VK_ENABLED"`
|
||||
Token SecureString `json:"token,omitzero" yaml:"token,omitempty" env:"PICOCLAW_CHANNELS_VK_TOKEN"`
|
||||
GroupID int `json:"group_id" yaml:"-" env:"PICOCLAW_CHANNELS_VK_GROUP_ID"`
|
||||
AllowFrom FlexibleStringSlice `json:"allow_from" yaml:"-" env:"PICOCLAW_CHANNELS_VK_ALLOW_FROM"`
|
||||
GroupTrigger GroupTriggerConfig `json:"group_trigger,omitempty" yaml:"-"`
|
||||
Typing TypingConfig `json:"typing,omitempty" yaml:"-"`
|
||||
Placeholder PlaceholderConfig `json:"placeholder,omitempty" yaml:"-"`
|
||||
ReasoningChannelID string `json:"reasoning_channel_id" yaml:"-" env:"PICOCLAW_CHANNELS_VK_REASONING_CHANNEL_ID"`
|
||||
}
|
||||
|
||||
func (c *VKConfig) SetToken(token string) {
|
||||
c.Token = *NewSecureString(token)
|
||||
}
|
||||
|
||||
type HeartbeatConfig struct {
|
||||
Enabled bool `json:"enabled" env:"PICOCLAW_HEARTBEAT_ENABLED"`
|
||||
Interval int `json:"interval" env:"PICOCLAW_HEARTBEAT_INTERVAL"` // minutes, min 5
|
||||
|
|
@ -598,6 +617,8 @@ type ModelConfig struct {
|
|||
// existing configs, the field is inferred during load: models with API keys
|
||||
// or the reserved "local-model" name are auto-enabled.
|
||||
Enabled bool `json:"enabled,omitempty" yaml:"enabled,omitempty"`
|
||||
// UserAgent is the user agent string to use for HTTP requests.
|
||||
UserAgent string `json:"user_agent,omitempty" yaml:"-"`
|
||||
|
||||
// isVirtual marks this model as a virtual model generated from multi-key expansion.
|
||||
// Virtual models should not be persisted to config files.
|
||||
|
|
@ -760,13 +781,13 @@ type WebToolsConfig struct {
|
|||
// the client-side web_search tool is hidden to avoid duplicate search surfaces,
|
||||
// and the provider's built-in search is used instead. Falls back to client-side
|
||||
// search when the provider does not support native search.
|
||||
PreferNative bool `json:"prefer_native" yaml:"-" env:"PICOCLAW_TOOLS_WEB_PREFER_NATIVE"`
|
||||
PreferNative bool `yaml:"-" json:"prefer_native" env:"PICOCLAW_TOOLS_WEB_PREFER_NATIVE"`
|
||||
// Proxy is an optional proxy URL for web tools (http/https/socks5/socks5h).
|
||||
// For authenticated proxies, prefer HTTP_PROXY/HTTPS_PROXY env vars instead of embedding credentials in config.
|
||||
Proxy string `json:"proxy,omitempty" yaml:"-" env:"PICOCLAW_TOOLS_WEB_PROXY"`
|
||||
FetchLimitBytes int64 `json:"fetch_limit_bytes,omitempty" yaml:"-" env:"PICOCLAW_TOOLS_WEB_FETCH_LIMIT_BYTES"`
|
||||
Format string `json:"format,omitempty" yaml:"-" env:"PICOCLAW_TOOLS_WEB_FORMAT"`
|
||||
PrivateHostWhitelist FlexibleStringSlice `json:"private_host_whitelist,omitempty" yaml:"-" env:"PICOCLAW_TOOLS_WEB_PRIVATE_HOST_WHITELIST"`
|
||||
Proxy string `yaml:"-" json:"proxy,omitempty" env:"PICOCLAW_TOOLS_WEB_PROXY"`
|
||||
FetchLimitBytes int64 `yaml:"-" json:"fetch_limit_bytes,omitempty" env:"PICOCLAW_TOOLS_WEB_FETCH_LIMIT_BYTES"`
|
||||
Format string `yaml:"-" json:"format,omitempty" env:"PICOCLAW_TOOLS_WEB_FORMAT"`
|
||||
PrivateHostWhitelist FlexibleStringSlice `yaml:"-" json:"private_host_whitelist,omitempty" env:"PICOCLAW_TOOLS_WEB_PRIVATE_HOST_WHITELIST"`
|
||||
}
|
||||
|
||||
type CronToolsConfig struct {
|
||||
|
|
@ -800,9 +821,26 @@ type MediaCleanupConfig struct {
|
|||
|
||||
type ReadFileToolConfig struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
Mode string `json:"mode"`
|
||||
MaxReadFileSize int `json:"max_read_file_size"`
|
||||
}
|
||||
|
||||
const (
|
||||
ReadFileModeBytes = "bytes"
|
||||
ReadFileModeLines = "lines"
|
||||
)
|
||||
|
||||
func (c ReadFileToolConfig) EffectiveMode() string {
|
||||
switch strings.ToLower(strings.TrimSpace(c.Mode)) {
|
||||
case ReadFileModeLines:
|
||||
return ReadFileModeLines
|
||||
case "", ReadFileModeBytes:
|
||||
return ReadFileModeBytes
|
||||
default:
|
||||
return ReadFileModeBytes
|
||||
}
|
||||
}
|
||||
|
||||
type ToolsConfig struct {
|
||||
AllowReadPaths []string `json:"allow_read_paths" yaml:"-" env:"PICOCLAW_TOOLS_ALLOW_READ_PATHS"`
|
||||
AllowWritePaths []string `json:"allow_write_paths" yaml:"-" env:"PICOCLAW_TOOLS_ALLOW_WRITE_PATHS"`
|
||||
|
|
@ -905,10 +943,21 @@ type MCPServerConfig struct {
|
|||
type MCPConfig struct {
|
||||
ToolConfig ` envPrefix:"PICOCLAW_TOOLS_MCP_"`
|
||||
Discovery ToolDiscoveryConfig ` json:"discovery"`
|
||||
// MaxInlineTextChars controls how much MCP text stays inline before it is saved as an artifact.
|
||||
MaxInlineTextChars int `json:"max_inline_text_chars,omitempty" env:"PICOCLAW_TOOLS_MCP_MAX_INLINE_TEXT_CHARS"`
|
||||
// Servers is a map of server name to server configuration
|
||||
Servers map[string]MCPServerConfig `json:"servers,omitempty"`
|
||||
}
|
||||
|
||||
const DefaultMCPMaxInlineTextChars = 16 * 1024
|
||||
|
||||
func (c *MCPConfig) GetMaxInlineTextChars() int {
|
||||
if c.MaxInlineTextChars > 0 {
|
||||
return c.MaxInlineTextChars
|
||||
}
|
||||
return DefaultMCPMaxInlineTextChars
|
||||
}
|
||||
|
||||
func LoadConfig(path string) (*Config, error) {
|
||||
logger.Debugf("loading config from %s", path)
|
||||
|
||||
|
|
@ -917,7 +966,10 @@ func LoadConfig(path string) (*Config, error) {
|
|||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
logger.WarnF("config file not found, using default config", map[string]any{"path": path})
|
||||
logger.WarnF(
|
||||
"config file not found, using default config",
|
||||
map[string]any{"path": path},
|
||||
)
|
||||
return DefaultConfig(), nil
|
||||
}
|
||||
logger.Errorf("failed to read config file: %v", err)
|
||||
|
|
@ -940,7 +992,10 @@ func LoadConfig(path string) (*Config, error) {
|
|||
var cfg *Config
|
||||
switch versionInfo.Version {
|
||||
case 0:
|
||||
logger.InfoF("config migrate start", map[string]any{"from": versionInfo.Version, "to": CurrentVersion})
|
||||
logger.InfoF(
|
||||
"config migrate start",
|
||||
map[string]any{"from": versionInfo.Version, "to": CurrentVersion},
|
||||
)
|
||||
// Legacy config (no version field)
|
||||
v, e := loadConfigV0(data)
|
||||
if e != nil {
|
||||
|
|
@ -948,10 +1003,16 @@ func LoadConfig(path string) (*Config, error) {
|
|||
}
|
||||
cfg, e = v.Migrate()
|
||||
if e != nil {
|
||||
logger.ErrorF("config migrate fail", map[string]any{"from": versionInfo.Version, "to": CurrentVersion})
|
||||
logger.ErrorF(
|
||||
"config migrate fail",
|
||||
map[string]any{"from": versionInfo.Version, "to": CurrentVersion},
|
||||
)
|
||||
return nil, e
|
||||
}
|
||||
logger.InfoF("config migrate success", map[string]any{"from": versionInfo.Version, "to": CurrentVersion})
|
||||
logger.InfoF(
|
||||
"config migrate success",
|
||||
map[string]any{"from": versionInfo.Version, "to": CurrentVersion},
|
||||
)
|
||||
err = makeBackup(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
|
@ -959,7 +1020,10 @@ func LoadConfig(path string) (*Config, error) {
|
|||
// Load existing security config and merge with migrated one to prevent data loss
|
||||
secErr := loadSecurityConfig(cfg, securityPath(path))
|
||||
if secErr != nil && !os.IsNotExist(secErr) {
|
||||
logger.WarnF("failed to load existing security config during migration", map[string]any{"error": secErr})
|
||||
logger.WarnF(
|
||||
"failed to load existing security config during migration",
|
||||
map[string]any{"error": secErr},
|
||||
)
|
||||
return nil, fmt.Errorf("failed to load existing security config: %w", secErr)
|
||||
}
|
||||
defer func(cfg *Config) {
|
||||
|
|
@ -967,7 +1031,10 @@ func LoadConfig(path string) (*Config, error) {
|
|||
}(cfg)
|
||||
case 1:
|
||||
// V1→V2 migration: infer Enabled and migrate channel config fields
|
||||
logger.InfoF("config migrate start", map[string]any{"from": versionInfo.Version, "to": CurrentVersion})
|
||||
logger.InfoF(
|
||||
"config migrate start",
|
||||
map[string]any{"from": versionInfo.Version, "to": CurrentVersion},
|
||||
)
|
||||
cfg, err = loadConfig(data)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
|
@ -981,7 +1048,10 @@ func LoadConfig(path string) (*Config, error) {
|
|||
oldCfg := &configV1{Config: *cfg}
|
||||
cfg, err = oldCfg.Migrate()
|
||||
if err != nil {
|
||||
logger.ErrorF("config migrate fail", map[string]any{"from": versionInfo.Version, "to": CurrentVersion})
|
||||
logger.ErrorF(
|
||||
"config migrate fail",
|
||||
map[string]any{"from": versionInfo.Version, "to": CurrentVersion},
|
||||
)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
|
|
@ -993,7 +1063,10 @@ func LoadConfig(path string) (*Config, error) {
|
|||
defer func(cfg *Config) {
|
||||
_ = SaveConfig(path, cfg)
|
||||
}(cfg)
|
||||
logger.InfoF("config migrate success", map[string]any{"from": versionInfo.Version, "to": CurrentVersion})
|
||||
logger.InfoF(
|
||||
"config migrate success",
|
||||
map[string]any{"from": versionInfo.Version, "to": CurrentVersion},
|
||||
)
|
||||
case CurrentVersion:
|
||||
// Current version
|
||||
cfg, err = loadConfig(data)
|
||||
|
|
|
|||
|
|
@ -198,6 +198,41 @@ func TestAgentConfig_FullParse(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestDefaultConfig_MCPMaxInlineTextChars(t *testing.T) {
|
||||
cfg := DefaultConfig()
|
||||
if cfg.Tools.MCP.GetMaxInlineTextChars() != DefaultMCPMaxInlineTextChars {
|
||||
t.Fatalf(
|
||||
"DefaultConfig().Tools.MCP.GetMaxInlineTextChars() = %d, want %d",
|
||||
cfg.Tools.MCP.GetMaxInlineTextChars(),
|
||||
DefaultMCPMaxInlineTextChars,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadConfig_MCPMaxInlineTextChars(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
configPath := filepath.Join(dir, "config.json")
|
||||
raw := `{
|
||||
"tools": {
|
||||
"mcp": {
|
||||
"enabled": true,
|
||||
"max_inline_text_chars": 2048
|
||||
}
|
||||
}
|
||||
}`
|
||||
if err := os.WriteFile(configPath, []byte(raw), 0o644); err != nil {
|
||||
t.Fatalf("WriteFile(configPath): %v", err)
|
||||
}
|
||||
|
||||
cfg, err := LoadConfig(configPath)
|
||||
if err != nil {
|
||||
t.Fatalf("LoadConfig() error: %v", err)
|
||||
}
|
||||
if got := cfg.Tools.MCP.GetMaxInlineTextChars(); got != 2048 {
|
||||
t.Fatalf("cfg.Tools.MCP.GetMaxInlineTextChars() = %d, want 2048", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConfig_BackwardCompat_NoAgentsList(t *testing.T) {
|
||||
jsonData := `{
|
||||
"agents": {
|
||||
|
|
@ -317,6 +352,13 @@ func TestDefaultConfig_WebTools(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestDefaultConfig_ReadFileMode(t *testing.T) {
|
||||
cfg := DefaultConfig()
|
||||
if cfg.Tools.ReadFile.EffectiveMode() != ReadFileModeBytes {
|
||||
t.Fatalf("expected default read_file mode %q, got %q", ReadFileModeBytes, cfg.Tools.ReadFile.EffectiveMode())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSaveConfig_FilePermissions(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("file permission bits are not enforced on Windows")
|
||||
|
|
|
|||
|
|
@ -185,6 +185,13 @@ func DefaultConfig() *Config {
|
|||
APIBase: "https://api.deepseek.com/v1",
|
||||
},
|
||||
|
||||
// Venice AI - https://venice.ai
|
||||
{
|
||||
ModelName: "venice-uncensored",
|
||||
Model: "venice/venice-uncensored",
|
||||
APIBase: "https://api.venice.ai/api/v1",
|
||||
},
|
||||
|
||||
// Google Gemini - https://ai.google.dev/
|
||||
{
|
||||
ModelName: "gemini-2.0-flash",
|
||||
|
|
@ -335,6 +342,13 @@ func DefaultConfig() *Config {
|
|||
APIBase: "http://localhost:8000/v1",
|
||||
},
|
||||
|
||||
// LM Studio (local) - http://localhost:1234
|
||||
{
|
||||
ModelName: "lmstudio-local",
|
||||
Model: "lmstudio/openai/gpt-oss-20b",
|
||||
APIBase: "http://localhost:1234/v1",
|
||||
},
|
||||
|
||||
// Azure OpenAI - https://portal.azure.com
|
||||
// model_name is a user-friendly alias; the model field's path after "azure/" is your deployment name
|
||||
{
|
||||
|
|
@ -448,6 +462,7 @@ func DefaultConfig() *Config {
|
|||
UseBM25: true,
|
||||
UseRegex: false,
|
||||
},
|
||||
MaxInlineTextChars: DefaultMCPMaxInlineTextChars,
|
||||
Servers: map[string]MCPServerConfig{},
|
||||
},
|
||||
AppendFile: ToolConfig{
|
||||
|
|
@ -473,6 +488,7 @@ func DefaultConfig() *Config {
|
|||
},
|
||||
ReadFile: ReadFileToolConfig{
|
||||
Enabled: true,
|
||||
Mode: ReadFileModeBytes,
|
||||
MaxReadFileSize: 64 * 1024, // 64KB
|
||||
},
|
||||
Spawn: ToolConfig{
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ import (
|
|||
_ "github.com/sipeed/picoclaw/pkg/channels/qq"
|
||||
_ "github.com/sipeed/picoclaw/pkg/channels/slack"
|
||||
_ "github.com/sipeed/picoclaw/pkg/channels/telegram"
|
||||
_ "github.com/sipeed/picoclaw/pkg/channels/vk"
|
||||
_ "github.com/sipeed/picoclaw/pkg/channels/wecom"
|
||||
_ "github.com/sipeed/picoclaw/pkg/channels/weixin"
|
||||
_ "github.com/sipeed/picoclaw/pkg/channels/whatsapp"
|
||||
|
|
|
|||
|
|
@ -2,12 +2,15 @@ package logger
|
|||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime/debug"
|
||||
"time"
|
||||
)
|
||||
|
||||
var panicWriter io.WriteCloser
|
||||
|
||||
func InitPanic(filePath string) (func(), error) {
|
||||
if err := os.MkdirAll(filepath.Dir(filePath), 0o755); err != nil {
|
||||
return nil, fmt.Errorf("failed to create log directory: %w", err)
|
||||
|
|
@ -16,9 +19,28 @@ func InitPanic(filePath string) (func(), error) {
|
|||
if writer == nil {
|
||||
return nil, fmt.Errorf("failed to create log file: %s", filePath)
|
||||
}
|
||||
if panicWriter != nil {
|
||||
_ = panicWriter.Close()
|
||||
}
|
||||
panicWriter = writer
|
||||
return func() {
|
||||
defer writer.Close()
|
||||
defer func() {
|
||||
writer.Close()
|
||||
panicWriter = nil
|
||||
}()
|
||||
if err := recover(); err != nil {
|
||||
RecoverPanicNoExit(err)
|
||||
|
||||
os.Exit(1)
|
||||
}
|
||||
}, nil
|
||||
}
|
||||
|
||||
func RecoverPanicNoExit(err any) {
|
||||
if panicWriter == nil {
|
||||
Errorf("panicWriter is nil, should not happen")
|
||||
return
|
||||
}
|
||||
now := time.Now().Format("2006-01-02 15:04:05")
|
||||
stack := debug.Stack()
|
||||
logMsg := "\n\n====================\n[" + now + "] PANIC OCCURRED: " + fmt.Sprintf(
|
||||
|
|
@ -28,9 +50,5 @@ func InitPanic(filePath string) (func(), error) {
|
|||
stack,
|
||||
)
|
||||
|
||||
writer.Write([]byte(logMsg))
|
||||
|
||||
os.Exit(1)
|
||||
}
|
||||
}, nil
|
||||
panicWriter.Write([]byte(logMsg))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -455,6 +455,33 @@ func (s *JSONLStore) rewriteJSONL(
|
|||
return fileutil.WriteFileAtomic(s.jsonlPath(sessionKey), buf.Bytes(), 0o644)
|
||||
}
|
||||
|
||||
// ListSessions returns all known session keys by reading .meta.json files.
|
||||
func (s *JSONLStore) ListSessions() []string {
|
||||
entries, err := os.ReadDir(s.dir)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
var keys []string
|
||||
for _, entry := range entries {
|
||||
if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".meta.json") {
|
||||
continue
|
||||
}
|
||||
// Read the meta file to get the original key
|
||||
data, err := os.ReadFile(filepath.Join(s.dir, entry.Name()))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
var meta sessionMeta
|
||||
if err := json.Unmarshal(data, &meta); err != nil {
|
||||
continue
|
||||
}
|
||||
if meta.Key != "" {
|
||||
keys = append(keys, meta.Key)
|
||||
}
|
||||
}
|
||||
return keys
|
||||
}
|
||||
|
||||
func (s *JSONLStore) Close() error {
|
||||
return nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -37,6 +37,9 @@ type Store interface {
|
|||
// data. Backends that do not accumulate dead data may return nil.
|
||||
Compact(ctx context.Context, sessionKey string) error
|
||||
|
||||
// ListSessions returns all known session keys.
|
||||
ListSessions() []string
|
||||
|
||||
// Close releases any resources held by the store.
|
||||
Close() error
|
||||
}
|
||||
|
|
|
|||
|
|
@ -41,15 +41,16 @@ type Provider struct {
|
|||
apiKey string
|
||||
apiBase string
|
||||
httpClient *http.Client
|
||||
userAgent string
|
||||
}
|
||||
|
||||
// NewProvider creates a new Anthropic Messages API provider.
|
||||
func NewProvider(apiKey, apiBase string) *Provider {
|
||||
return NewProviderWithTimeout(apiKey, apiBase, 0)
|
||||
func NewProvider(apiKey, apiBase, userAgent string) *Provider {
|
||||
return NewProviderWithTimeout(apiKey, apiBase, userAgent, 0)
|
||||
}
|
||||
|
||||
// NewProviderWithTimeout creates a provider with custom request timeout.
|
||||
func NewProviderWithTimeout(apiKey, apiBase string, timeoutSeconds int) *Provider {
|
||||
func NewProviderWithTimeout(apiKey, apiBase, userAgent string, timeoutSeconds int) *Provider {
|
||||
baseURL := normalizeBaseURL(apiBase)
|
||||
timeout := defaultRequestTimeout
|
||||
if timeoutSeconds > 0 {
|
||||
|
|
@ -59,6 +60,7 @@ func NewProviderWithTimeout(apiKey, apiBase string, timeoutSeconds int) *Provide
|
|||
return &Provider{
|
||||
apiKey: apiKey,
|
||||
apiBase: baseURL,
|
||||
userAgent: userAgent,
|
||||
httpClient: &http.Client{
|
||||
Timeout: timeout,
|
||||
},
|
||||
|
|
@ -105,6 +107,9 @@ func (p *Provider) Chat(
|
|||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("X-API-Key", p.apiKey) //nolint:canonicalheader // Anthropic API requires exact header name
|
||||
req.Header.Set("Anthropic-Version", defaultAPIVersion)
|
||||
if p.userAgent != "" {
|
||||
req.Header.Set("User-Agent", p.userAgent)
|
||||
}
|
||||
|
||||
// Execute request
|
||||
resp, err := p.httpClient.Do(req)
|
||||
|
|
|
|||
|
|
@ -411,7 +411,7 @@ func TestNormalizeBaseURL(t *testing.T) {
|
|||
}
|
||||
|
||||
func TestNewProvider(t *testing.T) {
|
||||
provider := NewProvider("test-key", "https://api.example.com")
|
||||
provider := NewProvider("test-key", "https://api.example.com", "")
|
||||
if provider == nil {
|
||||
t.Fatal("NewProvider() returned nil")
|
||||
}
|
||||
|
|
@ -424,7 +424,7 @@ func TestNewProvider(t *testing.T) {
|
|||
}
|
||||
|
||||
func TestGetDefaultModel(t *testing.T) {
|
||||
provider := NewProvider("test-key", "")
|
||||
provider := NewProvider("test-key", "", "")
|
||||
got := provider.GetDefaultModel()
|
||||
expected := "claude-sonnet-4.6"
|
||||
if got != expected {
|
||||
|
|
@ -743,7 +743,7 @@ func TestProviderChatErrors(t *testing.T) {
|
|||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Create provider using constructor to ensure proper initialization
|
||||
provider := NewProvider(tt.apiKey, "https://api.example.com")
|
||||
provider := NewProvider(tt.apiKey, "https://api.example.com", "")
|
||||
|
||||
_, err := provider.Chat(context.Background(), tt.messages, nil, "test-model", nil)
|
||||
if err == nil {
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ type Provider struct {
|
|||
apiKey string
|
||||
apiBase string
|
||||
httpClient *http.Client
|
||||
userAgent string
|
||||
}
|
||||
|
||||
// Option configures the Azure Provider.
|
||||
|
|
@ -50,11 +51,19 @@ func WithRequestTimeout(timeout time.Duration) Option {
|
|||
}
|
||||
}
|
||||
|
||||
// WithUserAgent sets the User-Agent header for requests.
|
||||
func WithUserAgent(userAgent string) Option {
|
||||
return func(p *Provider) {
|
||||
p.userAgent = userAgent
|
||||
}
|
||||
}
|
||||
|
||||
// NewProvider creates a new Azure OpenAI provider.
|
||||
func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider {
|
||||
func NewProvider(apiKey, apiBase, proxy, userAgent string, opts ...Option) *Provider {
|
||||
p := &Provider{
|
||||
apiKey: apiKey,
|
||||
apiBase: strings.TrimRight(apiBase, "/"),
|
||||
userAgent: userAgent,
|
||||
httpClient: common.NewHTTPClient(proxy),
|
||||
}
|
||||
|
||||
|
|
@ -68,9 +77,9 @@ func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider {
|
|||
}
|
||||
|
||||
// NewProviderWithTimeout creates a new Azure OpenAI provider with a custom request timeout in seconds.
|
||||
func NewProviderWithTimeout(apiKey, apiBase, proxy string, requestTimeoutSeconds int) *Provider {
|
||||
func NewProviderWithTimeout(apiKey, apiBase, proxy, userAgent string, requestTimeoutSeconds int) *Provider {
|
||||
return NewProvider(
|
||||
apiKey, apiBase, proxy,
|
||||
apiKey, apiBase, proxy, userAgent,
|
||||
WithRequestTimeout(time.Duration(requestTimeoutSeconds)*time.Second),
|
||||
)
|
||||
}
|
||||
|
|
@ -141,6 +150,9 @@ func (p *Provider) Chat(
|
|||
if p.apiKey != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+p.apiKey)
|
||||
}
|
||||
if p.userAgent != "" {
|
||||
req.Header.Set("User-Agent", p.userAgent)
|
||||
}
|
||||
|
||||
resp, err := p.httpClient.Do(req)
|
||||
if err != nil {
|
||||
|
|
|
|||
|
|
@ -46,7 +46,7 @@ func TestProviderChat_AzureURLConstruction(t *testing.T) {
|
|||
}))
|
||||
defer server.Close()
|
||||
|
||||
p := NewProvider("test-key", server.URL, "")
|
||||
p := NewProvider("test-key", server.URL, "", "")
|
||||
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "my-gpt5-deployment", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("Chat() error = %v", err)
|
||||
|
|
@ -69,7 +69,7 @@ func TestProviderChat_AzureAuthHeader(t *testing.T) {
|
|||
}))
|
||||
defer server.Close()
|
||||
|
||||
p := NewProvider("test-azure-key", server.URL, "")
|
||||
p := NewProvider("test-azure-key", server.URL, "", "")
|
||||
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("Chat() error = %v", err)
|
||||
|
|
@ -92,7 +92,7 @@ func TestProviderChat_AzureRequestBodyContainsModel(t *testing.T) {
|
|||
}))
|
||||
defer server.Close()
|
||||
|
||||
p := NewProvider("test-key", server.URL, "")
|
||||
p := NewProvider("test-key", server.URL, "", "")
|
||||
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "my-deployment", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("Chat() error = %v", err)
|
||||
|
|
@ -112,7 +112,7 @@ func TestProviderChat_AzureUsesMaxOutputTokens(t *testing.T) {
|
|||
}))
|
||||
defer server.Close()
|
||||
|
||||
p := NewProvider("test-key", server.URL, "")
|
||||
p := NewProvider("test-key", server.URL, "", "")
|
||||
_, err := p.Chat(
|
||||
t.Context(),
|
||||
[]Message{{Role: "user", Content: "hi"}},
|
||||
|
|
@ -144,7 +144,7 @@ func TestProviderChat_AzureStoreIsFalse(t *testing.T) {
|
|||
}))
|
||||
defer server.Close()
|
||||
|
||||
p := NewProvider("test-key", server.URL, "")
|
||||
p := NewProvider("test-key", server.URL, "", "")
|
||||
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("Chat() error = %v", err)
|
||||
|
|
@ -161,7 +161,7 @@ func TestProviderChat_AzureHTTPError(t *testing.T) {
|
|||
}))
|
||||
defer server.Close()
|
||||
|
||||
p := NewProvider("bad-key", server.URL, "")
|
||||
p := NewProvider("bad-key", server.URL, "", "")
|
||||
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil)
|
||||
if err == nil {
|
||||
t.Fatal("expected error, got nil")
|
||||
|
|
@ -176,7 +176,7 @@ func TestProviderChat_AzureRateLimitError(t *testing.T) {
|
|||
}))
|
||||
defer server.Close()
|
||||
|
||||
p := NewProvider("test-key", server.URL, "")
|
||||
p := NewProvider("test-key", server.URL, "", "")
|
||||
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for 429, got nil")
|
||||
|
|
@ -194,7 +194,7 @@ func TestProviderChat_AzureServerError(t *testing.T) {
|
|||
}))
|
||||
defer server.Close()
|
||||
|
||||
p := NewProvider("test-key", server.URL, "")
|
||||
p := NewProvider("test-key", server.URL, "", "")
|
||||
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for 500, got nil")
|
||||
|
|
@ -229,7 +229,7 @@ func TestProviderChat_AzureParseTextOutput(t *testing.T) {
|
|||
}))
|
||||
defer server.Close()
|
||||
|
||||
p := NewProvider("test-key", server.URL, "")
|
||||
p := NewProvider("test-key", server.URL, "", "")
|
||||
out, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("Chat() error = %v", err)
|
||||
|
|
@ -270,7 +270,7 @@ func TestProviderChat_AzureParseToolCalls(t *testing.T) {
|
|||
}))
|
||||
defer server.Close()
|
||||
|
||||
p := NewProvider("test-key", server.URL, "")
|
||||
p := NewProvider("test-key", server.URL, "", "")
|
||||
out, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "weather?"}}, nil, "deployment", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("Chat() error = %v", err)
|
||||
|
|
@ -287,7 +287,7 @@ func TestProviderChat_AzureParseToolCalls(t *testing.T) {
|
|||
}
|
||||
|
||||
func TestProvider_AzureEmptyAPIBase(t *testing.T) {
|
||||
p := NewProvider("test-key", "", "")
|
||||
p := NewProvider("test-key", "", "", "")
|
||||
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for empty API base")
|
||||
|
|
@ -295,21 +295,21 @@ func TestProvider_AzureEmptyAPIBase(t *testing.T) {
|
|||
}
|
||||
|
||||
func TestProvider_AzureRequestTimeoutDefault(t *testing.T) {
|
||||
p := NewProvider("test-key", "https://example.com", "")
|
||||
p := NewProvider("test-key", "https://example.com", "", "")
|
||||
if p.httpClient.Timeout != defaultRequestTimeout {
|
||||
t.Errorf("timeout = %v, want %v", p.httpClient.Timeout, defaultRequestTimeout)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProvider_AzureRequestTimeoutOverride(t *testing.T) {
|
||||
p := NewProvider("test-key", "https://example.com", "", WithRequestTimeout(300*time.Second))
|
||||
p := NewProvider("test-key", "https://example.com", "", "", WithRequestTimeout(300*time.Second))
|
||||
if p.httpClient.Timeout != 300*time.Second {
|
||||
t.Errorf("timeout = %v, want %v", p.httpClient.Timeout, 300*time.Second)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProvider_AzureNewProviderWithTimeout(t *testing.T) {
|
||||
p := NewProviderWithTimeout("test-key", "https://example.com", "", 180)
|
||||
p := NewProviderWithTimeout("test-key", "https://example.com", "", "", 180)
|
||||
if p.httpClient.Timeout != 180*time.Second {
|
||||
t.Errorf("timeout = %v, want %v", p.httpClient.Timeout, 180*time.Second)
|
||||
}
|
||||
|
|
@ -343,7 +343,7 @@ func TestProviderChat_AzureNativeWebSearchInjection(t *testing.T) {
|
|||
},
|
||||
}
|
||||
|
||||
p := NewProvider("test-key", server.URL, "")
|
||||
p := NewProvider("test-key", server.URL, "", "")
|
||||
|
||||
// With native_search=true: user-defined web_search should be replaced by built-in
|
||||
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, tools, "deployment",
|
||||
|
|
@ -393,7 +393,7 @@ func TestProviderChat_AzureNoNativeWebSearch(t *testing.T) {
|
|||
},
|
||||
}
|
||||
|
||||
p := NewProvider("test-key", server.URL, "")
|
||||
p := NewProvider("test-key", server.URL, "", "")
|
||||
|
||||
// Without native_search: user-defined web_search should be kept as-is
|
||||
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, tools, "deployment", nil)
|
||||
|
|
|
|||
|
|
@ -262,6 +262,22 @@ func TestDecodeToolCallArguments_StringJSON(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestDecodeToolCallArguments_StringJSON_NewlineEscape(t *testing.T) {
|
||||
raw := json.RawMessage(`"{\"content\":\"line1\\nline2\"}"`)
|
||||
args := DecodeToolCallArguments(raw, "write_file")
|
||||
if args["content"] != "line1\nline2" {
|
||||
t.Errorf("content = %q, want newline-expanded string", args["content"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeToolCallArguments_StringJSON_LiteralBackslashN(t *testing.T) {
|
||||
raw := json.RawMessage(`"{\"content\":\"line1\\\\nline2\"}"`)
|
||||
args := DecodeToolCallArguments(raw, "write_file")
|
||||
if args["content"] != `line1\nline2` {
|
||||
t.Errorf("content = %q, want literal backslash-n", args["content"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeToolCallArguments_EmptyInput(t *testing.T) {
|
||||
args := DecodeToolCallArguments(nil, "test")
|
||||
if len(args) != 0 {
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ type protocolMeta struct {
|
|||
|
||||
var protocolMetaByName = map[string]protocolMeta{
|
||||
"openai": {defaultAPIBase: "https://api.openai.com/v1"},
|
||||
"venice": {defaultAPIBase: "https://api.venice.ai/api/v1"},
|
||||
"openrouter": {defaultAPIBase: "https://openrouter.ai/api/v1"},
|
||||
"litellm": {defaultAPIBase: "http://localhost:4000/v1"},
|
||||
"lmstudio": {defaultAPIBase: "http://localhost:1234/v1", emptyAPIKeyAllowed: true},
|
||||
|
|
@ -128,6 +129,11 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
|
|||
|
||||
protocol, modelID := ExtractProtocol(cfg.Model)
|
||||
|
||||
userAgent := cfg.UserAgent
|
||||
if userAgent == "" {
|
||||
userAgent = fmt.Sprintf("PicoClaw/%s", config.Version)
|
||||
}
|
||||
|
||||
switch protocol {
|
||||
case "openai":
|
||||
// OpenAI with OAuth/token auth (Codex-style)
|
||||
|
|
@ -151,6 +157,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
|
|||
apiBase,
|
||||
cfg.Proxy,
|
||||
cfg.MaxTokensField,
|
||||
userAgent,
|
||||
cfg.RequestTimeout,
|
||||
cfg.ExtraBody,
|
||||
), modelID, nil
|
||||
|
|
@ -170,6 +177,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
|
|||
cfg.APIKey(),
|
||||
cfg.APIBase,
|
||||
cfg.Proxy,
|
||||
userAgent,
|
||||
cfg.RequestTimeout,
|
||||
), modelID, nil
|
||||
|
||||
|
|
@ -209,7 +217,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
|
|||
}
|
||||
return provider, modelID, nil
|
||||
|
||||
case "litellm", "lmstudio", "openrouter", "groq", "zhipu", "gemini", "nvidia",
|
||||
case "litellm", "lmstudio", "openrouter", "groq", "zhipu", "gemini", "nvidia", "venice",
|
||||
"ollama", "moonshot", "shengsuanyun", "deepseek", "cerebras",
|
||||
"vivgrid", "volcengine", "vllm", "qwen", "qwen-intl", "qwen-international", "dashscope-intl",
|
||||
"qwen-us", "dashscope-us", "mistral", "avian", "longcat", "modelscope", "novita",
|
||||
|
|
@ -227,6 +235,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
|
|||
apiBase,
|
||||
cfg.Proxy,
|
||||
cfg.MaxTokensField,
|
||||
userAgent,
|
||||
cfg.RequestTimeout,
|
||||
cfg.ExtraBody,
|
||||
), modelID, nil
|
||||
|
|
@ -252,6 +261,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
|
|||
apiBase,
|
||||
cfg.Proxy,
|
||||
cfg.MaxTokensField,
|
||||
userAgent,
|
||||
cfg.RequestTimeout,
|
||||
extraBody,
|
||||
), modelID, nil
|
||||
|
|
@ -278,6 +288,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
|
|||
apiBase,
|
||||
cfg.Proxy,
|
||||
cfg.MaxTokensField,
|
||||
userAgent,
|
||||
cfg.RequestTimeout,
|
||||
cfg.ExtraBody,
|
||||
), modelID, nil
|
||||
|
|
@ -294,6 +305,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
|
|||
return anthropicmessages.NewProviderWithTimeout(
|
||||
cfg.APIKey(),
|
||||
apiBase,
|
||||
userAgent,
|
||||
cfg.RequestTimeout,
|
||||
), modelID, nil
|
||||
|
||||
|
|
@ -309,6 +321,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
|
|||
return anthropicmessages.NewProviderWithTimeout(
|
||||
cfg.APIKey(),
|
||||
apiBase,
|
||||
userAgent,
|
||||
cfg.RequestTimeout,
|
||||
), modelID, nil
|
||||
|
||||
|
|
|
|||
|
|
@ -112,6 +112,7 @@ func TestCreateProviderFromConfig_DefaultAPIBase(t *testing.T) {
|
|||
protocol string
|
||||
}{
|
||||
{"openai", "openai"},
|
||||
{"venice", "venice"},
|
||||
{"groq", "groq"},
|
||||
{"novita", "novita"},
|
||||
{"openrouter", "openrouter"},
|
||||
|
|
@ -160,6 +161,12 @@ func TestGetDefaultAPIBase_LMStudio(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestGetDefaultAPIBase_Venice(t *testing.T) {
|
||||
if got := getDefaultAPIBase("venice"); got != "https://api.venice.ai/api/v1" {
|
||||
t.Fatalf("getDefaultAPIBase(%q) = %q, want %q", "venice", got, "https://api.venice.ai/api/v1")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateProviderFromConfig_LiteLLM(t *testing.T) {
|
||||
cfg := &config.ModelConfig{
|
||||
ModelName: "test-litellm",
|
||||
|
|
@ -362,6 +369,28 @@ func TestCreateProviderFromConfig_Mimo(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestCreateProviderFromConfig_Venice(t *testing.T) {
|
||||
cfg := &config.ModelConfig{
|
||||
ModelName: "test-venice",
|
||||
Model: "venice/venice-uncensored",
|
||||
}
|
||||
cfg.SetAPIKey("test-key")
|
||||
|
||||
provider, modelID, err := CreateProviderFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("CreateProviderFromConfig() error = %v", err)
|
||||
}
|
||||
if provider == nil {
|
||||
t.Fatal("CreateProviderFromConfig() returned nil provider")
|
||||
}
|
||||
if modelID != "venice-uncensored" {
|
||||
t.Errorf("modelID = %q, want %q", modelID, "venice-uncensored")
|
||||
}
|
||||
if _, ok := provider.(*HTTPProvider); !ok {
|
||||
t.Fatalf("expected *HTTPProvider, got %T", provider)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetDefaultAPIBase_Mimo(t *testing.T) {
|
||||
if got := getDefaultAPIBase("mimo"); got != "https://api.xiaomimimo.com/v1" {
|
||||
t.Fatalf("getDefaultAPIBase(%q) = %q, want %q", "mimo", got, "https://api.xiaomimimo.com/v1")
|
||||
|
|
@ -817,6 +846,107 @@ func TestCreateProviderFromConfig_MinimaxPreservesUserExtraBody(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
// openaiCompatResponse is the JSON response used by OpenAI-compatible providers.
|
||||
const openaiCompatResponse = `{"choices":[{"message":{"content":"ok"},"finish_reason":"stop"}]}`
|
||||
|
||||
// anthropicResponse is the JSON response used by Anthropic providers.
|
||||
const anthropicResponse = `{"content":[{"type":"text","text":"ok"}],"stop_reason":"end_turn","model":"claude-sonnet-4-20250514","usage":{"input_tokens":10,"output_tokens":5}}`
|
||||
|
||||
func TestCreateProviderFromConfig_UserAgent(t *testing.T) {
|
||||
defaultUA := "PicoClaw/" + config.Version
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
model string
|
||||
userAgent string
|
||||
apiKey string
|
||||
response string
|
||||
wantUA string
|
||||
chatOpts map[string]any
|
||||
}{
|
||||
{
|
||||
name: "openai default user agent",
|
||||
model: "openai/gpt-4o",
|
||||
apiKey: "test-key",
|
||||
response: openaiCompatResponse,
|
||||
wantUA: defaultUA,
|
||||
},
|
||||
{
|
||||
name: "openai custom user agent",
|
||||
model: "openai/gpt-4o",
|
||||
apiKey: "test-key",
|
||||
userAgent: "MyAgent/1.2.3",
|
||||
response: openaiCompatResponse,
|
||||
wantUA: "MyAgent/1.2.3",
|
||||
},
|
||||
{
|
||||
name: "anthropic default user agent",
|
||||
model: "anthropic/claude-sonnet-4-20250514",
|
||||
apiKey: "test-key",
|
||||
response: anthropicResponse,
|
||||
wantUA: defaultUA,
|
||||
},
|
||||
{
|
||||
name: "anthropic-messages default user agent",
|
||||
model: "anthropic-messages/claude-sonnet-4-20250514",
|
||||
apiKey: "test-key",
|
||||
response: anthropicResponse,
|
||||
wantUA: defaultUA,
|
||||
chatOpts: map[string]any{"max_tokens": 1024},
|
||||
},
|
||||
{
|
||||
name: "azure default user agent",
|
||||
model: "azure/my-deployment",
|
||||
apiKey: "test-azure-key",
|
||||
response: openaiCompatResponse,
|
||||
wantUA: defaultUA,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
var receivedUA string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
receivedUA = r.Header.Get("User-Agent")
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(tt.response))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
cfg := &config.ModelConfig{
|
||||
ModelName: "test-ua-" + tt.name,
|
||||
Model: tt.model,
|
||||
APIBase: server.URL,
|
||||
UserAgent: tt.userAgent,
|
||||
}
|
||||
cfg.SetAPIKey(tt.apiKey)
|
||||
|
||||
provider, modelID, err := CreateProviderFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("CreateProviderFromConfig() error = %v", err)
|
||||
}
|
||||
if provider == nil {
|
||||
t.Fatal("CreateProviderFromConfig() returned nil provider")
|
||||
}
|
||||
|
||||
_, err = provider.Chat(
|
||||
t.Context(),
|
||||
[]Message{{Role: "user", Content: "hi"}},
|
||||
nil,
|
||||
modelID,
|
||||
tt.chatOpts,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("Chat() error = %v", err)
|
||||
}
|
||||
|
||||
if receivedUA != tt.wantUA {
|
||||
t.Errorf("User-Agent = %q, want %q", receivedUA, tt.wantUA)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateProviderFromConfig_Bedrock(t *testing.T) {
|
||||
// Set dummy AWS env vars to make test deterministic
|
||||
t.Setenv("AWS_ACCESS_KEY_ID", "test-key")
|
||||
|
|
|
|||
|
|
@ -10,12 +10,24 @@ import (
|
|||
// FallbackChain orchestrates model fallback across multiple candidates.
|
||||
type FallbackChain struct {
|
||||
cooldown *CooldownTracker
|
||||
rl *RateLimiterRegistry
|
||||
}
|
||||
|
||||
// FallbackCandidate represents one model/provider to try.
|
||||
type FallbackCandidate struct {
|
||||
Provider string
|
||||
Model string
|
||||
RPM int // requests per minute; 0 means unrestricted
|
||||
IdentityKey string // optional stable config identity for cooldown/rate limiting
|
||||
}
|
||||
|
||||
// StableKey returns the candidate's config-level identity when available,
|
||||
// otherwise it falls back to the runtime provider/model key.
|
||||
func (c FallbackCandidate) StableKey() string {
|
||||
if key := strings.TrimSpace(c.IdentityKey); key != "" {
|
||||
return key
|
||||
}
|
||||
return ModelKey(c.Provider, c.Model)
|
||||
}
|
||||
|
||||
// FallbackResult contains the successful response and metadata about all attempts.
|
||||
|
|
@ -36,9 +48,10 @@ type FallbackAttempt struct {
|
|||
Skipped bool // true if skipped due to cooldown
|
||||
}
|
||||
|
||||
// NewFallbackChain creates a new fallback chain with the given cooldown tracker.
|
||||
func NewFallbackChain(cooldown *CooldownTracker) *FallbackChain {
|
||||
return &FallbackChain{cooldown: cooldown}
|
||||
// NewFallbackChain creates a new fallback chain with the given cooldown tracker
|
||||
// and rate limiter registry.
|
||||
func NewFallbackChain(cooldown *CooldownTracker, rl *RateLimiterRegistry) *FallbackChain {
|
||||
return &FallbackChain{cooldown: cooldown, rl: rl}
|
||||
}
|
||||
|
||||
// ResolveCandidates parses model config into a deduplicated candidate list.
|
||||
|
|
@ -117,9 +130,9 @@ func (fc *FallbackChain) Execute(
|
|||
return nil, context.Canceled
|
||||
}
|
||||
|
||||
// Check cooldown (per provider/model, not just provider).
|
||||
// This allows multi-key failover where different keys use different model names.
|
||||
cooldownKey := ModelKey(candidate.Provider, candidate.Model)
|
||||
// Check cooldown per stable candidate identity, not just provider/model.
|
||||
// This allows aliases and multi-key configs to fail over independently.
|
||||
cooldownKey := candidate.StableKey()
|
||||
if !fc.cooldown.IsAvailable(cooldownKey) {
|
||||
remaining := fc.cooldown.CooldownRemaining(cooldownKey)
|
||||
result.Attempts = append(result.Attempts, FallbackAttempt{
|
||||
|
|
@ -136,6 +149,33 @@ func (fc *FallbackChain) Execute(
|
|||
continue
|
||||
}
|
||||
|
||||
// Enforce per-candidate rate limit before calling the provider.
|
||||
// If this candidate is locally saturated, try other candidates first.
|
||||
if fc.rl != nil {
|
||||
if !fc.rl.TryAcquire(cooldownKey) {
|
||||
if i < len(candidates)-1 {
|
||||
result.Attempts = append(result.Attempts, FallbackAttempt{
|
||||
Provider: candidate.Provider,
|
||||
Model: candidate.Model,
|
||||
Skipped: true,
|
||||
Reason: FailoverRateLimit,
|
||||
Error: fmt.Errorf("%s waiting for local rate limit token", cooldownKey),
|
||||
})
|
||||
continue
|
||||
}
|
||||
if waitErr := fc.rl.Wait(ctx, cooldownKey); waitErr != nil {
|
||||
result.Attempts = append(result.Attempts, FallbackAttempt{
|
||||
Provider: candidate.Provider,
|
||||
Model: candidate.Model,
|
||||
Skipped: true,
|
||||
Reason: FailoverRateLimit,
|
||||
Error: waitErr,
|
||||
})
|
||||
return nil, waitErr
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Execute the run function.
|
||||
start := time.Now()
|
||||
resp, err := run(ctx, candidate.Provider, candidate.Model)
|
||||
|
|
@ -229,6 +269,34 @@ func (fc *FallbackChain) ExecuteImage(
|
|||
return nil, context.Canceled
|
||||
}
|
||||
|
||||
// Enforce per-candidate rate limit before calling the provider.
|
||||
// If this candidate is locally saturated, try other candidates first.
|
||||
imageKey := candidate.StableKey()
|
||||
if fc.rl != nil {
|
||||
if !fc.rl.TryAcquire(imageKey) {
|
||||
if i < len(candidates)-1 {
|
||||
result.Attempts = append(result.Attempts, FallbackAttempt{
|
||||
Provider: candidate.Provider,
|
||||
Model: candidate.Model,
|
||||
Skipped: true,
|
||||
Reason: FailoverRateLimit,
|
||||
Error: fmt.Errorf("%s waiting for local rate limit token", imageKey),
|
||||
})
|
||||
continue
|
||||
}
|
||||
if waitErr := fc.rl.Wait(ctx, imageKey); waitErr != nil {
|
||||
result.Attempts = append(result.Attempts, FallbackAttempt{
|
||||
Provider: candidate.Provider,
|
||||
Model: candidate.Model,
|
||||
Skipped: true,
|
||||
Reason: FailoverRateLimit,
|
||||
Error: waitErr,
|
||||
})
|
||||
return nil, waitErr
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
resp, err := run(ctx, candidate.Provider, candidate.Model)
|
||||
elapsed := time.Since(start)
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ func TestMultiKeyFailover(t *testing.T) {
|
|||
|
||||
// Create fallback chain
|
||||
cooldown := NewCooldownTracker()
|
||||
chain := NewFallbackChain(cooldown)
|
||||
chain := NewFallbackChain(cooldown, nil)
|
||||
|
||||
// Mock run function: first call fails with 429, second succeeds
|
||||
callCount := 0
|
||||
|
|
@ -82,7 +82,7 @@ func TestMultiKeyFailoverAllFail(t *testing.T) {
|
|||
candidates := ResolveCandidates(cfg, "zhipu")
|
||||
|
||||
cooldown := NewCooldownTracker()
|
||||
chain := NewFallbackChain(cooldown)
|
||||
chain := NewFallbackChain(cooldown, nil)
|
||||
|
||||
// Mock run function: all calls fail with rate limit
|
||||
callCount := 0
|
||||
|
|
@ -127,7 +127,7 @@ func TestMultiKeyFailoverCooldown(t *testing.T) {
|
|||
candidates := ResolveCandidates(cfg, "zhipu")
|
||||
|
||||
cooldown := NewCooldownTracker()
|
||||
chain := NewFallbackChain(cooldown)
|
||||
chain := NewFallbackChain(cooldown, nil)
|
||||
|
||||
// Put the first model in cooldown (using ModelKey now, not just provider)
|
||||
cooldownKey := ModelKey(candidates[0].Provider, candidates[0].Model)
|
||||
|
|
@ -183,7 +183,7 @@ func TestMultiKeyFailoverWithFormatError(t *testing.T) {
|
|||
candidates := ResolveCandidates(cfg, "zhipu")
|
||||
|
||||
cooldown := NewCooldownTracker()
|
||||
chain := NewFallbackChain(cooldown)
|
||||
chain := NewFallbackChain(cooldown, nil)
|
||||
|
||||
// Mock run function: first call fails with format error (bad request)
|
||||
callCount := 0
|
||||
|
|
@ -263,7 +263,7 @@ func TestMultiKeyWithModelFallback(t *testing.T) {
|
|||
}
|
||||
|
||||
cooldown := NewCooldownTracker()
|
||||
chain := NewFallbackChain(cooldown)
|
||||
chain := NewFallbackChain(cooldown, nil)
|
||||
|
||||
// Mock run function: first two fail, third succeeds (model fallback)
|
||||
callCount := 0
|
||||
|
|
@ -337,7 +337,7 @@ func TestMultiKeyFailoverMixedErrors(t *testing.T) {
|
|||
candidates := ResolveCandidates(cfg, "zhipu")
|
||||
|
||||
cooldown := NewCooldownTracker()
|
||||
chain := NewFallbackChain(cooldown)
|
||||
chain := NewFallbackChain(cooldown, nil)
|
||||
|
||||
// Mock run function: different errors for each key
|
||||
callCount := 0
|
||||
|
|
|
|||
|
|
@ -19,7 +19,7 @@ func successRun(content string) func(ctx context.Context, provider, model string
|
|||
|
||||
func TestFallback_SingleCandidate_Success(t *testing.T) {
|
||||
ct := NewCooldownTracker()
|
||||
fc := NewFallbackChain(ct)
|
||||
fc := NewFallbackChain(ct, nil)
|
||||
|
||||
candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4")}
|
||||
result, err := fc.Execute(context.Background(), candidates, successRun("hello"))
|
||||
|
|
@ -36,7 +36,7 @@ func TestFallback_SingleCandidate_Success(t *testing.T) {
|
|||
|
||||
func TestFallback_SecondCandidateSuccess(t *testing.T) {
|
||||
ct := NewCooldownTracker()
|
||||
fc := NewFallbackChain(ct)
|
||||
fc := NewFallbackChain(ct, nil)
|
||||
|
||||
candidates := []FallbackCandidate{
|
||||
makeCandidate("openai", "gpt-4"),
|
||||
|
|
@ -69,7 +69,7 @@ func TestFallback_SecondCandidateSuccess(t *testing.T) {
|
|||
|
||||
func TestFallback_AllFail(t *testing.T) {
|
||||
ct := NewCooldownTracker()
|
||||
fc := NewFallbackChain(ct)
|
||||
fc := NewFallbackChain(ct, nil)
|
||||
|
||||
candidates := []FallbackCandidate{
|
||||
makeCandidate("openai", "gpt-4"),
|
||||
|
|
@ -96,7 +96,7 @@ func TestFallback_AllFail(t *testing.T) {
|
|||
|
||||
func TestFallback_ContextCanceled(t *testing.T) {
|
||||
ct := NewCooldownTracker()
|
||||
fc := NewFallbackChain(ct)
|
||||
fc := NewFallbackChain(ct, nil)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
candidates := []FallbackCandidate{
|
||||
|
|
@ -123,7 +123,7 @@ func TestFallback_ContextCanceled(t *testing.T) {
|
|||
|
||||
func TestFallback_NonRetriableError(t *testing.T) {
|
||||
ct := NewCooldownTracker()
|
||||
fc := NewFallbackChain(ct)
|
||||
fc := NewFallbackChain(ct, nil)
|
||||
|
||||
candidates := []FallbackCandidate{
|
||||
makeCandidate("openai", "gpt-4"),
|
||||
|
|
@ -155,7 +155,7 @@ func TestFallback_NonRetriableError(t *testing.T) {
|
|||
func TestFallback_CooldownSkip(t *testing.T) {
|
||||
now := time.Now()
|
||||
ct, _ := newTestTracker(now)
|
||||
fc := NewFallbackChain(ct)
|
||||
fc := NewFallbackChain(ct, nil)
|
||||
|
||||
// Put openai/gpt-4 in cooldown (using ModelKey now)
|
||||
ct.MarkFailure(ModelKey("openai", "gpt-4"), FailoverRateLimit)
|
||||
|
|
@ -193,7 +193,7 @@ func TestFallback_CooldownSkip(t *testing.T) {
|
|||
|
||||
func TestFallback_AllInCooldown(t *testing.T) {
|
||||
ct := NewCooldownTracker()
|
||||
fc := NewFallbackChain(ct)
|
||||
fc := NewFallbackChain(ct, nil)
|
||||
|
||||
// Put all models in cooldown (using ModelKey now)
|
||||
ct.MarkFailure(ModelKey("openai", "gpt-4"), FailoverRateLimit)
|
||||
|
|
@ -221,7 +221,7 @@ func TestFallback_AllInCooldown(t *testing.T) {
|
|||
|
||||
func TestFallback_NoCandidates(t *testing.T) {
|
||||
ct := NewCooldownTracker()
|
||||
fc := NewFallbackChain(ct)
|
||||
fc := NewFallbackChain(ct, nil)
|
||||
|
||||
_, err := fc.Execute(context.Background(), nil, successRun("ok"))
|
||||
if err == nil {
|
||||
|
|
@ -232,7 +232,7 @@ func TestFallback_NoCandidates(t *testing.T) {
|
|||
func TestFallback_EmptyFallbacks(t *testing.T) {
|
||||
// Single primary, no fallbacks: should work like direct call
|
||||
ct := NewCooldownTracker()
|
||||
fc := NewFallbackChain(ct)
|
||||
fc := NewFallbackChain(ct, nil)
|
||||
|
||||
candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4")}
|
||||
result, err := fc.Execute(context.Background(), candidates, successRun("ok"))
|
||||
|
|
@ -246,7 +246,7 @@ func TestFallback_EmptyFallbacks(t *testing.T) {
|
|||
|
||||
func TestFallback_UnclassifiedError(t *testing.T) {
|
||||
ct := NewCooldownTracker()
|
||||
fc := NewFallbackChain(ct)
|
||||
fc := NewFallbackChain(ct, nil)
|
||||
|
||||
candidates := []FallbackCandidate{
|
||||
makeCandidate("openai", "gpt-4"),
|
||||
|
|
@ -270,7 +270,7 @@ func TestFallback_UnclassifiedError(t *testing.T) {
|
|||
|
||||
func TestFallback_SuccessResetsCooldown(t *testing.T) {
|
||||
ct := NewCooldownTracker()
|
||||
fc := NewFallbackChain(ct)
|
||||
fc := NewFallbackChain(ct, nil)
|
||||
|
||||
candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4")}
|
||||
modelKey := ModelKey("openai", "gpt-4")
|
||||
|
|
@ -293,11 +293,78 @@ func TestFallback_SuccessResetsCooldown(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func assertLocalRateLimitSkipsToHealthyFallback(
|
||||
t *testing.T,
|
||||
primaryKey string,
|
||||
fallbackKey string,
|
||||
fallbackProvider string,
|
||||
fallbackModel string,
|
||||
execute func(context.Context, *FallbackChain, []FallbackCandidate,
|
||||
func(context.Context, string, string) (*LLMResponse, error),
|
||||
) (*FallbackResult, error),
|
||||
responseContent string,
|
||||
) {
|
||||
t.Helper()
|
||||
|
||||
ct := NewCooldownTracker()
|
||||
rl := NewRateLimiterRegistry()
|
||||
rl.Register(primaryKey, 1)
|
||||
if err := rl.Wait(context.Background(), primaryKey); err != nil {
|
||||
t.Fatalf("failed to pre-drain primary limiter: %v", err)
|
||||
}
|
||||
|
||||
fc := NewFallbackChain(ct, rl)
|
||||
candidates := []FallbackCandidate{
|
||||
{Provider: "openai", Model: "gpt-4o", IdentityKey: primaryKey},
|
||||
{Provider: fallbackProvider, Model: fallbackModel, IdentityKey: fallbackKey},
|
||||
}
|
||||
|
||||
run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
|
||||
if provider != fallbackProvider || model != fallbackModel {
|
||||
t.Fatalf("expected fallback candidate to run, got %s/%s", provider, model)
|
||||
}
|
||||
return &LLMResponse{Content: responseContent, FinishReason: "stop"}, nil
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
result, err := execute(ctx, fc, candidates, run)
|
||||
if err != nil {
|
||||
t.Fatalf("expected fallback success, got error: %v", err)
|
||||
}
|
||||
if result.Provider != fallbackProvider || result.Model != fallbackModel {
|
||||
t.Fatalf("result = %s/%s, want %s/%s", result.Provider, result.Model, fallbackProvider, fallbackModel)
|
||||
}
|
||||
if len(result.Attempts) != 1 || !result.Attempts[0].Skipped {
|
||||
t.Fatalf("expected one skipped primary attempt, got %+v", result.Attempts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFallback_LocalRateLimitSkipsToHealthyFallback(t *testing.T) {
|
||||
assertLocalRateLimitSkipsToHealthyFallback(
|
||||
t,
|
||||
"model_name:primary",
|
||||
"model_name:fallback",
|
||||
"anthropic",
|
||||
"claude",
|
||||
func(
|
||||
ctx context.Context,
|
||||
fc *FallbackChain,
|
||||
candidates []FallbackCandidate,
|
||||
run func(context.Context, string, string) (*LLMResponse, error),
|
||||
) (*FallbackResult, error) {
|
||||
return fc.Execute(ctx, candidates, run)
|
||||
},
|
||||
"fallback ok",
|
||||
)
|
||||
}
|
||||
|
||||
// --- Image Fallback Tests ---
|
||||
|
||||
func TestImageFallback_Success(t *testing.T) {
|
||||
ct := NewCooldownTracker()
|
||||
fc := NewFallbackChain(ct)
|
||||
fc := NewFallbackChain(ct, nil)
|
||||
|
||||
candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4o")}
|
||||
result, err := fc.ExecuteImage(context.Background(), candidates, successRun("image result"))
|
||||
|
|
@ -311,7 +378,7 @@ func TestImageFallback_Success(t *testing.T) {
|
|||
|
||||
func TestImageFallback_DimensionError(t *testing.T) {
|
||||
ct := NewCooldownTracker()
|
||||
fc := NewFallbackChain(ct)
|
||||
fc := NewFallbackChain(ct, nil)
|
||||
|
||||
candidates := []FallbackCandidate{
|
||||
makeCandidate("openai", "gpt-4o"),
|
||||
|
|
@ -335,7 +402,7 @@ func TestImageFallback_DimensionError(t *testing.T) {
|
|||
|
||||
func TestImageFallback_SizeError(t *testing.T) {
|
||||
ct := NewCooldownTracker()
|
||||
fc := NewFallbackChain(ct)
|
||||
fc := NewFallbackChain(ct, nil)
|
||||
|
||||
candidates := []FallbackCandidate{
|
||||
makeCandidate("openai", "gpt-4o"),
|
||||
|
|
@ -359,7 +426,7 @@ func TestImageFallback_SizeError(t *testing.T) {
|
|||
|
||||
func TestImageFallback_RetryOnOtherErrors(t *testing.T) {
|
||||
ct := NewCooldownTracker()
|
||||
fc := NewFallbackChain(ct)
|
||||
fc := NewFallbackChain(ct, nil)
|
||||
|
||||
candidates := []FallbackCandidate{
|
||||
makeCandidate("openai", "gpt-4o"),
|
||||
|
|
@ -384,9 +451,28 @@ func TestImageFallback_RetryOnOtherErrors(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestImageFallback_LocalRateLimitSkipsToHealthyFallback(t *testing.T) {
|
||||
assertLocalRateLimitSkipsToHealthyFallback(
|
||||
t,
|
||||
"model_name:primary-image",
|
||||
"model_name:fallback-image",
|
||||
"anthropic",
|
||||
"claude-sonnet",
|
||||
func(
|
||||
ctx context.Context,
|
||||
fc *FallbackChain,
|
||||
candidates []FallbackCandidate,
|
||||
run func(context.Context, string, string) (*LLMResponse, error),
|
||||
) (*FallbackResult, error) {
|
||||
return fc.ExecuteImage(ctx, candidates, run)
|
||||
},
|
||||
"image fallback ok",
|
||||
)
|
||||
}
|
||||
|
||||
func TestImageFallback_NoCandidates(t *testing.T) {
|
||||
ct := NewCooldownTracker()
|
||||
fc := NewFallbackChain(ct)
|
||||
fc := NewFallbackChain(ct, nil)
|
||||
|
||||
_, err := fc.ExecuteImage(context.Background(), nil, successRun("ok"))
|
||||
if err == nil {
|
||||
|
|
|
|||
|
|
@ -24,11 +24,11 @@ func NewHTTPProvider(apiKey, apiBase, proxy string) *HTTPProvider {
|
|||
}
|
||||
|
||||
func NewHTTPProviderWithMaxTokensField(apiKey, apiBase, proxy, maxTokensField string) *HTTPProvider {
|
||||
return NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(apiKey, apiBase, proxy, maxTokensField, 0, nil)
|
||||
return NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(apiKey, apiBase, proxy, maxTokensField, "", 0, nil)
|
||||
}
|
||||
|
||||
func NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(
|
||||
apiKey, apiBase, proxy, maxTokensField string,
|
||||
apiKey, apiBase, proxy, maxTokensField, userAgent string,
|
||||
requestTimeoutSeconds int,
|
||||
extraBody map[string]any,
|
||||
) *HTTPProvider {
|
||||
|
|
@ -40,6 +40,7 @@ func NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(
|
|||
openai_compat.WithMaxTokensField(maxTokensField),
|
||||
openai_compat.WithRequestTimeout(time.Duration(requestTimeoutSeconds)*time.Second),
|
||||
openai_compat.WithExtraBody(extraBody),
|
||||
openai_compat.WithUserAgent(userAgent),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ type Provider struct {
|
|||
maxTokensField string // Field name for max tokens (e.g., "max_completion_tokens" for o1/glm models)
|
||||
httpClient *http.Client
|
||||
extraBody map[string]any // Additional fields to inject into request body
|
||||
userAgent string
|
||||
}
|
||||
|
||||
type Option func(*Provider)
|
||||
|
|
@ -44,6 +45,7 @@ const defaultRequestTimeout = common.DefaultRequestTimeout
|
|||
|
||||
var stripModelPrefixProviders = map[string]struct{}{
|
||||
"litellm": {},
|
||||
"venice": {},
|
||||
"moonshot": {},
|
||||
"nvidia": {},
|
||||
"groq": {},
|
||||
|
|
@ -65,6 +67,12 @@ func WithMaxTokensField(maxTokensField string) Option {
|
|||
}
|
||||
}
|
||||
|
||||
func WithUserAgent(userAgent string) Option {
|
||||
return func(p *Provider) {
|
||||
p.userAgent = userAgent
|
||||
}
|
||||
}
|
||||
|
||||
func WithRequestTimeout(timeout time.Duration) Option {
|
||||
return func(p *Provider) {
|
||||
if timeout > 0 {
|
||||
|
|
@ -197,6 +205,9 @@ func (p *Provider) Chat(
|
|||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if p.userAgent != "" {
|
||||
req.Header.Set("User-Agent", p.userAgent)
|
||||
}
|
||||
if p.apiKey != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+p.apiKey)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -479,6 +479,11 @@ func TestProviderChat_StripsKnownProviderPrefixes(t *testing.T) {
|
|||
input: "lmstudio/openai/gpt-oss-20b",
|
||||
wantModel: "openai/gpt-oss-20b",
|
||||
},
|
||||
{
|
||||
name: "strips venice prefix",
|
||||
input: "venice/venice-uncensored",
|
||||
wantModel: "venice-uncensored",
|
||||
},
|
||||
{
|
||||
name: "strips deepseek prefix",
|
||||
input: "deepseek/deepseek-chat",
|
||||
|
|
@ -587,6 +592,9 @@ func TestNormalizeModel_UsesAPIBase(t *testing.T) {
|
|||
if got := normalizeModel("lmstudio/openai/gpt-oss-20b", "http://localhost:1234/v1"); got != "openai/gpt-oss-20b" {
|
||||
t.Fatalf("normalizeModel(lmstudio) = %q, want %q", got, "openai/gpt-oss-20b")
|
||||
}
|
||||
if got := normalizeModel("venice/venice-uncensored", "https://api.venice.ai/api/v1"); got != "venice-uncensored" {
|
||||
t.Fatalf("normalizeModel(venice) = %q, want %q", got, "venice-uncensored")
|
||||
}
|
||||
if got := normalizeModel("openrouter/auto", "https://openrouter.ai/api/v1"); got != "openrouter/auto" {
|
||||
t.Fatalf("normalizeModel(openrouter) = %q, want %q", got, "openrouter/auto")
|
||||
}
|
||||
|
|
|
|||
144
pkg/providers/ratelimiter.go
Normal file
144
pkg/providers/ratelimiter.go
Normal file
|
|
@ -0,0 +1,144 @@
|
|||
package providers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// RateLimiter implements a token-bucket rate limiter for a single key.
|
||||
// Allows up to RPM requests per minute with a burst equal to RPM.
|
||||
// Thread-safe.
|
||||
type RateLimiter struct {
|
||||
mu sync.Mutex
|
||||
rpm int
|
||||
tokens float64
|
||||
maxBurst float64
|
||||
lastTick time.Time
|
||||
nowFunc func() time.Time // for testing
|
||||
}
|
||||
|
||||
func (rl *RateLimiter) refillLocked(now time.Time) {
|
||||
elapsed := now.Sub(rl.lastTick).Seconds()
|
||||
rl.lastTick = now
|
||||
|
||||
// Refill tokens proportional to elapsed time.
|
||||
refill := elapsed * float64(rl.rpm) / 60.0
|
||||
rl.tokens = min(rl.maxBurst, rl.tokens+refill)
|
||||
}
|
||||
|
||||
// newRateLimiter creates a RateLimiter that allows rpm requests/minute.
|
||||
func newRateLimiter(rpm int) *RateLimiter {
|
||||
return &RateLimiter{
|
||||
rpm: rpm,
|
||||
tokens: float64(rpm), // start full
|
||||
maxBurst: float64(rpm),
|
||||
lastTick: time.Now(),
|
||||
nowFunc: time.Now,
|
||||
}
|
||||
}
|
||||
|
||||
// Wait blocks until a token is available or ctx is canceled.
|
||||
// Returns ctx.Err() if canceled while waiting.
|
||||
func (rl *RateLimiter) Wait(ctx context.Context) error {
|
||||
for {
|
||||
rl.mu.Lock()
|
||||
now := rl.nowFunc()
|
||||
rl.refillLocked(now)
|
||||
|
||||
if rl.tokens >= 1.0 {
|
||||
rl.tokens--
|
||||
rl.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
// Calculate how long until a token is available.
|
||||
deficit := 1.0 - rl.tokens
|
||||
waitSec := deficit / (float64(rl.rpm) / 60.0)
|
||||
rl.mu.Unlock()
|
||||
|
||||
timer := time.NewTimer(time.Duration(waitSec * float64(time.Second)))
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
if !timer.Stop() {
|
||||
<-timer.C
|
||||
}
|
||||
return ctx.Err()
|
||||
case <-timer.C:
|
||||
// Loop to re-check (another goroutine may have consumed the token).
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TryAcquire attempts to consume a token without blocking.
|
||||
func (rl *RateLimiter) TryAcquire() bool {
|
||||
rl.mu.Lock()
|
||||
defer rl.mu.Unlock()
|
||||
|
||||
rl.refillLocked(rl.nowFunc())
|
||||
if rl.tokens < 1.0 {
|
||||
return false
|
||||
}
|
||||
rl.tokens--
|
||||
return true
|
||||
}
|
||||
|
||||
// RateLimiterRegistry holds per-candidate rate limiters.
|
||||
// Candidates with RPM=0 are unrestricted.
|
||||
// Thread-safe for concurrent reads/writes.
|
||||
type RateLimiterRegistry struct {
|
||||
mu sync.RWMutex
|
||||
limiters map[string]*RateLimiter
|
||||
}
|
||||
|
||||
// NewRateLimiterRegistry creates an empty registry.
|
||||
func NewRateLimiterRegistry() *RateLimiterRegistry {
|
||||
return &RateLimiterRegistry{
|
||||
limiters: make(map[string]*RateLimiter),
|
||||
}
|
||||
}
|
||||
|
||||
// Register adds a rate limiter for the given key at the given RPM.
|
||||
// If rpm <= 0, no limiter is registered (unrestricted).
|
||||
func (r *RateLimiterRegistry) Register(key string, rpm int) {
|
||||
if rpm <= 0 {
|
||||
return
|
||||
}
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.limiters[key] = newRateLimiter(rpm)
|
||||
}
|
||||
|
||||
// Wait acquires a token for the given key, blocking if needed.
|
||||
// If no limiter is registered for key, returns immediately.
|
||||
func (r *RateLimiterRegistry) Wait(ctx context.Context, key string) error {
|
||||
r.mu.RLock()
|
||||
rl := r.limiters[key]
|
||||
r.mu.RUnlock()
|
||||
if rl == nil {
|
||||
return nil
|
||||
}
|
||||
return rl.Wait(ctx)
|
||||
}
|
||||
|
||||
// TryAcquire attempts to consume a token for the given key without blocking.
|
||||
// If no limiter is registered for key, it returns true.
|
||||
func (r *RateLimiterRegistry) TryAcquire(key string) bool {
|
||||
r.mu.RLock()
|
||||
rl := r.limiters[key]
|
||||
r.mu.RUnlock()
|
||||
if rl == nil {
|
||||
return true
|
||||
}
|
||||
return rl.TryAcquire()
|
||||
}
|
||||
|
||||
// RegisterCandidates registers rate limiters for all candidates that have RPM > 0.
|
||||
// Candidates with RPM == 0 are ignored (no restriction).
|
||||
func (r *RateLimiterRegistry) RegisterCandidates(candidates []FallbackCandidate) {
|
||||
for _, c := range candidates {
|
||||
if c.RPM > 0 {
|
||||
r.Register(c.StableKey(), c.RPM)
|
||||
}
|
||||
}
|
||||
}
|
||||
209
pkg/providers/ratelimiter_test.go
Normal file
209
pkg/providers/ratelimiter_test.go
Normal file
|
|
@ -0,0 +1,209 @@
|
|||
package providers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TestRateLimiter_AllowsUpToRPM verifies that up to RPM requests pass immediately
|
||||
// (burst capacity) and the (RPM+1)-th request is delayed.
|
||||
func TestRateLimiter_AllowsUpToRPM(t *testing.T) {
|
||||
rpm := 5
|
||||
rl := newRateLimiter(rpm)
|
||||
|
||||
// All rpm tokens should be available immediately (bucket starts full).
|
||||
for i := 0; i < rpm; i++ {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
||||
if err := rl.Wait(ctx); err != nil {
|
||||
t.Fatalf("request %d should pass immediately, got: %v", i+1, err)
|
||||
}
|
||||
cancel()
|
||||
}
|
||||
|
||||
// The next request must wait; cancel it to confirm it blocks.
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancel()
|
||||
err := rl.Wait(ctx)
|
||||
if err == nil {
|
||||
t.Fatal("expected request beyond RPM to block, but it passed immediately")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRateLimiter_ContextCancellation verifies that a blocked Wait respects cancellation.
|
||||
func TestRateLimiter_ContextCancellation(t *testing.T) {
|
||||
rl := newRateLimiter(1)
|
||||
|
||||
// Drain the one token.
|
||||
ctx := context.Background()
|
||||
if err := rl.Wait(ctx); err != nil {
|
||||
t.Fatalf("first request failed: %v", err)
|
||||
}
|
||||
|
||||
// Second request should block; cancel it.
|
||||
cancelCtx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
|
||||
defer cancel()
|
||||
err := rl.Wait(cancelCtx)
|
||||
if err == nil {
|
||||
t.Fatal("expected cancellation error, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRateLimiter_TokenRefill verifies that tokens refill over time.
|
||||
func TestRateLimiter_TokenRefill(t *testing.T) {
|
||||
rpm := 60 // 1 token per second
|
||||
rl := newRateLimiter(rpm)
|
||||
|
||||
// Drain all tokens.
|
||||
for i := 0; i < rpm; i++ {
|
||||
rl.Wait(context.Background()) //nolint:errcheck
|
||||
}
|
||||
|
||||
// Advance time via nowFunc: simulate 2 seconds passing (should give 2 tokens).
|
||||
start := time.Now()
|
||||
rl.nowFunc = func() time.Time { return start.Add(2 * time.Second) }
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
||||
defer cancel()
|
||||
if err := rl.Wait(ctx); err != nil {
|
||||
t.Fatalf("expected refilled token to be available: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRateLimiterRegistry_NoLimiter verifies that keys without a registered limiter pass freely.
|
||||
func TestRateLimiterRegistry_NoLimiter(t *testing.T) {
|
||||
r := NewRateLimiterRegistry()
|
||||
ctx := context.Background()
|
||||
for i := 0; i < 100; i++ {
|
||||
if err := r.Wait(ctx, "unregistered/key"); err != nil {
|
||||
t.Fatalf("unregistered key should not block: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestRateLimiterRegistry_ZeroRPM verifies that RPM=0 means no limiter is registered.
|
||||
func TestRateLimiterRegistry_ZeroRPM(t *testing.T) {
|
||||
r := NewRateLimiterRegistry()
|
||||
r.Register("some/key", 0)
|
||||
ctx := context.Background()
|
||||
for i := 0; i < 50; i++ {
|
||||
if err := r.Wait(ctx, "some/key"); err != nil {
|
||||
t.Fatalf("zero-RPM key should not block: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestRateLimiterRegistry_Enforcement verifies the registry enforces RPM per key.
|
||||
func TestRateLimiterRegistry_Enforcement(t *testing.T) {
|
||||
r := NewRateLimiterRegistry()
|
||||
r.Register("openai/gpt-4o", 3)
|
||||
|
||||
// First 3 calls should pass (burst = RPM).
|
||||
for i := 0; i < 3; i++ {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
||||
if err := r.Wait(ctx, "openai/gpt-4o"); err != nil {
|
||||
t.Fatalf("call %d should pass: %v", i+1, err)
|
||||
}
|
||||
cancel()
|
||||
}
|
||||
|
||||
// 4th call should block.
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancel()
|
||||
if err := r.Wait(ctx, "openai/gpt-4o"); err == nil {
|
||||
t.Fatal("4th call should have been rate-limited")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRateLimiterRegistry_RegisterCandidates verifies that RegisterCandidates
|
||||
// correctly picks up RPM from FallbackCandidate.
|
||||
func TestRateLimiterRegistry_RegisterCandidates(t *testing.T) {
|
||||
r := NewRateLimiterRegistry()
|
||||
candidates := []FallbackCandidate{
|
||||
{Provider: "openai", Model: "gpt-4o", RPM: 2},
|
||||
{Provider: "anthropic", Model: "claude-3", RPM: 0}, // no limit
|
||||
}
|
||||
r.RegisterCandidates(candidates)
|
||||
|
||||
// openai/gpt-4o: 2 tokens burst, 3rd should block.
|
||||
for i := 0; i < 2; i++ {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
||||
if err := r.Wait(ctx, "openai/gpt-4o"); err != nil {
|
||||
t.Fatalf("openai call %d should pass: %v", i+1, err)
|
||||
}
|
||||
cancel()
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancel()
|
||||
if err := r.Wait(ctx, "openai/gpt-4o"); err == nil {
|
||||
t.Fatal("openai 3rd call should have been limited")
|
||||
}
|
||||
|
||||
// anthropic/claude-3: no limit, should always pass.
|
||||
for i := 0; i < 10; i++ {
|
||||
if err := r.Wait(context.Background(), "anthropic/claude-3"); err != nil {
|
||||
t.Fatalf("anthropic call should not be limited: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimiterRegistry_RegisterCandidatesUsesStableIdentity(t *testing.T) {
|
||||
r := NewRateLimiterRegistry()
|
||||
candidates := []FallbackCandidate{
|
||||
{Provider: "openai", Model: "gpt-4o", RPM: 1, IdentityKey: "model_name:primary"},
|
||||
{Provider: "openai", Model: "gpt-4o", RPM: 2, IdentityKey: "model_name:fallback"},
|
||||
}
|
||||
r.RegisterCandidates(candidates)
|
||||
|
||||
if err := r.Wait(context.Background(), "model_name:primary"); err != nil {
|
||||
t.Fatalf("primary first call should pass: %v", err)
|
||||
}
|
||||
if err := r.Wait(context.Background(), "model_name:fallback"); err != nil {
|
||||
t.Fatalf("fallback first call should pass: %v", err)
|
||||
}
|
||||
if err := r.Wait(context.Background(), "model_name:fallback"); err != nil {
|
||||
t.Fatalf("fallback second call should pass: %v", err)
|
||||
}
|
||||
|
||||
ctxPrimary, cancelPrimary := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancelPrimary()
|
||||
if err := r.Wait(ctxPrimary, "model_name:primary"); err == nil {
|
||||
t.Fatal("primary second call should have been limited")
|
||||
}
|
||||
|
||||
ctxFallback, cancelFallback := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancelFallback()
|
||||
if err := r.Wait(ctxFallback, "model_name:fallback"); err == nil {
|
||||
t.Fatal("fallback third call should have been limited")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRateLimiter_Concurrency verifies thread safety under concurrent access.
|
||||
func TestRateLimiter_Concurrency(t *testing.T) {
|
||||
rpm := 20
|
||||
rl := newRateLimiter(rpm)
|
||||
var passed atomic.Int64
|
||||
var wg sync.WaitGroup
|
||||
|
||||
// Launch 30 goroutines; only ~20 should pass immediately.
|
||||
for i := 0; i < 30; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
|
||||
defer cancel()
|
||||
if rl.Wait(ctx) == nil {
|
||||
passed.Add(1)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
got := passed.Load()
|
||||
// Allow small timing slack: between rpm-2 and rpm+2.
|
||||
if got < int64(rpm-2) || got > int64(rpm+2) {
|
||||
t.Fatalf("expected ~%d immediate passes, got %d", rpm, got)
|
||||
}
|
||||
}
|
||||
|
|
@ -23,6 +23,12 @@ func buildCLIToolsPrompt(tools []ToolDefinition) string {
|
|||
)
|
||||
sb.WriteString("\n```\n\n")
|
||||
sb.WriteString("CRITICAL: The 'arguments' field MUST be a JSON-encoded STRING.\n\n")
|
||||
sb.WriteString("Escaping rules (what to type in `function.arguments`):\n")
|
||||
sb.WriteString("- Use `\\n` to represent a real newline character.\n")
|
||||
sb.WriteString("- Use `\\\\n` to represent a literal backslash+n sequence (`\\n`).\n")
|
||||
sb.WriteString(
|
||||
"- `function.arguments` is a JSON-encoded string, so quotes/backslashes must be escaped in the outer payload.\n\n",
|
||||
)
|
||||
sb.WriteString("### Tool Definitions:\n\n")
|
||||
|
||||
for _, tool := range tools {
|
||||
|
|
|
|||
7
pkg/seahorse/.omc/state/last-tool-error.json
Normal file
7
pkg/seahorse/.omc/state/last-tool-error.json
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
{
|
||||
"tool_name": "Bash",
|
||||
"tool_input_preview": "{\"command\":\"cd /home/yliu/repos/picoclaw && make lint 2>&1\",\"timeout\":120000}",
|
||||
"error": "Exit code 2\npkg/agent/context_seahorse_test.go:1027:1: File is not properly formatted (gci)\n\t\t\tEarliestAt: &now,\n^\n1 issues:\n* gci: 1\nmake: *** [Makefile:264: lint] Error 1",
|
||||
"timestamp": "2026-04-04T02:38:32.067Z",
|
||||
"retry_count": 6
|
||||
}
|
||||
58
pkg/seahorse/compact_until_under_test.go
Normal file
58
pkg/seahorse/compact_until_under_test.go
Normal file
|
|
@ -0,0 +1,58 @@
|
|||
package seahorse
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// =============================================================================
|
||||
// CompactUntilUnder iteration cap
|
||||
// =============================================================================
|
||||
|
||||
func TestCompactUntilUnderIterationCap(t *testing.T) {
|
||||
// Setup: create a conversation with so many tokens that compaction
|
||||
// will never reach the budget. The iteration cap prevents infinite loops.
|
||||
//
|
||||
// We use a mock CompleteFn that always returns the same content,
|
||||
// and a budget of 0 which tokens can never reach.
|
||||
// Without the cap, this would loop forever.
|
||||
|
||||
db := openTestDB(t)
|
||||
if err := runSchema(db); err != nil {
|
||||
t.Fatalf("migration: %v", err)
|
||||
}
|
||||
s := &Store{db: db}
|
||||
|
||||
conv, _ := s.GetOrCreateConversation(context.Background(), "agent:iter-cap")
|
||||
convID := conv.ConversationID
|
||||
|
||||
// Add many messages to ensure there's plenty to compact
|
||||
for i := 0; i < 40; i++ {
|
||||
m, _ := s.AddMessage(context.Background(), convID, "user",
|
||||
"this is a long message with lots of tokens to push context over budget", 100)
|
||||
s.AppendContextMessage(context.Background(), convID, m.ID)
|
||||
}
|
||||
|
||||
// A completeFn that always succeeds but returns non-reducing content
|
||||
mockComplete := func(ctx context.Context, prompt string, opts CompleteOptions) (string, error) {
|
||||
return "Summary that doesn't reduce tokens much.", nil
|
||||
}
|
||||
|
||||
ce, cancel := newTestCompactionEngineWithStore(s, mockComplete)
|
||||
defer cancel()
|
||||
|
||||
// Use budget=1 so tokens can never reach budget
|
||||
// (each message is 100 tokens, so 40 messages = 4000 tokens, budget 1 is unreachable)
|
||||
// The function should stop after maxCompactIterations, not loop forever
|
||||
ce.config = Config{} // ensure defaults
|
||||
|
||||
result, err := ce.CompactUntilUnder(context.Background(), convID, 1)
|
||||
if err != nil {
|
||||
// Should not error — should stop gracefully
|
||||
t.Fatalf("CompactUntilUnder with budget=0: %v", err)
|
||||
}
|
||||
|
||||
// The function should have completed within reasonable time
|
||||
// If it exceeded the cap, it would still return (not hang)
|
||||
_ = result
|
||||
}
|
||||
144
pkg/seahorse/parts_roundtrip_test.go
Normal file
144
pkg/seahorse/parts_roundtrip_test.go
Normal file
|
|
@ -0,0 +1,144 @@
|
|||
package seahorse
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// =============================================================================
|
||||
// Bug 1: formatMessagesForSummary ignores Parts
|
||||
// - formatMessagesForSummary only reads m.Content, empty for Part-based messages
|
||||
// - truncateSummary has same issue
|
||||
// =============================================================================
|
||||
|
||||
func TestFormatMessagesForSummaryIncludesParts(t *testing.T) {
|
||||
ts := time.Date(2025, 1, 1, 12, 0, 0, 0, time.UTC)
|
||||
|
||||
messages := []Message{
|
||||
{ID: 1, Role: "user", Content: "hello world", CreatedAt: ts},
|
||||
{
|
||||
ID: 2,
|
||||
Role: "assistant",
|
||||
Content: "", // empty — real content is in Parts
|
||||
Parts: []MessagePart{
|
||||
{Type: "text", Text: "I will run a command"},
|
||||
{Type: "tool_use", Name: "bash", Arguments: `{"command":"ls -la"}`, ToolCallID: "call_1"},
|
||||
},
|
||||
CreatedAt: ts.Add(time.Minute),
|
||||
},
|
||||
{
|
||||
ID: 3,
|
||||
Role: "tool",
|
||||
Content: "", // empty — real content is in Parts
|
||||
Parts: []MessagePart{
|
||||
{Type: "tool_result", Text: "file1.txt\nfile2.txt", ToolCallID: "call_1"},
|
||||
},
|
||||
CreatedAt: ts.Add(2 * time.Minute),
|
||||
},
|
||||
}
|
||||
|
||||
result := formatMessagesForSummary(messages)
|
||||
|
||||
// Must contain the plain text message
|
||||
if !contains(result, "hello world") {
|
||||
t.Error("formatMessagesForSummary: missing plain text content")
|
||||
}
|
||||
|
||||
// Must contain tool_use info (not blank)
|
||||
if !contains(result, "bash") || !contains(result, "ls -la") {
|
||||
t.Errorf("formatMessagesForSummary: tool_use info missing from Parts.\nGot:\n%s", result)
|
||||
}
|
||||
|
||||
// Must contain tool_result info (not blank)
|
||||
if !contains(result, "file1.txt") {
|
||||
t.Errorf("formatMessagesForSummary: tool_result text missing from Parts.\nGot:\n%s", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTruncateSummaryIncludesParts(t *testing.T) {
|
||||
messages := []Message{
|
||||
{ID: 1, Role: "user", Content: "run the tests", CreatedAt: time.Now()},
|
||||
{
|
||||
ID: 2,
|
||||
Role: "assistant",
|
||||
Content: "", // empty
|
||||
Parts: []MessagePart{
|
||||
{Type: "tool_use", Name: "bash", Arguments: `{"command":"go test ./..."}`, ToolCallID: "call_1"},
|
||||
},
|
||||
CreatedAt: time.Now(),
|
||||
},
|
||||
{
|
||||
ID: 3,
|
||||
Role: "tool",
|
||||
Content: "", // empty
|
||||
Parts: []MessagePart{
|
||||
{Type: "tool_result", Text: "PASS\nok 3.2s", ToolCallID: "call_1"},
|
||||
},
|
||||
CreatedAt: time.Now(),
|
||||
},
|
||||
}
|
||||
|
||||
result := truncateSummary(messages)
|
||||
|
||||
// Must contain plain text
|
||||
if !contains(result, "run the tests") {
|
||||
t.Error("truncateSummary: missing plain text content")
|
||||
}
|
||||
|
||||
// Must contain tool info from Parts (not blank)
|
||||
if !contains(result, "bash") || !contains(result, "go test") {
|
||||
t.Errorf("truncateSummary: tool_use info missing from Parts.\nGot:\n%s", result)
|
||||
}
|
||||
|
||||
// Must contain tool_result from Parts
|
||||
if !contains(result, "PASS") {
|
||||
t.Errorf("truncateSummary: tool_result text missing from Parts.\nGot:\n%s", result)
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Bug 2: SearchMessages cannot find Part-based messages
|
||||
// - FTS5 indexes empty content, LIKE queries empty content
|
||||
// =============================================================================
|
||||
|
||||
func TestSearchMessagesFindsPartBasedMessages(t *testing.T) {
|
||||
s := openTestStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
conv, _ := s.GetOrCreateConversation(ctx, "agent:search-parts")
|
||||
convID := conv.ConversationID
|
||||
|
||||
// Add a plain message (searchable)
|
||||
s.AddMessage(ctx, convID, "user", "list the files please", 5)
|
||||
|
||||
// Add a Part-based message (tool_use) — currently NOT searchable
|
||||
parts := []MessagePart{
|
||||
{Type: "tool_use", Name: "bash", Arguments: `{"command":"grep -r TODO ."}`, ToolCallID: "call_1"},
|
||||
}
|
||||
s.AddMessageWithParts(ctx, convID, "assistant", parts, 10)
|
||||
|
||||
// Add a Part-based message (tool_result) — currently NOT searchable
|
||||
resultParts := []MessagePart{
|
||||
{Type: "tool_result", Text: "main.go:42: TODO fix this bug", ToolCallID: "call_1"},
|
||||
}
|
||||
s.AddMessageWithParts(ctx, convID, "tool", resultParts, 10)
|
||||
|
||||
// Search for "grep" — should find the tool_use message
|
||||
results, err := s.SearchMessages(ctx, SearchInput{Pattern: "grep"})
|
||||
if err != nil {
|
||||
t.Fatalf("SearchMessages: %v", err)
|
||||
}
|
||||
if len(results) == 0 {
|
||||
t.Error("SearchMessages: 'grep' not found — Part-based messages are invisible to search")
|
||||
}
|
||||
|
||||
// Search for "TODO fix" — should find the tool_result message
|
||||
results2, err := s.SearchMessages(ctx, SearchInput{Pattern: "TODO fix"})
|
||||
if err != nil {
|
||||
t.Fatalf("SearchMessages: %v", err)
|
||||
}
|
||||
if len(results2) == 0 {
|
||||
t.Error("SearchMessages: 'TODO fix' not found — tool_result messages are invisible to search")
|
||||
}
|
||||
}
|
||||
185
pkg/seahorse/schema.go
Normal file
185
pkg/seahorse/schema.go
Normal file
|
|
@ -0,0 +1,185 @@
|
|||
package seahorse
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
)
|
||||
|
||||
// SQL statements for FTS5 tables with trigram tokenizer.
|
||||
const (
|
||||
sqlCreateSummariesFTS = `CREATE VIRTUAL TABLE IF NOT EXISTS summaries_fts USING fts5(
|
||||
summary_id,
|
||||
content,
|
||||
tokenize="trigram"
|
||||
)`
|
||||
sqlCreateMessagesFTS = `CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5(
|
||||
message_id,
|
||||
content,
|
||||
tokenize="trigram"
|
||||
)`
|
||||
sqlCheckFTS5Available = `CREATE VIRTUAL TABLE IF NOT EXISTS _fts5_check USING fts5(content)`
|
||||
sqlCheckTrigramAvailable = `CREATE VIRTUAL TABLE IF NOT EXISTS _trigram_check USING fts5(content, tokenize="trigram")`
|
||||
sqlDropFTS5Check = `DROP TABLE IF EXISTS _fts5_check`
|
||||
sqlDropTrigramCheck = `DROP TABLE IF EXISTS _trigram_check`
|
||||
)
|
||||
|
||||
// runSchema creates or upgrades the database schema.
|
||||
// All schemas are idempotent (safe to run multiple times).
|
||||
func runSchema(db *sql.DB) error {
|
||||
// Check FTS5 support before creating tables
|
||||
if err := checkFTS5Support(db); err != nil {
|
||||
return fmt.Errorf("FTS5 check: %w", err)
|
||||
}
|
||||
|
||||
stmts := []string{
|
||||
`CREATE TABLE IF NOT EXISTS conversations (
|
||||
conversation_id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
session_key TEXT NOT NULL UNIQUE,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
)`,
|
||||
|
||||
`CREATE TABLE IF NOT EXISTS messages (
|
||||
message_id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
conversation_id INTEGER NOT NULL REFERENCES conversations(conversation_id),
|
||||
role TEXT NOT NULL,
|
||||
content TEXT NOT NULL DEFAULT '',
|
||||
token_count INTEGER NOT NULL DEFAULT 0,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
)`,
|
||||
|
||||
`CREATE TABLE IF NOT EXISTS message_parts (
|
||||
part_id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
message_id INTEGER NOT NULL REFERENCES messages(message_id),
|
||||
type TEXT NOT NULL,
|
||||
text TEXT,
|
||||
name TEXT,
|
||||
arguments TEXT,
|
||||
tool_call_id TEXT,
|
||||
media_uri TEXT,
|
||||
mime_type TEXT,
|
||||
ordinal INTEGER NOT NULL DEFAULT 0
|
||||
)`,
|
||||
|
||||
`CREATE TABLE IF NOT EXISTS summaries (
|
||||
summary_id TEXT PRIMARY KEY,
|
||||
conversation_id INTEGER NOT NULL REFERENCES conversations(conversation_id),
|
||||
kind TEXT NOT NULL,
|
||||
depth INTEGER NOT NULL DEFAULT 0,
|
||||
content TEXT NOT NULL,
|
||||
token_count INTEGER NOT NULL DEFAULT 0,
|
||||
earliest_at TEXT,
|
||||
latest_at TEXT,
|
||||
descendant_count INTEGER NOT NULL DEFAULT 0,
|
||||
descendant_token_count INTEGER NOT NULL DEFAULT 0,
|
||||
source_message_token_count INTEGER NOT NULL DEFAULT 0,
|
||||
model TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
)`,
|
||||
|
||||
`CREATE TABLE IF NOT EXISTS summary_parents (
|
||||
summary_id TEXT NOT NULL,
|
||||
parent_summary_id TEXT NOT NULL,
|
||||
PRIMARY KEY (summary_id, parent_summary_id)
|
||||
)`,
|
||||
|
||||
`CREATE TABLE IF NOT EXISTS summary_messages (
|
||||
summary_id TEXT NOT NULL,
|
||||
message_id INTEGER NOT NULL,
|
||||
ordinal INTEGER NOT NULL DEFAULT 0,
|
||||
PRIMARY KEY (summary_id, message_id)
|
||||
)`,
|
||||
|
||||
`CREATE TABLE IF NOT EXISTS context_items (
|
||||
conversation_id INTEGER NOT NULL,
|
||||
ordinal INTEGER NOT NULL,
|
||||
item_type TEXT NOT NULL,
|
||||
summary_id TEXT,
|
||||
message_id INTEGER,
|
||||
token_count INTEGER NOT NULL DEFAULT 0,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
PRIMARY KEY (conversation_id, ordinal)
|
||||
)`,
|
||||
|
||||
// FTS5 virtual table with trigram tokenizer for CJK support
|
||||
sqlCreateSummariesFTS,
|
||||
|
||||
// FTS5 virtual table for message search with trigram tokenizer
|
||||
sqlCreateMessagesFTS,
|
||||
|
||||
// Indexes for common query patterns
|
||||
`CREATE INDEX IF NOT EXISTS idx_messages_conversation ON messages(conversation_id)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_messages_created ON messages(conversation_id, created_at)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_summaries_conversation ON summaries(conversation_id)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_summaries_kind_depth ON summaries(conversation_id, kind, depth)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_summary_parents_parent ON summary_parents(parent_summary_id)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_summary_messages_message ON summary_messages(message_id)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_context_items_conv ON context_items(conversation_id, ordinal)`,
|
||||
|
||||
// FTS5 triggers to keep summaries_fts in sync with summaries table
|
||||
`CREATE TRIGGER IF NOT EXISTS summaries_ai AFTER INSERT ON summaries BEGIN
|
||||
INSERT INTO summaries_fts (summary_id, content) VALUES (new.summary_id, new.content);
|
||||
END`,
|
||||
`CREATE TRIGGER IF NOT EXISTS summaries_ad AFTER DELETE ON summaries BEGIN
|
||||
INSERT INTO summaries_fts (summaries_fts, summary_id, content) VALUES ('delete', old.summary_id, old.content);
|
||||
END`,
|
||||
`CREATE TRIGGER IF NOT EXISTS summaries_au AFTER UPDATE ON summaries BEGIN
|
||||
INSERT INTO summaries_fts (summaries_fts, summary_id, content) VALUES ('delete', old.summary_id, old.content);
|
||||
INSERT INTO summaries_fts (summary_id, content) VALUES (new.summary_id, new.content);
|
||||
END`,
|
||||
|
||||
// FTS5 triggers to keep messages_fts in sync with messages table
|
||||
`CREATE TRIGGER IF NOT EXISTS messages_ai AFTER INSERT ON messages BEGIN
|
||||
INSERT INTO messages_fts (message_id, content) VALUES (new.message_id, new.content);
|
||||
END`,
|
||||
`CREATE TRIGGER IF NOT EXISTS messages_ad AFTER DELETE ON messages BEGIN
|
||||
DELETE FROM messages_fts WHERE message_id = old.message_id;
|
||||
END`,
|
||||
`CREATE TRIGGER IF NOT EXISTS messages_au AFTER UPDATE ON messages BEGIN
|
||||
DELETE FROM messages_fts WHERE message_id = old.message_id;
|
||||
INSERT INTO messages_fts (message_id, content) VALUES (new.message_id, new.content);
|
||||
END`,
|
||||
}
|
||||
|
||||
for _, s := range stmts {
|
||||
if _, err := db.Exec(s); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// checkFTS5Support verifies that SQLite has FTS5 with trigram tokenizer enabled.
|
||||
// This is required for full-text search with CJK (Chinese, Japanese, Korean) support.
|
||||
func checkFTS5Support(db *sql.DB) error {
|
||||
// Check if FTS5 is compiled in
|
||||
var fts5Enabled int
|
||||
err := db.QueryRow(`SELECT sqlite_compileoption_used('ENABLE_FTS5')`).Scan(&fts5Enabled)
|
||||
if err != nil {
|
||||
// sqlite_compileoption_used might not exist in older SQLite
|
||||
// Try a different approach: create a test FTS5 table
|
||||
_, testErr := db.Exec(sqlCheckFTS5Available)
|
||||
if testErr != nil {
|
||||
return fmt.Errorf("SQLite FTS5 not available: %w (required for full-text search)", testErr)
|
||||
}
|
||||
db.Exec(sqlDropFTS5Check)
|
||||
} else if fts5Enabled == 0 {
|
||||
return fmt.Errorf("SQLite was compiled without FTS5 support (required for full-text search)")
|
||||
}
|
||||
|
||||
// Check if trigram tokenizer is available by trying to create a test table
|
||||
// Not all SQLite builds include the trigram tokenizer
|
||||
_, err = db.Exec(sqlCheckTrigramAvailable)
|
||||
if err != nil {
|
||||
logger.WarnCF("seahorse", "SQLite trigram tokenizer not available, CJK search may be limited",
|
||||
map[string]any{"error": err.Error()})
|
||||
// Trigram is not strictly required, just better for CJK
|
||||
// Don't return error, just log warning
|
||||
} else {
|
||||
db.Exec(sqlDropTrigramCheck)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
211
pkg/seahorse/schema_test.go
Normal file
211
pkg/seahorse/schema_test.go
Normal file
|
|
@ -0,0 +1,211 @@
|
|||
package seahorse
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"testing"
|
||||
|
||||
_ "modernc.org/sqlite"
|
||||
)
|
||||
|
||||
func openTestDB(t *testing.T) *sql.DB {
|
||||
t.Helper()
|
||||
db, err := sql.Open("sqlite", ":memory:")
|
||||
if err != nil {
|
||||
t.Fatalf("open test db: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { db.Close() })
|
||||
return db
|
||||
}
|
||||
|
||||
func TestRunMigrations(t *testing.T) {
|
||||
db := openTestDB(t)
|
||||
|
||||
if err := runSchema(db); err != nil {
|
||||
t.Fatalf("runSchema: %v", err)
|
||||
}
|
||||
|
||||
// Verify all tables exist
|
||||
tables := []string{
|
||||
"conversations",
|
||||
"messages",
|
||||
"message_parts",
|
||||
"summaries",
|
||||
"summary_parents",
|
||||
"summary_messages",
|
||||
"context_items",
|
||||
}
|
||||
for _, tbl := range tables {
|
||||
var name string
|
||||
err := db.QueryRow(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' AND name=?", tbl,
|
||||
).Scan(&name)
|
||||
if err != nil {
|
||||
t.Errorf("table %q not found: %v", tbl, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Verify FTS5 virtual table exists
|
||||
var ftsName string
|
||||
err := db.QueryRow(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' AND name='summaries_fts'",
|
||||
).Scan(&ftsName)
|
||||
if err != nil {
|
||||
t.Errorf("FTS5 table summaries_fts not found: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunMigrationsIdempotent(t *testing.T) {
|
||||
db := openTestDB(t)
|
||||
|
||||
// Run migrations twice — should succeed both times
|
||||
if err := runSchema(db); err != nil {
|
||||
t.Fatalf("first migration: %v", err)
|
||||
}
|
||||
if err := runSchema(db); err != nil {
|
||||
t.Fatalf("second migration (idempotent): %v", err)
|
||||
}
|
||||
|
||||
// Verify we can still insert data after double migration
|
||||
res, err := db.Exec(
|
||||
"INSERT INTO conversations (session_key, created_at, updated_at) VALUES (?, datetime('now'), datetime('now'))",
|
||||
"test-session",
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("insert after double migration: %v", err)
|
||||
}
|
||||
id, _ := res.LastInsertId()
|
||||
if id == 0 {
|
||||
t.Error("expected non-zero conversation id")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMigrationConversationUnique(t *testing.T) {
|
||||
db := openTestDB(t)
|
||||
if err := runSchema(db); err != nil {
|
||||
t.Fatalf("migration: %v", err)
|
||||
}
|
||||
|
||||
// Insert first
|
||||
_, err := db.Exec(
|
||||
"INSERT INTO conversations (session_key, created_at, updated_at) VALUES (?, datetime('now'), datetime('now'))",
|
||||
"unique-key",
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("first insert: %v", err)
|
||||
}
|
||||
|
||||
// Duplicate should fail
|
||||
_, err = db.Exec(
|
||||
"INSERT INTO conversations (session_key, created_at, updated_at) VALUES (?, datetime('now'), datetime('now'))",
|
||||
"unique-key",
|
||||
)
|
||||
if err == nil {
|
||||
t.Error("expected unique constraint violation for duplicate session_key")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMigrationSummaryFTSInsert(t *testing.T) {
|
||||
db := openTestDB(t)
|
||||
if err := runSchema(db); err != nil {
|
||||
t.Fatalf("migration: %v", err)
|
||||
}
|
||||
|
||||
// Insert a conversation first
|
||||
_, err := db.Exec(
|
||||
"INSERT INTO conversations (session_key, created_at, updated_at) VALUES (?, datetime('now'), datetime('now'))",
|
||||
"fts-test",
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("insert conversation: %v", err)
|
||||
}
|
||||
|
||||
// Insert a summary
|
||||
_, err = db.Exec(
|
||||
`INSERT INTO summaries (summary_id, conversation_id, kind, depth, content, token_count, created_at)
|
||||
VALUES ('sum_test1', 1, 'leaf', 0, '你好世界 hello world', 10, datetime('now'))`)
|
||||
if err != nil {
|
||||
t.Fatalf("insert summary: %v", err)
|
||||
}
|
||||
|
||||
// FTS should find it — trigram tokenizer requires >= 3 chars
|
||||
rows, err := db.Query(
|
||||
"SELECT summary_id FROM summaries_fts WHERE summaries_fts MATCH ?",
|
||||
"你好世",
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("FTS query: %v", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var found string
|
||||
if rows.Next() {
|
||||
if err := rows.Scan(&found); err != nil {
|
||||
t.Fatalf("scan: %v", err)
|
||||
}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
t.Fatalf("rows.Err: %v", err)
|
||||
}
|
||||
if found != "sum_test1" {
|
||||
t.Errorf("FTS: expected 'sum_test1', got %q", found)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMigrationSummaryParentsPK(t *testing.T) {
|
||||
db := openTestDB(t)
|
||||
if err := runSchema(db); err != nil {
|
||||
t.Fatalf("migration: %v", err)
|
||||
}
|
||||
|
||||
// Insert two summaries
|
||||
for _, id := range []string{"sum_a", "sum_b"} {
|
||||
_, err := db.Exec(
|
||||
`INSERT INTO summaries (summary_id, conversation_id, kind, depth, content, token_count, created_at)
|
||||
VALUES (?, 1, 'leaf', 0, 'content', 5, datetime('now'))`, id)
|
||||
if err != nil {
|
||||
t.Fatalf("insert summary %s: %v", id, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Link child to parent
|
||||
_, err := db.Exec(
|
||||
"INSERT INTO summary_parents (summary_id, parent_summary_id) VALUES ('sum_a', 'sum_b')")
|
||||
if err != nil {
|
||||
t.Fatalf("link: %v", err)
|
||||
}
|
||||
|
||||
// Duplicate link should fail (composite PK)
|
||||
_, err = db.Exec(
|
||||
"INSERT INTO summary_parents (summary_id, parent_summary_id) VALUES ('sum_a', 'sum_b')")
|
||||
if err == nil {
|
||||
t.Error("expected unique constraint violation for duplicate summary_parents link")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFTS5SQLConstants(t *testing.T) {
|
||||
db := openTestDB(t)
|
||||
|
||||
// Verify FTS5 check SQL executes without error
|
||||
_, err := db.Exec(sqlCheckFTS5Available)
|
||||
if err != nil {
|
||||
t.Errorf("sqlCheckFTS5Available failed: %v", err)
|
||||
}
|
||||
|
||||
// Verify trigram check SQL executes without error
|
||||
_, err = db.Exec(sqlCheckTrigramAvailable)
|
||||
if err != nil {
|
||||
t.Errorf("sqlCheckTrigramAvailable failed: %v", err)
|
||||
}
|
||||
|
||||
// Verify summaries_fts SQL executes without error
|
||||
_, err = db.Exec(sqlCreateSummariesFTS)
|
||||
if err != nil {
|
||||
t.Errorf("sqlCreateSummariesFTS failed: %v", err)
|
||||
}
|
||||
|
||||
// Verify messages_fts SQL executes without error
|
||||
_, err = db.Exec(sqlCreateMessagesFTS)
|
||||
if err != nil {
|
||||
t.Errorf("sqlCreateMessagesFTS failed: %v", err)
|
||||
}
|
||||
}
|
||||
261
pkg/seahorse/short_assembler.go
Normal file
261
pkg/seahorse/short_assembler.go
Normal file
|
|
@ -0,0 +1,261 @@
|
|||
package seahorse
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
)
|
||||
|
||||
// escapeXML escapes special characters for safe inclusion in XML content.
|
||||
func escapeXML(s string) string {
|
||||
s = strings.ReplaceAll(s, "&", "&")
|
||||
s = strings.ReplaceAll(s, "<", "<")
|
||||
s = strings.ReplaceAll(s, ">", ">")
|
||||
s = strings.ReplaceAll(s, "\"", """)
|
||||
s = strings.ReplaceAll(s, "'", "'")
|
||||
return s
|
||||
}
|
||||
|
||||
// resolvedItem is a context item resolved to its full content with token count.
|
||||
type resolvedItem struct {
|
||||
ordinal int
|
||||
itemType string // "message" or "summary"
|
||||
message *Message
|
||||
summary *Summary
|
||||
tokenCount int
|
||||
}
|
||||
|
||||
// Assemble builds budget-constrained context from summaries + messages.
|
||||
//
|
||||
// Algorithm:
|
||||
// 1. Fetch context_items, resolve to full content
|
||||
// 2. Split into evictable prefix + protected fresh tail
|
||||
// 3. If evictable fits in remaining budget → include all
|
||||
// 4. Else walk evictable from newest to oldest, keep while fits
|
||||
func (a *Assembler) Assemble(ctx context.Context, convID int64, input AssembleInput) (*AssembleResult, error) {
|
||||
items, err := a.store.GetContextItems(ctx, convID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get context items: %w", err)
|
||||
}
|
||||
if len(items) == 0 {
|
||||
return &AssembleResult{}, nil
|
||||
}
|
||||
|
||||
// Resolve all items
|
||||
resolved := make([]resolvedItem, len(items))
|
||||
for i, item := range items {
|
||||
r, err := a.resolveItem(ctx, item)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
resolved[i] = r
|
||||
}
|
||||
|
||||
// Split into evictable prefix and protected fresh tail
|
||||
tailStart := len(resolved) - FreshTailCount
|
||||
if tailStart < 0 {
|
||||
tailStart = 0
|
||||
}
|
||||
evictable := resolved[:tailStart]
|
||||
freshTail := resolved[tailStart:]
|
||||
|
||||
// Calculate fresh tail tokens
|
||||
freshTailTokens := 0
|
||||
for _, r := range freshTail {
|
||||
freshTailTokens += r.tokenCount
|
||||
}
|
||||
|
||||
// Budget-aware selection of evictable items
|
||||
remainingBudget := input.Budget - freshTailTokens
|
||||
if remainingBudget < 0 {
|
||||
// Fresh tail alone exceeds budget - we keep it anyway (design decision)
|
||||
// Log for debugging retry/overflow issues
|
||||
logger.InfoCF("seahorse", "assemble: fresh tail exceeds budget", map[string]any{
|
||||
"budget": input.Budget,
|
||||
"fresh_tail_tokens": freshTailTokens,
|
||||
"fresh_tail_count": len(freshTail),
|
||||
"over_budget_by": freshTailTokens - input.Budget,
|
||||
})
|
||||
remainingBudget = 0
|
||||
}
|
||||
|
||||
var selected []resolvedItem
|
||||
evictableTokens := 0
|
||||
for _, r := range evictable {
|
||||
evictableTokens += r.tokenCount
|
||||
}
|
||||
|
||||
if evictableTokens <= remainingBudget {
|
||||
// All evictable fit
|
||||
selected = append(selected, evictable...)
|
||||
} else {
|
||||
// Walk from newest to oldest, keep while fits
|
||||
var kept []resolvedItem
|
||||
accum := 0
|
||||
for i := len(evictable) - 1; i >= 0; i-- {
|
||||
if accum+evictable[i].tokenCount <= remainingBudget {
|
||||
kept = append(kept, evictable[i])
|
||||
accum += evictable[i].tokenCount
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
// Reverse to restore chronological order
|
||||
for i, j := 0, len(kept)-1; i < j; i, j = i+1, j-1 {
|
||||
kept[i], kept[j] = kept[j], kept[i]
|
||||
}
|
||||
selected = append(selected, kept...)
|
||||
}
|
||||
|
||||
// Combine: selected evictable + fresh tail
|
||||
final := append(selected, freshTail...)
|
||||
|
||||
// Build result
|
||||
var messages []Message
|
||||
var summaries []Summary
|
||||
var sourceIDs []string
|
||||
totalTokens := 0
|
||||
maxDepth := 0
|
||||
condensedCount := 0
|
||||
|
||||
for _, r := range final {
|
||||
totalTokens += r.tokenCount
|
||||
if r.itemType == "message" && r.message != nil {
|
||||
messages = append(messages, *r.message)
|
||||
sourceIDs = append(sourceIDs, fmt.Sprintf("msg:%d", r.message.ID))
|
||||
} else if r.itemType == "summary" && r.summary != nil {
|
||||
summaries = append(summaries, *r.summary)
|
||||
if r.summary.Depth > maxDepth {
|
||||
maxDepth = r.summary.Depth
|
||||
}
|
||||
if r.summary.Kind == SummaryKindCondensed {
|
||||
condensedCount++
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Build depth-aware system prompt addition
|
||||
systemPromptAddition := ""
|
||||
if len(summaries) > 0 {
|
||||
if maxDepth >= 2 || condensedCount >= 2 {
|
||||
systemPromptAddition = "Your context has been heavily compressed through multi-level summarization.\n" +
|
||||
"- Do NOT assert specific facts (commands, SHAs, paths, timestamps) from summaries without expanding.\n" +
|
||||
"- When uncertain, use expand to recover original detail before making claims.\n" +
|
||||
"- Tool escalation: grep \xe2\x86\x92 describe \xe2\x86\x92 expand"
|
||||
} else {
|
||||
systemPromptAddition = "Some earlier messages have been summarized. Use expand tools to recover details if needed."
|
||||
}
|
||||
}
|
||||
|
||||
// Build Summary field: all XML summaries + system prompt addition
|
||||
var summaryParts []string
|
||||
for _, sum := range summaries {
|
||||
if sum.Content == "" {
|
||||
continue
|
||||
}
|
||||
// Load parent IDs for XML formatting
|
||||
parentSummaries, err := a.store.GetSummaryParents(ctx, sum.SummaryID)
|
||||
if err != nil {
|
||||
logger.WarnCF("seahorse", "assemble: get summary parents", map[string]any{
|
||||
"summary_id": sum.SummaryID,
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
var parentIDs []string
|
||||
for _, ps := range parentSummaries {
|
||||
parentIDs = append(parentIDs, ps.SummaryID)
|
||||
}
|
||||
summaryParts = append(summaryParts, FormatSummaryXML(&sum, parentIDs))
|
||||
}
|
||||
summary := strings.Join(summaryParts, "\n\n")
|
||||
if systemPromptAddition != "" {
|
||||
if summary != "" {
|
||||
summary += "\n\n"
|
||||
}
|
||||
summary += systemPromptAddition
|
||||
}
|
||||
|
||||
return &AssembleResult{
|
||||
Messages: messages,
|
||||
Summary: summary,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// resolveItem loads the full message or summary for a context item.
|
||||
func (a *Assembler) resolveItem(ctx context.Context, item ContextItem) (resolvedItem, error) {
|
||||
if item.ItemType == "message" {
|
||||
msg, err := a.store.GetMessageByID(ctx, item.MessageID)
|
||||
if err != nil {
|
||||
return resolvedItem{}, err
|
||||
}
|
||||
tokens := item.TokenCount
|
||||
if tokens == 0 {
|
||||
tokens = msg.TokenCount
|
||||
}
|
||||
return resolvedItem{
|
||||
ordinal: item.Ordinal,
|
||||
itemType: "message",
|
||||
message: msg,
|
||||
tokenCount: tokens,
|
||||
}, nil
|
||||
}
|
||||
|
||||
if item.ItemType == "summary" {
|
||||
sum, err := a.store.GetSummary(ctx, item.SummaryID)
|
||||
if err != nil {
|
||||
return resolvedItem{}, err
|
||||
}
|
||||
tokens := item.TokenCount
|
||||
if tokens == 0 {
|
||||
tokens = sum.TokenCount
|
||||
}
|
||||
return resolvedItem{
|
||||
ordinal: item.Ordinal,
|
||||
itemType: "summary",
|
||||
summary: sum,
|
||||
tokenCount: tokens,
|
||||
}, nil
|
||||
}
|
||||
|
||||
return resolvedItem{
|
||||
ordinal: item.Ordinal,
|
||||
itemType: item.ItemType,
|
||||
tokenCount: item.TokenCount,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// FormatSummaryXML formats a summary as XML for LLM context.
|
||||
// This is exported so context managers can format summaries consistently.
|
||||
func FormatSummaryXML(s *Summary, parentIDs []string) string {
|
||||
// Build time attributes if available
|
||||
var attrs string
|
||||
if s.EarliestAt != nil {
|
||||
attrs += fmt.Sprintf(` earliest_at="%s"`, s.EarliestAt.Format(time.RFC3339))
|
||||
}
|
||||
if s.LatestAt != nil {
|
||||
attrs += fmt.Sprintf(` latest_at="%s"`, s.LatestAt.Format(time.RFC3339))
|
||||
}
|
||||
|
||||
var parentsSection string
|
||||
if s.Kind == SummaryKindCondensed && len(parentIDs) > 0 {
|
||||
parents := "<parents>\n"
|
||||
for _, pid := range parentIDs {
|
||||
parents += fmt.Sprintf(" <summary_ref id=\"%s\" />\n", pid)
|
||||
}
|
||||
parents += " </parents>\n"
|
||||
parentsSection = parents
|
||||
}
|
||||
return fmt.Sprintf(
|
||||
"<summary id=\"%s\" kind=\"%s\" depth=\"%d\" descendant_count=\"%d\"%s>\n <content>\n %s\n </content>\n%s</summary>",
|
||||
s.SummaryID,
|
||||
string(s.Kind),
|
||||
s.Depth,
|
||||
s.DescendantCount,
|
||||
attrs,
|
||||
escapeXML(s.Content),
|
||||
parentsSection,
|
||||
)
|
||||
}
|
||||
536
pkg/seahorse/short_assembler_test.go
Normal file
536
pkg/seahorse/short_assembler_test.go
Normal file
|
|
@ -0,0 +1,536 @@
|
|||
package seahorse
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// --- Assembler Tests ---
|
||||
|
||||
// helper: create a store with messages and summaries for assembly tests
|
||||
func setupAssemblerStore(t *testing.T) (*Store, int64) {
|
||||
t.Helper()
|
||||
s := openTestStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
conv, err := s.GetOrCreateConversation(ctx, "test:assemble")
|
||||
if err != nil {
|
||||
t.Fatalf("create conversation: %v", err)
|
||||
}
|
||||
|
||||
return s, conv.ConversationID
|
||||
}
|
||||
|
||||
func TestAssemblerAssembleEmpty(t *testing.T) {
|
||||
s, convID := setupAssemblerStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
a := &Assembler{store: s, config: Config{}}
|
||||
result, err := a.Assemble(ctx, convID, AssembleInput{Budget: 1000})
|
||||
if err != nil {
|
||||
t.Fatalf("Assemble: %v", err)
|
||||
}
|
||||
if len(result.Messages) != 0 {
|
||||
t.Errorf("Messages = %d, want 0", len(result.Messages))
|
||||
}
|
||||
if result.Summary != "" {
|
||||
t.Errorf("Summary = %q, want empty", result.Summary)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssemblerAssembleMessagesOnly(t *testing.T) {
|
||||
s, convID := setupAssemblerStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create messages
|
||||
msg1, _ := s.AddMessage(ctx, convID, "user", "hello", 5)
|
||||
msg2, _ := s.AddMessage(ctx, convID, "assistant", "world", 5)
|
||||
|
||||
// Create context items
|
||||
s.UpsertContextItems(ctx, convID, []ContextItem{
|
||||
{Ordinal: 100, ItemType: "message", MessageID: msg1.ID, TokenCount: 5},
|
||||
{Ordinal: 200, ItemType: "message", MessageID: msg2.ID, TokenCount: 5},
|
||||
})
|
||||
|
||||
a := &Assembler{store: s, config: Config{}}
|
||||
result, err := a.Assemble(ctx, convID, AssembleInput{Budget: 100})
|
||||
if err != nil {
|
||||
t.Fatalf("Assemble: %v", err)
|
||||
}
|
||||
|
||||
if len(result.Messages) != 2 {
|
||||
t.Fatalf("Messages = %d, want 2", len(result.Messages))
|
||||
}
|
||||
if result.Messages[0].Content != "hello" {
|
||||
t.Errorf("Messages[0].Content = %q, want 'hello'", result.Messages[0].Content)
|
||||
}
|
||||
if result.Messages[1].Content != "world" {
|
||||
t.Errorf("Messages[1].Content = %q, want 'world'", result.Messages[1].Content)
|
||||
}
|
||||
// No summaries, so Summary should be empty
|
||||
if result.Summary != "" {
|
||||
t.Errorf("Summary = %q, want empty", result.Summary)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssemblerAssembleWithSummary(t *testing.T) {
|
||||
s, convID := setupAssemblerStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a summary
|
||||
summary, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: "summary of early messages",
|
||||
TokenCount: 50,
|
||||
})
|
||||
|
||||
// Create recent messages
|
||||
msg1, _ := s.AddMessage(ctx, convID, "user", "recent", 5)
|
||||
msg2, _ := s.AddMessage(ctx, convID, "assistant", "reply", 5)
|
||||
|
||||
// Context: summary + recent messages
|
||||
s.UpsertContextItems(ctx, convID, []ContextItem{
|
||||
{Ordinal: 100, ItemType: "summary", SummaryID: summary.SummaryID, TokenCount: 50},
|
||||
{Ordinal: 200, ItemType: "message", MessageID: msg1.ID, TokenCount: 5},
|
||||
{Ordinal: 300, ItemType: "message", MessageID: msg2.ID, TokenCount: 5},
|
||||
})
|
||||
|
||||
a := &Assembler{store: s, config: Config{}}
|
||||
result, err := a.Assemble(ctx, convID, AssembleInput{Budget: 1000})
|
||||
if err != nil {
|
||||
t.Fatalf("Assemble: %v", err)
|
||||
}
|
||||
|
||||
// Messages = 2 raw messages (summaries are in Summary field, not Messages)
|
||||
if len(result.Messages) != 2 {
|
||||
t.Errorf("Messages = %d, want 2 (raw messages only)", len(result.Messages))
|
||||
}
|
||||
// Summary should contain XML with summary content
|
||||
if result.Summary == "" {
|
||||
t.Error("Summary should not be empty when summary exists")
|
||||
}
|
||||
if !strings.Contains(result.Summary, summary.Content) {
|
||||
t.Errorf("Summary should contain summary content %q", summary.Content)
|
||||
}
|
||||
if !strings.Contains(result.Summary, "<summary") {
|
||||
t.Error("Summary should contain <summary XML tag")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssemblerBudgetEvictsOldest(t *testing.T) {
|
||||
s, convID := setupAssemblerStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create 40 messages, each with 10 tokens = 400 total
|
||||
msgs := make([]*Message, 40)
|
||||
for i := 0; i < 40; i++ {
|
||||
m, _ := s.AddMessage(ctx, convID, "user", "msg", 10)
|
||||
msgs[i] = m
|
||||
}
|
||||
|
||||
// Context items for all messages
|
||||
items := make([]ContextItem, 40)
|
||||
for i := 0; i < 40; i++ {
|
||||
items[i] = ContextItem{
|
||||
Ordinal: (i + 1) * 100,
|
||||
ItemType: "message",
|
||||
MessageID: msgs[i].ID,
|
||||
TokenCount: 10,
|
||||
}
|
||||
}
|
||||
s.UpsertContextItems(ctx, convID, items)
|
||||
|
||||
// Budget of 200 tokens with FreshTailCount=32
|
||||
// Fresh tail = last 32 messages (320 tokens, over budget, but always included)
|
||||
// Evictable = first 8 messages (80 tokens)
|
||||
// Budget after tail: max(0, 200-320) = 0 → no evictable items included
|
||||
a := &Assembler{store: s, config: Config{}}
|
||||
result, err := a.Assemble(ctx, convID, AssembleInput{Budget: 200})
|
||||
if err != nil {
|
||||
t.Fatalf("Assemble: %v", err)
|
||||
}
|
||||
|
||||
// Should only include the 32-item fresh tail
|
||||
if len(result.Messages) != 32 {
|
||||
t.Errorf("Messages = %d, want 32 (fresh tail)", len(result.Messages))
|
||||
}
|
||||
// Should be the LAST 32 messages
|
||||
if result.Messages[0].ID != msgs[8].ID {
|
||||
t.Errorf("first message ID = %d, want %d (msgs[8])", result.Messages[0].ID, msgs[8].ID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssemblerBudgetFitsAll(t *testing.T) {
|
||||
s, convID := setupAssemblerStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
msgs := make([]*Message, 5)
|
||||
for i := 0; i < 5; i++ {
|
||||
m, _ := s.AddMessage(ctx, convID, "user", "msg", 10)
|
||||
msgs[i] = m
|
||||
}
|
||||
|
||||
items := make([]ContextItem, 5)
|
||||
for i := 0; i < 5; i++ {
|
||||
items[i] = ContextItem{
|
||||
Ordinal: (i + 1) * 100,
|
||||
ItemType: "message",
|
||||
MessageID: msgs[i].ID,
|
||||
TokenCount: 10,
|
||||
}
|
||||
}
|
||||
s.UpsertContextItems(ctx, convID, items)
|
||||
|
||||
// Budget = 100, total = 50, FreshTailCount=32 → all items in tail
|
||||
a := &Assembler{store: s, config: Config{}}
|
||||
result, err := a.Assemble(ctx, convID, AssembleInput{Budget: 100})
|
||||
if err != nil {
|
||||
t.Fatalf("Assemble: %v", err)
|
||||
}
|
||||
|
||||
if len(result.Messages) != 5 {
|
||||
t.Errorf("Messages = %d, want 5", len(result.Messages))
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssemblerSummaryXMLFormat(t *testing.T) {
|
||||
s, convID := setupAssemblerStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
summary, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: "test summary content",
|
||||
TokenCount: 20,
|
||||
})
|
||||
|
||||
msg, _ := s.AddMessage(ctx, convID, "user", "hello", 5)
|
||||
|
||||
s.UpsertContextItems(ctx, convID, []ContextItem{
|
||||
{Ordinal: 100, ItemType: "summary", SummaryID: summary.SummaryID, TokenCount: 20},
|
||||
{Ordinal: 200, ItemType: "message", MessageID: msg.ID, TokenCount: 5},
|
||||
})
|
||||
|
||||
a := &Assembler{store: s, config: Config{}}
|
||||
result, err := a.Assemble(ctx, convID, AssembleInput{Budget: 1000})
|
||||
if err != nil {
|
||||
t.Fatalf("Assemble: %v", err)
|
||||
}
|
||||
|
||||
// Messages should only contain raw messages (no XML summary in Messages)
|
||||
if len(result.Messages) != 1 {
|
||||
t.Errorf("Messages = %d, want 1 (raw message only)", len(result.Messages))
|
||||
}
|
||||
// Summary should contain XML with summary content
|
||||
if result.Summary == "" {
|
||||
t.Fatal("Summary should not be empty")
|
||||
}
|
||||
if !contains(result.Summary, "<summary") {
|
||||
t.Errorf("Summary missing <summary tag: %q", result.Summary)
|
||||
}
|
||||
if !contains(result.Summary, summary.SummaryID) {
|
||||
t.Errorf("Summary missing summary ID: %q", result.Summary)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssemblerSummaryXMLEscaping(t *testing.T) {
|
||||
// Summary content with special XML characters should be properly escaped
|
||||
s, convID := setupAssemblerStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create summary with content containing XML special characters
|
||||
summary, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: `User said: "hello" & asked about <tags>`,
|
||||
TokenCount: 20,
|
||||
})
|
||||
|
||||
s.UpsertContextItems(ctx, convID, []ContextItem{
|
||||
{Ordinal: 100, ItemType: "summary", SummaryID: summary.SummaryID, TokenCount: 20},
|
||||
})
|
||||
|
||||
a := &Assembler{store: s, config: Config{}}
|
||||
result, err := a.Assemble(ctx, convID, AssembleInput{Budget: 1000})
|
||||
if err != nil {
|
||||
t.Fatalf("Assemble: %v", err)
|
||||
}
|
||||
|
||||
// Summary field should contain XML with escaped special characters
|
||||
if result.Summary == "" {
|
||||
t.Fatal("Summary should not be empty")
|
||||
}
|
||||
|
||||
// Check that special characters are escaped
|
||||
if strings.Contains(result.Summary, "<tags>") {
|
||||
t.Errorf("BUG: unescaped < in summary content: %q", result.Summary)
|
||||
}
|
||||
if strings.Contains(result.Summary, `"hello"`) {
|
||||
t.Errorf("BUG: unescaped \" in summary content: %q", result.Summary)
|
||||
}
|
||||
// & should be escaped as &
|
||||
if strings.Contains(result.Summary, " & ") {
|
||||
t.Errorf("BUG: unescaped & in summary content: %q", result.Summary)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssemblerSummaryXMLWithParents(t *testing.T) {
|
||||
s, convID := setupAssemblerStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a leaf and a condensed summary (condensed has parent)
|
||||
leaf, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: "leaf content",
|
||||
TokenCount: 20,
|
||||
})
|
||||
condensed, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindCondensed,
|
||||
Depth: 1,
|
||||
Content: "condensed content",
|
||||
TokenCount: 15,
|
||||
ParentIDs: []string{leaf.SummaryID},
|
||||
})
|
||||
|
||||
msg, _ := s.AddMessage(ctx, convID, "user", "fresh", 5)
|
||||
|
||||
s.UpsertContextItems(ctx, convID, []ContextItem{
|
||||
{Ordinal: 100, ItemType: "summary", SummaryID: condensed.SummaryID, TokenCount: 15},
|
||||
{Ordinal: 200, ItemType: "message", MessageID: msg.ID, TokenCount: 5},
|
||||
})
|
||||
|
||||
a := &Assembler{store: s, config: Config{}}
|
||||
result, err := a.Assemble(ctx, convID, AssembleInput{Budget: 1000})
|
||||
if err != nil {
|
||||
t.Fatalf("Assemble: %v", err)
|
||||
}
|
||||
|
||||
// Summary field should contain XML with parent information
|
||||
if result.Summary == "" {
|
||||
t.Fatal("Summary should not be empty")
|
||||
}
|
||||
xmlContent := result.Summary
|
||||
|
||||
// Should contain <parents> section with parent ID
|
||||
if !contains(xmlContent, "<parents>") {
|
||||
t.Errorf("condensed summary XML missing <parents> section: %q", xmlContent)
|
||||
}
|
||||
if !contains(xmlContent, leaf.SummaryID) {
|
||||
t.Errorf("condensed summary XML missing parent ID %q: %q", leaf.SummaryID, xmlContent)
|
||||
}
|
||||
|
||||
// Should contain kind="condensed"
|
||||
if !contains(xmlContent, `kind="condensed"`) {
|
||||
t.Errorf("condensed summary XML missing kind attribute: %q", xmlContent)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssemblerSummaryXMLIncludesDescendantCount(t *testing.T) {
|
||||
s, convID := setupAssemblerStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a leaf summary with specific descendant count
|
||||
leaf, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: "leaf content",
|
||||
TokenCount: 20,
|
||||
DescendantCount: 8,
|
||||
DescendantTokenCount: 1200,
|
||||
})
|
||||
|
||||
msg, _ := s.AddMessage(ctx, convID, "user", "fresh", 5)
|
||||
|
||||
s.UpsertContextItems(ctx, convID, []ContextItem{
|
||||
{Ordinal: 100, ItemType: "summary", SummaryID: leaf.SummaryID, TokenCount: 20},
|
||||
{Ordinal: 200, ItemType: "message", MessageID: msg.ID, TokenCount: 5},
|
||||
})
|
||||
|
||||
a := &Assembler{store: s, config: Config{}}
|
||||
result, err := a.Assemble(ctx, convID, AssembleInput{Budget: 1000})
|
||||
if err != nil {
|
||||
t.Fatalf("Assemble: %v", err)
|
||||
}
|
||||
|
||||
if result.Summary == "" {
|
||||
t.Fatal("Summary should not be empty")
|
||||
}
|
||||
xmlContent := result.Summary
|
||||
|
||||
// Should contain descendant_count="8"
|
||||
if !contains(xmlContent, `descendant_count="8"`) {
|
||||
t.Errorf("summary XML missing descendant_count attribute: %q", xmlContent)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssemblerLeafSummaryNoParents(t *testing.T) {
|
||||
s, convID := setupAssemblerStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Leaf summary has no parents
|
||||
leaf, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: "leaf content",
|
||||
TokenCount: 20,
|
||||
})
|
||||
|
||||
msg, _ := s.AddMessage(ctx, convID, "user", "fresh", 5)
|
||||
|
||||
s.UpsertContextItems(ctx, convID, []ContextItem{
|
||||
{Ordinal: 100, ItemType: "summary", SummaryID: leaf.SummaryID, TokenCount: 20},
|
||||
{Ordinal: 200, ItemType: "message", MessageID: msg.ID, TokenCount: 5},
|
||||
})
|
||||
|
||||
a := &Assembler{store: s, config: Config{}}
|
||||
result, err := a.Assemble(ctx, convID, AssembleInput{Budget: 1000})
|
||||
if err != nil {
|
||||
t.Fatalf("Assemble: %v", err)
|
||||
}
|
||||
|
||||
if result.Summary == "" {
|
||||
t.Fatal("Summary should not be empty")
|
||||
}
|
||||
xmlContent := result.Summary
|
||||
|
||||
// Leaf summary should NOT have <parents> section
|
||||
if contains(xmlContent, "<parents>") {
|
||||
t.Errorf("leaf summary XML should not have <parents> section: %q", xmlContent)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssemblerDepthAwarePrompt(t *testing.T) {
|
||||
s, convID := setupAssemblerStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a condensed summary (depth >= 2) to trigger full guidance
|
||||
now := time.Now().UTC()
|
||||
leaf, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: "leaf summary",
|
||||
TokenCount: 20,
|
||||
EarliestAt: &now,
|
||||
LatestAt: &now,
|
||||
})
|
||||
condensed, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindCondensed,
|
||||
Depth: 2,
|
||||
Content: "condensed summary",
|
||||
TokenCount: 15,
|
||||
ParentIDs: []string{leaf.SummaryID},
|
||||
DescendantCount: 1,
|
||||
DescendantTokenCount: 20,
|
||||
})
|
||||
|
||||
msg, _ := s.AddMessage(ctx, convID, "user", "fresh", 5)
|
||||
|
||||
s.UpsertContextItems(ctx, convID, []ContextItem{
|
||||
{Ordinal: 100, ItemType: "summary", SummaryID: condensed.SummaryID, TokenCount: 15},
|
||||
{Ordinal: 200, ItemType: "message", MessageID: msg.ID, TokenCount: 5},
|
||||
})
|
||||
|
||||
a := &Assembler{store: s, config: Config{}}
|
||||
result, err := a.Assemble(ctx, convID, AssembleInput{Budget: 1000})
|
||||
if err != nil {
|
||||
t.Fatalf("Assemble: %v", err)
|
||||
}
|
||||
|
||||
// Should have a depth-aware prompt in Summary field
|
||||
if result.Summary == "" {
|
||||
t.Error("expected non-empty Summary when depth >= 2")
|
||||
}
|
||||
// SystemPromptAddition is embedded in Summary field
|
||||
if !strings.Contains(result.Summary, "multi-level summarization") {
|
||||
t.Error("Summary should contain system prompt addition about multi-level summarization")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatSummaryXMLUsesSummaryRef(t *testing.T) {
|
||||
// Spec: condensed summaries use <summary_ref id="parentId" /> not <parent>parentId</parent>
|
||||
now := time.Now().UTC()
|
||||
s := Summary{
|
||||
SummaryID: "sum_condensed1",
|
||||
Kind: SummaryKindCondensed,
|
||||
Depth: 1,
|
||||
Content: "condensed content",
|
||||
TokenCount: 50,
|
||||
DescendantCount: 2,
|
||||
EarliestAt: &now,
|
||||
LatestAt: &now,
|
||||
}
|
||||
parentIDs := []string{"sum_leaf1", "sum_leaf2"}
|
||||
|
||||
xml := FormatSummaryXML(&s, parentIDs)
|
||||
|
||||
// Must use <summary_ref id="..." /> per spec
|
||||
if !contains(xml, `<summary_ref id="sum_leaf1" />`) {
|
||||
t.Errorf("expected <summary_ref id=\"sum_leaf1\" />, got: %s", xml)
|
||||
}
|
||||
if !contains(xml, `<summary_ref id="sum_leaf2" />`) {
|
||||
t.Errorf("expected <summary_ref id=\"sum_leaf2\" />, got: %s", xml)
|
||||
}
|
||||
// Must NOT use old <parent> tag
|
||||
if contains(xml, "<parent>") {
|
||||
t.Errorf("should not use <parent> tag, got: %s", xml)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatSummaryXMLIncludesTimestamps(t *testing.T) {
|
||||
// Spec: summary XML includes earliest_at and latest_at attributes
|
||||
earliest := time.Date(2026, 3, 15, 10, 0, 0, 0, time.UTC)
|
||||
latest := time.Date(2026, 3, 15, 14, 30, 0, 0, time.UTC)
|
||||
s := Summary{
|
||||
SummaryID: "sum_leaf1",
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: "leaf content",
|
||||
TokenCount: 30,
|
||||
DescendantCount: 0,
|
||||
EarliestAt: &earliest,
|
||||
LatestAt: &latest,
|
||||
}
|
||||
|
||||
xml := FormatSummaryXML(&s, nil)
|
||||
|
||||
if !contains(xml, `earliest_at="2026-03-15T10:00:00Z"`) {
|
||||
t.Errorf("missing earliest_at attribute, got: %s", xml)
|
||||
}
|
||||
if !contains(xml, `latest_at="2026-03-15T14:30:00Z"`) {
|
||||
t.Errorf("missing latest_at attribute, got: %s", xml)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatSummaryXMLNoTimestampsWhenNil(t *testing.T) {
|
||||
// When EarliestAt/LatestAt are nil, attributes should be omitted
|
||||
s := Summary{
|
||||
SummaryID: "sum_leaf1",
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: "leaf content",
|
||||
TokenCount: 30,
|
||||
DescendantCount: 0,
|
||||
}
|
||||
|
||||
xml := FormatSummaryXML(&s, nil)
|
||||
|
||||
if contains(xml, "earliest_at=") {
|
||||
t.Errorf("should not have earliest_at when nil, got: %s", xml)
|
||||
}
|
||||
if contains(xml, "latest_at=") {
|
||||
t.Errorf("should not have latest_at when nil, got: %s", xml)
|
||||
}
|
||||
}
|
||||
336
pkg/seahorse/short_bench_test.go
Normal file
336
pkg/seahorse/short_bench_test.go
Normal file
|
|
@ -0,0 +1,336 @@
|
|||
package seahorse
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
_ "modernc.org/sqlite"
|
||||
)
|
||||
|
||||
// newBenchStore creates a test store for benchmarks.
|
||||
func newBenchStore(b *testing.B) (*Store, func()) {
|
||||
b.Helper()
|
||||
db, err := sql.Open("sqlite", ":memory:")
|
||||
if err != nil {
|
||||
b.Fatalf("open test db: %v", err)
|
||||
}
|
||||
if err := runSchema(db); err != nil {
|
||||
db.Close()
|
||||
b.Fatalf("migration: %v", err)
|
||||
}
|
||||
return &Store{db: db}, func() { db.Close() }
|
||||
}
|
||||
|
||||
// --- Ingest benchmarks ---
|
||||
|
||||
func BenchmarkIngest_SingleMessage(b *testing.B) {
|
||||
s, cleanup := newBenchStore(b)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
conv, _ := s.GetOrCreateConversation(ctx, "bench:ingest")
|
||||
convID := conv.ConversationID
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, err := s.AddMessage(ctx, convID, "user", "Test message content", 15)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkIngest_BatchMessages(b *testing.B) {
|
||||
s, cleanup := newBenchStore(b)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
conv, _ := s.GetOrCreateConversation(ctx, fmt.Sprintf("bench:ingest-batch:%d", i))
|
||||
convID := conv.ConversationID
|
||||
|
||||
for j := 0; j < 10; j++ {
|
||||
added, err := s.AddMessage(ctx, convID, "user",
|
||||
fmt.Sprintf("Message %d in batch", j), 10)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
s.AppendContextMessage(ctx, convID, added.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- Assemble benchmarks ---
|
||||
|
||||
func BenchmarkAssemble_MessagesOnly(b *testing.B) {
|
||||
s, cleanup := newBenchStore(b)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
conv, _ := s.GetOrCreateConversation(ctx, "bench:assemble-msgs")
|
||||
convID := conv.ConversationID
|
||||
|
||||
// Add 100 messages
|
||||
for i := 0; i < 100; i++ {
|
||||
m, _ := s.AddMessage(ctx, convID, "user",
|
||||
fmt.Sprintf("Message content %d with some text", i), 10)
|
||||
s.AppendContextMessage(ctx, convID, m.ID)
|
||||
}
|
||||
|
||||
a := &Assembler{store: s}
|
||||
input := AssembleInput{Budget: 50000}
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, err := a.Assemble(ctx, convID, input)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkAssemble_WithSummaries(b *testing.B) {
|
||||
s, cleanup := newBenchStore(b)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
conv, _ := s.GetOrCreateConversation(ctx, "bench:assemble-sums")
|
||||
convID := conv.ConversationID
|
||||
|
||||
now := time.Now().UTC()
|
||||
|
||||
// Add 10 leaf summaries
|
||||
for i := 0; i < 10; i++ {
|
||||
sum, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: fmt.Sprintf("Leaf summary %d", i),
|
||||
TokenCount: 500,
|
||||
EarliestAt: &now,
|
||||
LatestAt: &now,
|
||||
})
|
||||
s.AppendContextSummary(ctx, convID, sum.SummaryID)
|
||||
}
|
||||
|
||||
// Add 20 fresh messages
|
||||
for i := 0; i < 20; i++ {
|
||||
m, _ := s.AddMessage(ctx, convID, "user", fmt.Sprintf("Fresh message %d", i), 10)
|
||||
s.AppendContextMessage(ctx, convID, m.ID)
|
||||
}
|
||||
|
||||
a := &Assembler{store: s}
|
||||
input := AssembleInput{Budget: 10000}
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, err := a.Assemble(ctx, convID, input)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkAssemble_BudgetEviction(b *testing.B) {
|
||||
s, cleanup := newBenchStore(b)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
conv, _ := s.GetOrCreateConversation(ctx, "bench:assemble-evict")
|
||||
convID := conv.ConversationID
|
||||
|
||||
now := time.Now().UTC()
|
||||
|
||||
// Add 50 leaf summaries (more than budget can hold)
|
||||
for i := 0; i < 50; i++ {
|
||||
sum, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: fmt.Sprintf("Summary %d", i),
|
||||
TokenCount: 300,
|
||||
EarliestAt: &now,
|
||||
LatestAt: &now,
|
||||
})
|
||||
s.AppendContextSummary(ctx, convID, sum.SummaryID)
|
||||
}
|
||||
|
||||
// Add fresh tail
|
||||
for i := 0; i < FreshTailCount; i++ {
|
||||
m, _ := s.AddMessage(ctx, convID, "user", "fresh", 10)
|
||||
s.AppendContextMessage(ctx, convID, m.ID)
|
||||
}
|
||||
|
||||
a := &Assembler{store: s}
|
||||
input := AssembleInput{Budget: 5000} // Force eviction
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, err := a.Assemble(ctx, convID, input)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- Search (FTS5) benchmarks ---
|
||||
|
||||
// benchSeedSummaries adds n summaries to a conversation for search benchmarks.
|
||||
func benchSeedSummaries(b *testing.B, s *Store, convID int64, n int, contentTpl string) {
|
||||
b.Helper()
|
||||
now := time.Now().UTC()
|
||||
for i := 0; i < n; i++ {
|
||||
sum, err := s.CreateSummary(context.Background(), CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: fmt.Sprintf(contentTpl, i),
|
||||
TokenCount: 200,
|
||||
EarliestAt: &now,
|
||||
LatestAt: &now,
|
||||
})
|
||||
if err != nil {
|
||||
b.Fatalf("create summary: %v", err)
|
||||
}
|
||||
s.AppendContextSummary(context.Background(), convID, sum.SummaryID)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkSearchSummaries_FTS5(b *testing.B) {
|
||||
s, cleanup := newBenchStore(b)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
conv, _ := s.GetOrCreateConversation(ctx, "bench:search-fts")
|
||||
convID := conv.ConversationID
|
||||
|
||||
benchSeedSummaries(b, s, convID, 100, "Summary about database configuration and API endpoints %d")
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, err := s.SearchSummaries(ctx, SearchInput{
|
||||
Pattern: "database",
|
||||
Mode: "full_text",
|
||||
ConversationID: convID,
|
||||
})
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkSearchSummaries_Like(b *testing.B) {
|
||||
s, cleanup := newBenchStore(b)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
conv, _ := s.GetOrCreateConversation(ctx, "bench:search-like")
|
||||
convID := conv.ConversationID
|
||||
|
||||
benchSeedSummaries(b, s, convID, 100, "Summary about configuration %d")
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, err := s.SearchSummaries(ctx, SearchInput{
|
||||
Pattern: "config",
|
||||
Mode: "like",
|
||||
ConversationID: convID,
|
||||
})
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkSearchMessages_FTS5(b *testing.B) {
|
||||
s, cleanup := newBenchStore(b)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
conv, _ := s.GetOrCreateConversation(ctx, "bench:search-msg-fts")
|
||||
convID := conv.ConversationID
|
||||
|
||||
// Add 500 messages
|
||||
for i := 0; i < 500; i++ {
|
||||
m, _ := s.AddMessage(ctx, convID, "user",
|
||||
fmt.Sprintf("User message about API and database integration %d", i), 20)
|
||||
s.AppendContextMessage(ctx, convID, m.ID)
|
||||
}
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, err := s.SearchMessages(ctx, SearchInput{
|
||||
Pattern: "API database",
|
||||
Mode: "full_text",
|
||||
ConversationID: convID,
|
||||
})
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- Bootstrap benchmarks ---
|
||||
|
||||
func BenchmarkBootstrap_Empty(b *testing.B) {
|
||||
s, cleanup := newBenchStore(b)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
conv, _ := s.GetOrCreateConversation(ctx, fmt.Sprintf("bench:bootstrap-empty:%d", i))
|
||||
convID := conv.ConversationID
|
||||
_ = convID // Bootstrap with empty history
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkBootstrap_100Messages(b *testing.B) {
|
||||
s, cleanup := newBenchStore(b)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
// Prepare 100 messages
|
||||
msgs := make([]Message, 100)
|
||||
for i := 0; i < 100; i++ {
|
||||
msgs[i] = Message{
|
||||
Role: "user",
|
||||
Content: fmt.Sprintf("Bootstrap message %d", i),
|
||||
TokenCount: 15,
|
||||
}
|
||||
}
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
conv, _ := s.GetOrCreateConversation(ctx, fmt.Sprintf("bench:bootstrap-100:%d", i))
|
||||
convID := conv.ConversationID
|
||||
|
||||
for _, m := range msgs {
|
||||
added, _ := s.AddMessage(ctx, convID, m.Role, m.Content, m.TokenCount)
|
||||
s.AppendContextMessage(ctx, convID, added.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkBootstrap_500Messages(b *testing.B) {
|
||||
s, cleanup := newBenchStore(b)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
msgs := make([]Message, 500)
|
||||
for i := 0; i < 500; i++ {
|
||||
msgs[i] = Message{
|
||||
Role: "user",
|
||||
Content: fmt.Sprintf("Bootstrap message %d", i),
|
||||
TokenCount: 15,
|
||||
}
|
||||
}
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
conv, _ := s.GetOrCreateConversation(ctx, fmt.Sprintf("bench:bootstrap-500:%d", i))
|
||||
convID := conv.ConversationID
|
||||
|
||||
for _, m := range msgs {
|
||||
added, _ := s.AddMessage(ctx, convID, m.Role, m.Content, m.TokenCount)
|
||||
s.AppendContextMessage(ctx, convID, added.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
898
pkg/seahorse/short_compaction.go
Normal file
898
pkg/seahorse/short_compaction.go
Normal file
|
|
@ -0,0 +1,898 @@
|
|||
package seahorse
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sort"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
"github.com/sipeed/picoclaw/pkg/tokenizer"
|
||||
)
|
||||
|
||||
// CompactInput controls compaction behavior.
|
||||
type CompactInput struct {
|
||||
Budget *int // Token budget override
|
||||
Force bool // Force compaction even if below threshold
|
||||
}
|
||||
|
||||
// CompactResult describes what was compacted.
|
||||
type CompactResult struct {
|
||||
SummariesCreated []string `json:"summariesCreated"`
|
||||
TokensSaved int `json:"tokensSaved"`
|
||||
LeafSummaries int `json:"leafSummaries"`
|
||||
CondensedSummaries int `json:"condensedSummaries"`
|
||||
}
|
||||
|
||||
// NeedsCompaction returns true if context tokens >= ContextThreshold × contextWindow.
|
||||
func (e *CompactionEngine) NeedsCompaction(ctx context.Context, convID int64, contextWindow int) (bool, error) {
|
||||
tokens, err := e.store.GetContextTokenCount(ctx, convID)
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("get token count: %w", err)
|
||||
}
|
||||
threshold := int(float64(contextWindow) * ContextThreshold)
|
||||
return tokens >= threshold, nil
|
||||
}
|
||||
|
||||
// Close cancels the shutdown context, stopping async goroutines.
|
||||
func (e *CompactionEngine) Close() {
|
||||
if e.shutdownCancel != nil {
|
||||
e.shutdownCancel()
|
||||
}
|
||||
}
|
||||
|
||||
// Compact runs leaf compaction (sync) and optionally condensed compaction.
|
||||
func (e *CompactionEngine) Compact(ctx context.Context, convID int64, input CompactInput) (*CompactResult, error) {
|
||||
result := &CompactResult{}
|
||||
|
||||
// Phase 1: leaf compaction (synchronous, every turn)
|
||||
summaryID, err := e.compactLeaf(ctx, convID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("compact leaf: %w", err)
|
||||
}
|
||||
if summaryID != nil {
|
||||
result.SummariesCreated = append(result.SummariesCreated, *summaryID)
|
||||
result.LeafSummaries++
|
||||
logger.InfoCF("seahorse", "compact: leaf", map[string]any{
|
||||
"conv_id": convID,
|
||||
"summary_id": *summaryID,
|
||||
})
|
||||
}
|
||||
|
||||
// Phase 2: condensed compaction if over threshold
|
||||
tokensBefore, _ := e.store.GetContextTokenCount(ctx, convID)
|
||||
var budget int
|
||||
if input.Budget != nil {
|
||||
budget = *input.Budget
|
||||
if budget == 0 {
|
||||
logger.ErrorCF("seahorse", "Compact: budget is 0, this should not happen", map[string]any{
|
||||
"conv_id": convID,
|
||||
})
|
||||
}
|
||||
} else {
|
||||
budget = int(float64(tokensBefore) * ContextThreshold)
|
||||
}
|
||||
|
||||
if input.Force || (tokensBefore > budget && budget > 0) {
|
||||
// Launch async condensed compaction with dedup
|
||||
if _, loaded := e.condensing.LoadOrStore(convID, struct{}{}); !loaded {
|
||||
go func() {
|
||||
defer e.condensing.Delete(convID)
|
||||
e.runCondensedLoop(e.shutdownCtx, convID)
|
||||
}()
|
||||
}
|
||||
}
|
||||
|
||||
tokensAfter, _ := e.store.GetContextTokenCount(ctx, convID)
|
||||
if tokensAfter < tokensBefore {
|
||||
result.TokensSaved = tokensBefore - tokensAfter
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// CompactUntilUnder aggressively compacts until context is under budget.
|
||||
func (e *CompactionEngine) CompactUntilUnder(ctx context.Context, convID int64, budget int) (*CompactResult, error) {
|
||||
result := &CompactResult{}
|
||||
prevTokens := 0
|
||||
logger.InfoCF("seahorse", "compact_until_under: start", map[string]any{"conv_id": convID, "budget": budget})
|
||||
|
||||
for iter := 0; iter < MaxCompactIterations; iter++ {
|
||||
tokens, err := e.store.GetContextTokenCount(ctx, convID)
|
||||
if err != nil {
|
||||
return result, fmt.Errorf("get tokens: %w", err)
|
||||
}
|
||||
if tokens <= budget {
|
||||
logger.InfoCF("seahorse", "compact_until_under: done", map[string]any{
|
||||
"conv_id": convID,
|
||||
"budget": budget,
|
||||
"tokens": tokens,
|
||||
"leaf": result.LeafSummaries,
|
||||
"condensed": result.CondensedSummaries,
|
||||
})
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// Try leaf first
|
||||
summaryID, err := e.compactLeaf(ctx, convID, true)
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
if summaryID != nil {
|
||||
result.SummariesCreated = append(result.SummariesCreated, *summaryID)
|
||||
result.LeafSummaries++
|
||||
logger.InfoCF("seahorse", "compact_until_under: leaf", map[string]any{
|
||||
"conv_id": convID,
|
||||
"summary_id": *summaryID,
|
||||
})
|
||||
continue
|
||||
}
|
||||
|
||||
// Try condensed with forced fanout
|
||||
condensedID, err := e.compactCondensed(ctx, convID)
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
if condensedID != nil {
|
||||
result.SummariesCreated = append(result.SummariesCreated, *condensedID)
|
||||
result.CondensedSummaries++
|
||||
logger.InfoCF("seahorse", "compact_until_under: condensed", map[string]any{
|
||||
"conv_id": convID,
|
||||
"summary_id": *condensedID,
|
||||
})
|
||||
continue
|
||||
}
|
||||
|
||||
// No progress
|
||||
newTokens, _ := e.store.GetContextTokenCount(ctx, convID)
|
||||
if newTokens >= prevTokens {
|
||||
logger.WarnCF("seahorse", "compact_until_under: no progress", map[string]any{
|
||||
"conv_id": convID,
|
||||
"tokens": newTokens,
|
||||
})
|
||||
return result, nil
|
||||
}
|
||||
prevTokens = newTokens
|
||||
}
|
||||
|
||||
// Safety cap exceeded — see MaxCompactIterations doc for rationale.
|
||||
logger.WarnCF("seahorse", "compact_until_under: exceeded max iterations", map[string]any{
|
||||
"conv_id": convID,
|
||||
"budget": budget,
|
||||
"iterations": MaxCompactIterations,
|
||||
"tokens": prevTokens,
|
||||
})
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// compactLeaf compresses the oldest contiguous message chunk into a leaf summary.
|
||||
// When force is true, FreshTailCount protection is bypassed (used by CompactUntilUnder).
|
||||
func (e *CompactionEngine) compactLeaf(ctx context.Context, convID int64, force ...bool) (*string, error) {
|
||||
items, err := e.store.GetContextItems(ctx, convID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Find oldest contiguous message chunk outside fresh tail
|
||||
msgCount := 0
|
||||
msgTokens := 0
|
||||
for _, item := range items {
|
||||
if item.ItemType == "message" {
|
||||
msgCount++
|
||||
msgTokens += item.TokenCount
|
||||
}
|
||||
}
|
||||
|
||||
// Trigger if either message count or token threshold is met
|
||||
if msgCount < LeafMinFanout && msgTokens < LeafChunkTokens {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Calculate fresh tail boundary (bypass when forced)
|
||||
useForce := len(force) > 0 && force[0]
|
||||
tailStartIdx := len(items) - FreshTailCount
|
||||
if useForce {
|
||||
tailStartIdx = len(items) // allow compacting everything
|
||||
}
|
||||
if tailStartIdx < 0 {
|
||||
tailStartIdx = 0
|
||||
}
|
||||
|
||||
// Find oldest contiguous message chunk, accumulating up to LeafChunkTokens
|
||||
var chunk []ContextItem
|
||||
chunkStart := -1
|
||||
chunkEnd := -1
|
||||
accumTokens := 0
|
||||
for i := 0; i < tailStartIdx; i++ {
|
||||
if items[i].ItemType == "message" {
|
||||
if chunkStart == -1 {
|
||||
chunkStart = i
|
||||
}
|
||||
chunkEnd = i
|
||||
accumTokens += items[i].TokenCount
|
||||
// Stop accumulating once we reach the token budget
|
||||
if accumTokens >= LeafChunkTokens {
|
||||
break
|
||||
}
|
||||
} else {
|
||||
// Non-message breaks the chunk
|
||||
if chunkStart != -1 && (chunkEnd-chunkStart+1) >= LeafMinFanout {
|
||||
break
|
||||
}
|
||||
chunkStart = -1
|
||||
chunkEnd = -1
|
||||
accumTokens = 0
|
||||
}
|
||||
}
|
||||
|
||||
if chunkStart == -1 || (chunkEnd-chunkStart+1) < LeafMinFanout {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
chunk = items[chunkStart : chunkEnd+1]
|
||||
|
||||
// Collect messages for the chunk
|
||||
var messages []Message
|
||||
for _, item := range chunk {
|
||||
msg, innerErr := e.store.GetMessageByID(ctx, item.MessageID)
|
||||
if innerErr != nil {
|
||||
return nil, innerErr
|
||||
}
|
||||
messages = append(messages, *msg)
|
||||
}
|
||||
|
||||
// Get prior summaries for context
|
||||
priorSummary := ""
|
||||
priorCount := 0
|
||||
for i := chunkStart - 1; i >= 0 && priorCount < 2; i-- {
|
||||
if items[i].ItemType == "summary" {
|
||||
sum, innerErr2 := e.store.GetSummary(ctx, items[i].SummaryID)
|
||||
if innerErr2 == nil {
|
||||
priorSummary = sum.Content + "\n" + priorSummary
|
||||
priorCount++
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Generate summary
|
||||
content, err := e.generateLeafSummary(ctx, messages, priorSummary)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Create summary in store
|
||||
tokenCount := tokenizer.EstimateMessageTokens(providers.Message{Content: content})
|
||||
|
||||
var earliestAt, latestAt *time.Time
|
||||
if len(messages) > 0 {
|
||||
earliestAt = &messages[0].CreatedAt
|
||||
latestAt = &messages[len(messages)-1].CreatedAt
|
||||
}
|
||||
|
||||
summary, err := e.store.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: content,
|
||||
TokenCount: tokenCount,
|
||||
EarliestAt: earliestAt,
|
||||
LatestAt: latestAt,
|
||||
SourceMessageTokens: sumMessageTokens(messages),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Link to source messages
|
||||
msgIDs := make([]int64, len(messages))
|
||||
for i, m := range messages {
|
||||
msgIDs[i] = m.ID
|
||||
}
|
||||
if err := e.store.LinkSummaryToMessages(ctx, summary.SummaryID, msgIDs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Replace context range with summary
|
||||
if err := e.store.ReplaceContextRangeWithSummary(
|
||||
ctx, convID, chunk[0].Ordinal, chunk[len(chunk)-1].Ordinal, summary.SummaryID,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &summary.SummaryID, nil
|
||||
}
|
||||
|
||||
// compactCondensed compresses multiple summaries into one higher-level summary.
|
||||
func (e *CompactionEngine) compactCondensed(ctx context.Context, convID int64) (*string, error) {
|
||||
// Try ordinal-aware selection first (respects consecutive ordering)
|
||||
var candidates []Summary
|
||||
|
||||
depths, err := e.store.GetDistinctDepthsInContext(ctx, convID, 0)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, depth := range depths {
|
||||
var chunkAtDepth []Summary
|
||||
var err2 error
|
||||
chunkAtDepth, err2 = e.selectOldestChunkAtDepth(ctx, convID, depth)
|
||||
if err2 != nil {
|
||||
continue
|
||||
}
|
||||
if len(chunkAtDepth) > 0 {
|
||||
candidates = chunkAtDepth
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to depth-grouping selection
|
||||
if len(candidates) == 0 {
|
||||
candidates, err = e.selectShallowestCondensationCandidate(ctx, convID, false)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if len(candidates) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Generate condensed summary
|
||||
content, err := e.generateCondensedSummary(ctx, candidates)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Merge metadata
|
||||
maxDepth := 0
|
||||
descendantCount := 0
|
||||
descendantTokenCount := 0
|
||||
sourceMessageTokens := 0
|
||||
var earliestAt, latestAt *time.Time
|
||||
|
||||
parentIDs := make([]string, len(candidates))
|
||||
for i, c := range candidates {
|
||||
parentIDs[i] = c.SummaryID
|
||||
if c.Depth > maxDepth {
|
||||
maxDepth = c.Depth
|
||||
}
|
||||
descendantCount += c.DescendantCount + 1
|
||||
descendantTokenCount += c.TokenCount + c.DescendantTokenCount
|
||||
sourceMessageTokens += c.SourceMessageTokenCount
|
||||
if c.EarliestAt != nil {
|
||||
if earliestAt == nil || c.EarliestAt.Before(*earliestAt) {
|
||||
earliestAt = c.EarliestAt
|
||||
}
|
||||
}
|
||||
if c.LatestAt != nil {
|
||||
if latestAt == nil || c.LatestAt.After(*latestAt) {
|
||||
latestAt = c.LatestAt
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tokenCount := tokenizer.EstimateMessageTokens(providers.Message{Content: content})
|
||||
|
||||
summary, err := e.store.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindCondensed,
|
||||
Depth: maxDepth + 1,
|
||||
Content: content,
|
||||
TokenCount: tokenCount,
|
||||
EarliestAt: earliestAt,
|
||||
LatestAt: latestAt,
|
||||
DescendantCount: descendantCount,
|
||||
DescendantTokenCount: descendantTokenCount,
|
||||
SourceMessageTokens: sourceMessageTokens,
|
||||
ParentIDs: parentIDs,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Find the ordinal range for the candidate summaries in context
|
||||
items, err := e.store.GetContextItems(ctx, convID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
candidateSet := make(map[string]bool)
|
||||
for _, c := range candidates {
|
||||
candidateSet[c.SummaryID] = true
|
||||
}
|
||||
|
||||
startOrd := -1
|
||||
endOrd := -1
|
||||
hasNonCandidate := false
|
||||
for _, item := range items {
|
||||
if item.ItemType == "summary" && candidateSet[item.SummaryID] {
|
||||
if startOrd == -1 {
|
||||
startOrd, endOrd = item.Ordinal, item.Ordinal
|
||||
} else {
|
||||
// Check for non-candidate items between endOrd and current ordinal
|
||||
for _, it := range items {
|
||||
if it.Ordinal > endOrd && it.Ordinal <= item.Ordinal {
|
||||
if it.ItemType != "summary" || !candidateSet[it.SummaryID] {
|
||||
hasNonCandidate = true
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if hasNonCandidate {
|
||||
break
|
||||
}
|
||||
if item.Ordinal < startOrd {
|
||||
startOrd = item.Ordinal
|
||||
}
|
||||
if item.Ordinal > endOrd {
|
||||
endOrd = item.Ordinal
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if startOrd == -1 || endOrd == -1 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Collect candidate summary IDs
|
||||
candidateIDs := make([]string, 0, len(candidates))
|
||||
for _, c := range candidates {
|
||||
candidateIDs = append(candidateIDs, c.SummaryID)
|
||||
}
|
||||
|
||||
if hasNonCandidate {
|
||||
// Use safe per-item deletion to avoid deleting non-candidate items
|
||||
if err := e.store.ReplaceContextItemsWithSummary(ctx, convID, candidateIDs, summary.SummaryID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
} else {
|
||||
// Candidates are consecutive, use efficient range deletion
|
||||
if err := e.store.ReplaceContextRangeWithSummary(ctx, convID, startOrd, endOrd, summary.SummaryID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
return &summary.SummaryID, nil
|
||||
}
|
||||
|
||||
// selectShallowestCondensationCandidate finds the shallowest consecutive summary group.
|
||||
func (e *CompactionEngine) selectShallowestCondensationCandidate(
|
||||
ctx context.Context, convID int64, forced bool,
|
||||
) ([]Summary, error) {
|
||||
items, err := e.store.GetContextItems(ctx, convID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Group by depth, find consecutive runs
|
||||
tailStartIdx := len(items) - FreshTailCount
|
||||
if tailStartIdx < 0 {
|
||||
tailStartIdx = 0
|
||||
}
|
||||
|
||||
minFanout := CondensedMinFanout
|
||||
if forced {
|
||||
minFanout = CondensedMinFanoutHard
|
||||
}
|
||||
|
||||
// Track depth groups
|
||||
depthGroups := make(map[int][]ContextItem)
|
||||
for i := 0; i < tailStartIdx; i++ {
|
||||
item := items[i]
|
||||
if item.ItemType != "summary" {
|
||||
continue
|
||||
}
|
||||
sum, err := e.store.GetSummary(ctx, item.SummaryID)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
depthGroups[sum.Depth] = append(depthGroups[sum.Depth], item)
|
||||
}
|
||||
|
||||
// Find shallowest depth with enough candidates
|
||||
// Collect all depths and sort to handle non-consecutive depths
|
||||
var depths []int
|
||||
for depth := range depthGroups {
|
||||
depths = append(depths, depth)
|
||||
}
|
||||
sort.Ints(depths)
|
||||
|
||||
for _, depth := range depths {
|
||||
group := depthGroups[depth]
|
||||
if len(group) >= minFanout {
|
||||
// Load summaries
|
||||
var result []Summary
|
||||
for _, item := range group[:minFanout] {
|
||||
sum, err := e.store.GetSummary(ctx, item.SummaryID)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
result = append(result, *sum)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
}
|
||||
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// selectOldestChunkAtDepth scans context_items from oldest ordinal, collecting consecutive
|
||||
// summaries at the given depth. Stops at non-summary items, different depth, fresh tail, or
|
||||
// token overflow. Returns contiguous chunk of summaries.
|
||||
func (e *CompactionEngine) selectOldestChunkAtDepth(
|
||||
ctx context.Context, convID int64, targetDepth int,
|
||||
) ([]Summary, error) {
|
||||
items, err := e.store.GetContextItems(ctx, convID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
tailStartIdx := len(items) - FreshTailCount
|
||||
if tailStartIdx < 0 {
|
||||
tailStartIdx = 0
|
||||
}
|
||||
|
||||
var chunk []Summary
|
||||
accumTokens := 0
|
||||
|
||||
for i := 0; i < tailStartIdx; i++ {
|
||||
item := items[i]
|
||||
if item.ItemType != "summary" {
|
||||
// Non-summary breaks the chunk
|
||||
break
|
||||
}
|
||||
sum, err := e.store.GetSummary(ctx, item.SummaryID)
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
if sum.Depth != targetDepth {
|
||||
// Different depth breaks the chunk
|
||||
break
|
||||
}
|
||||
if accumTokens+sum.TokenCount > LeafChunkTokens {
|
||||
// Token overflow stops collection
|
||||
break
|
||||
}
|
||||
chunk = append(chunk, *sum)
|
||||
accumTokens += sum.TokenCount
|
||||
}
|
||||
|
||||
// Min tokens check: spec line 808
|
||||
// chunk tokens must be >= max(CondensedTargetTokens, LeafChunkTokens × 0.1) = 2000
|
||||
minTokens := CondensedTargetTokens // 2000
|
||||
if accumTokens < minTokens {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
return chunk, nil
|
||||
}
|
||||
|
||||
// generateLeafSummary calls the LLM to generate a leaf summary with 3-level escalation.
|
||||
// Level 1: normal LLM prompt. Level 2: aggressive prompt. Level 3: deterministic truncation.
|
||||
func (e *CompactionEngine) generateLeafSummary(
|
||||
ctx context.Context,
|
||||
messages []Message,
|
||||
previousSummary string,
|
||||
) (string, error) {
|
||||
if e.complete == nil {
|
||||
return truncateSummary(messages), nil
|
||||
}
|
||||
|
||||
sourceText := formatMessagesForSummary(messages)
|
||||
inputTokens := sumMessageTokens(messages)
|
||||
targetTokens := minInt(LeafTargetTokens, int(float64(inputTokens)*0.35))
|
||||
|
||||
// Level 1: normal prompt
|
||||
prompt := buildLeafSummaryPrompt(sourceText, previousSummary, targetTokens)
|
||||
content, err := e.complete(ctx, prompt, CompleteOptions{
|
||||
MaxTokens: LeafTargetTokens * 2,
|
||||
Temperature: 0.3,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if content == "" {
|
||||
// Retry with temperature=0
|
||||
content, err = e.complete(ctx, prompt, CompleteOptions{
|
||||
MaxTokens: LeafTargetTokens * 2,
|
||||
Temperature: 0,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
|
||||
// Check if level 1 succeeded
|
||||
if content != "" && tokenizer.EstimateMessageTokens(providers.Message{Content: content}) < inputTokens {
|
||||
return content, nil
|
||||
}
|
||||
|
||||
// Level 2: aggressive prompt
|
||||
aggressiveTarget := minInt(640, int(float64(inputTokens)*0.20))
|
||||
aggressivePrompt := buildAggressiveLeafSummaryPrompt(sourceText, previousSummary, aggressiveTarget)
|
||||
content, err = e.complete(ctx, aggressivePrompt, CompleteOptions{
|
||||
MaxTokens: aggressiveTarget * 2,
|
||||
Temperature: 0.3,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if content == "" {
|
||||
// Retry with temperature=0
|
||||
content, err = e.complete(ctx, aggressivePrompt, CompleteOptions{
|
||||
MaxTokens: aggressiveTarget * 2,
|
||||
Temperature: 0,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
if content != "" && tokenizer.EstimateMessageTokens(providers.Message{Content: content}) < inputTokens {
|
||||
return content, nil
|
||||
}
|
||||
|
||||
// Level 3: deterministic truncation
|
||||
return truncateSummary(messages), nil
|
||||
}
|
||||
|
||||
// generateCondensedSummary calls the LLM to generate a condensed summary with 3-level escalation.
|
||||
func (e *CompactionEngine) generateCondensedSummary(ctx context.Context, summaries []Summary) (string, error) {
|
||||
if e.complete == nil {
|
||||
return truncateCondensedSummaries(summaries), nil
|
||||
}
|
||||
|
||||
sourceText := formatSummariesForCondensation(summaries)
|
||||
inputTokens := sumSummaryTokens(summaries)
|
||||
targetTokens := minInt(CondensedTargetTokens, int(float64(inputTokens)*0.35))
|
||||
|
||||
// Level 1: normal prompt
|
||||
prompt := buildCondensedSummaryPrompt(sourceText, targetTokens)
|
||||
content, err := e.complete(ctx, prompt, CompleteOptions{
|
||||
MaxTokens: CondensedTargetTokens * 2,
|
||||
Temperature: 0.3,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if content == "" {
|
||||
content, err = e.complete(ctx, prompt, CompleteOptions{
|
||||
MaxTokens: CondensedTargetTokens * 2,
|
||||
Temperature: 0,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
if content != "" {
|
||||
return content, nil
|
||||
}
|
||||
|
||||
// Level 2: aggressive prompt
|
||||
aggressiveTarget := minInt(640, int(float64(inputTokens)*0.20))
|
||||
aggressivePrompt := buildCondensedSummaryPrompt(sourceText, aggressiveTarget)
|
||||
content, err = e.complete(ctx, aggressivePrompt, CompleteOptions{
|
||||
MaxTokens: aggressiveTarget * 2,
|
||||
Temperature: 0.3,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if content != "" {
|
||||
return content, nil
|
||||
}
|
||||
|
||||
// Level 3: deterministic fallback
|
||||
return truncateCondensedSummaries(summaries), nil
|
||||
}
|
||||
|
||||
// runCondensedLoop runs condensed compaction in a loop until:
|
||||
// a) context tokens <= threshold (success), OR
|
||||
// b) No candidate found (nothing to condense), OR
|
||||
// c) tokensAfter >= tokensBefore (no progress this iteration), OR
|
||||
// d) tokensAfter >= previousTokens (no improvement over last iteration)
|
||||
func (e *CompactionEngine) runCondensedLoop(ctx context.Context, convID int64) {
|
||||
var prevTokens int
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
tokensBefore, err := e.store.GetContextTokenCount(ctx, convID)
|
||||
if err != nil {
|
||||
logger.ErrorCF("seahorse", "condensed: get tokens", map[string]any{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
condensedID, err := e.compactCondensed(ctx, convID)
|
||||
if err != nil {
|
||||
logger.ErrorCF("seahorse", "condensed: compact", map[string]any{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if condensedID == nil {
|
||||
// No candidate found
|
||||
logger.DebugCF("seahorse", "condensed: no candidate", map[string]any{"conv_id": convID})
|
||||
return
|
||||
}
|
||||
|
||||
tokensAfter, _ := e.store.GetContextTokenCount(ctx, convID)
|
||||
|
||||
if tokensAfter >= tokensBefore {
|
||||
// No progress this iteration
|
||||
logger.DebugCF(
|
||||
"seahorse",
|
||||
"condensed: no progress",
|
||||
map[string]any{"conv_id": convID, "tokens_before": tokensBefore, "tokens_after": tokensAfter},
|
||||
)
|
||||
return
|
||||
}
|
||||
if tokensAfter >= prevTokens && prevTokens > 0 {
|
||||
// No improvement over last iteration
|
||||
logger.DebugCF(
|
||||
"seahorse",
|
||||
"condensed: no improvement",
|
||||
map[string]any{"conv_id": convID, "tokens": tokensAfter},
|
||||
)
|
||||
return
|
||||
}
|
||||
|
||||
prevTokens = tokensAfter
|
||||
}
|
||||
}
|
||||
|
||||
// --- Helper functions ---
|
||||
|
||||
func formatMessagesForSummary(messages []Message) string {
|
||||
var result string
|
||||
for _, m := range messages {
|
||||
ts := m.CreatedAt.Format("2006-01-02 15:04 MST")
|
||||
content := m.Content
|
||||
if content == "" && len(m.Parts) > 0 {
|
||||
content = partsToReadableContent(m.Parts)
|
||||
}
|
||||
result += fmt.Sprintf("[%s]\n%s\n\n", ts, content)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func formatSummariesForCondensation(summaries []Summary) string {
|
||||
var result string
|
||||
for _, s := range summaries {
|
||||
earliest := ""
|
||||
if s.EarliestAt != nil {
|
||||
earliest = s.EarliestAt.Format("2006-01-02")
|
||||
}
|
||||
latest := ""
|
||||
if s.LatestAt != nil {
|
||||
latest = s.LatestAt.Format("2006-01-02")
|
||||
}
|
||||
result += fmt.Sprintf("[%s - %s]\n%s\n\n", earliest, latest, s.Content)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func buildLeafSummaryPrompt(sourceText, previousSummary string, targetTokens int) string {
|
||||
prev := "(none)"
|
||||
if previousSummary != "" {
|
||||
prev = previousSummary
|
||||
}
|
||||
return fmt.Sprintf(`You summarize a SEGMENT of a conversation for future model turns.
|
||||
Treat this as incremental memory compaction input, not a full-conversation summary.
|
||||
|
||||
Normal summary policy:
|
||||
- Preserve key decisions, rationale, constraints, and active tasks.
|
||||
- Keep essential technical details needed to continue work safely.
|
||||
- Remove obvious repetition and conversational filler.
|
||||
|
||||
Output requirements:
|
||||
- Plain text only.
|
||||
- No preamble, headings, or markdown formatting.
|
||||
- Track file operations (created, modified, deleted, renamed) with file paths and current status.
|
||||
- If no file operations appear, include exactly: "Files: none".
|
||||
- End with exactly: "Expand for details about: <comma-separated list of what was dropped or compressed>".
|
||||
- Target length: about %d tokens or less.
|
||||
|
||||
<previous_context>
|
||||
%s
|
||||
</previous_context>
|
||||
|
||||
<conversation_segment>
|
||||
%s
|
||||
</conversation_segment>`, targetTokens, prev, sourceText)
|
||||
}
|
||||
|
||||
func buildCondensedSummaryPrompt(sourceText string, targetTokens int) string {
|
||||
return fmt.Sprintf(`You condense multiple summaries into a single higher-level summary.
|
||||
Preserve all important decisions, constraints, and outcomes.
|
||||
Merge overlapping topics. Keep technical details intact.
|
||||
|
||||
Output requirements:
|
||||
- Plain text only.
|
||||
- No preamble, headings, or markdown formatting.
|
||||
- End with exactly: "Expand for details about: <comma-separated list>".
|
||||
- Target length: about %d tokens or less.
|
||||
|
||||
<summaries>
|
||||
%s
|
||||
</summaries>`, targetTokens, sourceText)
|
||||
}
|
||||
|
||||
func buildAggressiveLeafSummaryPrompt(sourceText, previousSummary string, targetTokens int) string {
|
||||
prev := "(none)"
|
||||
if previousSummary != "" {
|
||||
prev = previousSummary
|
||||
}
|
||||
return fmt.Sprintf(`You summarize a SEGMENT of a conversation for future model turns.
|
||||
Aggressive summary policy:
|
||||
- Keep only durable facts and current task state.
|
||||
- Remove examples, repetition, and low-value narrative details.
|
||||
- Preserve explicit TODOs, blockers, decisions, and constraints.
|
||||
|
||||
Output requirements:
|
||||
- Plain text only.
|
||||
- No preamble, headings, or markdown formatting.
|
||||
- Track file operations (created, modified, deleted, renamed) with file paths and current status.
|
||||
- If no file operations appear, include exactly: "Files: none".
|
||||
- End with exactly: "Expand for details about: <comma-separated list of what was dropped or compressed>".
|
||||
- Target length: about %d tokens or less.
|
||||
|
||||
<previous_context>
|
||||
%s
|
||||
</previous_context>
|
||||
|
||||
<conversation_segment>
|
||||
%s
|
||||
</conversation_segment>`, targetTokens, prev, sourceText)
|
||||
}
|
||||
|
||||
func truncateSummary(messages []Message) string {
|
||||
content := ""
|
||||
for _, m := range messages {
|
||||
c := m.Content
|
||||
if c == "" && len(m.Parts) > 0 {
|
||||
c = partsToReadableContent(m.Parts)
|
||||
}
|
||||
content += c + "\n"
|
||||
}
|
||||
if len(content) > 2048 {
|
||||
content = content[:2048]
|
||||
}
|
||||
content += fmt.Sprintf("\n[Truncated from %d messages]", len(messages))
|
||||
return content
|
||||
}
|
||||
|
||||
func truncateCondensedSummaries(summaries []Summary) string {
|
||||
content := ""
|
||||
for _, s := range summaries {
|
||||
content += s.Content + "\n"
|
||||
}
|
||||
if len(content) > 2048 {
|
||||
content = content[:2048]
|
||||
}
|
||||
content += fmt.Sprintf("\n[Condensed from %d summaries]", len(summaries))
|
||||
return content
|
||||
}
|
||||
|
||||
func sumMessageTokens(messages []Message) int {
|
||||
total := 0
|
||||
for _, m := range messages {
|
||||
total += m.TokenCount
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
func sumSummaryTokens(summaries []Summary) int {
|
||||
total := 0
|
||||
for _, s := range summaries {
|
||||
total += s.TokenCount
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
func minInt(a, b int) int {
|
||||
if a < b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
974
pkg/seahorse/short_compaction_test.go
Normal file
974
pkg/seahorse/short_compaction_test.go
Normal file
|
|
@ -0,0 +1,974 @@
|
|||
package seahorse
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// --- Test Helpers ---
|
||||
|
||||
// waitForCondensed blocks until the async condensed goroutine for convID finishes.
|
||||
// Returns false if timeout is reached.
|
||||
func waitForCondensed(ce *CompactionEngine, convID int64, timeout time.Duration) bool {
|
||||
deadline := time.Now().Add(timeout)
|
||||
for time.Now().Before(deadline) {
|
||||
if _, exists := ce.condensing.Load(convID); !exists {
|
||||
return true
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// --- Compaction Tests ---
|
||||
|
||||
func newTestCompactionEngine(t *testing.T) (*CompactionEngine, *Store, int64) {
|
||||
t.Helper()
|
||||
db := openTestDB(t)
|
||||
if err := runSchema(db); err != nil {
|
||||
t.Fatalf("migration: %v", err)
|
||||
}
|
||||
s := &Store{db: db}
|
||||
ctx := context.Background()
|
||||
conv, _ := s.GetOrCreateConversation(ctx, "test:compact")
|
||||
shutdownCtx, shutdownCancel := context.WithCancel(context.Background())
|
||||
ce := &CompactionEngine{
|
||||
store: s,
|
||||
config: Config{},
|
||||
complete: mockCompleteFn,
|
||||
shutdownCtx: shutdownCtx,
|
||||
shutdownCancel: shutdownCancel,
|
||||
}
|
||||
convID := conv.ConversationID
|
||||
// Ensure async goroutines are stopped before database is closed.
|
||||
// Register cleanup here (after openTestDB) so it runs BEFORE openTestDB's db.Close().
|
||||
t.Cleanup(func() {
|
||||
shutdownCancel()
|
||||
// Wait for async condensed goroutine to finish (poll condensing map)
|
||||
deadline := time.Now().Add(2 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if _, exists := ce.condensing.Load(convID); !exists {
|
||||
break
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
})
|
||||
return ce, s, conv.ConversationID
|
||||
}
|
||||
|
||||
// newTestCompactionEngineWithStore creates a CompactionEngine with existing store.
|
||||
// Note: Caller is responsible for calling shutdownCancel when test ends.
|
||||
func newTestCompactionEngineWithStore(
|
||||
s *Store, complete CompleteFn,
|
||||
) (ce *CompactionEngine, shutdownCancel context.CancelFunc) {
|
||||
shutdownCtx, cancel := context.WithCancel(context.Background())
|
||||
return &CompactionEngine{
|
||||
store: s,
|
||||
config: Config{},
|
||||
complete: complete,
|
||||
shutdownCtx: shutdownCtx,
|
||||
shutdownCancel: cancel,
|
||||
}, cancel
|
||||
}
|
||||
|
||||
// mockCompleteFn returns a simple summary for testing
|
||||
var mockCompleteFn CompleteFn = func(ctx context.Context, prompt string, opts CompleteOptions) (string, error) {
|
||||
return "Mock summary of the conversation segment.", nil
|
||||
}
|
||||
|
||||
func TestNeedsCompaction(t *testing.T) {
|
||||
ce, s, convID := newTestCompactionEngine(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Empty context — no compaction needed
|
||||
needed, err := ce.NeedsCompaction(ctx, convID, 10000)
|
||||
if err != nil {
|
||||
t.Fatalf("NeedsCompaction: %v", err)
|
||||
}
|
||||
if needed {
|
||||
t.Error("expected no compaction for empty context")
|
||||
}
|
||||
|
||||
// Add messages to context, total tokens = 8000
|
||||
for i := 0; i < 8; i++ {
|
||||
m, _ := s.AddMessage(ctx, convID, "user", "test message content", 1000)
|
||||
s.AppendContextMessage(ctx, convID, m.ID)
|
||||
}
|
||||
|
||||
// Threshold = 0.75 × 10000 = 7500. We have 8000 tokens → needs compaction
|
||||
needed, err = ce.NeedsCompaction(ctx, convID, 10000)
|
||||
if err != nil {
|
||||
t.Fatalf("NeedsCompaction: %v", err)
|
||||
}
|
||||
if !needed {
|
||||
t.Error("expected compaction needed at 8000/10000 tokens (threshold 75%)")
|
||||
}
|
||||
|
||||
// Below threshold: 5000 / 10000 → no compaction
|
||||
s.UpsertContextItems(ctx, convID, nil) // clear
|
||||
for i := 0; i < 5; i++ {
|
||||
m, _ := s.AddMessage(ctx, convID, "user", "test", 1000)
|
||||
s.AppendContextMessage(ctx, convID, m.ID)
|
||||
}
|
||||
needed, _ = ce.NeedsCompaction(ctx, convID, 10000)
|
||||
if needed {
|
||||
t.Error("expected no compaction at 5000/10000 tokens")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactLeaf(t *testing.T) {
|
||||
ce, s, convID := newTestCompactionEngine(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create enough messages to trigger leaf compaction:
|
||||
// Need > FreshTailCount(32) evictable messages with >= LeafMinFanout(8) contiguous
|
||||
for i := 0; i < 40; i++ {
|
||||
m, _ := s.AddMessage(ctx, convID, "user", "message content for compaction test", 100)
|
||||
s.AppendContextMessage(ctx, convID, m.ID)
|
||||
}
|
||||
|
||||
// Compact
|
||||
result, err := ce.Compact(ctx, convID, CompactInput{})
|
||||
if err != nil {
|
||||
t.Fatalf("Compact: %v", err)
|
||||
}
|
||||
if result == nil {
|
||||
t.Fatal("expected non-nil result")
|
||||
}
|
||||
|
||||
// Should have created at least one leaf summary
|
||||
if result.LeafSummaries == 0 {
|
||||
t.Error("expected at least 1 leaf summary")
|
||||
}
|
||||
|
||||
// Context should now contain a summary item
|
||||
items, _ := s.GetContextItems(ctx, convID)
|
||||
foundSummary := false
|
||||
for _, item := range items {
|
||||
if item.ItemType == "summary" {
|
||||
foundSummary = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !foundSummary {
|
||||
t.Error("expected a summary in context_items after leaf compaction")
|
||||
}
|
||||
|
||||
// Some messages should have been replaced
|
||||
if len(result.SummariesCreated) == 0 {
|
||||
t.Error("expected at least 1 summary created")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactLeafNoCandidate(t *testing.T) {
|
||||
ce, _, convID := newTestCompactionEngine(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Too few messages to trigger leaf compaction
|
||||
m, _ := ce.store.AddMessage(ctx, convID, "user", "short", 10)
|
||||
ce.store.AppendContextMessage(ctx, convID, m.ID)
|
||||
|
||||
result, err := ce.Compact(ctx, convID, CompactInput{})
|
||||
if err != nil {
|
||||
t.Fatalf("Compact: %v", err)
|
||||
}
|
||||
if result == nil {
|
||||
t.Fatal("expected non-nil result even with no candidate")
|
||||
}
|
||||
if result.LeafSummaries != 0 {
|
||||
t.Errorf("LeafSummaries = %d, want 0 (too few messages)", result.LeafSummaries)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactCondensed(t *testing.T) {
|
||||
ce, s, convID := newTestCompactionEngine(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create enough leaf summaries and fresh messages to enable condensation
|
||||
leafIDs := make([]string, CondensedMinFanout)
|
||||
for i := 0; i < CondensedMinFanout; i++ {
|
||||
now := time.Now().UTC()
|
||||
summary, err := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: "leaf summary content " + time.Now().String(),
|
||||
TokenCount: 500,
|
||||
EarliestAt: &now,
|
||||
LatestAt: &now,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("CreateSummary %d: %v", i, err)
|
||||
}
|
||||
leafIDs[i] = summary.SummaryID
|
||||
s.AppendContextSummary(ctx, convID, summary.SummaryID)
|
||||
}
|
||||
|
||||
// Add enough fresh messages to have a fresh tail (>= FreshTailCount)
|
||||
for i := 0; i < FreshTailCount; i++ {
|
||||
m, _ := s.AddMessage(ctx, convID, "user", "fresh message", 10)
|
||||
s.AppendContextMessage(ctx, convID, m.ID)
|
||||
}
|
||||
|
||||
// Compact with force to trigger condensation
|
||||
_, err := ce.Compact(ctx, convID, CompactInput{Force: true})
|
||||
if err != nil {
|
||||
t.Fatalf("Compact: %v", err)
|
||||
}
|
||||
|
||||
// Wait for async condensed goroutine to complete
|
||||
if !waitForCondensed(ce, convID, 2*time.Second) {
|
||||
t.Fatal("timeout waiting for condensed compaction")
|
||||
}
|
||||
|
||||
// Should have created a condensed summary in the DB
|
||||
summaries, _ := s.GetSummariesByConversation(ctx, convID)
|
||||
foundCondensed := false
|
||||
for _, sum := range summaries {
|
||||
if sum.Kind == SummaryKindCondensed {
|
||||
foundCondensed = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !foundCondensed {
|
||||
t.Error("expected at least 1 condensed summary")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactCondensedDoesNotOrphanSummaryWhenCandidatesRemovedConcurrently(t *testing.T) {
|
||||
// Reproduce orphan bug: candidates found by selectOldestChunkAtDepth are removed
|
||||
// from context_items between candidate selection and ordinal range scan.
|
||||
// Use a slow CompleteFn with barrier sync to control timing.
|
||||
s := openTestStore(t)
|
||||
ctx := context.Background()
|
||||
conv, _ := s.GetOrCreateConversation(ctx, "test:orphan-race")
|
||||
convID := conv.ConversationID
|
||||
|
||||
// Create leaf summaries with enough tokens for condensation
|
||||
var leafIDs []string
|
||||
for i := 0; i < CondensedMinFanout; i++ {
|
||||
now := time.Now().UTC()
|
||||
sum, err := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: fmt.Sprintf("leaf summary %d", i),
|
||||
TokenCount: 500,
|
||||
EarliestAt: &now,
|
||||
LatestAt: &now,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("CreateSummary: %v", err)
|
||||
}
|
||||
leafIDs = append(leafIDs, sum.SummaryID)
|
||||
s.AppendContextSummary(ctx, convID, sum.SummaryID)
|
||||
}
|
||||
|
||||
// Add fresh tail so leaf summaries are in evictable range
|
||||
for i := 0; i < FreshTailCount+1; i++ {
|
||||
m, _ := s.AddMessage(ctx, convID, "user", "fresh", 10)
|
||||
s.AppendContextMessage(ctx, convID, m.ID)
|
||||
}
|
||||
|
||||
// Barrier: CompleteFn waits until test removes context_items, then returns
|
||||
var barrier1, barrier2 sync.WaitGroup
|
||||
barrier1.Add(1) // CompleteFn signals when called
|
||||
barrier2.Add(1) // test signals when context_items removed
|
||||
|
||||
slowComplete := func(ctx context.Context, prompt string, opts CompleteOptions) (string, error) {
|
||||
barrier1.Done() // signal: LLM called, candidates selected
|
||||
barrier2.Wait() // wait: test removes context_items
|
||||
return "Condensed summary.", nil
|
||||
}
|
||||
|
||||
ce, cancel := newTestCompactionEngineWithStore(s, slowComplete)
|
||||
t.Cleanup(func() {
|
||||
cancel()
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
})
|
||||
|
||||
// Run compactCondensed in background
|
||||
type compactResult struct {
|
||||
summaryID *string
|
||||
err error
|
||||
}
|
||||
resultCh := make(chan compactResult, 1)
|
||||
go func() {
|
||||
sid, err := ce.compactCondensed(context.Background(), convID)
|
||||
resultCh <- compactResult{summaryID: sid, err: err}
|
||||
}()
|
||||
|
||||
// Wait for CompleteFn to be called (candidates selected)
|
||||
barrier1.Wait()
|
||||
|
||||
// Remove leaf summaries from context_items (simulating concurrent replacement)
|
||||
items, _ := s.GetContextItems(ctx, convID)
|
||||
var preserved []ContextItem
|
||||
for _, item := range items {
|
||||
isLeaf := false
|
||||
for _, lid := range leafIDs {
|
||||
if item.SummaryID == lid {
|
||||
isLeaf = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !isLeaf {
|
||||
preserved = append(preserved, item)
|
||||
}
|
||||
}
|
||||
s.UpsertContextItems(ctx, convID, preserved)
|
||||
|
||||
// Let CompleteFn return
|
||||
barrier2.Done()
|
||||
|
||||
// Get result
|
||||
res := <-resultCh
|
||||
if res.err != nil {
|
||||
t.Fatalf("compactCondensed: %v", res.err)
|
||||
}
|
||||
|
||||
// With the bug: returns non-nil summaryID even though context_items has no matching ordinals
|
||||
// The fix: should return nil when startOrd == -1
|
||||
if res.summaryID != nil {
|
||||
t.Errorf("compactCondensed returned summaryID=%s, want nil (orphan created)", *res.summaryID)
|
||||
|
||||
// Verify the orphan exists in DB
|
||||
summary, _ := s.GetSummary(context.Background(), *res.summaryID)
|
||||
if summary != nil && summary.Kind == SummaryKindCondensed {
|
||||
// Check it's NOT in context_items (orphan)
|
||||
items2, _ := s.GetContextItems(context.Background(), convID)
|
||||
found := false
|
||||
for _, item := range items2 {
|
||||
if item.SummaryID == *res.summaryID {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Error("condensed summary exists in DB but not in context_items — orphan confirmed")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactUntilUnder(t *testing.T) {
|
||||
ce, s, convID := newTestCompactionEngine(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create many leaf summaries to ensure we can condense
|
||||
for i := 0; i < 8; i++ {
|
||||
now := time.Now().UTC()
|
||||
summary, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: "leaf summary for condensation test",
|
||||
TokenCount: 500,
|
||||
EarliestAt: &now,
|
||||
LatestAt: &now,
|
||||
})
|
||||
s.AppendContextSummary(ctx, convID, summary.SummaryID)
|
||||
}
|
||||
|
||||
// Force compact until under budget
|
||||
result, err := ce.CompactUntilUnder(ctx, convID, 2000)
|
||||
if err != nil {
|
||||
t.Fatalf("CompactUntilUnder: %v", err)
|
||||
}
|
||||
|
||||
if result == nil {
|
||||
t.Fatal("expected non-nil result")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectShallowestCondensationCandidate(t *testing.T) {
|
||||
ce, s, convID := newTestCompactionEngine(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create enough leaf summaries + fresh messages for candidates
|
||||
for i := 0; i < LeafMinFanout; i++ {
|
||||
summary, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: "leaf",
|
||||
TokenCount: 100,
|
||||
})
|
||||
s.AppendContextSummary(ctx, convID, summary.SummaryID)
|
||||
}
|
||||
|
||||
// Add fresh tail messages so summaries are in evictable range
|
||||
for i := 0; i < FreshTailCount+1; i++ {
|
||||
m, _ := s.AddMessage(ctx, convID, "user", "fresh", 5)
|
||||
s.AppendContextMessage(ctx, convID, m.ID)
|
||||
}
|
||||
|
||||
candidates, err := ce.selectShallowestCondensationCandidate(ctx, convID, false)
|
||||
if err != nil {
|
||||
t.Fatalf("selectShallowestCondensationCandidate: %v", err)
|
||||
}
|
||||
|
||||
// Should find leaf summaries at depth 0
|
||||
if len(candidates) < CondensedMinFanout {
|
||||
t.Errorf("candidates = %d, want >= %d", len(candidates), CondensedMinFanout)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectShallowestCondensationCandidateEmpty(t *testing.T) {
|
||||
ce, _, convID := newTestCompactionEngine(t)
|
||||
ctx := context.Background()
|
||||
|
||||
candidates, err := ce.selectShallowestCondensationCandidate(ctx, convID, false)
|
||||
if err != nil {
|
||||
t.Fatalf("selectShallowestCondensationCandidate: %v", err)
|
||||
}
|
||||
if len(candidates) != 0 {
|
||||
t.Errorf("candidates = %d, want 0 for empty context", len(candidates))
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactCondensedUsesSelectOldestChunk(t *testing.T) {
|
||||
// Verify that compactCondensed prefers ordinal-ordered chunks via selectOldestChunkAtDepth
|
||||
// rather than just grouping by depth without regard to order
|
||||
ce, s, convID := newTestCompactionEngine(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create interleaved summaries at depth 0 with a message in between:
|
||||
// sum1 (ordinal 100), msg (ordinal 200), sum2 (ordinal 300)
|
||||
|
||||
for i := 0; i < LeafMinFanout+2; i++ {
|
||||
now := time.Now().UTC()
|
||||
|
||||
s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: fmt.Sprintf("leaf summary %d", i),
|
||||
TokenCount: 100,
|
||||
EarliestAt: &now,
|
||||
LatestAt: &now,
|
||||
})
|
||||
}
|
||||
|
||||
// Insert a message between first two summaries to break contiguity
|
||||
// for selectShallowestCondensationCandidate but would still find all 3
|
||||
// but selectOldestChunkAtDepth should only find sum1 + sum2 (not sum3)
|
||||
|
||||
msg, _ := s.AddMessage(ctx, convID, "user", "interrupting message", 5)
|
||||
s.AppendContextMessage(ctx, convID, msg.ID)
|
||||
|
||||
// Run compactCondensed
|
||||
result, err := ce.compactCondensed(ctx, convID)
|
||||
if err != nil {
|
||||
t.Fatalf("compactCondensed: %v", err)
|
||||
}
|
||||
|
||||
// The result should have merged the two summaries at the start
|
||||
// (skipping the message in between), This proves ordinal-aware selection works.
|
||||
|
||||
_ = result // verify summary was created
|
||||
|
||||
if result != nil {
|
||||
summaries, _ := s.GetSummariesByConversation(ctx, convID)
|
||||
found := false
|
||||
for _, sum := range summaries {
|
||||
if sum.Kind == SummaryKindCondensed {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Error("expected condensed summary to be created via ordinal-aware selection")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactCondensedUsesOrdinalAwareSelection(t *testing.T) {
|
||||
ce, s, convID := newTestCompactionEngine(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create leaf summaries at depth 0 (total tokens >= CondensedTargetTokens)
|
||||
for i := 0; i < 5; i++ {
|
||||
summary, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: fmt.Sprintf("leaf summary %d", i),
|
||||
TokenCount: 500, // 5 × 500 = 2500 >= CondensedTargetTokens (2000)
|
||||
})
|
||||
s.AppendContextSummary(ctx, convID, summary.SummaryID)
|
||||
}
|
||||
|
||||
// Add fresh tail
|
||||
for i := 0; i < FreshTailCount+1; i++ {
|
||||
m, _ := s.AddMessage(ctx, convID, "user", "fresh", 5)
|
||||
s.AppendContextMessage(ctx, convID, m.ID)
|
||||
}
|
||||
|
||||
chunk, err := ce.selectOldestChunkAtDepth(ctx, convID, 0)
|
||||
if err != nil {
|
||||
t.Fatalf("selectOldestChunkAtDepth: %v", err)
|
||||
}
|
||||
if len(chunk) < 2 {
|
||||
t.Errorf("chunk length = %d, want >= 2 contiguous summaries", len(chunk))
|
||||
}
|
||||
for _, s := range chunk {
|
||||
if s.Depth != 0 {
|
||||
t.Errorf("got depth %d, want 0", s.Depth)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectOldestChunkAtDepthBreaksOnMessage(t *testing.T) {
|
||||
ce, s, convID := newTestCompactionEngine(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create 3 summaries, then a message, then 3 more summaries
|
||||
for i := 0; i < 3; i++ {
|
||||
summary, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: fmt.Sprintf("leaf %d", i),
|
||||
TokenCount: 100,
|
||||
})
|
||||
s.AppendContextSummary(ctx, convID, summary.SummaryID)
|
||||
}
|
||||
msg, _ := s.AddMessage(ctx, convID, "user", "break", 10)
|
||||
s.AppendContextMessage(ctx, convID, msg.ID)
|
||||
for i := 0; i < 3; i++ {
|
||||
summary, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: fmt.Sprintf("leaf-after %d", i),
|
||||
TokenCount: 100,
|
||||
})
|
||||
s.AppendContextSummary(ctx, convID, summary.SummaryID)
|
||||
}
|
||||
for i := 0; i < FreshTailCount+1; i++ {
|
||||
m, _ := s.AddMessage(ctx, convID, "user", "fresh", 5)
|
||||
s.AppendContextMessage(ctx, convID, m.ID)
|
||||
}
|
||||
|
||||
chunk, _ := ce.selectOldestChunkAtDepth(ctx, convID, 0)
|
||||
if len(chunk) > 3 {
|
||||
t.Errorf("chunk length = %d, want <= 3 (message breaks chain)", len(chunk))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectOldestChunkAtDepthMinTokens(t *testing.T) {
|
||||
ce, s, convID := newTestCompactionEngine(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create summaries with very low token counts (total < 2000)
|
||||
for i := 0; i < 5; i++ {
|
||||
summary, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: fmt.Sprintf("tiny summary %d", i),
|
||||
TokenCount: 50, // very small
|
||||
})
|
||||
s.AppendContextSummary(ctx, convID, summary.SummaryID)
|
||||
}
|
||||
|
||||
// Add fresh tail to protect from compaction
|
||||
for i := 0; i < FreshTailCount+1; i++ {
|
||||
m, _ := s.AddMessage(ctx, convID, "user", fmt.Sprintf("tail %d", i), 10)
|
||||
s.AppendContextMessage(ctx, convID, m.ID)
|
||||
}
|
||||
|
||||
// Should return nil because total tokens (250) < 2000 minimum
|
||||
chunk, err := ce.selectOldestChunkAtDepth(ctx, convID, 0)
|
||||
if err != nil {
|
||||
t.Fatalf("selectOldestChunkAtDepth: %v", err)
|
||||
}
|
||||
if len(chunk) > 0 {
|
||||
t.Errorf("expected empty chunk when tokens < 2000, got %d summaries", len(chunk))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectOldestChunkAtDepthPassesMinTokens(t *testing.T) {
|
||||
ce, s, convID := newTestCompactionEngine(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create summaries with enough tokens (total >= 2000)
|
||||
for i := 0; i < 5; i++ {
|
||||
summary, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: fmt.Sprintf(
|
||||
"substantial summary with enough content to meet minimum token threshold for condensation candidate %d",
|
||||
i,
|
||||
),
|
||||
TokenCount: 500, // 5 × 500 = 2500 >= 2000
|
||||
})
|
||||
s.AppendContextSummary(ctx, convID, summary.SummaryID)
|
||||
}
|
||||
|
||||
// Add fresh tail
|
||||
for i := 0; i < FreshTailCount+1; i++ {
|
||||
m, _ := s.AddMessage(ctx, convID, "user", fmt.Sprintf("tail %d", i), 10)
|
||||
s.AppendContextMessage(ctx, convID, m.ID)
|
||||
}
|
||||
|
||||
// Should return chunk because total tokens (2500) >= 2000
|
||||
chunk, err := ce.selectOldestChunkAtDepth(ctx, convID, 0)
|
||||
if err != nil {
|
||||
t.Fatalf("selectOldestChunkAtDepth: %v", err)
|
||||
}
|
||||
if len(chunk) == 0 {
|
||||
t.Error("expected non-empty chunk when tokens >= 2000")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateLeafSummary(t *testing.T) {
|
||||
ce, _, _ := newTestCompactionEngine(t)
|
||||
ctx := context.Background()
|
||||
|
||||
msgs := []Message{
|
||||
{Role: "user", Content: "hello world", TokenCount: 5},
|
||||
{Role: "assistant", Content: "hi there", TokenCount: 5},
|
||||
}
|
||||
|
||||
content, err := ce.generateLeafSummary(ctx, msgs, "")
|
||||
if err != nil {
|
||||
t.Fatalf("generateLeafSummary: %v", err)
|
||||
}
|
||||
if content == "" {
|
||||
t.Error("expected non-empty summary content")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateLeafSummaryEscalationToAggressive(t *testing.T) {
|
||||
// Level 1 returns summary that's too large (tokens >= input), should escalate to level 2
|
||||
var calls []string
|
||||
escalateComplete := func(ctx context.Context, prompt string, opts CompleteOptions) (string, error) {
|
||||
if contains(prompt, "Aggressive summary policy") {
|
||||
calls = append(calls, "aggressive")
|
||||
return "Short aggressive summary.", nil
|
||||
}
|
||||
calls = append(calls, "normal")
|
||||
// Return a very long summary to trigger escalation
|
||||
longContent := make([]byte, 5000)
|
||||
for i := range longContent {
|
||||
longContent[i] = 'x'
|
||||
}
|
||||
return string(longContent), nil
|
||||
}
|
||||
|
||||
s := openTestStore(t)
|
||||
ce, _ := newTestCompactionEngineWithStore(s, escalateComplete)
|
||||
|
||||
msgs := []Message{
|
||||
{Role: "user", Content: "hello world", TokenCount: 10},
|
||||
{Role: "assistant", Content: "response", TokenCount: 10},
|
||||
}
|
||||
|
||||
content, err := ce.generateLeafSummary(context.Background(), msgs, "")
|
||||
if err != nil {
|
||||
t.Fatalf("generateLeafSummary: %v", err)
|
||||
}
|
||||
if content == "" {
|
||||
t.Error("expected non-empty summary content")
|
||||
}
|
||||
// Should have called both normal and aggressive
|
||||
foundNormal := false
|
||||
foundAggressive := false
|
||||
for _, c := range calls {
|
||||
if c == "normal" {
|
||||
foundNormal = true
|
||||
}
|
||||
if c == "aggressive" {
|
||||
foundAggressive = true
|
||||
}
|
||||
}
|
||||
if !foundNormal {
|
||||
t.Error("expected normal LLM call")
|
||||
}
|
||||
if !foundAggressive {
|
||||
t.Error("expected aggressive LLM call (level 2 escalation)")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateLeafSummaryEscalationToTruncation(t *testing.T) {
|
||||
// Both normal and aggressive return empty, should escalate to level 3 truncation
|
||||
emptyComplete := func(ctx context.Context, prompt string, opts CompleteOptions) (string, error) {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
s := openTestStore(t)
|
||||
ce, _ := newTestCompactionEngineWithStore(s, emptyComplete)
|
||||
|
||||
msgs := []Message{
|
||||
{Role: "user", Content: "hello world from test", TokenCount: 10},
|
||||
{Role: "assistant", Content: "response text here", TokenCount: 10},
|
||||
}
|
||||
|
||||
content, err := ce.generateLeafSummary(context.Background(), msgs, "")
|
||||
if err != nil {
|
||||
t.Fatalf("generateLeafSummary: %v", err)
|
||||
}
|
||||
// Level 3 truncation should have produced something
|
||||
if content == "" {
|
||||
t.Error("expected non-empty content from level 3 truncation fallback")
|
||||
}
|
||||
if !contains(content, "Truncated from") {
|
||||
t.Errorf("expected truncation marker in content: %q", content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateCondensedSummary(t *testing.T) {
|
||||
ce, _, _ := newTestCompactionEngine(t)
|
||||
ctx := context.Background()
|
||||
|
||||
summaries := []Summary{
|
||||
{SummaryID: "sum_a", Content: "first summary", TokenCount: 100},
|
||||
{SummaryID: "sum_b", Content: "second summary", TokenCount: 100},
|
||||
}
|
||||
|
||||
content, err := ce.generateCondensedSummary(ctx, summaries)
|
||||
if err != nil {
|
||||
t.Fatalf("generateCondensedSummary: %v", err)
|
||||
}
|
||||
if content == "" {
|
||||
t.Error("expected non-empty condensed summary content")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateCondensedSummaryEscalation(t *testing.T) {
|
||||
// When LLM returns empty, should fall back to deterministic concatenation
|
||||
emptyComplete := func(ctx context.Context, prompt string, opts CompleteOptions) (string, error) {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
s := openTestStore(t)
|
||||
ce, _ := newTestCompactionEngineWithStore(s, emptyComplete)
|
||||
|
||||
summaries := []Summary{
|
||||
{SummaryID: "sum_a", Content: "first summary text", TokenCount: 50},
|
||||
{SummaryID: "sum_b", Content: "second summary text", TokenCount: 50},
|
||||
}
|
||||
|
||||
content, err := ce.generateCondensedSummary(context.Background(), summaries)
|
||||
if err != nil {
|
||||
t.Fatalf("generateCondensedSummary: %v", err)
|
||||
}
|
||||
// Should fall back to concatenation
|
||||
if content == "" {
|
||||
t.Error("expected non-empty content from fallback")
|
||||
}
|
||||
}
|
||||
|
||||
// --- Async Condensed Compaction (Phase 2) ---
|
||||
|
||||
func TestCompactAsyncReturnsBeforeCondensed(t *testing.T) {
|
||||
// Use a slow CompleteFn to verify Compact returns before condensed finishes
|
||||
var callCount int32
|
||||
slowComplete := func(ctx context.Context, prompt string, opts CompleteOptions) (string, error) {
|
||||
atomic.AddInt32(&callCount, 1)
|
||||
time.Sleep(500 * time.Millisecond) // simulate slow LLM
|
||||
return "Slow condensed summary.", nil
|
||||
}
|
||||
|
||||
s := openTestStore(t)
|
||||
ctx := context.Background()
|
||||
conv, _ := s.GetOrCreateConversation(ctx, "test:async")
|
||||
convID := conv.ConversationID
|
||||
|
||||
ce, cancel := newTestCompactionEngineWithStore(s, slowComplete)
|
||||
t.Cleanup(func() {
|
||||
cancel()
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
})
|
||||
|
||||
// Create enough leaf summaries for condensation + fresh tail
|
||||
for i := 0; i < CondensedMinFanout; i++ {
|
||||
now := time.Now().UTC()
|
||||
summary, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: "leaf for async test",
|
||||
TokenCount: 500,
|
||||
EarliestAt: &now,
|
||||
LatestAt: &now,
|
||||
})
|
||||
s.AppendContextSummary(ctx, convID, summary.SummaryID)
|
||||
}
|
||||
for i := 0; i < FreshTailCount; i++ {
|
||||
m, _ := s.AddMessage(ctx, convID, "user", "fresh", 10)
|
||||
s.AppendContextMessage(ctx, convID, m.ID)
|
||||
}
|
||||
|
||||
// Compact with force — should return quickly, condensed runs async
|
||||
start := time.Now()
|
||||
result, err := ce.Compact(ctx, convID, CompactInput{Force: true})
|
||||
elapsed := time.Since(start)
|
||||
|
||||
if err != nil {
|
||||
t.Fatalf("Compact: %v", err)
|
||||
}
|
||||
if result == nil {
|
||||
t.Fatal("expected non-nil result")
|
||||
}
|
||||
|
||||
// Should return well before the 500ms LLM call
|
||||
if elapsed > 200*time.Millisecond {
|
||||
t.Errorf("Compact took %v, should return before async condensed finishes", elapsed)
|
||||
}
|
||||
|
||||
// Wait for async to complete
|
||||
time.Sleep(800 * time.Millisecond)
|
||||
|
||||
// Verify condensed summary was created by background goroutine
|
||||
summaries, _ := s.GetSummariesByConversation(ctx, convID)
|
||||
foundCondensed := false
|
||||
for _, sum := range summaries {
|
||||
if sum.Kind == SummaryKindCondensed {
|
||||
foundCondensed = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !foundCondensed {
|
||||
t.Error("expected at least one condensed summary from async Phase 2")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactAsyncDedup(t *testing.T) {
|
||||
var callCount int32
|
||||
slowComplete := func(ctx context.Context, prompt string, opts CompleteOptions) (string, error) {
|
||||
atomic.AddInt32(&callCount, 1)
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
return "Slow condensed summary.", nil
|
||||
}
|
||||
|
||||
s := openTestStore(t)
|
||||
ctx := context.Background()
|
||||
conv, _ := s.GetOrCreateConversation(ctx, "test:dedup")
|
||||
convID := conv.ConversationID
|
||||
|
||||
ce, cancel := newTestCompactionEngineWithStore(s, slowComplete)
|
||||
t.Cleanup(func() {
|
||||
cancel()
|
||||
waitForCondensed(ce, convID, 2*time.Second)
|
||||
})
|
||||
|
||||
// Create conditions for condensed compaction
|
||||
for i := 0; i < CondensedMinFanout; i++ {
|
||||
now := time.Now().UTC()
|
||||
summary, _ := s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: "leaf for dedup",
|
||||
TokenCount: 500,
|
||||
EarliestAt: &now,
|
||||
LatestAt: &now,
|
||||
})
|
||||
s.AppendContextSummary(ctx, convID, summary.SummaryID)
|
||||
}
|
||||
for i := 0; i < FreshTailCount; i++ {
|
||||
m, _ := s.AddMessage(ctx, convID, "user", "fresh", 10)
|
||||
s.AppendContextMessage(ctx, convID, m.ID)
|
||||
}
|
||||
|
||||
// Call Compact twice rapidly
|
||||
ce.Compact(ctx, convID, CompactInput{Force: true})
|
||||
ce.Compact(ctx, convID, CompactInput{Force: true})
|
||||
|
||||
// Wait for async to finish
|
||||
time.Sleep(600 * time.Millisecond)
|
||||
|
||||
// LLM should only be called once for condensed (dedup)
|
||||
// callCount may be 0 if no leaf was created (only condensed in goroutine)
|
||||
// The key is that we don't get 2+ condensed calls
|
||||
if atomic.LoadInt32(&callCount) > 1 {
|
||||
t.Errorf("LLM called %d times, expected at most 1 (dedup)", callCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactLeafForceBypassesFreshTail(t *testing.T) {
|
||||
// Spec: compactLeaf with force=true should bypass FreshTailCount protection
|
||||
// so CompactUntilUnder can compress messages inside the fresh tail
|
||||
ce, s, convID := newTestCompactionEngine(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create exactly FreshTailCount+4 messages (36 total)
|
||||
// Without force: all messages are in fresh tail → no candidate
|
||||
// With force: should compact the oldest messages
|
||||
total := FreshTailCount + 4
|
||||
for i := 0; i < total; i++ {
|
||||
m, _ := s.AddMessage(ctx, convID, "user", fmt.Sprintf("message %d for force test", i), 100)
|
||||
s.AppendContextMessage(ctx, convID, m.ID)
|
||||
}
|
||||
|
||||
// Without force: should return nil (all in fresh tail)
|
||||
summaryID, err := ce.compactLeaf(ctx, convID)
|
||||
if err != nil {
|
||||
t.Fatalf("compactLeaf no-force: %v", err)
|
||||
}
|
||||
if summaryID != nil {
|
||||
t.Error("expected nil without force (all messages in fresh tail)")
|
||||
}
|
||||
|
||||
// With force: should compact despite fresh tail protection
|
||||
summaryID, err = ce.compactLeaf(ctx, convID, true)
|
||||
if err != nil {
|
||||
t.Fatalf("compactLeaf force: %v", err)
|
||||
}
|
||||
if summaryID == nil {
|
||||
t.Error("expected summary with force=true (bypasses fresh tail)")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactLeafAccumulatesUpToLeafChunkTokens(t *testing.T) {
|
||||
// Spec: compactLeaf should accumulate messages up to LeafChunkTokens before stopping
|
||||
// It should NOT take the entire contiguous chunk regardless of token count
|
||||
ce, s, convID := newTestCompactionEngine(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create messages totaling far more than LeafChunkTokens (20000)
|
||||
// Each message is ~500 tokens, create 80 messages = 40000 tokens
|
||||
for i := 0; i < 80; i++ {
|
||||
m, _ := s.AddMessage(
|
||||
ctx,
|
||||
convID,
|
||||
"user",
|
||||
fmt.Sprintf(
|
||||
"message %d with lots of content to make it big enough for token counting purposes and this should be a substantial message body that represents a meaningful conversation turn",
|
||||
i,
|
||||
),
|
||||
500,
|
||||
)
|
||||
s.AppendContextMessage(ctx, convID, m.ID)
|
||||
}
|
||||
|
||||
summaryID, err := ce.compactLeaf(ctx, convID)
|
||||
if err != nil {
|
||||
t.Fatalf("compactLeaf: %v", err)
|
||||
}
|
||||
if summaryID == nil {
|
||||
t.Fatal("expected a summary to be created")
|
||||
}
|
||||
|
||||
// The source messages that were compacted should total roughly LeafChunkTokens (20000),
|
||||
// not the entire 40000 tokens worth of messages
|
||||
summary, _ := s.GetSummary(ctx, *summaryID)
|
||||
if summary == nil {
|
||||
t.Fatal("summary not found")
|
||||
}
|
||||
|
||||
// Source message tokens should be roughly <= LeafChunkTokens (20000)
|
||||
// Spec says: "Stop when accumulated tokens >= LeafChunkTokens"
|
||||
if summary.SourceMessageTokenCount > LeafChunkTokens {
|
||||
t.Errorf("source tokens = %d, should be <= LeafChunkTokens (%d)",
|
||||
summary.SourceMessageTokenCount, LeafChunkTokens)
|
||||
}
|
||||
}
|
||||
30
pkg/seahorse/short_constants.go
Normal file
30
pkg/seahorse/short_constants.go
Normal file
|
|
@ -0,0 +1,30 @@
|
|||
package seahorse
|
||||
|
||||
// Short-term memory configuration constants — all are experience-based defaults.
|
||||
|
||||
const (
|
||||
// OrdinalStep is the gap between ordinals in context_items.
|
||||
// Insert at midpoint; resequence only when precision exhausted.
|
||||
OrdinalStep = 100
|
||||
|
||||
// ContextThreshold is the compaction trigger for the context window.
|
||||
ContextThreshold float64 = 0.75 // Compact at 75% of context window
|
||||
FreshTailCount int = 32 // Recent messages protected from compaction
|
||||
|
||||
// LeafMinFanout is the fanout parameter.
|
||||
LeafMinFanout int = 8 // Min messages per leaf summary
|
||||
CondensedMinFanout int = 4 // Min summaries per condensed
|
||||
CondensedMinFanoutHard int = 2 // Min for forced compaction
|
||||
|
||||
// LeafChunkTokens is the token target.
|
||||
LeafChunkTokens int = 20000 // Max tokens per leaf chunk
|
||||
LeafTargetTokens int = 1200 // Target tokens for leaf summaries
|
||||
CondensedTargetTokens int = 2000 // Target tokens for condensed summaries
|
||||
MaxExpandTokens int = 4000 // Token cap for expansion queries
|
||||
|
||||
// MaxCompactIterations caps CompactUntilUnder to prevent infinite loops.
|
||||
// Each iteration reduces ~4x tokens via leaf (8:1) or condensed (4:1) compaction.
|
||||
// With a 200k token context window and 75% threshold, ~20 iterations is enough
|
||||
// for any realistic scenario. If exceeded, the issue is logged as a warning.
|
||||
MaxCompactIterations int = 20
|
||||
)
|
||||
568
pkg/seahorse/short_engine.go
Normal file
568
pkg/seahorse/short_engine.go
Normal file
|
|
@ -0,0 +1,568 @@
|
|||
package seahorse
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
_ "modernc.org/sqlite"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
)
|
||||
|
||||
// Config holds engine configuration.
|
||||
type Config struct {
|
||||
DBPath string `json:"dbPath"`
|
||||
IgnoreSessionPatterns []string `json:"ignoreSessionPatterns,omitempty"`
|
||||
StatelessSessionPatterns []string `json:"statelessSessionPatterns,omitempty"`
|
||||
}
|
||||
|
||||
// CompleteFn is the LLM completion function type.
|
||||
type CompleteFn func(ctx context.Context, prompt string, opts CompleteOptions) (string, error)
|
||||
|
||||
// CompleteOptions holds LLM completion parameters.
|
||||
type CompleteOptions struct {
|
||||
Model string
|
||||
MaxTokens int
|
||||
Temperature float64
|
||||
}
|
||||
|
||||
// IngestResult is the result of message ingestion.
|
||||
type IngestResult struct {
|
||||
MessageCount int `json:"messageCount"`
|
||||
TokenCount int `json:"tokenCount"`
|
||||
}
|
||||
|
||||
// AssembleInput controls context assembly.
|
||||
type AssembleInput struct {
|
||||
Budget int `json:"budget"`
|
||||
Query string `json:"query,omitempty"`
|
||||
}
|
||||
|
||||
// AssembleResult contains assembled context.
|
||||
type AssembleResult struct {
|
||||
Messages []Message `json:"messages"`
|
||||
Summary string `json:"summary"` // formatted XML summaries + system prompt addition
|
||||
}
|
||||
|
||||
const numSessionShards = 256
|
||||
|
||||
// Engine is the main short-term memory engine.
|
||||
type Engine struct {
|
||||
store *Store
|
||||
compaction *CompactionEngine
|
||||
compactionMu sync.Mutex
|
||||
assembler *Assembler
|
||||
assemblerMu sync.Mutex
|
||||
retrieval *RetrievalEngine
|
||||
config Config
|
||||
complete CompleteFn
|
||||
ignorePatterns []*regexp.Regexp
|
||||
statelessPatterns []*regexp.Regexp
|
||||
sessionShards [numSessionShards]struct {
|
||||
mu sync.Mutex
|
||||
}
|
||||
}
|
||||
|
||||
// CompactionEngine handles LLM-based summarization (defined in short_compaction.go).
|
||||
type CompactionEngine struct {
|
||||
store *Store
|
||||
config Config
|
||||
complete CompleteFn
|
||||
condensing sync.Map // map[int64]struct{} — dedup for async condensed goroutines
|
||||
shutdownCtx context.Context
|
||||
shutdownCancel context.CancelFunc
|
||||
}
|
||||
|
||||
// Assembler handles budget-aware context assembly (defined in short_assembler.go).
|
||||
type Assembler struct {
|
||||
store *Store
|
||||
config Config
|
||||
}
|
||||
|
||||
// RetrievalEngine handles search and expansion (defined in short_retrieval.go).
|
||||
type RetrievalEngine struct {
|
||||
store *Store
|
||||
config Config
|
||||
}
|
||||
|
||||
// Store returns the underlying store for direct access.
|
||||
func (r *RetrievalEngine) Store() *Store {
|
||||
return r.store
|
||||
}
|
||||
|
||||
// NewEngine creates a new short-term memory engine.
|
||||
func NewEngine(config Config, completeFn CompleteFn) (*Engine, error) {
|
||||
dir := filepath.Dir(config.DBPath)
|
||||
if dir != "" && dir != "." {
|
||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||
return nil, fmt.Errorf("create db directory: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
db, err := sql.Open("sqlite", config.DBPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open db: %w", err)
|
||||
}
|
||||
|
||||
// Configure SQLite for concurrent access
|
||||
if _, err := db.Exec("PRAGMA journal_mode = WAL;"); err != nil {
|
||||
db.Close()
|
||||
return nil, fmt.Errorf("enable WAL: %w", err)
|
||||
}
|
||||
if _, err := db.Exec("PRAGMA busy_timeout = 5000;"); err != nil {
|
||||
db.Close()
|
||||
return nil, fmt.Errorf("set busy_timeout: %w", err)
|
||||
}
|
||||
if _, err := db.Exec("PRAGMA synchronous = NORMAL;"); err != nil {
|
||||
db.Close()
|
||||
return nil, fmt.Errorf("set synchronous: %w", err)
|
||||
}
|
||||
|
||||
if err := runSchema(db); err != nil {
|
||||
db.Close()
|
||||
return nil, fmt.Errorf("migrations: %w", err)
|
||||
}
|
||||
|
||||
store := &Store{db: db}
|
||||
|
||||
// Prepend hardcoded ignore patterns (spec lines 1326-1328)
|
||||
ignorePatterns := make([]string, 0, 1+len(config.IgnoreSessionPatterns))
|
||||
ignorePatterns = append(ignorePatterns, "heartbeat")
|
||||
ignorePatterns = append(ignorePatterns, config.IgnoreSessionPatterns...)
|
||||
|
||||
retrieval := &RetrievalEngine{store: store, config: config}
|
||||
|
||||
return &Engine{
|
||||
store: store,
|
||||
compaction: nil,
|
||||
assembler: nil,
|
||||
retrieval: retrieval,
|
||||
config: config,
|
||||
complete: completeFn,
|
||||
ignorePatterns: compileSessionPatterns(ignorePatterns),
|
||||
statelessPatterns: compileSessionPatterns(config.StatelessSessionPatterns),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// compileSessionPattern converts a glob pattern to a compiled regex.
|
||||
// Pattern rules:
|
||||
// - * matches any sequence of non-colon characters ([^:]*)
|
||||
// - ** matches any sequence of characters including colons (.*)
|
||||
// - All other characters are treated literally
|
||||
// - Pattern is anchored (^...$)
|
||||
func compileSessionPattern(pattern string) *regexp.Regexp {
|
||||
var b strings.Builder
|
||||
b.WriteByte('^')
|
||||
|
||||
i := 0
|
||||
for i < len(pattern) {
|
||||
if i+1 < len(pattern) && pattern[i] == '*' && pattern[i+1] == '*' {
|
||||
b.WriteString(".*")
|
||||
i += 2
|
||||
continue
|
||||
}
|
||||
if pattern[i] == '*' {
|
||||
b.WriteString("[^:]*")
|
||||
i++
|
||||
continue
|
||||
}
|
||||
b.WriteString(regexp.QuoteMeta(string(pattern[i])))
|
||||
i++
|
||||
}
|
||||
|
||||
b.WriteByte('$')
|
||||
return regexp.MustCompile(b.String())
|
||||
}
|
||||
|
||||
// compileSessionPatterns compiles multiple glob patterns into regex patterns.
|
||||
func compileSessionPatterns(patterns []string) []*regexp.Regexp {
|
||||
result := make([]*regexp.Regexp, 0, len(patterns))
|
||||
for _, p := range patterns {
|
||||
if p == "" {
|
||||
continue
|
||||
}
|
||||
result = append(result, compileSessionPattern(p))
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// shouldIgnoreSession returns true if the session key matches any ignore pattern.
|
||||
func (e *Engine) shouldIgnoreSession(sessionKey string) bool {
|
||||
for _, p := range e.ignorePatterns {
|
||||
if p.MatchString(sessionKey) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// isStatelessSession returns true if the session key matches any stateless pattern.
|
||||
func (e *Engine) isStatelessSession(sessionKey string) bool {
|
||||
for _, p := range e.statelessPatterns {
|
||||
if p.MatchString(sessionKey) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// fnv32 computes FNV-1a 32-bit hash for session key sharding.
|
||||
func fnv32(key string) uint32 {
|
||||
h := uint32(2166136261)
|
||||
for _, c := range key {
|
||||
h ^= uint32(c)
|
||||
h *= 16777619
|
||||
}
|
||||
return h
|
||||
}
|
||||
|
||||
// getSessionMutex returns the sharded mutex for a session key.
|
||||
func (e *Engine) getSessionMutex(sessionKey string) *sync.Mutex {
|
||||
h := fnv32(sessionKey)
|
||||
shard := h % numSessionShards
|
||||
return &e.sessionShards[shard].mu
|
||||
}
|
||||
|
||||
// Ingest adds messages to a conversation identified by sessionKey.
|
||||
func (e *Engine) Ingest(ctx context.Context, sessionKey string, messages []Message) (*IngestResult, error) {
|
||||
if e.shouldIgnoreSession(sessionKey) {
|
||||
return nil, nil
|
||||
}
|
||||
if e.isStatelessSession(sessionKey) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
mu := e.getSessionMutex(sessionKey)
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
|
||||
conv, err := e.store.GetOrCreateConversation(ctx, sessionKey)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get conversation: %w", err)
|
||||
}
|
||||
|
||||
var totalTokens int
|
||||
var msgIDs []int64
|
||||
for _, msg := range messages {
|
||||
var added *Message
|
||||
var err error
|
||||
if len(msg.Parts) > 0 {
|
||||
added, err = e.store.AddMessageWithParts(ctx, conv.ConversationID, msg.Role, msg.Parts, msg.TokenCount)
|
||||
} else {
|
||||
added, err = e.store.AddMessage(ctx, conv.ConversationID, msg.Role, msg.Content, msg.TokenCount)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("add message: %w", err)
|
||||
}
|
||||
totalTokens += msg.TokenCount
|
||||
msgIDs = append(msgIDs, added.ID)
|
||||
}
|
||||
|
||||
// Append to context_items using actual inserted IDs
|
||||
if err := e.store.AppendContextMessages(ctx, conv.ConversationID, msgIDs); err != nil {
|
||||
return nil, fmt.Errorf("append context: %w", err)
|
||||
}
|
||||
|
||||
logger.InfoCF("seahorse", "ingest", map[string]any{
|
||||
"conv_id": conv.ConversationID,
|
||||
"messages": len(messages),
|
||||
"tokens": totalTokens,
|
||||
})
|
||||
return &IngestResult{
|
||||
MessageCount: len(messages),
|
||||
TokenCount: totalTokens,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Close releases resources.
|
||||
func (e *Engine) Close() error {
|
||||
// Signal compaction goroutines to stop
|
||||
if e.compaction != nil {
|
||||
e.compaction.Close()
|
||||
}
|
||||
if e.store != nil && e.store.db != nil {
|
||||
return e.store.db.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetRetrieval returns the retrieval engine for tool implementations.
|
||||
func (e *Engine) GetRetrieval() *RetrievalEngine {
|
||||
return e.retrieval
|
||||
}
|
||||
|
||||
// Assemble builds budget-constrained context for a session.
|
||||
func (e *Engine) Assemble(ctx context.Context, sessionKey string, input AssembleInput) (*AssembleResult, error) {
|
||||
if e.shouldIgnoreSession(sessionKey) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
conv, err := e.store.GetOrCreateConversation(ctx, sessionKey)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get conversation: %w", err)
|
||||
}
|
||||
|
||||
e.initAssemblerOnce()
|
||||
return e.assembler.Assemble(ctx, conv.ConversationID, input)
|
||||
}
|
||||
|
||||
// Compact compresses conversation history for a session.
|
||||
func (e *Engine) Compact(ctx context.Context, sessionKey string, input CompactInput) (*CompactResult, error) {
|
||||
if e.shouldIgnoreSession(sessionKey) || e.isStatelessSession(sessionKey) {
|
||||
return &CompactResult{}, nil
|
||||
}
|
||||
|
||||
conv, err := e.store.GetOrCreateConversation(ctx, sessionKey)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get conversation: %w", err)
|
||||
}
|
||||
|
||||
e.initCompactionOnce()
|
||||
return e.compaction.Compact(ctx, conv.ConversationID, input)
|
||||
}
|
||||
|
||||
// CompactUntilUnder aggressively compacts until context is under budget.
|
||||
// Used for emergency compaction after LLM overflow (retry reason).
|
||||
func (e *Engine) CompactUntilUnder(ctx context.Context, sessionKey string, budget int) (*CompactResult, error) {
|
||||
if e.shouldIgnoreSession(sessionKey) || e.isStatelessSession(sessionKey) {
|
||||
return &CompactResult{}, nil
|
||||
}
|
||||
|
||||
conv, err := e.store.GetOrCreateConversation(ctx, sessionKey)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get conversation: %w", err)
|
||||
}
|
||||
|
||||
e.initCompactionOnce()
|
||||
return e.compaction.CompactUntilUnder(ctx, conv.ConversationID, budget)
|
||||
}
|
||||
|
||||
// initCompactionOnce lazily initializes the compaction engine.
|
||||
func (e *Engine) initCompactionOnce() {
|
||||
if e.compaction == nil {
|
||||
e.compactionMu.Lock()
|
||||
defer e.compactionMu.Unlock()
|
||||
if e.compaction == nil {
|
||||
shutdownCtx, shutdownCancel := context.WithCancel(context.Background())
|
||||
e.compaction = &CompactionEngine{
|
||||
store: e.store,
|
||||
config: e.config,
|
||||
complete: e.complete,
|
||||
shutdownCtx: shutdownCtx,
|
||||
shutdownCancel: shutdownCancel,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// initAssemblerOnce lazily initializes the assembler.
|
||||
func (e *Engine) initAssemblerOnce() {
|
||||
if e.assembler == nil {
|
||||
e.assemblerMu.Lock()
|
||||
defer e.assemblerMu.Unlock()
|
||||
if e.assembler == nil {
|
||||
e.assembler = &Assembler{store: e.store, config: e.config}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// IngestMessages is an alias for Ingest.
|
||||
func (e *Engine) IngestMessages(ctx context.Context, sessionKey string, messages []Message) (*IngestResult, error) {
|
||||
return e.Ingest(ctx, sessionKey, messages)
|
||||
}
|
||||
|
||||
// Bootstrap reconciles a session's messages with the database.
|
||||
// Called once at startup for each known session.
|
||||
// Bootstrap reconciles JSONL history with SQLite by ingesting only the delta.
|
||||
// Simple approach: find longest matching prefix and append delta.
|
||||
// If any mismatch is detected, clear and rebuild.
|
||||
func (e *Engine) Bootstrap(ctx context.Context, sessionKey string, messages []Message) error {
|
||||
if e.shouldIgnoreSession(sessionKey) {
|
||||
return nil
|
||||
}
|
||||
if e.isStatelessSession(sessionKey) {
|
||||
return nil
|
||||
}
|
||||
if len(messages) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
conv, err := e.store.GetOrCreateConversation(ctx, sessionKey)
|
||||
if err != nil {
|
||||
return fmt.Errorf("bootstrap: get conversation: %w", err)
|
||||
}
|
||||
|
||||
// Get messages already in DB
|
||||
dbMsgs, err := e.store.GetMessages(ctx, conv.ConversationID, len(messages), 0)
|
||||
if err != nil {
|
||||
return fmt.Errorf("bootstrap: get messages: %w", err)
|
||||
}
|
||||
|
||||
// Fast path: DB has same count and exact match → no-op
|
||||
if len(dbMsgs) == len(messages) {
|
||||
matched := true
|
||||
for i := 0; i < len(messages); i++ {
|
||||
if !messageMatches(dbMsgs[i], messages[i]) {
|
||||
matched = false
|
||||
break
|
||||
}
|
||||
}
|
||||
if matched {
|
||||
return nil // DB is up to date
|
||||
}
|
||||
}
|
||||
|
||||
// Find longest matching prefix from the start
|
||||
anchor := -1
|
||||
compareLen := len(dbMsgs)
|
||||
if compareLen > len(messages) {
|
||||
compareLen = len(messages)
|
||||
}
|
||||
|
||||
for i := 0; i < compareLen; i++ {
|
||||
if messageMatches(dbMsgs[i], messages[i]) {
|
||||
anchor = i
|
||||
} else {
|
||||
// Mismatch detected - log details and rebuild
|
||||
logger.InfoCF("seahorse", "bootstrap: mismatch detected", map[string]any{
|
||||
"conv_id": conv.ConversationID,
|
||||
"index": i,
|
||||
"db_role": dbMsgs[i].Role,
|
||||
"db_content": truncate(dbMsgs[i].Content, 50),
|
||||
"db_parts": len(dbMsgs[i].Parts),
|
||||
"msg_role": messages[i].Role,
|
||||
"msg_content": truncate(messages[i].Content, 50),
|
||||
"msg_parts": len(messages[i].Parts),
|
||||
})
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// If we hit a mismatch before reaching the end of DB messages, delete delta and re-ingest
|
||||
// Note: anchor can be -1 if first message didn't match (history completely changed)
|
||||
if anchor >= 0 && anchor < len(dbMsgs)-1 && len(dbMsgs) > 0 {
|
||||
anchorID := dbMsgs[anchor].ID
|
||||
logger.InfoCF("seahorse", "bootstrap: history edit detected", map[string]any{
|
||||
"conv_id": conv.ConversationID,
|
||||
"db_count": len(dbMsgs),
|
||||
"anchor": anchor,
|
||||
"anchor_id": anchorID,
|
||||
"msg_count": len(messages),
|
||||
"delta_start": anchor + 1,
|
||||
})
|
||||
|
||||
// Delete messages after anchor (also clears context_items)
|
||||
if err := e.store.DeleteMessagesAfterID(ctx, conv.ConversationID, anchorID); err != nil {
|
||||
return fmt.Errorf("bootstrap: delete messages: %w", err)
|
||||
}
|
||||
|
||||
// Re-ingest from anchor+1 to end
|
||||
delta := messages[anchor+1:]
|
||||
if len(delta) > 0 {
|
||||
_, err := e.Ingest(ctx, sessionKey, delta)
|
||||
if err != nil {
|
||||
return fmt.Errorf("bootstrap: re-ingest: %w", err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Normal case: append delta after anchor
|
||||
if anchor >= 0 && anchor < len(messages)-1 {
|
||||
delta := messages[anchor+1:]
|
||||
if len(delta) > 0 {
|
||||
_, err := e.Ingest(ctx, sessionKey, delta)
|
||||
if err != nil {
|
||||
return fmt.Errorf("bootstrap: ingest delta: %w", err)
|
||||
}
|
||||
}
|
||||
} else if anchor == -1 && len(dbMsgs) > 0 {
|
||||
// First message changed (history completely different) - rebuild from scratch
|
||||
logger.InfoCF("seahorse", "bootstrap: history replaced, rebuilding", map[string]any{
|
||||
"conv_id": conv.ConversationID,
|
||||
"db_count": len(dbMsgs),
|
||||
"msg_count": len(messages),
|
||||
})
|
||||
// Delete all existing messages
|
||||
if err := e.store.DeleteMessagesAfterID(ctx, conv.ConversationID, 0); err != nil {
|
||||
return fmt.Errorf("bootstrap: delete all messages: %w", err)
|
||||
}
|
||||
// Re-ingest everything
|
||||
if len(messages) > 0 {
|
||||
_, err := e.Ingest(ctx, sessionKey, messages)
|
||||
if err != nil {
|
||||
return fmt.Errorf("bootstrap: re-ingest all: %w", err)
|
||||
}
|
||||
}
|
||||
} else if anchor == -1 && len(dbMsgs) == 0 {
|
||||
// DB is empty, ingest everything
|
||||
_, err := e.Ingest(ctx, sessionKey, messages)
|
||||
if err != nil {
|
||||
return fmt.Errorf("bootstrap: ingest all: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// truncate shortens a string for logging.
|
||||
func truncate(s string, maxLen int) string {
|
||||
if len(s) <= maxLen {
|
||||
return s
|
||||
}
|
||||
return s[:maxLen] + "..."
|
||||
}
|
||||
|
||||
// messageMatches compares two messages using (role, content) or (role, parts).
|
||||
// TokenCount is NOT compared because it may be re-estimated differently
|
||||
// during bootstrap (e.g., via tokenizer.EstimateMessageTokens).
|
||||
// For messages with Parts (tool_use, tool_result), compare Parts instead of Content
|
||||
// since AddMessageWithParts stores empty Content in DB.
|
||||
func messageMatches(a, b Message) bool {
|
||||
if a.Role != b.Role {
|
||||
return false
|
||||
}
|
||||
// If either message has Parts, compare Parts
|
||||
if len(a.Parts) > 0 || len(b.Parts) > 0 {
|
||||
return partsMatch(a.Parts, b.Parts)
|
||||
}
|
||||
// Simple text messages: compare Content
|
||||
return a.Content == b.Content
|
||||
}
|
||||
|
||||
// partsMatch compares two slices of MessagePart for equality.
|
||||
func partsMatch(a, b []MessagePart) bool {
|
||||
if len(a) != len(b) {
|
||||
return false
|
||||
}
|
||||
for i := range a {
|
||||
if a[i].Type != b[i].Type {
|
||||
return false
|
||||
}
|
||||
switch a[i].Type {
|
||||
case "text":
|
||||
if a[i].Text != b[i].Text {
|
||||
return false
|
||||
}
|
||||
case "tool_use":
|
||||
if a[i].Name != b[i].Name || a[i].Arguments != b[i].Arguments || a[i].ToolCallID != b[i].ToolCallID {
|
||||
return false
|
||||
}
|
||||
case "tool_result":
|
||||
if a[i].ToolCallID != b[i].ToolCallID || a[i].Text != b[i].Text {
|
||||
return false
|
||||
}
|
||||
case "media":
|
||||
if a[i].MediaURI != b[i].MediaURI || a[i].MimeType != b[i].MimeType {
|
||||
return false
|
||||
}
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
1448
pkg/seahorse/short_engine_test.go
Normal file
1448
pkg/seahorse/short_engine_test.go
Normal file
File diff suppressed because it is too large
Load diff
212
pkg/seahorse/short_retrieval.go
Normal file
212
pkg/seahorse/short_retrieval.go
Normal file
|
|
@ -0,0 +1,212 @@
|
|||
package seahorse
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ParseLastDuration parses a "last" duration string like "6h", "7d", "2w", "1m".
|
||||
// Returns the duration and nil error, or zero and error if invalid.
|
||||
func ParseLastDuration(s string) (time.Duration, error) {
|
||||
if s == "" {
|
||||
return 0, fmt.Errorf("empty duration")
|
||||
}
|
||||
|
||||
re := regexp.MustCompile(`^(\d+)([hdwm])$`)
|
||||
matches := re.FindStringSubmatch(s)
|
||||
if matches == nil {
|
||||
return 0, fmt.Errorf("invalid duration format: %q (use format like 6h, 7d, 2w, 1m)", s)
|
||||
}
|
||||
|
||||
value, _ := strconv.Atoi(matches[1])
|
||||
unit := matches[2]
|
||||
|
||||
switch unit {
|
||||
case "h":
|
||||
return time.Duration(value) * time.Hour, nil
|
||||
case "d":
|
||||
return time.Duration(value) * 24 * time.Hour, nil
|
||||
case "w":
|
||||
return time.Duration(value) * 7 * 24 * time.Hour, nil
|
||||
case "m":
|
||||
return time.Duration(value) * 30 * 24 * time.Hour, nil
|
||||
default:
|
||||
return 0, fmt.Errorf("unknown unit: %q", unit)
|
||||
}
|
||||
}
|
||||
|
||||
// GrepInput controls search across summaries and messages.
|
||||
type GrepInput struct {
|
||||
Pattern string `json:"pattern"`
|
||||
Scope string `json:"scope,omitempty"` // "both" (default), "summary", or "message"
|
||||
Role string `json:"role,omitempty"` // "user", "assistant", or "" (all)
|
||||
AllConversations bool `json:"allConversations,omitempty"`
|
||||
Since *time.Time `json:"since,omitempty"`
|
||||
Before *time.Time `json:"before,omitempty"`
|
||||
Last string `json:"last,omitempty"` // shortcut: "6h", "7d", "2w", "1m"
|
||||
Limit int `json:"limit,omitempty"`
|
||||
}
|
||||
|
||||
// GrepResult contains search results.
|
||||
type GrepResult struct {
|
||||
Success bool `json:"success"`
|
||||
Summaries []GrepSummaryResult `json:"summaries"`
|
||||
Messages []GrepMessageResult `json:"messages"`
|
||||
TotalSummaries int `json:"totalSummaries"`
|
||||
TotalMessages int `json:"totalMessages"`
|
||||
Hint string `json:"hint,omitempty"`
|
||||
}
|
||||
|
||||
// GrepSummaryResult is a summary match from grep.
|
||||
type GrepSummaryResult struct {
|
||||
ID string `json:"id"`
|
||||
Content string `json:"content"`
|
||||
Depth int `json:"depth"`
|
||||
Kind SummaryKind `json:"kind"`
|
||||
ConversationID int64 `json:"conversationId"`
|
||||
// Rank is the bm25 relevance score (negative value, closer to 0 = better match).
|
||||
// Examples: -0.5 = excellent match, -2.0 = good match, -10.0 = partial match.
|
||||
Rank float64 `json:"rank,omitempty"`
|
||||
}
|
||||
|
||||
// GrepMessageResult is a message match from grep.
|
||||
type GrepMessageResult struct {
|
||||
ID int64 `json:"id,string"`
|
||||
Snippet string `json:"snippet"`
|
||||
Role string `json:"role"`
|
||||
ConversationID int64 `json:"conversationId"`
|
||||
Rank float64 `json:"rank,omitempty"` // Relevance score (lower = better match)
|
||||
}
|
||||
|
||||
// ExpandMessagesResult contains expanded messages.
|
||||
type ExpandMessagesResult struct {
|
||||
Messages []Message `json:"messages"`
|
||||
TokenCount int `json:"tokenCount"`
|
||||
}
|
||||
|
||||
// Grep searches summaries and messages for matching content.
|
||||
func (r *RetrievalEngine) Grep(ctx context.Context, input GrepInput) (*GrepResult, error) {
|
||||
if input.Pattern == "" {
|
||||
return nil, fmt.Errorf("grep: pattern is required")
|
||||
}
|
||||
|
||||
limit := input.Limit
|
||||
if limit == 0 {
|
||||
limit = 20
|
||||
}
|
||||
|
||||
// Handle Last parameter: convert to Since
|
||||
since := input.Since
|
||||
if input.Last != "" {
|
||||
dur, err := ParseLastDuration(input.Last)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("grep: invalid last: %w", err)
|
||||
}
|
||||
t := time.Now().UTC().Add(-dur)
|
||||
since = &t
|
||||
}
|
||||
|
||||
// Auto-detect mode: use LIKE if pattern contains %, otherwise full-text
|
||||
mode := ""
|
||||
if strings.Contains(input.Pattern, "%") {
|
||||
mode = "like"
|
||||
}
|
||||
|
||||
searchInput := SearchInput{
|
||||
Pattern: input.Pattern,
|
||||
Mode: mode,
|
||||
Role: input.Role,
|
||||
AllConversations: input.AllConversations,
|
||||
Since: since,
|
||||
Before: input.Before,
|
||||
Limit: limit,
|
||||
}
|
||||
|
||||
result := &GrepResult{
|
||||
Success: true,
|
||||
Summaries: make([]GrepSummaryResult, 0),
|
||||
Messages: make([]GrepMessageResult, 0),
|
||||
TotalSummaries: 0,
|
||||
TotalMessages: 0,
|
||||
}
|
||||
|
||||
// Determine scope
|
||||
scope := input.Scope
|
||||
if scope == "" {
|
||||
scope = "both"
|
||||
}
|
||||
|
||||
// Search summaries if requested
|
||||
if scope == "both" || scope == "summary" {
|
||||
sumResults, err := r.store.SearchSummaries(ctx, searchInput)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("search summaries: %w", err)
|
||||
}
|
||||
for _, sr := range sumResults {
|
||||
if sr.SummaryID != "" {
|
||||
result.Summaries = append(result.Summaries, GrepSummaryResult{
|
||||
ID: sr.SummaryID,
|
||||
Content: sr.Content,
|
||||
Depth: sr.Depth,
|
||||
Kind: sr.Kind,
|
||||
ConversationID: sr.ConversationID,
|
||||
Rank: sr.Rank,
|
||||
})
|
||||
}
|
||||
}
|
||||
if len(sumResults) > 0 {
|
||||
result.TotalSummaries = sumResults[0].TotalCount
|
||||
}
|
||||
}
|
||||
|
||||
// Search messages if requested
|
||||
if scope == "both" || scope == "message" {
|
||||
msgResults, err := r.store.SearchMessages(ctx, searchInput)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("search messages: %w", err)
|
||||
}
|
||||
for _, sr := range msgResults {
|
||||
if sr.MessageID > 0 {
|
||||
result.Messages = append(result.Messages, GrepMessageResult{
|
||||
ID: sr.MessageID,
|
||||
Snippet: sr.Snippet,
|
||||
Role: sr.Role,
|
||||
ConversationID: sr.ConversationID,
|
||||
Rank: sr.Rank,
|
||||
})
|
||||
}
|
||||
}
|
||||
if len(msgResults) > 0 {
|
||||
result.TotalMessages = msgResults[0].TotalCount
|
||||
}
|
||||
}
|
||||
|
||||
// Add hint if no results
|
||||
if len(result.Summaries) == 0 && len(result.Messages) == 0 {
|
||||
result.Hint = "No matches. Try: %keyword% for fuzzy search, or all_conversations: true"
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// ExpandMessages retrieves full message content by IDs.
|
||||
func (r *RetrievalEngine) ExpandMessages(ctx context.Context, messageIDs []int64) (*ExpandMessagesResult, error) {
|
||||
result := &ExpandMessagesResult{
|
||||
Messages: make([]Message, 0, len(messageIDs)),
|
||||
}
|
||||
|
||||
for _, msgID := range messageIDs {
|
||||
msg, err := r.store.GetMessageByID(ctx, msgID)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
result.Messages = append(result.Messages, *msg)
|
||||
result.TokenCount += msg.TokenCount
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
362
pkg/seahorse/short_retrieval_test.go
Normal file
362
pkg/seahorse/short_retrieval_test.go
Normal file
|
|
@ -0,0 +1,362 @@
|
|||
package seahorse
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// --- Retrieval Tests ---
|
||||
|
||||
func newTestRetrieval(t *testing.T) (*RetrievalEngine, *Store, int64) {
|
||||
t.Helper()
|
||||
s := openTestStore(t)
|
||||
ctx := context.Background()
|
||||
conv, _ := s.GetOrCreateConversation(ctx, "test:retrieval")
|
||||
return &RetrievalEngine{store: s}, s, conv.ConversationID
|
||||
}
|
||||
|
||||
func TestRetrievalGrepSummaries(t *testing.T) {
|
||||
r, s, convID := newTestRetrieval(t)
|
||||
ctx := context.Background()
|
||||
|
||||
s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: "数据库连接配置说明",
|
||||
TokenCount: 50,
|
||||
})
|
||||
s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: "API endpoint documentation",
|
||||
TokenCount: 50,
|
||||
})
|
||||
|
||||
// FTS5 search (trigram, needs >= 3 chars)
|
||||
results, err := r.Grep(ctx, GrepInput{
|
||||
Pattern: "数据库连",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Grep: %v", err)
|
||||
}
|
||||
if len(results.Summaries) == 0 {
|
||||
t.Error("expected at least 1 FTS result")
|
||||
}
|
||||
|
||||
// LIKE search with wildcard
|
||||
results, err = r.Grep(ctx, GrepInput{
|
||||
Pattern: "%endpoint%",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Grep LIKE: %v", err)
|
||||
}
|
||||
if len(results.Summaries) == 0 {
|
||||
t.Error("expected at least 1 LIKE result")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetrievalGrepMessages(t *testing.T) {
|
||||
r, s, convID := newTestRetrieval(t)
|
||||
ctx := context.Background()
|
||||
|
||||
s.AddMessage(ctx, convID, "user", "find this message about testing", 5)
|
||||
s.AddMessage(ctx, convID, "user", "unrelated content here", 5)
|
||||
|
||||
results, err := r.Grep(ctx, GrepInput{
|
||||
Pattern: "testing",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Grep: %v", err)
|
||||
}
|
||||
if len(results.Messages) == 0 {
|
||||
t.Error("expected at least 1 result for 'testing'")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetrievalExpandMessages(t *testing.T) {
|
||||
r, s, convID := newTestRetrieval(t)
|
||||
ctx := context.Background()
|
||||
|
||||
msg, _ := s.AddMessage(ctx, convID, "user", "expand this message", 10)
|
||||
|
||||
result, err := r.ExpandMessages(ctx, []int64{msg.ID})
|
||||
if err != nil {
|
||||
t.Fatalf("ExpandMessages: %v", err)
|
||||
}
|
||||
if len(result.Messages) != 1 {
|
||||
t.Errorf("Messages = %d, want 1", len(result.Messages))
|
||||
}
|
||||
if result.Messages[0].Content != "expand this message" {
|
||||
t.Errorf("Content = %q, want 'expand this message'", result.Messages[0].Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetrievalExpandMultipleMessages(t *testing.T) {
|
||||
r, s, convID := newTestRetrieval(t)
|
||||
ctx := context.Background()
|
||||
|
||||
msg1, _ := s.AddMessage(ctx, convID, "user", "first message", 10)
|
||||
msg2, _ := s.AddMessage(ctx, convID, "assistant", "second message", 10)
|
||||
msg3, _ := s.AddMessage(ctx, convID, "user", "third message", 10)
|
||||
|
||||
result, err := r.ExpandMessages(ctx, []int64{msg1.ID, msg2.ID, msg3.ID})
|
||||
if err != nil {
|
||||
t.Fatalf("ExpandMessages: %v", err)
|
||||
}
|
||||
if len(result.Messages) != 3 {
|
||||
t.Errorf("Messages = %d, want 3", len(result.Messages))
|
||||
}
|
||||
if result.TokenCount != 30 {
|
||||
t.Errorf("TokenCount = %d, want 30", result.TokenCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetrievalGrepWithTimeFilter(t *testing.T) {
|
||||
r, s, convID := newTestRetrieval(t)
|
||||
ctx := context.Background()
|
||||
|
||||
now := time.Now().UTC()
|
||||
before := now.Add(-2 * time.Hour)
|
||||
|
||||
// Create messages at different times
|
||||
s.AddMessage(ctx, convID, "user", "old message about auth", 5)
|
||||
s.AddMessage(ctx, convID, "user", "recent message about auth", 5)
|
||||
|
||||
// Search with time filter
|
||||
results, err := r.Grep(ctx, GrepInput{
|
||||
Pattern: "auth",
|
||||
Since: &before,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Grep: %v", err)
|
||||
}
|
||||
_ = results // Just verify no error
|
||||
}
|
||||
|
||||
func TestRetrievalGrepAllConversations(t *testing.T) {
|
||||
r, s, _ := newTestRetrieval(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create another conversation
|
||||
conv2, _ := s.GetOrCreateConversation(ctx, "test:retrieval2")
|
||||
|
||||
// Add messages to both
|
||||
s.AddMessage(ctx, conv2.ConversationID, "user", "unique keyword xyz", 5)
|
||||
|
||||
// Search all conversations
|
||||
results, err := r.Grep(ctx, GrepInput{
|
||||
Pattern: "xyz",
|
||||
AllConversations: true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Grep: %v", err)
|
||||
}
|
||||
if len(results.Messages) == 0 {
|
||||
t.Error("expected to find message in other conversation")
|
||||
}
|
||||
}
|
||||
|
||||
// --- Last Duration Parsing Tests ---
|
||||
|
||||
func TestParseLastDuration(t *testing.T) {
|
||||
tests := []struct {
|
||||
input string
|
||||
wantDur time.Duration
|
||||
wantErr bool
|
||||
}{
|
||||
{"6h", 6 * time.Hour, false},
|
||||
{"1d", 24 * time.Hour, false},
|
||||
{"7d", 7 * 24 * time.Hour, false},
|
||||
{"2w", 14 * 24 * time.Hour, false},
|
||||
{"1m", 30 * 24 * time.Hour, false}, // month = 30 days
|
||||
{"3m", 90 * 24 * time.Hour, false},
|
||||
{"", 0, true},
|
||||
{"invalid", 0, true},
|
||||
{"5x", 0, true}, // unknown unit
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.input, func(t *testing.T) {
|
||||
got, err := ParseLastDuration(tt.input)
|
||||
if tt.wantErr {
|
||||
if err == nil {
|
||||
t.Error("expected error, got nil")
|
||||
}
|
||||
} else {
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got != tt.wantDur {
|
||||
t.Errorf("ParseLastDuration(%q) = %v, want %v", tt.input, got, tt.wantDur)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// --- Role Filter Tests ---
|
||||
|
||||
func TestRetrievalGrepRoleFilter(t *testing.T) {
|
||||
r, s, convID := newTestRetrieval(t)
|
||||
ctx := context.Background()
|
||||
|
||||
s.AddMessage(ctx, convID, "user", "user message about alpha", 5)
|
||||
s.AddMessage(ctx, convID, "assistant", "assistant reply about alpha", 5)
|
||||
s.AddMessage(ctx, convID, "user", "another user message", 5)
|
||||
|
||||
// Search all roles
|
||||
allResults, err := r.Grep(ctx, GrepInput{
|
||||
Pattern: "alpha",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Grep: %v", err)
|
||||
}
|
||||
if len(allResults.Messages) != 2 {
|
||||
t.Errorf("expected 2 messages, got %d", len(allResults.Messages))
|
||||
}
|
||||
|
||||
// Search user only
|
||||
userResults, err := r.Grep(ctx, GrepInput{
|
||||
Pattern: "alpha",
|
||||
Role: "user",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Grep: %v", err)
|
||||
}
|
||||
if len(userResults.Messages) != 1 {
|
||||
t.Errorf("expected 1 user message, got %d", len(userResults.Messages))
|
||||
}
|
||||
if userResults.Messages[0].Role != "user" {
|
||||
t.Errorf("expected role=user, got %s", userResults.Messages[0].Role)
|
||||
}
|
||||
|
||||
// Search assistant only
|
||||
assistantResults, err := r.Grep(ctx, GrepInput{
|
||||
Pattern: "alpha",
|
||||
Role: "assistant",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Grep: %v", err)
|
||||
}
|
||||
if len(assistantResults.Messages) != 1 {
|
||||
t.Errorf("expected 1 assistant message, got %d", len(assistantResults.Messages))
|
||||
}
|
||||
}
|
||||
|
||||
// --- Last Parameter Tests ---
|
||||
|
||||
func TestRetrievalGrepWithLast(t *testing.T) {
|
||||
r, s, convID := newTestRetrieval(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Add messages (we can't control timestamps in SQLite easily,
|
||||
// but we can verify the parameter is parsed correctly)
|
||||
s.AddMessage(ctx, convID, "user", "recent message about testing", 5)
|
||||
|
||||
// Test that Last parameter is converted to Since
|
||||
results, err := r.Grep(ctx, GrepInput{
|
||||
Pattern: "testing",
|
||||
Last: "1d", // last 1 day
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Grep: %v", err)
|
||||
}
|
||||
// Should still find the message since it's recent
|
||||
if len(results.Messages) == 0 {
|
||||
t.Error("expected to find recent message")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRetrievalGrepRoleFilterWithSummaries tests that role filter works when
|
||||
// searching both summaries and messages (summaries don't have role column).
|
||||
func TestRetrievalGrepRoleFilterWithSummaries(t *testing.T) {
|
||||
r, s, convID := newTestRetrieval(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a summary (no role column)
|
||||
s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: "summary about testing",
|
||||
TokenCount: 50,
|
||||
})
|
||||
|
||||
// Add messages with different roles
|
||||
s.AddMessage(ctx, convID, "user", "user message about testing", 5)
|
||||
s.AddMessage(ctx, convID, "assistant", "assistant reply about testing", 5)
|
||||
|
||||
// Search with role filter and scope=both (default), using LIKE mode (%)
|
||||
// This should NOT error even though summaries don't have role column
|
||||
bothResults, err := r.Grep(ctx, GrepInput{
|
||||
Pattern: "%testing%", // LIKE mode to trigger the bug
|
||||
Role: "user",
|
||||
Scope: "both",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Grep with role and scope=both: %v", err)
|
||||
}
|
||||
|
||||
// Should only return user messages, not summaries or assistant messages
|
||||
if len(bothResults.Messages) != 1 {
|
||||
t.Errorf("expected 1 user message, got %d", len(bothResults.Messages))
|
||||
}
|
||||
if len(bothResults.Messages) > 0 && bothResults.Messages[0].Role != "user" {
|
||||
t.Errorf("expected role=user, got %s", bothResults.Messages[0].Role)
|
||||
}
|
||||
|
||||
// Summaries should be empty since they don't have roles to filter
|
||||
// (or we could return all summaries - either is acceptable)
|
||||
}
|
||||
|
||||
// TestRetrievalGrepTotalCounts tests that grep returns total counts.
|
||||
func TestRetrievalGrepTotalCounts(t *testing.T) {
|
||||
r, s, convID := newTestRetrieval(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Create 3 summaries
|
||||
for i := 0; i < 3; i++ {
|
||||
s.CreateSummary(ctx, CreateSummaryInput{
|
||||
ConversationID: convID,
|
||||
Kind: SummaryKindLeaf,
|
||||
Depth: 0,
|
||||
Content: fmt.Sprintf("summary about testing %d", i),
|
||||
TokenCount: 50,
|
||||
})
|
||||
}
|
||||
|
||||
// Add 5 messages
|
||||
for i := 0; i < 5; i++ {
|
||||
s.AddMessage(ctx, convID, "user", fmt.Sprintf("message about testing %d", i), 5)
|
||||
}
|
||||
|
||||
// Search with limit smaller than total
|
||||
results, err := r.Grep(ctx, GrepInput{
|
||||
Pattern: "%testing%", // LIKE mode
|
||||
Scope: "both",
|
||||
Limit: 2,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Grep: %v", err)
|
||||
}
|
||||
|
||||
// Should return limited results
|
||||
if len(results.Summaries) > 2 {
|
||||
t.Errorf("expected at most 2 summaries, got %d", len(results.Summaries))
|
||||
}
|
||||
if len(results.Messages) > 2 {
|
||||
t.Errorf("expected at most 2 messages, got %d", len(results.Messages))
|
||||
}
|
||||
|
||||
// But total counts should reflect all matches
|
||||
if results.TotalSummaries != 3 {
|
||||
t.Errorf("expected TotalSummaries=3, got %d", results.TotalSummaries)
|
||||
}
|
||||
if results.TotalMessages != 5 {
|
||||
t.Errorf("expected TotalMessages=5, got %d", results.TotalMessages)
|
||||
}
|
||||
}
|
||||
1532
pkg/seahorse/store.go
Normal file
1532
pkg/seahorse/store.go
Normal file
File diff suppressed because it is too large
Load diff
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue