From 99f99c2fbfffe1d2b9dc1568fe6a00135ad8f596 Mon Sep 17 00:00:00 2001 From: Guillaume Simonneau Date: Tue, 18 Aug 2026 14:17:54 +0200 Subject: [PATCH] RLM fallback --- .env.example | 21 ++ CHANGELOG.md | 11 + action.yml | 33 ++ codespy.yaml | 22 ++ docs/architecture.md | 23 +- docs/configuration.md | 4 + src/codespy/agents/context_safe.py | 68 +++- .../hippocampus/modules/cartographer.py | 6 +- .../memory/hippocampus/modules/distiller.py | 6 +- .../agents/reviewer/modules/auditor.py | 1 + .../agents/reviewer/modules/code_reviewer.py | 1 + .../agents/reviewer/modules/doc_reviewer.py | 1 + .../agents/reviewer/modules/scope_resolver.py | 1 + .../agents/reviewer/modules/summarizer.py | 1 + .../reviewer/modules/supply_chain_auditor.py | 1 + src/codespy/config.py | 26 +- src/codespy/config_dspy.py | 48 +++ tests/test_context_safe.py | 303 ++++++++++++++++++ 18 files changed, 548 insertions(+), 29 deletions(-) create mode 100644 tests/test_context_safe.py diff --git a/.env.example b/.env.example index 3b10e3d..52a98b3 100644 --- a/.env.example +++ b/.env.example @@ -155,6 +155,27 @@ DEFAULT_MODEL=anthropic/claude-opus-4-6 # Minimum confidence threshold for reported issues (0.0-1.0) # MIN_CONFIDENCE=0.81 +# ============================================================================= +# RLM Fallback (Context Rot Prevention) +# ============================================================================= +# Proactive fallback to RLM when input tokens exceed a ratio of the model's +# max_input_tokens. Prevents quality degradation ("context rot") that occurs +# well before the hard context window limit. The hard overflow safety net +# remains active regardless of these settings. +# +# Thresholds are per module type: +# react: Multi-step tool-using agents (code_review, scope, supply_chain) +# chain_of_thought: Single-pass reasoning (doc, summary, audit) +# predict: Basic completion (not currently used) +# +# Minimum floor: proactive threshold never fires below 8192 input tokens. + +# RLM_FALLBACK_ENABLED=true +# RLM_FALLBACK_REACT_THRESHOLD=0.30 +# RLM_FALLBACK_CHAIN_OF_THOUGHT_THRESHOLD=0.40 +# RLM_FALLBACK_PREDICT_THRESHOLD=0.50 + + # ============================================================================= # Memory (Hippocampus) # ============================================================================= diff --git a/CHANGELOG.md b/CHANGELOG.md index 8960943..0d6bc12 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,16 @@ # Changelog +## [1.0.5] - 2026-08-18 + +### Added +- Proactive RLM fallback with configurable context rot thresholds per module type (ReAct: 0.30, ChainOfThought: 0.40, Predict: 0.50) +- New `rlm_fallback` config section with `enabled`, `react_threshold`, `chain_of_thought_threshold`, `predict_threshold` +- GitHub Action inputs: `rlm-fallback-enabled`, `rlm-fallback-react-threshold`, `rlm-fallback-chain-of-thought-threshold`, `rlm-fallback-predict-threshold` + +### Changed +- `ContextSafe` now checks proactive context rot threshold before checking hard overflow (existing overflow detection retained as safety net) +- `ContextSafe._would_overflow` renamed to `_should_use_rlm` with expanded three-layer logic + ## [1.0.4] - 2026-08-18 ### Fixed diff --git a/action.yml b/action.yml index 79b3101..8d02ec4 100644 --- a/action.yml +++ b/action.yml @@ -99,6 +99,27 @@ inputs: required: false default: '30' + # RLM fallback (context rot prevention) + rlm-fallback-enabled: + description: 'Enable proactive RLM fallback before context window overflow' + required: false + default: 'true' + + rlm-fallback-react-threshold: + description: 'Context ratio for ReAct modules (0.0-1.0). RLM used when input exceeds this ratio of max_input_tokens.' + required: false + default: '0.30' + + rlm-fallback-chain-of-thought-threshold: + description: 'Context ratio for ChainOfThought modules (0.0-1.0). RLM used when input exceeds this ratio.' + required: false + default: '0.40' + + rlm-fallback-predict-threshold: + description: 'Context ratio for Predict modules (0.0-1.0). RLM used when input exceeds this ratio.' + required: false + default: '0.50' + min-confidence: description: 'Minimum confidence threshold (0.0-1.0) for reported issues' required: false @@ -390,6 +411,12 @@ runs: DEFAULT_MAX_LLM_CALLS: ${{ inputs.default-max-llm-calls }} MIN_CONFIDENCE: ${{ inputs.min-confidence }} + # RLM fallback + RLM_FALLBACK_ENABLED: ${{ inputs.rlm-fallback-enabled }} + RLM_FALLBACK_REACT_THRESHOLD: ${{ inputs.rlm-fallback-react-threshold }} + RLM_FALLBACK_CHAIN_OF_THOUGHT_THRESHOLD: ${{ inputs.rlm-fallback-chain-of-thought-threshold }} + RLM_FALLBACK_PREDICT_THRESHOLD: ${{ inputs.rlm-fallback-predict-threshold }} + # Scope identification signature SCOPE_ENABLED: ${{ inputs.scope-enabled }} SCOPE_MODEL: ${{ inputs.scope-model }} @@ -476,6 +503,12 @@ runs: [ -n "$DEFAULT_MAX_LLM_CALLS" ] && DOCKER_ARGS="$DOCKER_ARGS -e DEFAULT_MAX_LLM_CALLS" [ -n "$MIN_CONFIDENCE" ] && DOCKER_ARGS="$DOCKER_ARGS -e MIN_CONFIDENCE" + # RLM fallback + [ -n "$RLM_FALLBACK_ENABLED" ] && DOCKER_ARGS="$DOCKER_ARGS -e RLM_FALLBACK_ENABLED" + [ -n "$RLM_FALLBACK_REACT_THRESHOLD" ] && DOCKER_ARGS="$DOCKER_ARGS -e RLM_FALLBACK_REACT_THRESHOLD" + [ -n "$RLM_FALLBACK_CHAIN_OF_THOUGHT_THRESHOLD" ] && DOCKER_ARGS="$DOCKER_ARGS -e RLM_FALLBACK_CHAIN_OF_THOUGHT_THRESHOLD" + [ -n "$RLM_FALLBACK_PREDICT_THRESHOLD" ] && DOCKER_ARGS="$DOCKER_ARGS -e RLM_FALLBACK_PREDICT_THRESHOLD" + # Scope identification [ -n "$SCOPE_ENABLED" ] && DOCKER_ARGS="$DOCKER_ARGS -e SCOPE_ENABLED" [ -n "$SCOPE_MODEL" ] && DOCKER_ARGS="$DOCKER_ARGS -e SCOPE_MODEL" diff --git a/codespy.yaml b/codespy.yaml index 73106be..e2f6954 100644 --- a/codespy.yaml +++ b/codespy.yaml @@ -144,6 +144,28 @@ memory: max_tokens: null # MEMORY_CARTOGRAPHER_MAX_TOKENS +# ============================================================================ +# RLM FALLBACK (CONTEXT ROT PREVENTION) +# ============================================================================ +# Proactive fallback to RLM when input tokens exceed a ratio of the model's +# max_input_tokens. Prevents quality degradation that occurs well before the +# hard context window limit. The hard overflow safety net remains active +# regardless of these settings. +# +# Thresholds are per DSPy module type: +# - react: Multi-step tool-using agents (code_review, scope, supply_chain) +# - chain_of_thought: Single-pass reasoning (doc, summary, audit) +# - predict: Basic completion (not currently used) +# +# Minimum floor: proactive threshold never fires below 8192 input tokens. + +rlm_fallback: + enabled: true # RLM_FALLBACK_ENABLED + react_threshold: 0.30 # RLM_FALLBACK_REACT_THRESHOLD + chain_of_thought_threshold: 0.40 # RLM_FALLBACK_CHAIN_OF_THOUGHT_THRESHOLD + predict_threshold: 0.50 # RLM_FALLBACK_PREDICT_THRESHOLD + + # ============================================================================ # SIGNATURES # ============================================================================ diff --git a/docs/architecture.md b/docs/architecture.md index 7bccbf7..514473f 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -90,21 +90,26 @@ See [Configuration](configuration.md) for per-signature settings. ### Context Window Overflow Resilience -All signature modules are wrapped in `ContextSafe`, which provides transparent -fallback to `dspy.RLM` (Recursive Language Model) when the input exceeds the -model's context window. - -- **Pre-flight**: For models in litellm's database, estimates input tokens - before calling the module. If overflow is predicted, skips directly to RLM. -- **Try/catch**: For models not in litellm's database, catches the provider - error and retries with RLM automatically. +All signature modules are wrapped in `ContextSafe`, which provides three-layer +defense against context-related quality degradation: + +1. **Proactive threshold (context rot)**: Estimates input tokens and compares + against a configurable ratio of `max_input_tokens`. When the threshold is + exceeded (and input >= 8192 tokens), skips the inner module and uses RLM + directly — preventing quality degradation that occurs well before the hard + context limit. Thresholds differ by module type: ReAct (0.30), + ChainOfThought (0.40), Predict (0.50). +2. **Pre-flight overflow**: For models in litellm's database, detects when + input + output budget + safety margin would exceed the hard context limit. +3. **Try/catch**: For models not in litellm's database, catches the provider + error and retries with RLM automatically. RLM puts inputs into a sandboxed Python interpreter rather than the LLM prompt. The model writes code to access and process variables (e.g., chunking large patches via `llm_query()`), eliminating context window limits entirely. The composition order is: `Hippocampus(ContextSafe(Module))` — Hippocampus -handles memory, ContextSafe handles overflow, each with single responsibility. +handles memory, ContextSafe handles context quality, each with single responsibility. ## Hippocampus Memory diff --git a/docs/configuration.md b/docs/configuration.md index b88f9fc..09d3e3a 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -112,6 +112,10 @@ AUTO_DISCOVER_GEMINI=false | Temperature | `DEFAULT_TEMPERATURE` | `0.2` | Default temperature for LLM calls | | Max iterations | `DEFAULT_MAX_ITERS` | `10` | Maximum ReAct iterations for tool-using agents | | Prompt caching | `ENABLE_PROMPT_CACHING` | `true` | Provider-side prompt caching (Anthropic, OpenAI, Bedrock) | +| RLM fallback | `RLM_FALLBACK_ENABLED` | `true` | Proactive RLM fallback for context rot prevention | +| RLM react threshold | `RLM_FALLBACK_REACT_THRESHOLD` | `0.30` | Context ratio triggering RLM for ReAct modules | +| RLM CoT threshold | `RLM_FALLBACK_CHAIN_OF_THOUGHT_THRESHOLD` | `0.40` | Context ratio triggering RLM for ChainOfThought modules | +| RLM predict threshold | `RLM_FALLBACK_PREDICT_THRESHOLD` | `0.50` | Context ratio triggering RLM for Predict modules | ## Recommended Model Strategy diff --git a/src/codespy/agents/context_safe.py b/src/codespy/agents/context_safe.py index 4b565af..64af235 100644 --- a/src/codespy/agents/context_safe.py +++ b/src/codespy/agents/context_safe.py @@ -15,6 +15,10 @@ # Regex for detecting context window overflow in error messages _RE_CONTEXT_LENGTH = re.compile(r"maximum context length is \d+ tokens", re.IGNORECASE) +# Minimum input tokens before proactive threshold applies. +# Below this, context rot is not a concern regardless of ratio. +_MIN_RLM_THRESHOLD_TOKENS = 8192 + def estimate_context_overflow(model: str, max_tokens: int, input_text: str) -> bool: """Estimate whether input + max_tokens would exceed the model's context window. @@ -67,6 +71,7 @@ def __init__( name: str = "", max_iters: int | None = None, max_llm_calls: int | None = None, + rlm_threshold: float = 1.0, ): super().__init__() self.module = module @@ -75,6 +80,7 @@ def __init__( self._name = name or signature.__name__ self._max_iters = max_iters self._max_llm_calls = max_llm_calls + self._rlm_threshold = rlm_threshold @property def signature(self): @@ -85,12 +91,11 @@ def signature(self, value): self.module.signature = value def forward(self, **kwargs) -> dspy.Prediction: - if self._would_overflow(kwargs): + should_fallback, reason = self._should_use_rlm(kwargs) + if should_fallback: logger.warning( - "ContextSafe[%s]: pre-flight detected context overflow for model=%s; " - "falling back to RLM", - self._name, - getattr(dspy.settings.lm, "model", "unknown"), + "ContextSafe[%s]: %s for model=%s; falling back to RLM", + self._name, reason, getattr(dspy.settings.lm, "model", "unknown"), ) return self._create_rlm_fallback()(**kwargs) @@ -110,12 +115,11 @@ def forward(self, **kwargs) -> dspy.Prediction: async def aforward(self, **kwargs) -> dspy.Prediction: """Async path — used by code_review, scope, supply_chain via Hippocampus.aforward.""" - if self._would_overflow(kwargs): + should_fallback, reason = self._should_use_rlm(kwargs) + if should_fallback: logger.warning( - "ContextSafe[%s]: pre-flight detected context overflow for model=%s; " - "falling back to RLM", - self._name, - getattr(dspy.settings.lm, "model", "unknown"), + "ContextSafe[%s]: %s for model=%s; falling back to RLM", + self._name, reason, getattr(dspy.settings.lm, "model", "unknown"), ) return await self._create_rlm_fallback().acall(**kwargs) @@ -133,20 +137,50 @@ async def aforward(self, **kwargs) -> dspy.Prediction: ) return await self._create_rlm_fallback().acall(**kwargs) - def _would_overflow(self, kwargs: dict) -> bool: - """Pre-flight: estimate overflow from current LM config and input size.""" + def _should_use_rlm(self, kwargs: dict) -> tuple[bool, str]: + """Check if RLM should be used (proactive threshold or hard overflow). + + Returns (should_fallback, reason) for logging. + """ try: lm = dspy.settings.lm if lm is None: - return False + return False, "" model = lm.model max_tokens = lm.kwargs.get("max_tokens") or 0 - if not max_tokens: - return False + + info = litellm.get_model_info(model) + max_input = info.get("max_input_tokens") or 0 + if not max_input: + return False, "" + input_text = "\n".join(str(v) for v in kwargs.values()) - return estimate_context_overflow(model, max_tokens, input_text) + estimated_input = litellm.token_counter(model=model, text=input_text) + + # Layer 1: Proactive threshold (context rot prevention) + # Only applies when input exceeds the minimum floor (8192 tokens) — + # below that, context rot is not a concern. + if (self._rlm_threshold < 1.0 + and estimated_input >= _MIN_RLM_THRESHOLD_TOKENS): + threshold_tokens = int(max_input * self._rlm_threshold) + if estimated_input > threshold_tokens: + return True, ( + f"context rot threshold exceeded " + f"({estimated_input} > {threshold_tokens} = " + f"{self._rlm_threshold:.0%} of {max_input})" + ) + + # Layer 2: Hard overflow check (existing safety net) + SAFETY_MARGIN = 4096 + if max_tokens and (estimated_input + max_tokens + SAFETY_MARGIN) > max_input: + return True, ( + f"context overflow predicted " + f"({estimated_input} + {max_tokens} + {SAFETY_MARGIN} > {max_input})" + ) + + return False, "" except Exception: - return False + return False, "" def _get_current_signature(self): """Get current signature (may include context_memory if Hippocampus modified it).""" diff --git a/src/codespy/agents/memory/hippocampus/modules/cartographer.py b/src/codespy/agents/memory/hippocampus/modules/cartographer.py index a9626bc..518cfee 100644 --- a/src/codespy/agents/memory/hippocampus/modules/cartographer.py +++ b/src/codespy/agents/memory/hippocampus/modules/cartographer.py @@ -142,7 +142,11 @@ class Cartographer(dspy.Module): def __init__(self): super().__init__() - self.predict = ContextSafe(dspy.ChainOfThought(CartographerSig), CartographerSig, name="cartographer") + from codespy.config import settings # lazy to avoid circular import + self.predict = ContextSafe( + dspy.ChainOfThought(CartographerSig), CartographerSig, name="cartographer", + rlm_threshold=settings.get_rlm_threshold("chain_of_thought"), + ) def forward(self, diagnosis, item_tags, cache_candidates, current_map, question, token_budget, current_tokens, max_context_item_tokens): diff --git a/src/codespy/agents/memory/hippocampus/modules/distiller.py b/src/codespy/agents/memory/hippocampus/modules/distiller.py index 440bbea..612f46f 100644 --- a/src/codespy/agents/memory/hippocampus/modules/distiller.py +++ b/src/codespy/agents/memory/hippocampus/modules/distiller.py @@ -151,7 +151,11 @@ class Distiller(dspy.Module): def __init__(self): super().__init__() - self.predict = ContextSafe(dspy.ChainOfThought(DistillerSig), DistillerSig, name="distiller") + from codespy.config import settings # lazy to avoid circular import + self.predict = ContextSafe( + dspy.ChainOfThought(DistillerSig), DistillerSig, name="distiller", + rlm_threshold=settings.get_rlm_threshold("chain_of_thought"), + ) def forward( self, diff --git a/src/codespy/agents/reviewer/modules/auditor.py b/src/codespy/agents/reviewer/modules/auditor.py index 1ebbde1..475d984 100644 --- a/src/codespy/agents/reviewer/modules/auditor.py +++ b/src/codespy/agents/reviewer/modules/auditor.py @@ -151,6 +151,7 @@ def forward( AuditSignature, name="audit", max_llm_calls=self._settings.get_max_llm_calls("audit"), + rlm_threshold=self._settings.get_rlm_threshold("chain_of_thought"), ) logger.info("Running audit...") diff --git a/src/codespy/agents/reviewer/modules/code_reviewer.py b/src/codespy/agents/reviewer/modules/code_reviewer.py index e4c1249..29323e4 100644 --- a/src/codespy/agents/reviewer/modules/code_reviewer.py +++ b/src/codespy/agents/reviewer/modules/code_reviewer.py @@ -240,6 +240,7 @@ async def aforward( name="code_review", max_iters=max_iters, max_llm_calls=self._settings.get_max_llm_calls("code_review"), + rlm_threshold=self._settings.get_rlm_threshold("react"), ) scoped = make_scope_relative(scope) logger.info( diff --git a/src/codespy/agents/reviewer/modules/doc_reviewer.py b/src/codespy/agents/reviewer/modules/doc_reviewer.py index 6ce861f..d0bebd2 100644 --- a/src/codespy/agents/reviewer/modules/doc_reviewer.py +++ b/src/codespy/agents/reviewer/modules/doc_reviewer.py @@ -177,6 +177,7 @@ async def aforward( DocReviewSignature, name="doc", max_llm_calls=self._settings.get_max_llm_calls("doc"), + rlm_threshold=self._settings.get_rlm_threshold("chain_of_thought"), ) logger.info( f" Doc review: scope {scope.subroot} " diff --git a/src/codespy/agents/reviewer/modules/scope_resolver.py b/src/codespy/agents/reviewer/modules/scope_resolver.py index 3c58835..1e826c0 100644 --- a/src/codespy/agents/reviewer/modules/scope_resolver.py +++ b/src/codespy/agents/reviewer/modules/scope_resolver.py @@ -952,6 +952,7 @@ async def _refine_scopes( name="scope", max_iters=max_iters, max_llm_calls=self._settings.get_max_llm_calls("scope"), + rlm_threshold=self._settings.get_rlm_threshold("react"), ) mem: Hippocampus | None = None diff --git a/src/codespy/agents/reviewer/modules/summarizer.py b/src/codespy/agents/reviewer/modules/summarizer.py index 91f6278..0f0548c 100644 --- a/src/codespy/agents/reviewer/modules/summarizer.py +++ b/src/codespy/agents/reviewer/modules/summarizer.py @@ -103,6 +103,7 @@ def forward( PRSummarySignature, name="summary", max_llm_calls=self._settings.get_max_llm_calls("summary"), + rlm_threshold=self._settings.get_rlm_threshold("chain_of_thought"), ) logger.info("Generating PR summary...") diff --git a/src/codespy/agents/reviewer/modules/supply_chain_auditor.py b/src/codespy/agents/reviewer/modules/supply_chain_auditor.py index d735f42..45c7550 100644 --- a/src/codespy/agents/reviewer/modules/supply_chain_auditor.py +++ b/src/codespy/agents/reviewer/modules/supply_chain_auditor.py @@ -325,6 +325,7 @@ async def aforward( name="supply_chain", max_iters=supply_chain_max_iters, max_llm_calls=self._settings.get_max_llm_calls("supply_chain"), + rlm_threshold=self._settings.get_rlm_threshold("react"), ) logger.debug( diff --git a/src/codespy/config.py b/src/codespy/config.py index de7a170..0e329aa 100644 --- a/src/codespy/config.py +++ b/src/codespy/config.py @@ -10,7 +10,9 @@ from codespy.config_dspy import ( ReasoningEffort, + RLMFallbackConfig, SignatureConfig, + apply_rlm_fallback_env_overrides, apply_signature_env_overrides, ) from codespy.config_git import ( @@ -115,6 +117,7 @@ class Settings(BaseSettings): github: GitHubConfig = Field(default_factory=GitHubConfig) gitlab: GitLabConfig = Field(default_factory=GitLabConfig) memory: MemoryConfig = Field(default_factory=MemoryConfig) + rlm_fallback: RLMFallbackConfig = Field(default_factory=RLMFallbackConfig) # Flat signature configs (signature_name -> SignatureConfig) signatures: dict[str, SignatureConfig] = Field(default_factory=dict) @@ -282,8 +285,28 @@ def get_memory_budget(self, signature_name: str) -> "MemoryBudget": max_question_tokens=self.memory.default_max_question_tokens, ) + def get_rlm_threshold(self, module_type: str) -> float: + """Resolve RLM fallback threshold for a module type. + + Args: + module_type: "react" | "chain_of_thought" | "predict" + + Returns: + Threshold ratio (0.0-1.0), or 1.0 if disabled. + """ + if not self.rlm_fallback.enabled: + return 1.0 + return getattr(self.rlm_fallback, f"{module_type}_threshold", 1.0) + def log_signature_configs(self) -> None: """Log all signature and reflection module LLM configurations.""" + logger.info("RLM fallback configuration:") + logger.info( + f" enabled={self.rlm_fallback.enabled}, " + f"react_threshold={self.rlm_fallback.react_threshold}, " + f"chain_of_thought_threshold={self.rlm_fallback.chain_of_thought_threshold}, " + f"predict_threshold={self.rlm_fallback.predict_threshold}" + ) logger.info("Signature configurations:") for sig_name, sig_config in self.signatures.items(): status = "enabled" if sig_config.enabled else "disabled" @@ -316,7 +339,8 @@ def load_yaml_config(cls, values: dict[str, Any]) -> dict[str, Any]: # MEMORY_* env vars target the nested `memory` model, which # pydantic-settings cannot populate on its own (no env_nested_delimiter). yaml_config = apply_memory_env_overrides(yaml_config) - + # RLM_FALLBACK_* env vars target the nested `rlm_fallback` model. + yaml_config = apply_rlm_fallback_env_overrides(yaml_config) # Merge YAML config into values only if not already set (env vars take precedence) for key, val in yaml_config.items(): diff --git a/src/codespy/config_dspy.py b/src/codespy/config_dspy.py index 15fcf28..22dc069 100644 --- a/src/codespy/config_dspy.py +++ b/src/codespy/config_dspy.py @@ -15,6 +15,20 @@ ReasoningEffort = Literal["minimal", "low", "medium", "high"] +class RLMFallbackConfig(BaseModel): + """Proactive RLM fallback thresholds to avoid context rot. + + When input tokens exceed threshold * max_input_tokens, ContextSafe + switches to RLM before quality degrades. Thresholds differ by DSPy + module type because multi-step reasoning (ReAct) is more vulnerable + to context rot than single-pass (ChainOfThought/Predict). + """ + enabled: bool = True + react_threshold: float = Field(default=0.30, ge=0.0, le=1.0) + chain_of_thought_threshold: float = Field(default=0.40, ge=0.0, le=1.0) + predict_threshold: float = Field(default=0.50, ge=0.0, le=1.0) + + class MemorySignatureConfig(BaseModel): """Per-signature Hippocampus memory overrides. @@ -64,6 +78,40 @@ class SignatureConfig(BaseModel): MEMORY_SIGNATURE_SETTINGS = set(MemorySignatureConfig.model_fields) +# Env var name (without RLM_FALLBACK_ prefix) -> RLMFallbackConfig field name. +RLM_FALLBACK_ENV_SETTINGS = { + "ENABLED": "enabled", + "REACT_THRESHOLD": "react_threshold", + "CHAIN_OF_THOUGHT_THRESHOLD": "chain_of_thought_threshold", + "PREDICT_THRESHOLD": "predict_threshold", +} + + +def apply_rlm_fallback_env_overrides(config: dict[str, Any]) -> dict[str, Any]: + """Apply RLM_FALLBACK_* environment variable overrides. + + Maps: RLM_FALLBACK_REACT_THRESHOLD=0.30 -> rlm_fallback.react_threshold + """ + from dotenv import dotenv_values + env_vars = {**dotenv_values(".env"), **os.environ} + + for key, value in env_vars.items(): + if value is None: + continue + key_upper = key.upper() + if not key_upper.startswith("RLM_FALLBACK_"): + continue + remainder = key_upper[len("RLM_FALLBACK_"):] + field = RLM_FALLBACK_ENV_SETTINGS.get(remainder) + if field is None: + continue + rlm_config = config.setdefault("rlm_fallback", {}) + if not isinstance(rlm_config, dict): + continue + rlm_config[field] = convert_env_value(value) + + return config + def convert_env_value(value: str) -> Any: """Convert environment variable string to appropriate Python type.""" diff --git a/tests/test_context_safe.py b/tests/test_context_safe.py new file mode 100644 index 0000000..4009bf1 --- /dev/null +++ b/tests/test_context_safe.py @@ -0,0 +1,303 @@ +"""Tests for ContextSafe proactive RLM fallback thresholds.""" + +import os +from unittest.mock import MagicMock, patch + +import pytest + +from codespy.agents.context_safe import ContextSafe, _MIN_RLM_THRESHOLD_TOKENS +from codespy.config_dspy import ( + RLMFallbackConfig, + apply_rlm_fallback_env_overrides, + convert_env_value, +) + + +class TestRLMFallbackConfig: + """Tests for RLMFallbackConfig model.""" + + def test_default_values(self): + """Test default threshold values.""" + config = RLMFallbackConfig() + assert config.enabled is True + assert config.react_threshold == 0.30 + assert config.chain_of_thought_threshold == 0.40 + assert config.predict_threshold == 0.50 + + def test_threshold_bounds(self): + """Test that thresholds are bounded 0.0-1.0.""" + # Valid values + config = RLMFallbackConfig(react_threshold=0.5) + assert config.react_threshold == 0.5 + + # Invalid: below 0 + with pytest.raises(Exception): + RLMFallbackConfig(react_threshold=-0.1) + + # Invalid: above 1 + with pytest.raises(Exception): + RLMFallbackConfig(react_threshold=1.1) + + +class TestApplyRLMFallbackEnvOverrides: + """Tests for apply_rlm_fallback_env_overrides function.""" + + def test_react_threshold_override(self): + """Test RLM_FALLBACK_REACT_THRESHOLD env var.""" + config = {} + with patch.dict(os.environ, {"RLM_FALLBACK_REACT_THRESHOLD": "0.25"}): + result = apply_rlm_fallback_env_overrides(config) + assert result["rlm_fallback"]["react_threshold"] == "0.25" + + def test_enabled_override(self): + """Test RLM_FALLBACK_ENABLED env var.""" + config = {} + with patch.dict(os.environ, {"RLM_FALLBACK_ENABLED": "false"}): + result = apply_rlm_fallback_env_overrides(config) + assert result["rlm_fallback"]["enabled"] is False + + def test_all_thresholds(self): + """Test all threshold env vars.""" + config = {} + env_vars = { + "RLM_FALLBACK_REACT_THRESHOLD": "0.35", + "RLM_FALLBACK_CHAIN_OF_THOUGHT_THRESHOLD": "0.45", + "RLM_FALLBACK_PREDICT_THRESHOLD": "0.55", + } + with patch.dict(os.environ, env_vars): + result = apply_rlm_fallback_env_overrides(config) + assert result["rlm_fallback"]["react_threshold"] == "0.35" + assert result["rlm_fallback"]["chain_of_thought_threshold"] == "0.45" + assert result["rlm_fallback"]["predict_threshold"] == "0.55" + + def test_unrelated_env_vars_ignored(self): + """Test that unrelated env vars are ignored.""" + config = {} + with patch.dict(os.environ, {"OTHER_VAR": "value"}): + result = apply_rlm_fallback_env_overrides(config) + assert "rlm_fallback" not in result + + def test_existing_config_preserved(self): + """Test that existing config is preserved.""" + config = {"rlm_fallback": {"enabled": False, "react_threshold": 0.20}} + with patch.dict(os.environ, {"RLM_FALLBACK_REACT_THRESHOLD": "0.35"}): + result = apply_rlm_fallback_env_overrides(config) + assert result["rlm_fallback"]["enabled"] is False + assert result["rlm_fallback"]["react_threshold"] == "0.35" + + +class MockSignature: + """Mock signature class with __name__ attribute.""" + __name__ = "MockSignature" + + +class TestGetRLMThreshold: + """Tests for Settings.get_rlm_threshold() method.""" + + def test_returns_correct_threshold_per_type(self): + """Test that correct threshold is returned for each module type.""" + from codespy.config import Settings + + settings = Settings() + settings.rlm_fallback.enabled = True + settings.rlm_fallback.react_threshold = 0.30 + settings.rlm_fallback.chain_of_thought_threshold = 0.40 + settings.rlm_fallback.predict_threshold = 0.50 + + assert settings.get_rlm_threshold("react") == 0.30 + assert settings.get_rlm_threshold("chain_of_thought") == 0.40 + assert settings.get_rlm_threshold("predict") == 0.50 + + def test_returns_1_0_when_disabled(self): + """Test that 1.0 is returned when rlm_fallback is disabled.""" + from codespy.config import Settings + + settings = Settings() + settings.rlm_fallback.enabled = False + + assert settings.get_rlm_threshold("react") == 1.0 + assert settings.get_rlm_threshold("chain_of_thought") == 1.0 + assert settings.get_rlm_threshold("predict") == 1.0 + + def test_returns_1_0_for_unknown_type(self): + """Test that 1.0 is returned for unknown module types.""" + from codespy.config import Settings + + settings = Settings() + settings.rlm_fallback.enabled = True + + assert settings.get_rlm_threshold("unknown_type") == 1.0 + + +class TestContextSafeShouldUseRLM: + """Tests for ContextSafe._should_use_rlm() method.""" + + def _create_mock_lm(self, model="anthropic/claude-opus-4-6", max_tokens=64000): + """Create a mock LM for testing.""" + lm = MagicMock() + lm.model = model + lm.kwargs = {"max_tokens": max_tokens} + return lm + + @patch("codespy.agents.context_safe.litellm") + def test_proactive_threshold_triggered(self, mock_litellm): + """Test that proactive threshold triggers when input exceeds threshold.""" + # Mock model info: 200k max input tokens + mock_litellm.get_model_info.return_value = {"max_input_tokens": 200000} + # Mock token counter to return 70k tokens (above 30% of 200k = 60k, above 8192 floor) + mock_litellm.token_counter.return_value = 70000 + + # Create ContextSafe with 0.30 threshold + module = MagicMock() + cs = ContextSafe(module, MockSignature(), rlm_threshold=0.30) + + # Set up dspy.settings.lm + with patch("codespy.agents.context_safe.dspy.settings") as mock_settings: + mock_settings.lm = self._create_mock_lm() + should_fallback, reason = cs._should_use_rlm({"input": "test"}) + + assert should_fallback is True + assert "context rot threshold exceeded" in reason + assert "70000" in reason + + @patch("codespy.agents.context_safe.litellm") + def test_proactive_threshold_not_triggered_below_threshold(self, mock_litellm): + """Test that proactive threshold does not trigger below threshold.""" + mock_litellm.get_model_info.return_value = {"max_input_tokens": 200000} + # 50k tokens is below 30% of 200k (60k) + mock_litellm.token_counter.return_value = 50000 + + module = MagicMock() + cs = ContextSafe(module, MockSignature(), rlm_threshold=0.30) + + with patch("codespy.agents.context_safe.dspy.settings") as mock_settings: + mock_settings.lm = self._create_mock_lm() + should_fallback, reason = cs._should_use_rlm({"input": "test"}) + + assert should_fallback is False + + @patch("codespy.agents.context_safe.litellm") + def test_proactive_threshold_not_triggered_below_floor(self, mock_litellm): + """Test that proactive threshold does not trigger below 8192 floor.""" + mock_litellm.get_model_info.return_value = {"max_input_tokens": 200000} + # 4000 tokens is above 30% threshold (would be 6000) but below 8192 floor + mock_litellm.token_counter.return_value = 4000 + + module = MagicMock() + cs = ContextSafe(module, MockSignature(), rlm_threshold=0.30) + + with patch("codespy.agents.context_safe.dspy.settings") as mock_settings: + mock_settings.lm = self._create_mock_lm() + should_fallback, reason = cs._should_use_rlm({"input": "test"}) + + assert should_fallback is False + + @patch("codespy.agents.context_safe.litellm") + def test_proactive_threshold_disabled_at_1_0(self, mock_litellm): + """Test that proactive threshold is disabled when rlm_threshold=1.0.""" + mock_litellm.get_model_info.return_value = {"max_input_tokens": 300000} + # With 150k tokens, proactive would trigger at threshold < 0.75 + # but with threshold=1.0, proactive is disabled + # Hard overflow: 150k + 64k + 4k = 218k < 300k, so no overflow either + mock_litellm.token_counter.return_value = 150000 + + module = MagicMock() + cs = ContextSafe(module, MockSignature(), rlm_threshold=1.0) + + with patch("codespy.agents.context_safe.dspy.settings") as mock_settings: + mock_settings.lm = self._create_mock_lm() + should_fallback, reason = cs._should_use_rlm({"input": "test"}) + + # Should not trigger proactive fallback (disabled at 1.0) + # and should not trigger hard overflow (218k < 300k) + assert should_fallback is False + + @patch("codespy.agents.context_safe.litellm") + def test_hard_overflow_still_works(self, mock_litellm): + """Test that hard overflow check still works with rlm_threshold=1.0.""" + mock_litellm.get_model_info.return_value = {"max_input_tokens": 100000} + # Input + max_tokens (64000) + safety_margin (4096) > max_input + mock_litellm.token_counter.return_value = 50000 + + module = MagicMock() + cs = ContextSafe(module, MockSignature(), rlm_threshold=1.0) + + with patch("codespy.agents.context_safe.dspy.settings") as mock_settings: + mock_settings.lm = self._create_mock_lm(max_tokens=64000) + should_fallback, reason = cs._should_use_rlm({"input": "test"}) + + # 50000 + 64000 + 4096 = 118096 > 100000 + assert should_fallback is True + assert "context overflow predicted" in reason + + @patch("codespy.agents.context_safe.litellm") + def test_model_not_in_litellm_db(self, mock_litellm): + """Test behavior when model is not in litellm's database.""" + mock_litellm.get_model_info.return_value = {"max_input_tokens": 0} + + module = MagicMock() + cs = ContextSafe(module, MockSignature(), rlm_threshold=0.30) + + with patch("codespy.agents.context_safe.dspy.settings") as mock_settings: + mock_settings.lm = self._create_mock_lm() + should_fallback, reason = cs._should_use_rlm({"input": "test"}) + + # Should skip proactive check when max_input_tokens is 0 + assert should_fallback is False + + def test_no_lm_configured(self): + """Test behavior when no LM is configured.""" + module = MagicMock() + cs = ContextSafe(module, MockSignature(), rlm_threshold=0.30) + + with patch("codespy.agents.context_safe.dspy.settings") as mock_settings: + mock_settings.lm = None + should_fallback, reason = cs._should_use_rlm({"input": "test"}) + + assert should_fallback is False + + +class TestContextSafeThreeLayerDefense: + """Tests demonstrating the three-layer defense strategy.""" + + @patch("codespy.agents.context_safe.litellm") + def test_layer1_proactive_triggers_first(self, mock_litellm): + """Test that Layer 1 (proactive) triggers before Layer 2 (hard overflow).""" + mock_litellm.get_model_info.return_value = {"max_input_tokens": 200000} + # At 70k tokens with 0.30 threshold: proactive triggers (70k > 60k) + # But hard overflow: 70k + 64k + 4k = 138k < 200k, so no overflow + mock_litellm.token_counter.return_value = 70000 + + module = MagicMock() + cs = ContextSafe(module, MockSignature(), rlm_threshold=0.30) + + with patch("codespy.agents.context_safe.dspy.settings") as mock_settings: + mock_settings.lm = MagicMock() + mock_settings.lm.model = "anthropic/claude-opus-4-6" + mock_settings.lm.kwargs = {"max_tokens": 64000} + should_fallback, reason = cs._should_use_rlm({"input": "test"}) + + assert should_fallback is True + # Should be proactive, not overflow + assert "context rot threshold exceeded" in reason + + @patch("codespy.agents.context_safe.litellm") + def test_layer2_hard_overflow_as_fallback(self, mock_litellm): + """Test that Layer 2 (hard overflow) triggers when Layer 1 doesn't.""" + mock_litellm.get_model_info.return_value = {"max_input_tokens": 80000} + # At 20k tokens with 0.30 threshold: proactive doesn't trigger (20k < 24k threshold) + # But hard overflow: 20k + 64k + 4k = 88k > 80k + mock_litellm.token_counter.return_value = 20000 + + module = MagicMock() + cs = ContextSafe(module, MockSignature(), rlm_threshold=0.30) + + with patch("codespy.agents.context_safe.dspy.settings") as mock_settings: + mock_settings.lm = MagicMock() + mock_settings.lm.model = "anthropic/claude-opus-4-6" + mock_settings.lm.kwargs = {"max_tokens": 64000} + should_fallback, reason = cs._should_use_rlm({"input": "test"}) + + # 20k + 64k + 4k = 88k > 80k + assert should_fallback is True + assert "context overflow predicted" in reason