Python package

Session intelligence
for small LLMs

Deterministic JSON repair, confidence scoring, progressive tool disclosure, and session memory — built for 2B–4B on-device models that hallucinate tool calls.

164 tests passing · 25/25 benchmark cases · pydantic v2 · zero dependencies beyond pydantic

Small models make big messes

2B–4B parameter models are fast and cheap, but they output broken JSON, hallucinate tool names, repeat the same failed call, and eat your context window. This harness fixes all of it — deterministically, without re-prompting.

# What the model actually outputs:
{'tool': 'speak', 'args': {'text': 'hello', 'speed': '1.5'}}

# What the harness produces:
{"tool": "21labs.speak", "arguments":
  {"text": "hello", "engine": "kokoro", "speed": 1.5}}

What it does

Five modules, one import, zero config.

R

Deterministic Repair

Fixes malformed JSON without re-prompting. Trailing commas, single quotes, unquoted keys, markdown fences, type drift, camelCase/snake_case mismatch. Pipeline: rename → coerce → inject defaults.

repair_tool_call()
C

Confidence Scoring

Multi-signal scoring: schema match (35%), history (25%), completeness (25%), repair cost (15%). Decides whether to trust the local model or escalate to cloud.

score_tool_call()
D

Progressive Disclosure

Small models fail on tool selection with too many options. Ranks tools by query relevance, shows only the top N. Keyword overlap, historical success, category matching.

rank_tools() · build_compact_tool_prompt()
S

Session Intelligence

Tracks success/failure per tool, detects same-tool loops, injects steering hints into the system prompt, monitors context pressure, compacts responses to save tokens.

HarnessState · create_harness_session()

Benchmark results

25 test cases across 4 categories. 100% pass rate. 9.4ms total.

10/10
JSON Repair
5/5
Confidence
5/5
Disclosure
5/5
Session
$ python benchmarks/benchmark.py

--- Benchmark 1: JSON Repair ---
  [PASS] trailing comma          0.4ms
  [PASS] single quotes           0.1ms
  [PASS] missing quotes on keys  0.4ms
  [PASS] markdown fences         0.1ms
  [PASS] leading garbage         0.1ms
  [PASS] string number arg       0.1ms
  [PASS] boolean as string       0.1ms
  [PASS] camelCase keys          0.3ms
  [PASS] comments in JSON        0.1ms
  [PASS] empty arguments         0.1ms

--- Benchmark 2: Confidence Scoring ---
  [PASS] clean call              0.0ms
  [PASS] missing required        0.0ms
  [PASS] unknown tool            0.0ms
  [PASS] heavy repairs           0.0ms
  [PASS] avoided tool            0.0ms

BENCHMARK RESULTS: 25/25 passed (100%)
Total time: 9.4ms

Quick start

Three lines to repair, score, and rank. The harness handles the rest.

01 Create a session with context-aware budgeting
02 Repair malformed JSON from the model output
03 Score confidence and decide: trust or escalate
04 Track session state, detect loops, inject hints
from small_model_harness import (
    create_harness_session,
    compact_tool_response,
)
from small_model_harness.tool_repair import repair_tool_call
from small_model_harness.confidence import score_tool_call
from small_model_harness.tool_disclosure import rank_tools

# 1. Session with auto-budget from context window
harness = create_harness_session(n_ctx=4096)

# 2. Repair malformed model output
raw = \'{"tool": "speak", "args": {"text": "hi"}}\'
args, tool, fixes = repair_tool_call(raw, schemas)

# 3. Score confidence
score = score_tool_call(tool, args, schema, harness)
if score.should_escalate:
    escalate_to_cloud()

# 4. Rank tools for next turn
top = rank_tools("generate speech", schemas, harness)

Install

One dependency. Works everywhere Python runs.

$ pip install small-model-harness