basemode 0.1.0__py3-none-any.whl
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/__init__.py +5 -0
- basemode/cli.py +136 -0
- basemode/continue_.py +56 -0
- basemode/detect.py +64 -0
- basemode/models.py +28 -0
- basemode/params.py +9 -0
- basemode/settings.py +47 -0
- basemode/strategies/__init__.py +27 -0
- basemode/strategies/base.py +14 -0
- basemode/strategies/completion.py +27 -0
- basemode/strategies/few_shot.py +69 -0
- basemode/strategies/fim.py +44 -0
- basemode/strategies/prefill.py +45 -0
- basemode/strategies/system.py +69 -0
- basemode/strategies/utils.py +26 -0
- basemode-0.1.0.dist-info/METADATA +11 -0
- basemode-0.1.0.dist-info/RECORD +19 -0
- basemode-0.1.0.dist-info/WHEEL +4 -0
- basemode-0.1.0.dist-info/entry_points.txt +2 -0
basemode/__init__.py
ADDED
basemode/cli.py
ADDED
|
@@ -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]")
|
basemode/continue_.py
ADDED
|
@@ -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)
|
basemode/detect.py
ADDED
|
@@ -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()
|
basemode/models.py
ADDED
|
@@ -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())
|
basemode/params.py
ADDED
basemode/settings.py
ADDED
|
@@ -0,0 +1,47 @@
|
|
|
1
|
+
import os
|
|
2
|
+
from pathlib import Path
|
|
3
|
+
|
|
4
|
+
from dotenv import load_dotenv
|
|
5
|
+
from pydantic_settings import BaseSettings, SettingsConfigDict
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def _find_env() -> Path | None:
|
|
9
|
+
for p in [Path(".env"), Path(__file__).parent.parent.parent / ".env"]:
|
|
10
|
+
if p.exists():
|
|
11
|
+
return p
|
|
12
|
+
return None
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
# Load into os.environ so litellm can pick up keys directly.
|
|
16
|
+
_env_path = _find_env()
|
|
17
|
+
if _env_path:
|
|
18
|
+
load_dotenv(_env_path, override=False)
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class Settings(BaseSettings):
|
|
22
|
+
model_config = SettingsConfigDict(env_file=_env_path, extra="allow")
|
|
23
|
+
|
|
24
|
+
openai_api_key: str = ""
|
|
25
|
+
anthropic_api_key: str = ""
|
|
26
|
+
openrouter_api_key: str = ""
|
|
27
|
+
groq_api_key: str = ""
|
|
28
|
+
gemini_api_key: str = ""
|
|
29
|
+
together_api_key: str = ""
|
|
30
|
+
|
|
31
|
+
@property
|
|
32
|
+
def available_providers(self) -> list[str]:
|
|
33
|
+
return [
|
|
34
|
+
provider
|
|
35
|
+
for provider, key in [
|
|
36
|
+
("openai", self.openai_api_key),
|
|
37
|
+
("anthropic", self.anthropic_api_key),
|
|
38
|
+
("openrouter", self.openrouter_api_key),
|
|
39
|
+
("groq", self.groq_api_key),
|
|
40
|
+
("gemini", self.gemini_api_key),
|
|
41
|
+
("together_ai", self.together_api_key),
|
|
42
|
+
]
|
|
43
|
+
if key
|
|
44
|
+
]
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
settings = Settings()
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
from .base import ContinuationStrategy
|
|
2
|
+
from .completion import CompletionStrategy
|
|
3
|
+
from .few_shot import FewShotStrategy
|
|
4
|
+
from .fim import FIMStrategy
|
|
5
|
+
from .prefill import PrefillStrategy
|
|
6
|
+
from .system import SystemPromptStrategy
|
|
7
|
+
|
|
8
|
+
REGISTRY: dict[str, type[ContinuationStrategy]] = {
|
|
9
|
+
s.name: s
|
|
10
|
+
for s in [
|
|
11
|
+
CompletionStrategy,
|
|
12
|
+
PrefillStrategy,
|
|
13
|
+
SystemPromptStrategy,
|
|
14
|
+
FewShotStrategy,
|
|
15
|
+
FIMStrategy,
|
|
16
|
+
]
|
|
17
|
+
}
|
|
18
|
+
|
|
19
|
+
__all__ = [
|
|
20
|
+
"ContinuationStrategy",
|
|
21
|
+
"CompletionStrategy",
|
|
22
|
+
"FewShotStrategy",
|
|
23
|
+
"FIMStrategy",
|
|
24
|
+
"PrefillStrategy",
|
|
25
|
+
"SystemPromptStrategy",
|
|
26
|
+
"REGISTRY",
|
|
27
|
+
]
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
from abc import ABC, abstractmethod
|
|
2
|
+
from collections.abc import AsyncGenerator
|
|
3
|
+
|
|
4
|
+
from ..params import GenerationParams
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
class ContinuationStrategy(ABC):
|
|
8
|
+
name: str
|
|
9
|
+
|
|
10
|
+
@abstractmethod
|
|
11
|
+
def stream(self, prefix: str, params: GenerationParams) -> AsyncGenerator[str, None]: ...
|
|
12
|
+
|
|
13
|
+
def __repr__(self) -> str:
|
|
14
|
+
return f"{self.__class__.__name__}()"
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
"""OpenAI-compatible /completions endpoint — works natively with base models."""
|
|
2
|
+
from collections.abc import AsyncGenerator
|
|
3
|
+
|
|
4
|
+
import litellm
|
|
5
|
+
|
|
6
|
+
from ..params import GenerationParams
|
|
7
|
+
from .base import ContinuationStrategy
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class CompletionStrategy(ContinuationStrategy):
|
|
11
|
+
"""Uses the text completions API. Best for true base models (davinci, etc.)."""
|
|
12
|
+
|
|
13
|
+
name = "completion"
|
|
14
|
+
|
|
15
|
+
async def stream(self, prefix: str, params: GenerationParams) -> AsyncGenerator[str, None]:
|
|
16
|
+
response = await litellm.atext_completion(
|
|
17
|
+
model=params.model,
|
|
18
|
+
prompt=prefix,
|
|
19
|
+
max_tokens=params.max_tokens,
|
|
20
|
+
temperature=params.temperature,
|
|
21
|
+
stream=True,
|
|
22
|
+
**params.extra,
|
|
23
|
+
)
|
|
24
|
+
async for chunk in response:
|
|
25
|
+
token = chunk.choices[0].text or ""
|
|
26
|
+
if token:
|
|
27
|
+
yield token
|
|
@@ -0,0 +1,69 @@
|
|
|
1
|
+
"""Few-shot coercion — for stubborn models that ignore system prompts."""
|
|
2
|
+
from collections.abc import AsyncGenerator
|
|
3
|
+
|
|
4
|
+
import litellm
|
|
5
|
+
|
|
6
|
+
from ..params import GenerationParams
|
|
7
|
+
from .base import ContinuationStrategy
|
|
8
|
+
from .utils import needs_leading_space, normalize_prefix
|
|
9
|
+
|
|
10
|
+
# Varied examples: fiction, technical, poetry, dialogue
|
|
11
|
+
_EXAMPLES = [
|
|
12
|
+
(
|
|
13
|
+
"The experiment had been running for three days when Dr. Chen noticed the anomaly.",
|
|
14
|
+
" At first she thought it was a measurement error—the kind that comes from",
|
|
15
|
+
),
|
|
16
|
+
(
|
|
17
|
+
"To install the package, first ensure you have Python 3.11 or higher.",
|
|
18
|
+
" Then run the following command in your terminal:\n\n```\npip install",
|
|
19
|
+
),
|
|
20
|
+
(
|
|
21
|
+
"the rain comes down like static",
|
|
22
|
+
"\nbetween stations, the city\nblurs into signal",
|
|
23
|
+
),
|
|
24
|
+
(
|
|
25
|
+
"Look, I'm just saying—if you'd been there, you'd understand why I had to",
|
|
26
|
+
" make that call. There was no good option. Either way, someone was going to",
|
|
27
|
+
),
|
|
28
|
+
]
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def _build_system_prompt() -> str:
|
|
32
|
+
examples = "\n\n".join(
|
|
33
|
+
f'Input: """{inp}"""\nOutput: """{out}"""' for inp, out in _EXAMPLES
|
|
34
|
+
)
|
|
35
|
+
return (
|
|
36
|
+
"You continue text. Given input text, output only the natural continuation "
|
|
37
|
+
"with no preamble or acknowledgment. Examples:\n\n" + examples
|
|
38
|
+
)
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
_SYSTEM_PROMPT = _build_system_prompt()
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
class FewShotStrategy(ContinuationStrategy):
|
|
45
|
+
"""Few-shot examples in system prompt. For models that ignore plain instructions."""
|
|
46
|
+
|
|
47
|
+
name = "few_shot"
|
|
48
|
+
|
|
49
|
+
async def stream(self, prefix: str, params: GenerationParams) -> AsyncGenerator[str, None]:
|
|
50
|
+
response = await litellm.acompletion(
|
|
51
|
+
model=params.model,
|
|
52
|
+
messages=[
|
|
53
|
+
{"role": "system", "content": _SYSTEM_PROMPT},
|
|
54
|
+
{"role": "user", "content": normalize_prefix(prefix)},
|
|
55
|
+
],
|
|
56
|
+
max_tokens=params.max_tokens,
|
|
57
|
+
temperature=params.temperature,
|
|
58
|
+
stream=True,
|
|
59
|
+
**params.extra,
|
|
60
|
+
)
|
|
61
|
+
first = True
|
|
62
|
+
async for chunk in response:
|
|
63
|
+
token = chunk.choices[0].delta.content or ""
|
|
64
|
+
if not token:
|
|
65
|
+
continue
|
|
66
|
+
if first and needs_leading_space(prefix, token):
|
|
67
|
+
token = " " + token
|
|
68
|
+
first = False
|
|
69
|
+
yield token
|
|
@@ -0,0 +1,44 @@
|
|
|
1
|
+
"""Fill-in-the-middle — for models that support FIM tokens natively."""
|
|
2
|
+
from collections.abc import AsyncGenerator
|
|
3
|
+
|
|
4
|
+
import litellm
|
|
5
|
+
|
|
6
|
+
from ..params import GenerationParams
|
|
7
|
+
from .base import ContinuationStrategy
|
|
8
|
+
|
|
9
|
+
# Token formats by model family
|
|
10
|
+
_FIM_FORMATS = {
|
|
11
|
+
"deepseek": ("<|fim▁begin|>", "<|fim▁hole|>", "<|fim▁end|>"),
|
|
12
|
+
"starcoder": ("<fim_prefix>", "<fim_suffix>", "<fim_middle>"),
|
|
13
|
+
"codellama": ("▁<PRE>", "▁<SUF>", "▁<MID>"),
|
|
14
|
+
}
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
def _fim_prompt(prefix: str, model: str) -> str:
|
|
18
|
+
for key, (pre, suf, mid) in _FIM_FORMATS.items():
|
|
19
|
+
if key in model.lower():
|
|
20
|
+
return f"{pre}{prefix}{suf}{mid}"
|
|
21
|
+
# Generic fallback
|
|
22
|
+
pre, suf, mid = _FIM_FORMATS["starcoder"]
|
|
23
|
+
return f"{pre}{prefix}{suf}{mid}"
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class FIMStrategy(ContinuationStrategy):
|
|
27
|
+
"""Fill-in-the-middle via text completion. DeepSeek, StarCoder, CodeLlama."""
|
|
28
|
+
|
|
29
|
+
name = "fim"
|
|
30
|
+
|
|
31
|
+
async def stream(self, prefix: str, params: GenerationParams) -> AsyncGenerator[str, None]:
|
|
32
|
+
prompt = _fim_prompt(prefix, params.model)
|
|
33
|
+
response = await litellm.atext_completion(
|
|
34
|
+
model=params.model,
|
|
35
|
+
prompt=prompt,
|
|
36
|
+
max_tokens=params.max_tokens,
|
|
37
|
+
temperature=params.temperature,
|
|
38
|
+
stream=True,
|
|
39
|
+
**params.extra,
|
|
40
|
+
)
|
|
41
|
+
async for chunk in response:
|
|
42
|
+
token = chunk.choices[0].text or ""
|
|
43
|
+
if token:
|
|
44
|
+
yield token
|
|
@@ -0,0 +1,45 @@
|
|
|
1
|
+
"""Anthropic prefill trick — force continuation by seeding the assistant turn."""
|
|
2
|
+
from collections.abc import AsyncGenerator
|
|
3
|
+
|
|
4
|
+
import litellm
|
|
5
|
+
|
|
6
|
+
from ..params import GenerationParams
|
|
7
|
+
from .base import ContinuationStrategy
|
|
8
|
+
|
|
9
|
+
# How many trailing chars of the prefix to use as the assistant seed.
|
|
10
|
+
# Long enough that the model clearly understands it's mid-sentence.
|
|
11
|
+
SEED_LEN = 50
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
class PrefillStrategy(ContinuationStrategy):
|
|
15
|
+
"""
|
|
16
|
+
Works by splitting the prefix: the first part goes in the user turn,
|
|
17
|
+
the last SEED_LEN characters become the start of the assistant turn.
|
|
18
|
+
The model is forced to continue from exactly where the prefix ends.
|
|
19
|
+
"""
|
|
20
|
+
|
|
21
|
+
name = "prefill"
|
|
22
|
+
|
|
23
|
+
async def stream(self, prefix: str, params: GenerationParams) -> AsyncGenerator[str, None]:
|
|
24
|
+
if len(prefix) <= SEED_LEN:
|
|
25
|
+
user_content = "Continue:"
|
|
26
|
+
assistant_seed = prefix
|
|
27
|
+
else:
|
|
28
|
+
user_content = prefix[:-SEED_LEN]
|
|
29
|
+
assistant_seed = prefix[-SEED_LEN:]
|
|
30
|
+
|
|
31
|
+
response = await litellm.acompletion(
|
|
32
|
+
model=params.model,
|
|
33
|
+
messages=[
|
|
34
|
+
{"role": "user", "content": user_content},
|
|
35
|
+
{"role": "assistant", "content": assistant_seed},
|
|
36
|
+
],
|
|
37
|
+
max_tokens=params.max_tokens,
|
|
38
|
+
temperature=params.temperature,
|
|
39
|
+
stream=True,
|
|
40
|
+
**params.extra,
|
|
41
|
+
)
|
|
42
|
+
async for chunk in response:
|
|
43
|
+
token = chunk.choices[0].delta.content or ""
|
|
44
|
+
if token:
|
|
45
|
+
yield token
|
|
@@ -0,0 +1,69 @@
|
|
|
1
|
+
"""System prompt coercion — works on any chat model."""
|
|
2
|
+
from collections.abc import AsyncGenerator
|
|
3
|
+
|
|
4
|
+
import litellm
|
|
5
|
+
|
|
6
|
+
from ..params import GenerationParams
|
|
7
|
+
from .base import ContinuationStrategy
|
|
8
|
+
from .utils import needs_leading_space, normalize_prefix
|
|
9
|
+
|
|
10
|
+
SYSTEM_PROMPT = """\
|
|
11
|
+
You are a text continuation engine. Your only function is to extend the provided \
|
|
12
|
+
text naturally, as if you wrote it yourself. Rules:
|
|
13
|
+
- Output ONLY the continuation — no acknowledgment, preamble, or commentary
|
|
14
|
+
- Continue in the exact same voice, style, and register
|
|
15
|
+
- Begin immediately with the next character that naturally follows
|
|
16
|
+
- Never start with "Sure", "Of course", "Certainly", or any other acknowledgment"""
|
|
17
|
+
|
|
18
|
+
# Gemini 2.5+ models are "thinking" models that consume tokens on internal reasoning
|
|
19
|
+
# before producing output. Without a thinking budget, max_tokens is exhausted by
|
|
20
|
+
# thoughts and the visible output is empty or truncated.
|
|
21
|
+
_THINKING_MODELS = {"gemini-2.5-flash", "gemini-2.5-pro", "gemini-2.5-flash-lite"}
|
|
22
|
+
_THINKING_BUDGET = 1024
|
|
23
|
+
_THINKING_MIN_OUTPUT = 512
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def _is_thinking_model(model: str) -> bool:
|
|
27
|
+
m = model.lower().split("/")[-1]
|
|
28
|
+
return any(t in m for t in _THINKING_MODELS)
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def _thinking_kwargs(params: GenerationParams) -> dict:
|
|
32
|
+
if not _is_thinking_model(params.model):
|
|
33
|
+
return {}
|
|
34
|
+
return {
|
|
35
|
+
"thinking": {"type": "enabled", "budget_tokens": _THINKING_BUDGET},
|
|
36
|
+
"max_tokens": max(params.max_tokens, _THINKING_BUDGET + _THINKING_MIN_OUTPUT),
|
|
37
|
+
}
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
class SystemPromptStrategy(ContinuationStrategy):
|
|
41
|
+
"""Generic coercion via system prompt. Fallback for any chat model."""
|
|
42
|
+
|
|
43
|
+
name = "system"
|
|
44
|
+
|
|
45
|
+
async def stream(self, prefix: str, params: GenerationParams) -> AsyncGenerator[str, None]:
|
|
46
|
+
kwargs = {
|
|
47
|
+
"max_tokens": params.max_tokens,
|
|
48
|
+
"temperature": params.temperature,
|
|
49
|
+
**_thinking_kwargs(params), # may override max_tokens for thinking models
|
|
50
|
+
**params.extra,
|
|
51
|
+
}
|
|
52
|
+
response = await litellm.acompletion(
|
|
53
|
+
model=params.model,
|
|
54
|
+
messages=[
|
|
55
|
+
{"role": "system", "content": SYSTEM_PROMPT},
|
|
56
|
+
{"role": "user", "content": normalize_prefix(prefix)},
|
|
57
|
+
],
|
|
58
|
+
stream=True,
|
|
59
|
+
**kwargs,
|
|
60
|
+
)
|
|
61
|
+
first = True
|
|
62
|
+
async for chunk in response:
|
|
63
|
+
token = chunk.choices[0].delta.content or ""
|
|
64
|
+
if not token:
|
|
65
|
+
continue
|
|
66
|
+
if first and needs_leading_space(prefix, token):
|
|
67
|
+
token = " " + token
|
|
68
|
+
first = False
|
|
69
|
+
yield token
|
|
@@ -0,0 +1,26 @@
|
|
|
1
|
+
def normalize_prefix(prefix: str) -> str:
|
|
2
|
+
"""Ensure prefix ends with exactly one space for the model input.
|
|
3
|
+
|
|
4
|
+
Chat models respond without a leading space, so we strip trailing whitespace
|
|
5
|
+
and add exactly one space. This makes the model output tokens that join
|
|
6
|
+
correctly when we prepend a space to the first token if needed.
|
|
7
|
+
|
|
8
|
+
Not applied to completion/prefill strategies — they handle boundaries natively.
|
|
9
|
+
"""
|
|
10
|
+
return prefix.rstrip() + " "
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def needs_leading_space(prefix: str, first_token: str) -> bool:
|
|
14
|
+
"""Return True if a space must be injected between prefix and first_token.
|
|
15
|
+
|
|
16
|
+
After sending normalize_prefix(prefix) to the model, the model outputs
|
|
17
|
+
first_token without a leading space. If the original prefix didn't end
|
|
18
|
+
with whitespace, the space was consumed in the model input and must be
|
|
19
|
+
restored so that prefix + tokens is correct text.
|
|
20
|
+
"""
|
|
21
|
+
return (
|
|
22
|
+
bool(prefix)
|
|
23
|
+
and not prefix[-1].isspace()
|
|
24
|
+
and bool(first_token)
|
|
25
|
+
and not first_token[0].isspace()
|
|
26
|
+
)
|
|
@@ -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
|
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
basemode/__init__.py,sha256=fdPLFqO2bCXcZeTFmxhM3rJQ3_6ByBmcekSGQ01GXcg,206
|
|
2
|
+
basemode/cli.py,sha256=HYrPF3mT1Fkl2d_l04fDtMkLGWpsBypB540iQ4cqvcg,4861
|
|
3
|
+
basemode/continue_.py,sha256=iquyxU8beFpoEDSL_HO0nxavfoEsTX3U7875cmqEh-4,1659
|
|
4
|
+
basemode/detect.py,sha256=T1bx3NaPrMDEN-YBDWc8lydERkjQZE6cAgPN_5GA1hM,1762
|
|
5
|
+
basemode/models.py,sha256=l4rcGJQ4KrBeldtS9drBxfzIVo3jE_Z399u3-RXQoLI,738
|
|
6
|
+
basemode/params.py,sha256=_2tycQWR5vZmj2L3vXcIGil3zgZaJCEtCcX9HAU7dRo,194
|
|
7
|
+
basemode/settings.py,sha256=qlvrr_73O2erHxxmbx1HmpOwfop-qIu9mAeuMSqWXsc,1244
|
|
8
|
+
basemode/strategies/__init__.py,sha256=AvdKHyJKvl2UY2bovrVEOvy_JQwhxE8j5I3O7Dqrqyw,624
|
|
9
|
+
basemode/strategies/base.py,sha256=GPj7_hWm-XeaXVqpY193l_Dn-GSym3OcVuv0qYcHR8c,360
|
|
10
|
+
basemode/strategies/completion.py,sha256=7ywtAtylzyvSsrujNPM3WUJUbhW71KxQ07NMTtihYJc,870
|
|
11
|
+
basemode/strategies/few_shot.py,sha256=4qpEyMaXIWJ_8g650WPwLLQCsSqYqYwtiUrJqTOI6w4,2344
|
|
12
|
+
basemode/strategies/fim.py,sha256=EywjHVkQgUooUSs-qTG0pBm5V_SrUfNEmL4IfWIqxjs,1445
|
|
13
|
+
basemode/strategies/prefill.py,sha256=HqhlkQJXb571xY5MMoRmvrLUCPB8wkIAVOD31zlQpC0,1528
|
|
14
|
+
basemode/strategies/system.py,sha256=pXCtqDIAOqbAQeiXb8jhxBqQmDWongzZUoFaHGBaK-c,2587
|
|
15
|
+
basemode/strategies/utils.py,sha256=K7qgf3TN_2JRXclDFKPs5BT3hscdr2tOH_pJOx62e2s,1047
|
|
16
|
+
basemode-0.1.0.dist-info/METADATA,sha256=KeIcJaZTCNfsEmi6jHwTUXde3NRI0wL6lM3TWPOXtmE,315
|
|
17
|
+
basemode-0.1.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
|
|
18
|
+
basemode-0.1.0.dist-info/entry_points.txt,sha256=N9pXaPSXGtGOreQNEuSYXKSh7gK27Dc-3sKmOHCqLgQ,46
|
|
19
|
+
basemode-0.1.0.dist-info/RECORD,,
|