fix: replace litellm with tiktoken for token counting
- 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:
+8
-2
@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
|
||||
## [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
|
||||
|
||||
### 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
|
||||
- Lazy database initialization pattern
|
||||
- Context management infrastructure
|
||||
- Token counting utilities using `litellm`
|
||||
- Token counting utilities using `tiktoken`
|
||||
- Context summarization at 80% token threshold
|
||||
- XML-tagged context prompt building for agent injection
|
||||
- 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
|
||||
- `POST /conversations/{id}/messages` - Add message (triggers agent)
|
||||
- `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`
|
||||
- 19 conversation tests, 6 token counting tests (176 total tests passing)
|
||||
|
||||
|
||||
@@ -4,8 +4,9 @@ A Claude Code-inspired development assistant powered by local LLMs via Ollama.
|
||||
|
||||
## Features
|
||||
|
||||
- **Explore Agent** - Search, read, and understand codebases
|
||||
- **8 Tools** - File read/write, glob, grep, bash, web search
|
||||
- **3 Agents** - Explore (read-only), Plan (architecture), Task (orchestrator)
|
||||
- **8 Tools** - File read/write/edit, glob, grep, bash, web search
|
||||
- **Conversations** - Multi-turn memory with context summarization
|
||||
- **Streaming** - Real-time response display
|
||||
- **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 |
|
||||
| `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
|
||||
|
||||
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)
|
||||
|
||||
## Requirements
|
||||
|
||||
@@ -24,18 +24,30 @@ webber/
|
||||
├── src/
|
||||
│ ├── 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
|
||||
│ │ ├── base.py # BaseController, BaseSchema
|
||||
│ │ ├── config.py # Pydantic Settings
|
||||
│ │ ├── logging.py # @logged decorator, trace_span
|
||||
│ │ ├── exceptions.py # Custom exception hierarchy
|
||||
│ │ ├── auth.py # API key validation
|
||||
│ │ └── context.py # UserProvider singleton
|
||||
│ │ ├── context.py # UserProvider singleton
|
||||
│ │ └── tokens.py # Token counting utilities (litellm)
|
||||
│ │
|
||||
│ └── domains/ # Feature domains
|
||||
│ ├── router.py # Root router (composes all)
|
||||
│ ├── health/ # Health endpoints
|
||||
│ ├── 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
|
||||
│ │ ├── explore/ # Codebase navigation
|
||||
│ │ ├── plan/ # Implementation design
|
||||
@@ -201,6 +213,10 @@ All settings via environment variables or `.env`:
|
||||
| ALLOWED_PATHS | [] | Paths accessible to tools |
|
||||
| SESSION_TTL_HOURS | 24 | Session expiry |
|
||||
| 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
|
||||
|
||||
1. Client sends `X-API-Key` header
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "webber-api"
|
||||
version = "0.4.0"
|
||||
version = "0.4.1"
|
||||
description = "Webber API - Multi-Agent AI Development Server"
|
||||
authors = [
|
||||
{name = "jpmschweitzer"}
|
||||
|
||||
@@ -31,4 +31,4 @@ sqlalchemy[asyncio]~=2.0.36
|
||||
aiosqlite~=0.21.0 # SQLite async driver (dev)
|
||||
|
||||
# Token counting
|
||||
litellm~=1.57.0 # Multi-model token counting
|
||||
tiktoken>=0.12.0 # OpenAI tokenizer (used for estimation)
|
||||
|
||||
@@ -1,53 +1,64 @@
|
||||
"""
|
||||
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
|
||||
|
||||
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.
|
||||
|
||||
Args:
|
||||
text: Text to count tokens for
|
||||
model: Model identifier for tokenizer selection
|
||||
|
||||
Returns:
|
||||
Token count
|
||||
"""
|
||||
try:
|
||||
from litellm import token_counter
|
||||
return token_counter(model=model, text=text)
|
||||
encoding = _get_encoding()
|
||||
return len(encoding.encode(text))
|
||||
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}")
|
||||
return len(text) // 4
|
||||
|
||||
|
||||
def count_message_tokens(
|
||||
messages: list[dict[str, str]],
|
||||
model: str = DEFAULT_MODEL
|
||||
) -> int:
|
||||
def count_message_tokens(messages: list[dict[str, str]]) -> int:
|
||||
"""
|
||||
Count tokens for a list of chat messages.
|
||||
|
||||
Args:
|
||||
messages: List of message dicts with 'role' and 'content' keys
|
||||
model: Model identifier for tokenizer selection
|
||||
|
||||
Returns:
|
||||
Total token count including message overhead
|
||||
"""
|
||||
try:
|
||||
from litellm import token_counter
|
||||
return token_counter(model=model, messages=messages)
|
||||
encoding = _get_encoding()
|
||||
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:
|
||||
# Fallback to rough estimate
|
||||
logger.warning(f"Token counting failed, using estimate: {e}")
|
||||
|
||||
@@ -59,7 +59,7 @@ class TestTokenCounting:
|
||||
def test_count_message_tokens_empty_list(self):
|
||||
"""Test counting empty message list."""
|
||||
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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user