kami
968cbfa973
Found while live-QAing freestyle_planning on a 12B local model: - list_dir tool: recursive, .gitignore-aware listing so weak models stop flooding context with `ls -R` over node_modules/build/dist. Wired into the fileRead toggle + advertised to the planner (architect_freestyle). - ContextClassifier: assistantToolCall turns are STRUCTURED, so the token pruner never shreds the model's own tool-call history — that was causing amnesia loops (re-issuing calls it had already made). - Retire instruction-doc LLMLingua pruning (DOC_SOURCE_TYPES emptied): it fused load-bearing procedural text into unparseable soup. The static block stays small by dropping CLAUDE.md at the loader instead. - AgentInstructionsLoader: load only AGENTS.md, not CLAUDE.md — the latter targets the outer assistant and polluted the agent's stage context. - DefaultSessionReducer: WorkflowFailed flips session status to FAILED (was stuck ACTIVE forever, so clients/approval loops never saw a terminal). - ShellTool: run shell command lines (cd/&&/pipes) via `sh -c`; unrunnable program is recoverable instead of an uncaught IOException killing the stage; malformed argv (non-string/collapsed-array) rejected with guidance. - llmlingua sidecar: cap force_tokens to max_force_token (big docs blew the assert and 500'd, so doc pruning silently failed open). Tests added/updated across all of the above.
73 lines
2.4 KiB
Python
73 lines
2.4 KiB
Python
"""LLMLingua-2 token-pruning sidecar for correx.
|
|
|
|
A tiny HTTP service correx calls (via HttpTokenPruner) to prune low-perplexity tokens from
|
|
freeform prose before it's sent to the local LLM. Kept in Python because LLMLingua-2 is a
|
|
torch/BERT classifier with no JVM equivalent (pipeline §4).
|
|
|
|
Contract:
|
|
POST /prune {"text": str, "protected": [str], "rate": float} -> {"compressed": str}
|
|
rate = fraction of tokens to KEEP (0.55 keeps ~55%, i.e. ~45% compression).
|
|
`protected` substrings are force-kept verbatim (IDs, numbers, paths, code).
|
|
GET /health -> {"status": "ok"}
|
|
|
|
Run:
|
|
pip install -r requirements.txt
|
|
uvicorn server:app --host 127.0.0.1 --port 8199
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from fastapi import FastAPI
|
|
from pydantic import BaseModel
|
|
|
|
app = FastAPI(title="correx-llmlingua")
|
|
|
|
_MODEL = os.environ.get("LLMLINGUA_MODEL", "microsoft/llmlingua-2-xlm-roberta-large-meetingbank")
|
|
_compressor = None
|
|
|
|
|
|
def _get_compressor():
|
|
# Lazy-load so /health works (and the process starts fast) before torch spins up.
|
|
global _compressor
|
|
if _compressor is None:
|
|
from llmlingua import PromptCompressor
|
|
_compressor = PromptCompressor(model_name=_MODEL, use_llmlingua2=True)
|
|
return _compressor
|
|
|
|
|
|
class PruneRequest(BaseModel):
|
|
text: str
|
|
protected: list[str] = []
|
|
rate: float = 0.55
|
|
|
|
|
|
class PruneResponse(BaseModel):
|
|
compressed: str
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "ok"}
|
|
|
|
|
|
@app.post("/prune", response_model=PruneResponse)
|
|
def prune(req: PruneRequest):
|
|
text = req.text.strip()
|
|
if not text:
|
|
return PruneResponse(compressed=req.text)
|
|
compressor = _get_compressor()
|
|
# LLMLingua-2 asserts len(force_tokens) <= max_force_token (default 100). Large static docs
|
|
# (CLAUDE.md/AGENTS.md) yield hundreds of protected spans, which used to blow the assert and
|
|
# 500 -> the doc came back uncompressed. Dedup and keep the longest spans (most load-bearing:
|
|
# full paths, hashes, code fences beat bare numbers) up to the cap.
|
|
cap = getattr(compressor, "max_force_token", 100)
|
|
forced = sorted(set(req.protected), key=len, reverse=True)[:cap] or None
|
|
# rate is fraction to keep; LLMLingua-2 force_tokens keeps the protected spans verbatim.
|
|
result = compressor.compress_prompt(
|
|
text,
|
|
rate=max(0.1, min(1.0, req.rate)),
|
|
force_tokens=forced,
|
|
drop_consecutive=True,
|
|
)
|
|
return PruneResponse(compressed=result["compressed_prompt"])
|