basemode 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- basemode-0.1.0/.env.example +6 -0
- basemode-0.1.0/.gitignore +9 -0
- basemode-0.1.0/.pre-commit-config.yaml +15 -0
- basemode-0.1.0/Makefile +13 -0
- basemode-0.1.0/PKG-INFO +11 -0
- basemode-0.1.0/README.md +243 -0
- basemode-0.1.0/pyproject.toml +62 -0
- basemode-0.1.0/src/basemode/__init__.py +5 -0
- basemode-0.1.0/src/basemode/cli.py +136 -0
- basemode-0.1.0/src/basemode/continue_.py +56 -0
- basemode-0.1.0/src/basemode/detect.py +64 -0
- basemode-0.1.0/src/basemode/models.py +28 -0
- basemode-0.1.0/src/basemode/params.py +9 -0
- basemode-0.1.0/src/basemode/settings.py +47 -0
- basemode-0.1.0/src/basemode/strategies/__init__.py +27 -0
- basemode-0.1.0/src/basemode/strategies/base.py +14 -0
- basemode-0.1.0/src/basemode/strategies/completion.py +27 -0
- basemode-0.1.0/src/basemode/strategies/few_shot.py +69 -0
- basemode-0.1.0/src/basemode/strategies/fim.py +44 -0
- basemode-0.1.0/src/basemode/strategies/prefill.py +45 -0
- basemode-0.1.0/src/basemode/strategies/system.py +69 -0
- basemode-0.1.0/src/basemode/strategies/utils.py +26 -0
- basemode-0.1.0/tests/conftest.py +6 -0
- basemode-0.1.0/tests/test_detect.py +47 -0
- basemode-0.1.0/tests/test_integration.py +87 -0
- basemode-0.1.0/tests/test_models.py +32 -0
- basemode-0.1.0/tests/test_strategies.py +39 -0
- basemode-0.1.0/tests/test_utils.py +90 -0
- basemode-0.1.0/uv.lock +2184 -0
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
repos:
|
|
2
|
+
- repo: https://github.com/astral-sh/ruff-pre-commit
|
|
3
|
+
rev: v0.5.0
|
|
4
|
+
hooks:
|
|
5
|
+
- id: ruff
|
|
6
|
+
args: [--fix]
|
|
7
|
+
- id: ruff-format
|
|
8
|
+
- repo: https://github.com/pre-commit/pre-commit-hooks
|
|
9
|
+
rev: v4.6.0
|
|
10
|
+
hooks:
|
|
11
|
+
- id: trailing-whitespace
|
|
12
|
+
- id: end-of-file-fixer
|
|
13
|
+
- id: check-yaml
|
|
14
|
+
- id: check-toml
|
|
15
|
+
- id: check-added-large-files
|
basemode-0.1.0/Makefile
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
.PHONY: publish test test-integration lint
|
|
2
|
+
|
|
3
|
+
publish:
|
|
4
|
+
@export $$(grep UV_PUBLISH_TOKEN .env | xargs) && uv publish
|
|
5
|
+
|
|
6
|
+
test:
|
|
7
|
+
uv run pytest
|
|
8
|
+
|
|
9
|
+
test-integration:
|
|
10
|
+
uv run pytest -m integration
|
|
11
|
+
|
|
12
|
+
lint:
|
|
13
|
+
uv run ruff check src tests && uv run ruff format --check src tests
|
basemode-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: basemode
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Make any LLM do raw text continuation
|
|
5
|
+
Requires-Python: >=3.11
|
|
6
|
+
Requires-Dist: anyio>=4.0.0
|
|
7
|
+
Requires-Dist: litellm>=1.63.0
|
|
8
|
+
Requires-Dist: pydantic-settings>=2.0.0
|
|
9
|
+
Requires-Dist: python-dotenv>=1.0.0
|
|
10
|
+
Requires-Dist: rich>=13.0.0
|
|
11
|
+
Requires-Dist: typer>=0.12.0
|
basemode-0.1.0/README.md
ADDED
|
@@ -0,0 +1,243 @@
|
|
|
1
|
+
# basemode
|
|
2
|
+
|
|
3
|
+
**Make any LLM do raw text continuation.**
|
|
4
|
+
|
|
5
|
+
Most language models today are chat-tuned — they want to respond, acknowledge, and help. This is the opposite of what you need for loom-style tree exploration, creative writing, or any workflow that requires the model to simply *continue* text as if it wrote it. Feed a chat model a sentence mid-thought and you'll get "Sure! Here's a continuation:" instead of the next word.
|
|
6
|
+
|
|
7
|
+
`basemode` solves this. It wraps any LLM in the right coercion strategy for that provider — native completions API, assistant prefill, or a carefully-tuned system prompt — so you always get a clean continuation back.
|
|
8
|
+
|
|
9
|
+
```bash
|
|
10
|
+
basemode "The defendant, who had been sitting quietly throughout the proceedings, suddenly"
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
```
|
|
14
|
+
The defendant, who had been sitting quietly throughout the proceedings, suddenly
|
|
15
|
+
rose to his feet, overturning his chair. "You have no idea what actually happened
|
|
16
|
+
that night," he said, his voice barely above a whisper.
|
|
17
|
+
```
|
|
18
|
+
|
|
19
|
+
---
|
|
20
|
+
|
|
21
|
+
## Install
|
|
22
|
+
|
|
23
|
+
```bash
|
|
24
|
+
pip install basemode
|
|
25
|
+
```
|
|
26
|
+
|
|
27
|
+
Requires Python 3.11+. Uses [LiteLLM](https://docs.litellm.ai/) under the hood, so any model LiteLLM supports is available.
|
|
28
|
+
|
|
29
|
+
Set API keys as environment variables or in a `.env` file:
|
|
30
|
+
|
|
31
|
+
```bash
|
|
32
|
+
OPENAI_API_KEY=sk-...
|
|
33
|
+
ANTHROPIC_API_KEY=sk-ant-...
|
|
34
|
+
OPENROUTER_API_KEY=sk-or-...
|
|
35
|
+
GROQ_API_KEY=gsk_...
|
|
36
|
+
GEMINI_API_KEY=...
|
|
37
|
+
TOGETHER_API_KEY=...
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
---
|
|
41
|
+
|
|
42
|
+
## CLI
|
|
43
|
+
|
|
44
|
+
```bash
|
|
45
|
+
# Stream a single continuation (defaults to gpt-4o-mini)
|
|
46
|
+
basemode "The ship rounded the headland and"
|
|
47
|
+
|
|
48
|
+
# Choose a model
|
|
49
|
+
basemode "The ship rounded" --model anthropic/claude-3-haiku-20240307
|
|
50
|
+
|
|
51
|
+
# Generate 4 parallel branches, shown side by side
|
|
52
|
+
basemode "The ship rounded" --branches 4
|
|
53
|
+
|
|
54
|
+
# Pipe text in
|
|
55
|
+
cat chapter1.txt | basemode --model groq/llama-3.3-70b-versatile
|
|
56
|
+
|
|
57
|
+
# Show which coercion strategy is being used
|
|
58
|
+
basemode "text" --show-strategy
|
|
59
|
+
|
|
60
|
+
# Override the auto-detected strategy
|
|
61
|
+
basemode "text" --strategy prefill
|
|
62
|
+
|
|
63
|
+
# Explore available models
|
|
64
|
+
basemode models --available # only models with keys configured
|
|
65
|
+
basemode models --provider openai
|
|
66
|
+
basemode models --search claude
|
|
67
|
+
|
|
68
|
+
# Check strategy detection for any model
|
|
69
|
+
basemode info mistral-large-latest
|
|
70
|
+
```
|
|
71
|
+
|
|
72
|
+
---
|
|
73
|
+
|
|
74
|
+
## Python API
|
|
75
|
+
|
|
76
|
+
```python
|
|
77
|
+
from basemode import continue_text, branch_text
|
|
78
|
+
|
|
79
|
+
# Stream a single continuation
|
|
80
|
+
async for token in continue_text(
|
|
81
|
+
"She opened the letter with trembling hands.",
|
|
82
|
+
model="gpt-4o-mini",
|
|
83
|
+
max_tokens=200,
|
|
84
|
+
temperature=0.9,
|
|
85
|
+
):
|
|
86
|
+
print(token, end="", flush=True)
|
|
87
|
+
|
|
88
|
+
# Generate n parallel branches as (branch_idx, token) tuples
|
|
89
|
+
async for idx, token in branch_text(
|
|
90
|
+
"She opened the letter",
|
|
91
|
+
model="anthropic/claude-3-haiku-20240307",
|
|
92
|
+
n=4,
|
|
93
|
+
max_tokens=200,
|
|
94
|
+
):
|
|
95
|
+
print(f"[{idx}] {token}", end="", flush=True)
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
Token boundaries are handled correctly: `prefix + "".join(tokens)` always produces clean, properly-spaced text regardless of which strategy was used.
|
|
99
|
+
|
|
100
|
+
---
|
|
101
|
+
|
|
102
|
+
## How it works
|
|
103
|
+
|
|
104
|
+
Chat-tuned models need different tricks depending on the provider. `basemode` auto-detects the best strategy from the model name:
|
|
105
|
+
|
|
106
|
+
| Strategy | When used | How |
|
|
107
|
+
|----------|-----------|-----|
|
|
108
|
+
| `completion` | OpenAI base models (`davinci-002`, `gpt-3.5-turbo-instruct`) | Native `/completions` endpoint — no coercion needed |
|
|
109
|
+
| `prefill` | Anthropic (`claude-*`) | Splits the prefix: last 50 chars become the start of the assistant turn, forcing the model to continue from exactly that point |
|
|
110
|
+
| `system` | Everything else | System prompt instructs the model to output only continuation, plus trailing-space normalization to prevent word smashing |
|
|
111
|
+
| `few_shot` | Stubborn models | Four varied examples (fiction, technical, poetry, dialogue) in the system prompt |
|
|
112
|
+
| `fim` | DeepSeek Coder, StarCoder, CodeLlama | Fill-in-the-middle special tokens |
|
|
113
|
+
|
|
114
|
+
Override auto-detection with `--strategy` / `strategy=` parameter.
|
|
115
|
+
|
|
116
|
+
---
|
|
117
|
+
|
|
118
|
+
## Model compatibility
|
|
119
|
+
|
|
120
|
+
Tested results from the integration suite. Reliability is assessed on two axes: **does it produce output** and **is the output a clean continuation** (no "Sure!", no preamble, correct word boundaries).
|
|
121
|
+
|
|
122
|
+
### OpenAI
|
|
123
|
+
|
|
124
|
+
| Model | Strategy | Reliability | Notes |
|
|
125
|
+
|-------|----------|-------------|-------|
|
|
126
|
+
| `gpt-4o` | `system` | ⭐⭐⭐⭐⭐ | Excellent continuation quality, zero preamble |
|
|
127
|
+
| `gpt-4o-mini` | `system` | ⭐⭐⭐⭐⭐ | Fast, cheap, reliable. Recommended default |
|
|
128
|
+
| `gpt-3.5-turbo-instruct` | `completion` | ⭐⭐⭐⭐⭐ | Native completions endpoint, no coercion needed |
|
|
129
|
+
| `davinci-002` | `completion` | ⭐⭐⭐⭐⭐ | True base model, best raw continuation behavior |
|
|
130
|
+
|
|
131
|
+
### Anthropic
|
|
132
|
+
|
|
133
|
+
| Model | Strategy | Reliability | Notes |
|
|
134
|
+
|-------|----------|-------------|-------|
|
|
135
|
+
| `anthropic/claude-3-opus-20240229` | `prefill` | ⭐⭐⭐⭐⭐ | Superb quality, stays in voice perfectly |
|
|
136
|
+
| `anthropic/claude-3-5-sonnet-20241022` | `prefill` | ⭐⭐⭐⭐⭐ | Best balance of quality and speed |
|
|
137
|
+
| `anthropic/claude-3-haiku-20240307` | `prefill` | ⭐⭐⭐⭐⭐ | Fast and reliable, tested extensively here |
|
|
138
|
+
| `anthropic/claude-3-5-haiku-20241022` | `prefill` | ⭐⭐⭐⭐ | Use dated model ID; `claude-3-5-haiku-latest` alias not supported by all API tiers |
|
|
139
|
+
|
|
140
|
+
> **The prefill trick**: Anthropic's API lets you pre-fill the assistant's response. `basemode` puts the last 50 characters of your prefix into the assistant turn, so Claude is literally mid-sentence before it generates a single new token. This produces the cleanest continuation of any strategy tested.
|
|
141
|
+
|
|
142
|
+
### Groq
|
|
143
|
+
|
|
144
|
+
| Model | Strategy | Reliability | Notes |
|
|
145
|
+
|-------|----------|-------------|-------|
|
|
146
|
+
| `groq/llama-3.3-70b-versatile` | `system` | ⭐⭐⭐⭐⭐ | Very fast inference, high-quality output, no preamble |
|
|
147
|
+
| `groq/llama-3.1-70b-versatile` | `system` | ⭐⭐⭐⭐ | Slightly older, still reliable |
|
|
148
|
+
| `groq/mixtral-8x7b-32768` | `system` | ⭐⭐⭐⭐ | Good for longer context continuations |
|
|
149
|
+
|
|
150
|
+
Groq's speed (often under 1s for short continuations) makes it excellent for interactive loom exploration.
|
|
151
|
+
|
|
152
|
+
### Google Gemini
|
|
153
|
+
|
|
154
|
+
| Model | Strategy | Reliability | Notes |
|
|
155
|
+
|-------|----------|-------------|-------|
|
|
156
|
+
| `gemini/gemini-2.5-flash` | `system` | ⭐⭐⭐⭐ | **Thinking model** — `basemode` automatically allocates a thinking token budget. Do not use `max_tokens < 1536` |
|
|
157
|
+
| `gemini/gemini-2.5-pro` | `system` | ⭐⭐⭐⭐ | Same thinking model caveat |
|
|
158
|
+
| `gemini/gemini-flash-latest` | `system` | ⭐⭐⭐ | Non-thinking, faster, but output sometimes truncates at low token counts |
|
|
159
|
+
|
|
160
|
+
> **Gemini 2.5 thinking models**: These models spend tokens on internal reasoning before producing visible output. `basemode` detects Gemini 2.5 models and automatically sets a `thinking` budget, ensuring the visible output isn't starved. If you're passing `max_tokens` directly, use at least `1536`.
|
|
161
|
+
|
|
162
|
+
### Together AI
|
|
163
|
+
|
|
164
|
+
| Model | Strategy | Reliability | Notes |
|
|
165
|
+
|-------|----------|-------------|-------|
|
|
166
|
+
| `together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo` | `system` | ⭐⭐⭐⭐⭐ | Excellent quality and speed, well-behaved with system prompt |
|
|
167
|
+
| `together_ai/meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo` | `system` | ⭐⭐⭐⭐⭐ | Best open-weight quality available |
|
|
168
|
+
| `together_ai/mistralai/Mixtral-8x22B-Instruct-v0.1` | `system` | ⭐⭐⭐⭐ | Strong multilingual continuation |
|
|
169
|
+
|
|
170
|
+
### OpenRouter
|
|
171
|
+
|
|
172
|
+
OpenRouter routes to many providers. Any model available there works via `openrouter/<provider>/<model>`:
|
|
173
|
+
|
|
174
|
+
```bash
|
|
175
|
+
basemode "text" --model openrouter/openai/gpt-4o
|
|
176
|
+
basemode "text" --model openrouter/anthropic/claude-3-5-sonnet
|
|
177
|
+
basemode "text" --model openrouter/mistralai/mistral-large
|
|
178
|
+
```
|
|
179
|
+
|
|
180
|
+
| Reliability | Notes |
|
|
181
|
+
|-------------|-------|
|
|
182
|
+
| ⭐⭐⭐⭐ | Model availability varies; use `basemode models --provider openrouter` to check |
|
|
183
|
+
|
|
184
|
+
OpenRouter is particularly useful for accessing base model variants (e.g. `openrouter/mistralai/mistral-7b`) which don't require any coercion.
|
|
185
|
+
|
|
186
|
+
---
|
|
187
|
+
|
|
188
|
+
## Reliability ratings
|
|
189
|
+
|
|
190
|
+
| ⭐⭐⭐⭐⭐ | Zero preamble, correct spacing, stays in voice |
|
|
191
|
+
|-----------|---|
|
|
192
|
+
| ⭐⭐⭐⭐ | Occasionally needs strategy override or specific model ID |
|
|
193
|
+
| ⭐⭐⭐ | Works but has edge cases or quirks |
|
|
194
|
+
| ⭐⭐ | Usable with manual `--strategy` override |
|
|
195
|
+
| ⭐ | Unreliable, avoid |
|
|
196
|
+
|
|
197
|
+
---
|
|
198
|
+
|
|
199
|
+
## Development
|
|
200
|
+
|
|
201
|
+
```bash
|
|
202
|
+
git clone https://github.com/FergusFettes/basemode
|
|
203
|
+
cd basemode
|
|
204
|
+
uv sync
|
|
205
|
+
cp .env.example .env # add your API keys
|
|
206
|
+
|
|
207
|
+
# Unit tests (no API calls)
|
|
208
|
+
uv run pytest
|
|
209
|
+
|
|
210
|
+
# Integration tests (hits real APIs, costs ~$0.01)
|
|
211
|
+
uv run pytest -m integration
|
|
212
|
+
```
|
|
213
|
+
|
|
214
|
+
Pre-commit hooks (ruff format + lint):
|
|
215
|
+
|
|
216
|
+
```bash
|
|
217
|
+
uv run pre-commit install
|
|
218
|
+
```
|
|
219
|
+
|
|
220
|
+
---
|
|
221
|
+
|
|
222
|
+
## Strategies reference
|
|
223
|
+
|
|
224
|
+
```bash
|
|
225
|
+
basemode strategies # list all strategies
|
|
226
|
+
basemode info <model> # show detected strategy for any model
|
|
227
|
+
```
|
|
228
|
+
|
|
229
|
+
To force a specific strategy:
|
|
230
|
+
|
|
231
|
+
```bash
|
|
232
|
+
basemode "text" --strategy prefill # Anthropic prefill trick
|
|
233
|
+
basemode "text" --strategy system # system prompt coercion
|
|
234
|
+
basemode "text" --strategy few_shot # few-shot examples
|
|
235
|
+
basemode "text" --strategy completion # OpenAI /completions endpoint
|
|
236
|
+
basemode "text" --strategy fim # fill-in-the-middle
|
|
237
|
+
```
|
|
238
|
+
|
|
239
|
+
---
|
|
240
|
+
|
|
241
|
+
## Why this exists
|
|
242
|
+
|
|
243
|
+
This is the model layer for [loom](https://github.com/FergusFettes/loom) — a multiverse writing interface for human-AI collaboration. Every loom node needs raw continuation. Getting that reliably out of modern chat models is surprisingly hard, and the solution is different for every provider. `basemode` packages that knowledge into a single interface so loom (and anything else that needs it) doesn't have to think about it.
|
|
@@ -0,0 +1,62 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "basemode"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Make any LLM do raw text continuation"
|
|
9
|
+
requires-python = ">=3.11"
|
|
10
|
+
dependencies = [
|
|
11
|
+
"litellm>=1.63.0",
|
|
12
|
+
"typer>=0.12.0",
|
|
13
|
+
"rich>=13.0.0",
|
|
14
|
+
"pydantic-settings>=2.0.0",
|
|
15
|
+
"python-dotenv>=1.0.0",
|
|
16
|
+
"anyio>=4.0.0",
|
|
17
|
+
]
|
|
18
|
+
|
|
19
|
+
[project.scripts]
|
|
20
|
+
basemode = "basemode.cli:app"
|
|
21
|
+
|
|
22
|
+
[dependency-groups]
|
|
23
|
+
dev = [
|
|
24
|
+
"pytest>=8.0.0",
|
|
25
|
+
"pytest-asyncio>=0.23.0",
|
|
26
|
+
"pytest-cov>=5.0.0",
|
|
27
|
+
"ruff>=0.5.0",
|
|
28
|
+
"pre-commit>=3.7.0",
|
|
29
|
+
]
|
|
30
|
+
|
|
31
|
+
[tool.hatch.build.targets.wheel]
|
|
32
|
+
packages = ["src/basemode"]
|
|
33
|
+
|
|
34
|
+
[tool.ruff]
|
|
35
|
+
line-length = 88
|
|
36
|
+
src = ["src"]
|
|
37
|
+
target-version = "py311"
|
|
38
|
+
|
|
39
|
+
[tool.ruff.lint]
|
|
40
|
+
select = ["E", "F", "I", "UP", "B", "RUF"]
|
|
41
|
+
ignore = ["E501"]
|
|
42
|
+
|
|
43
|
+
[tool.ruff.lint.isort]
|
|
44
|
+
known-first-party = ["basemode"]
|
|
45
|
+
|
|
46
|
+
[tool.pytest.ini_options]
|
|
47
|
+
asyncio_mode = "auto"
|
|
48
|
+
addopts = "--cov=src/basemode --cov-report=term-missing -m 'not integration'"
|
|
49
|
+
testpaths = ["tests"]
|
|
50
|
+
markers = ["integration: hits real APIs, not run by default"]
|
|
51
|
+
|
|
52
|
+
[tool.coverage.run]
|
|
53
|
+
source = ["src/basemode"]
|
|
54
|
+
omit = ["*/tests/*"]
|
|
55
|
+
|
|
56
|
+
[tool.coverage.report]
|
|
57
|
+
exclude_lines = [
|
|
58
|
+
"pragma: no cover",
|
|
59
|
+
"def __repr__",
|
|
60
|
+
"raise NotImplementedError",
|
|
61
|
+
"if TYPE_CHECKING:",
|
|
62
|
+
]
|
|
@@ -0,0 +1,136 @@
|
|
|
1
|
+
import asyncio
|
|
2
|
+
import sys
|
|
3
|
+
from typing import Annotated
|
|
4
|
+
|
|
5
|
+
import typer
|
|
6
|
+
from rich.columns import Columns
|
|
7
|
+
from rich.console import Console
|
|
8
|
+
from rich.panel import Panel
|
|
9
|
+
from rich.table import Table
|
|
10
|
+
from rich.text import Text
|
|
11
|
+
|
|
12
|
+
from .continue_ import branch_text, continue_text
|
|
13
|
+
from .detect import detect_strategy, normalize_model
|
|
14
|
+
from .models import list_models, list_providers
|
|
15
|
+
from .strategies import REGISTRY
|
|
16
|
+
|
|
17
|
+
app = typer.Typer(help="Make any LLM do raw text continuation.")
|
|
18
|
+
console = Console()
|
|
19
|
+
|
|
20
|
+
_BRANCH_COLORS = ["green", "blue", "yellow", "magenta", "cyan"]
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def _read_prefix(prefix: str | None) -> str | None:
|
|
24
|
+
"""Return prefix from arg, stdin pipe, or None."""
|
|
25
|
+
if prefix is not None:
|
|
26
|
+
return prefix
|
|
27
|
+
if not sys.stdin.isatty():
|
|
28
|
+
return sys.stdin.read()
|
|
29
|
+
return None
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
@app.callback(invoke_without_command=True)
|
|
33
|
+
def main(
|
|
34
|
+
ctx: typer.Context,
|
|
35
|
+
prefix: Annotated[str | None, typer.Argument(help="Text to continue (or pipe via stdin)")] = None,
|
|
36
|
+
model: Annotated[str, typer.Option("-m", "--model")] = "gpt-4o-mini",
|
|
37
|
+
n: Annotated[int, typer.Option("-n", "--branches", help="Number of parallel continuations")] = 1,
|
|
38
|
+
max_tokens: Annotated[int, typer.Option("--max-tokens")] = 200,
|
|
39
|
+
temperature: Annotated[float, typer.Option("-t", "--temperature")] = 0.9,
|
|
40
|
+
strategy: Annotated[str | None, typer.Option("-s", "--strategy")] = None,
|
|
41
|
+
show_strategy: Annotated[bool, typer.Option("--show-strategy")] = False,
|
|
42
|
+
) -> None:
|
|
43
|
+
if ctx.invoked_subcommand is not None:
|
|
44
|
+
return
|
|
45
|
+
|
|
46
|
+
text = _read_prefix(prefix)
|
|
47
|
+
if text is None:
|
|
48
|
+
console.print(ctx.get_help())
|
|
49
|
+
return
|
|
50
|
+
|
|
51
|
+
prefix = text.rstrip("\n")
|
|
52
|
+
|
|
53
|
+
if show_strategy:
|
|
54
|
+
strat = detect_strategy(normalize_model(model), strategy)
|
|
55
|
+
console.print(f"[dim]strategy: {strat.name}[/dim]")
|
|
56
|
+
|
|
57
|
+
if n == 1:
|
|
58
|
+
asyncio.run(_stream_one(prefix, model, max_tokens, temperature, strategy))
|
|
59
|
+
else:
|
|
60
|
+
asyncio.run(_stream_branches(prefix, model, n, max_tokens, temperature, strategy)) # noqa: E501
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
async def _stream_one(prefix: str, model: str, max_tokens: int, temperature: float, strategy: str | None) -> None:
|
|
64
|
+
console.print(f"[dim]{prefix}[/dim]", end="")
|
|
65
|
+
async for token in continue_text(prefix, model, max_tokens=max_tokens, temperature=temperature, strategy=strategy):
|
|
66
|
+
console.print(token, end="")
|
|
67
|
+
console.print()
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
async def _stream_branches(
|
|
71
|
+
prefix: str, model: str, n: int, max_tokens: int, temperature: float, strategy: str | None
|
|
72
|
+
) -> None:
|
|
73
|
+
buffers: list[list[str]] = [[] for _ in range(n)]
|
|
74
|
+
console.print(f"[dim]{prefix}[/dim]\n")
|
|
75
|
+
|
|
76
|
+
async for idx, token in branch_text(
|
|
77
|
+
prefix, model, n=n, max_tokens=max_tokens, temperature=temperature, strategy=strategy
|
|
78
|
+
):
|
|
79
|
+
buffers[idx].append(token)
|
|
80
|
+
|
|
81
|
+
panels = []
|
|
82
|
+
for i, buf in enumerate(buffers):
|
|
83
|
+
color = _BRANCH_COLORS[i % len(_BRANCH_COLORS)]
|
|
84
|
+
text = Text(prefix, style="dim")
|
|
85
|
+
text.append("".join(buf), style=color)
|
|
86
|
+
panels.append(Panel(text, title=f"[{color}]Branch {i + 1}[/{color}]"))
|
|
87
|
+
|
|
88
|
+
console.print(Columns(panels, equal=True))
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
@app.command()
|
|
92
|
+
def models(
|
|
93
|
+
provider: Annotated[str | None, typer.Option("-p", "--provider")] = None,
|
|
94
|
+
search: Annotated[str | None, typer.Option("-s", "--search")] = None,
|
|
95
|
+
available: Annotated[bool, typer.Option("-a", "--available", help="Only show models with keys set")] = False,
|
|
96
|
+
) -> None:
|
|
97
|
+
"""List available models."""
|
|
98
|
+
results = list_models(provider=provider, search=search, available_only=available)
|
|
99
|
+
if not results:
|
|
100
|
+
console.print("[yellow]No models found.[/yellow]")
|
|
101
|
+
return
|
|
102
|
+
|
|
103
|
+
table = Table("Model", show_header=True, header_style="bold")
|
|
104
|
+
for m in results:
|
|
105
|
+
table.add_row(m)
|
|
106
|
+
console.print(table)
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
@app.command()
|
|
110
|
+
def providers() -> None:
|
|
111
|
+
"""List all known providers."""
|
|
112
|
+
for p in list_providers():
|
|
113
|
+
console.print(p)
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
@app.command()
|
|
117
|
+
def strategies() -> None:
|
|
118
|
+
"""List available continuation strategies."""
|
|
119
|
+
table = Table("Name", "Description", show_header=True, header_style="bold")
|
|
120
|
+
descriptions = {
|
|
121
|
+
"completion": "OpenAI /completions endpoint — for true base models",
|
|
122
|
+
"prefill": "Anthropic assistant prefill trick",
|
|
123
|
+
"system": "System prompt coercion — generic fallback for any chat model",
|
|
124
|
+
"few_shot": "Few-shot examples in system prompt — for stubborn models",
|
|
125
|
+
"fim": "Fill-in-the-middle tokens — DeepSeek, StarCoder, CodeLlama",
|
|
126
|
+
}
|
|
127
|
+
for name in REGISTRY:
|
|
128
|
+
table.add_row(name, descriptions.get(name, ""))
|
|
129
|
+
console.print(table)
|
|
130
|
+
|
|
131
|
+
|
|
132
|
+
@app.command()
|
|
133
|
+
def info(model: Annotated[str, typer.Argument(help="Model name to inspect")]) -> None:
|
|
134
|
+
"""Show which strategy would be used for a given model."""
|
|
135
|
+
strat = detect_strategy(model)
|
|
136
|
+
console.print(f"[bold]{model}[/bold] → [green]{strat.name}[/green]")
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
import asyncio
|
|
2
|
+
from collections.abc import AsyncGenerator
|
|
3
|
+
|
|
4
|
+
from .detect import detect_strategy, normalize_model
|
|
5
|
+
from .params import GenerationParams
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
async def continue_text(
|
|
9
|
+
prefix: str,
|
|
10
|
+
model: str = "gpt-4o-mini",
|
|
11
|
+
*,
|
|
12
|
+
max_tokens: int = 200,
|
|
13
|
+
temperature: float = 0.9,
|
|
14
|
+
strategy: str | None = None,
|
|
15
|
+
**extra,
|
|
16
|
+
) -> AsyncGenerator[str, None]:
|
|
17
|
+
"""Stream a single continuation."""
|
|
18
|
+
model = normalize_model(model)
|
|
19
|
+
params = GenerationParams(model=model, max_tokens=max_tokens, temperature=temperature, extra=extra)
|
|
20
|
+
strat = detect_strategy(model, strategy)
|
|
21
|
+
async for token in strat.stream(prefix, params):
|
|
22
|
+
yield token
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
async def branch_text(
|
|
26
|
+
prefix: str,
|
|
27
|
+
model: str = "gpt-4o-mini",
|
|
28
|
+
*,
|
|
29
|
+
n: int = 4,
|
|
30
|
+
max_tokens: int = 200,
|
|
31
|
+
temperature: float = 0.9,
|
|
32
|
+
strategy: str | None = None,
|
|
33
|
+
**extra,
|
|
34
|
+
) -> AsyncGenerator[tuple[int, str], None]:
|
|
35
|
+
"""Stream n parallel continuations as (branch_idx, token) tuples."""
|
|
36
|
+
model = normalize_model(model)
|
|
37
|
+
params = GenerationParams(model=model, max_tokens=max_tokens, temperature=temperature, extra=extra)
|
|
38
|
+
strat = detect_strategy(model, strategy)
|
|
39
|
+
|
|
40
|
+
queue: asyncio.Queue[tuple[int, str] | None] = asyncio.Queue()
|
|
41
|
+
|
|
42
|
+
async def run_branch(idx: int) -> None:
|
|
43
|
+
async for token in strat.stream(prefix, params):
|
|
44
|
+
await queue.put((idx, token))
|
|
45
|
+
await queue.put(None)
|
|
46
|
+
|
|
47
|
+
tasks = [asyncio.create_task(run_branch(i)) for i in range(n)]
|
|
48
|
+
done = 0
|
|
49
|
+
while done < n:
|
|
50
|
+
item = await queue.get()
|
|
51
|
+
if item is None:
|
|
52
|
+
done += 1
|
|
53
|
+
else:
|
|
54
|
+
yield item
|
|
55
|
+
|
|
56
|
+
await asyncio.gather(*tasks)
|
|
@@ -0,0 +1,64 @@
|
|
|
1
|
+
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
|
2
|
+
|
|
3
|
+
from .strategies import (
|
|
4
|
+
REGISTRY,
|
|
5
|
+
CompletionStrategy,
|
|
6
|
+
ContinuationStrategy,
|
|
7
|
+
FIMStrategy,
|
|
8
|
+
PrefillStrategy,
|
|
9
|
+
SystemPromptStrategy,
|
|
10
|
+
)
|
|
11
|
+
|
|
12
|
+
# Models that use the native completions API
|
|
13
|
+
_COMPLETION_MODELS = {
|
|
14
|
+
"gpt-3.5-turbo-instruct",
|
|
15
|
+
"davinci-002",
|
|
16
|
+
"babbage-002",
|
|
17
|
+
}
|
|
18
|
+
_COMPLETION_SUBSTRINGS = ["text-davinci", "text-curie", "text-babbage", "text-ada"]
|
|
19
|
+
|
|
20
|
+
# Models where FIM is the right move
|
|
21
|
+
_FIM_SUBSTRINGS = ["deepseek-coder", "starcoder", "codellama", "fim"]
|
|
22
|
+
|
|
23
|
+
# Provider prefix to add when litellm can't auto-detect from model name alone
|
|
24
|
+
_PREFIX_MAP = {
|
|
25
|
+
"claude": "anthropic",
|
|
26
|
+
"gemini": "gemini",
|
|
27
|
+
"command": "cohere",
|
|
28
|
+
}
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def normalize_model(model: str) -> str:
|
|
32
|
+
"""Add provider prefix if litellm can't resolve the model name."""
|
|
33
|
+
if "/" in model:
|
|
34
|
+
return model
|
|
35
|
+
try:
|
|
36
|
+
get_llm_provider(model)
|
|
37
|
+
return model
|
|
38
|
+
except Exception:
|
|
39
|
+
m = model.lower()
|
|
40
|
+
for fragment, provider in _PREFIX_MAP.items():
|
|
41
|
+
if fragment in m:
|
|
42
|
+
return f"{provider}/{model}"
|
|
43
|
+
return model
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def detect_strategy(model: str, override: str | None = None) -> ContinuationStrategy:
|
|
47
|
+
if override:
|
|
48
|
+
if override not in REGISTRY:
|
|
49
|
+
valid = ", ".join(REGISTRY)
|
|
50
|
+
raise ValueError(f"Unknown strategy {override!r}. Valid: {valid}")
|
|
51
|
+
return REGISTRY[override]()
|
|
52
|
+
|
|
53
|
+
m = model.lower()
|
|
54
|
+
|
|
55
|
+
if "claude" in m:
|
|
56
|
+
return PrefillStrategy()
|
|
57
|
+
|
|
58
|
+
if model in _COMPLETION_MODELS or any(s in m for s in _COMPLETION_SUBSTRINGS):
|
|
59
|
+
return CompletionStrategy()
|
|
60
|
+
|
|
61
|
+
if any(s in m for s in _FIM_SUBSTRINGS):
|
|
62
|
+
return FIMStrategy()
|
|
63
|
+
|
|
64
|
+
return SystemPromptStrategy()
|
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
import litellm
|
|
2
|
+
|
|
3
|
+
from .settings import settings
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
def list_models(
|
|
7
|
+
provider: str | None = None,
|
|
8
|
+
search: str | None = None,
|
|
9
|
+
available_only: bool = False,
|
|
10
|
+
) -> list[str]:
|
|
11
|
+
by_provider: dict[str, list[str]] = litellm.models_by_provider
|
|
12
|
+
|
|
13
|
+
if available_only:
|
|
14
|
+
providers = settings.available_providers
|
|
15
|
+
models = [m for p in providers for m in by_provider.get(p, [])]
|
|
16
|
+
elif provider:
|
|
17
|
+
models = by_provider.get(provider, [])
|
|
18
|
+
else:
|
|
19
|
+
models = [m for ms in by_provider.values() for m in ms]
|
|
20
|
+
|
|
21
|
+
if search:
|
|
22
|
+
models = [m for m in models if search.lower() in m.lower()]
|
|
23
|
+
|
|
24
|
+
return sorted(set(models))
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def list_providers() -> list[str]:
|
|
28
|
+
return sorted(litellm.models_by_provider.keys())
|