fix: replace litellm with tiktoken for token counting
Build and Push API / release (push) Successful in 3s
Build and Push API / build (push) Successful in 2m26s

- litellm had dependency conflicts with pydantic-ai
- tiktoken is lighter and already required by pydantic-ai
- Updated documentation (README.md, architecture.md)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
2026-01-11 23:05:11 +01:00
co-authored by Claude Opus 4.5
parent 2523db4da7
commit 2f97041aa9
7 changed files with 127 additions and 24 deletions
+8 -2
View File
@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
## [Unreleased] ## [Unreleased]
## [0.4.1] - 2026-01-11
### Fixed
- Replace `litellm` with `tiktoken` for token counting (dependency conflict with pydantic-ai)
- Update documentation (README.md, architecture.md) with conversation layer info
## [0.4.0] - 2026-01-11 ## [0.4.0] - 2026-01-11
### Added ### Added
@@ -15,7 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
- SQLite (dev) and PostgreSQL (prod) support via async engines - SQLite (dev) and PostgreSQL (prod) support via async engines
- Lazy database initialization pattern - Lazy database initialization pattern
- Context management infrastructure - Context management infrastructure
- Token counting utilities using `litellm` - Token counting utilities using `tiktoken`
- Context summarization at 80% token threshold - Context summarization at 80% token threshold
- XML-tagged context prompt building for agent injection - XML-tagged context prompt building for agent injection
- REST API for multi-turn conversations - REST API for multi-turn conversations
@@ -24,7 +30,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
- `GET /conversations/{id}` - Get conversation with history - `GET /conversations/{id}` - Get conversation with history
- `POST /conversations/{id}/messages` - Add message (triggers agent) - `POST /conversations/{id}/messages` - Add message (triggers agent)
- `DELETE /conversations/{id}` - Delete conversation - `DELETE /conversations/{id}` - Delete conversation
- New dependencies: `sqlalchemy[asyncio]~=2.0.36`, `aiosqlite~=0.21.0`, `litellm~=1.57.0` - New dependencies: `sqlalchemy[asyncio]~=2.0.36`, `aiosqlite~=0.21.0`, `tiktoken>=0.12.0`
- Config settings: `database_url`, `summarization_threshold`, `keep_recent_messages` - Config settings: `database_url`, `summarization_threshold`, `keep_recent_messages`
- 19 conversation tests, 6 token counting tests (176 total tests passing) - 19 conversation tests, 6 token counting tests (176 total tests passing)
+26 -3
View File
@@ -4,8 +4,9 @@ A Claude Code-inspired development assistant powered by local LLMs via Ollama.
## Features ## Features
- **Explore Agent** - Search, read, and understand codebases - **3 Agents** - Explore (read-only), Plan (architecture), Task (orchestrator)
- **8 Tools** - File read/write, glob, grep, bash, web search - **8 Tools** - File read/write/edit, glob, grep, bash, web search
- **Conversations** - Multi-turn memory with context summarization
- **Streaming** - Real-time response display - **Streaming** - Real-time response display
- **Self-hosted** - Runs on your own hardware with Ollama - **Self-hosted** - Runs on your own hardware with Ollama
@@ -74,11 +75,33 @@ webber-cli chat -d /path/to/project
| `bash` | Full bash with safety controls | | `bash` | Full bash with safety controls |
| `web_search` | Search web via SearXNG | | `web_search` | Search web via SearXNG |
## Agents
| Agent | Purpose | Tools |
|-------|---------|-------|
| **Explore** | Fast codebase navigation, search | Read-only (glob, grep, read, bash_readonly) |
| **Plan** | Design implementation strategies | Read-only (same as Explore) |
| **Task** | Autonomous multi-step execution | All tools + spawn_agent |
## API Endpoints
```bash
# Stateless agent execution
POST /agents/run # Execute agent, get response
POST /agents/stream # Execute with SSE streaming
GET /agents/ # List available agents
# Stateful conversations (multi-turn memory)
POST /conversations/ # Create conversation
GET /conversations/ # List conversations
POST /conversations/{id}/messages # Add message, get agent response
```
## Versioning ## Versioning
This project uses prefixed tags for independent release cycles: This project uses prefixed tags for independent release cycles:
- `api/v0.3.0` - Triggers API Docker build and deployment - `api/v0.4.0` - Triggers API Docker build and deployment
- `cli/v0.1.0` - Triggers CLI installer build (future) - `cli/v0.1.0` - Triggers CLI installer build (future)
## Requirements ## Requirements
+64 -1
View File
@@ -24,18 +24,30 @@ webber/
├── src/ ├── src/
│ ├── main.py # App entry point (NO routes) │ ├── main.py # App entry point (NO routes)
│ │ │ │
│ ├── db/ # Database layer
│ │ ├── __init__.py # Exports: Database, get_database, get_session
│ │ ├── database.py # SQLAlchemy async engine, session factory
│ │ └── models.py # Base declarative model
│ │
│ ├── shared/ # Cross-cutting concerns │ ├── shared/ # Cross-cutting concerns
│ │ ├── base.py # BaseController, BaseSchema │ │ ├── base.py # BaseController, BaseSchema
│ │ ├── config.py # Pydantic Settings │ │ ├── config.py # Pydantic Settings
│ │ ├── logging.py # @logged decorator, trace_span │ │ ├── logging.py # @logged decorator, trace_span
│ │ ├── exceptions.py # Custom exception hierarchy │ │ ├── exceptions.py # Custom exception hierarchy
│ │ ├── auth.py # API key validation │ │ ├── auth.py # API key validation
│ │ ── context.py # UserProvider singleton │ │ ── context.py # UserProvider singleton
│ │ └── tokens.py # Token counting utilities (litellm)
│ │ │ │
│ └── domains/ # Feature domains │ └── domains/ # Feature domains
│ ├── router.py # Root router (composes all) │ ├── router.py # Root router (composes all)
│ ├── health/ # Health endpoints │ ├── health/ # Health endpoints
│ ├── auth/ # Authentication │ ├── auth/ # Authentication
│ ├── conversations/ # Multi-turn conversation memory
│ │ ├── models.py # Conversation, Message SQLAlchemy models
│ │ ├── schemas.py # Pydantic request/response models
│ │ ├── service.py # ConversationService business logic
│ │ ├── router.py # REST API endpoints
│ │ └── summarize.py # Context summarization logic
│ ├── agents/ # Agent orchestration │ ├── agents/ # Agent orchestration
│ │ ├── explore/ # Codebase navigation │ │ ├── explore/ # Codebase navigation
│ │ ├── plan/ # Implementation design │ │ ├── plan/ # Implementation design
@@ -201,6 +213,10 @@ All settings via environment variables or `.env`:
| ALLOWED_PATHS | [] | Paths accessible to tools | | ALLOWED_PATHS | [] | Paths accessible to tools |
| SESSION_TTL_HOURS | 24 | Session expiry | | SESSION_TTL_HOURS | 24 | Session expiry |
| MAX_CONTEXT_TOKENS | 128000 | Max context window | | MAX_CONTEXT_TOKENS | 128000 | Max context window |
| DATABASE_URL | sqlite+aiosqlite:///./webber.db | Database connection URL |
| SUMMARIZATION_THRESHOLD | 0.8 | Summarize at N% of max tokens |
| SUMMARIZATION_TARGET_TOKENS | 500 | Target summary size |
| KEEP_RECENT_MESSAGES | 6 | Messages to keep unsummarized |
--- ---
@@ -250,6 +266,53 @@ Tools are sandboxed operations agents can invoke:
--- ---
## Database Layer
SQLAlchemy 2.0 async with lazy initialization pattern.
### Supported Databases
- **Development**: SQLite via `aiosqlite`
- **Production**: PostgreSQL via `asyncpg`
### Pattern
```python
from src.db import get_session
from sqlalchemy.ext.asyncio import AsyncSession
async def my_endpoint(session: AsyncSession = Depends(get_session)):
# Session auto-commits on success, rollbacks on exception
result = await session.execute(query)
```
Tables are created lazily on first `get_session()` call.
---
## Conversation API
Multi-turn conversation memory with automatic context summarization.
### Endpoints
| Endpoint | Method | Description |
|----------|--------|-------------|
| `/conversations/` | POST | Create new conversation |
| `/conversations/` | GET | List user's conversations |
| `/conversations/{id}` | GET | Get conversation with history |
| `/conversations/{id}/messages` | POST | Add message, triggers agent |
| `/conversations/{id}` | DELETE | Delete conversation |
### Models
- **Conversation**: User session with agent type, working directory
- **Message**: Individual messages with role, content, token count
### Context Summarization
When total tokens exceed 80% of `MAX_CONTEXT_TOKENS`:
1. Keep last 6 messages intact
2. Summarize older messages into a single summary message
3. Mark old messages as summarized (soft delete)
---
## Authentication Flow ## Authentication Flow
1. Client sends `X-API-Key` header 1. Client sends `X-API-Key` header
+1 -1
View File
@@ -1,6 +1,6 @@
[project] [project]
name = "webber-api" name = "webber-api"
version = "0.4.0" version = "0.4.1"
description = "Webber API - Multi-Agent AI Development Server" description = "Webber API - Multi-Agent AI Development Server"
authors = [ authors = [
{name = "jpmschweitzer"} {name = "jpmschweitzer"}
+1 -1
View File
@@ -31,4 +31,4 @@ sqlalchemy[asyncio]~=2.0.36
aiosqlite~=0.21.0 # SQLite async driver (dev) aiosqlite~=0.21.0 # SQLite async driver (dev)
# Token counting # Token counting
litellm~=1.57.0 # Multi-model token counting tiktoken>=0.12.0 # OpenAI tokenizer (used for estimation)
+26 -15
View File
@@ -1,53 +1,64 @@
""" """
Token counting utilities for context management. Token counting utilities for context management.
Uses litellm for accurate multi-model token counting. Uses tiktoken for token counting. While tiktoken is OpenAI's tokenizer,
cl100k_base encoding provides reasonable estimates for most LLMs.
""" """
from functools import lru_cache
from src.shared.logging import get_logger from src.shared.logging import get_logger
logger = get_logger(__name__) logger = get_logger(__name__)
# Default model for token counting (Mistral Nemo)
DEFAULT_MODEL = "mistral/mistral-nemo" @lru_cache(maxsize=1)
def _get_encoding():
"""Get tiktoken encoding (cached)."""
import tiktoken
# cl100k_base is used by GPT-4 and provides reasonable estimates for most models
return tiktoken.get_encoding("cl100k_base")
def count_tokens(text: str, model: str = DEFAULT_MODEL) -> int: def count_tokens(text: str) -> int:
""" """
Count tokens in a text string. Count tokens in a text string.
Args: Args:
text: Text to count tokens for text: Text to count tokens for
model: Model identifier for tokenizer selection
Returns: Returns:
Token count Token count
""" """
try: try:
from litellm import token_counter encoding = _get_encoding()
return token_counter(model=model, text=text) return len(encoding.encode(text))
except Exception as e: except Exception as e:
# Fallback to rough estimate if litellm fails # Fallback to rough estimate if tiktoken fails
logger.warning(f"Token counting failed, using estimate: {e}") logger.warning(f"Token counting failed, using estimate: {e}")
return len(text) // 4 return len(text) // 4
def count_message_tokens( def count_message_tokens(messages: list[dict[str, str]]) -> int:
messages: list[dict[str, str]],
model: str = DEFAULT_MODEL
) -> int:
""" """
Count tokens for a list of chat messages. Count tokens for a list of chat messages.
Args: Args:
messages: List of message dicts with 'role' and 'content' keys messages: List of message dicts with 'role' and 'content' keys
model: Model identifier for tokenizer selection
Returns: Returns:
Total token count including message overhead Total token count including message overhead
""" """
try: try:
from litellm import token_counter encoding = _get_encoding()
return token_counter(model=model, messages=messages) total = 0
for msg in messages:
# Each message has ~4 tokens overhead for role/formatting
total += 4
total += len(encoding.encode(msg.get("content", "")))
total += len(encoding.encode(msg.get("role", "")))
# Add 2 tokens for assistant response priming
total += 2
return total
except Exception as e: except Exception as e:
# Fallback to rough estimate # Fallback to rough estimate
logger.warning(f"Token counting failed, using estimate: {e}") logger.warning(f"Token counting failed, using estimate: {e}")
+1 -1
View File
@@ -59,7 +59,7 @@ class TestTokenCounting:
def test_count_message_tokens_empty_list(self): def test_count_message_tokens_empty_list(self):
"""Test counting empty message list.""" """Test counting empty message list."""
tokens = count_message_tokens([]) tokens = count_message_tokens([])
# litellm may return small overhead even for empty list # tiktoken returns small overhead for empty list (assistant priming)
assert tokens < 10 assert tokens < 10