diff --git a/app.py b/app.py index b24e1b4..35e6a96 100644 --- a/app.py +++ b/app.py @@ -318,6 +318,7 @@ def require_api_token() -> None: classification_model=CLASSIFICATION_MODEL, logger=logger, stats=state.classification_stats, + circuit=state.enrichment_circuit, ) diff --git a/automem/api/enrichment.py b/automem/api/enrichment.py index ba5b2ae..c8841ed 100644 --- a/automem/api/enrichment.py +++ b/automem/api/enrichment.py @@ -30,6 +30,7 @@ def enrichment_status() -> Any: "max_attempts": max_attempts, "stats": state.enrichment_stats.to_dict(), "classification": state.classification_stats.to_dict(), + "circuit": state.enrichment_circuit.to_dict(), } return jsonify(response) diff --git a/automem/api/memory.py b/automem/api/memory.py index 947df7b..4f9fd33 100644 --- a/automem/api/memory.py +++ b/automem/api/memory.py @@ -514,6 +514,7 @@ def store() -> Any: openai_client, CLASSIFICATION_MODEL, MEMORY_SUMMARY_TARGET_LENGTH, + state.enrichment_circuit, ) if summary: original_content = content @@ -1248,6 +1249,7 @@ def store_batch() -> Any: openai_client, CLASSIFICATION_MODEL, MEMORY_SUMMARY_TARGET_LENGTH, + state.enrichment_circuit, ) if summary: logger.info( diff --git a/automem/classification/memory_classifier.py b/automem/classification/memory_classifier.py index 17f6509..4869f96 100644 --- a/automem/classification/memory_classifier.py +++ b/automem/classification/memory_classifier.py @@ -97,6 +97,7 @@ def __init__( classification_model: str, logger: Any, stats: Any = None, + circuit: Any = None, ) -> None: self._normalize_memory_type = normalize_memory_type self._ensure_openai_client = ensure_openai_client @@ -104,6 +105,7 @@ def __init__( self._classification_model = classification_model self._logger = logger self._stats = stats + self._circuit = circuit def classify(self, content: str, *, use_llm: bool = True) -> tuple[str, float]: """Classify memory type and return confidence score.""" @@ -134,12 +136,17 @@ def classify(self, content: str, *, use_llm: bool = True) -> tuple[str, float]: except Exception as exc: self._logger.exception("LLM classification failed, using fallback") llm_error = str(exc) + if self._circuit is not None: + self._circuit.record_failure(llm_error) if self._stats is not None: self._stats.record_fallback(llm_error) return "Memory", 0.3 def _classify_with_llm(self, content: str) -> Optional[tuple[str, float]]: + if self._circuit is not None and not self._circuit.allow_request(): + self._logger.info("Skipping LLM classification while enrichment circuit is open") + return None client = self._get_openai_client() if client is None: self._ensure_openai_client() @@ -166,6 +173,8 @@ def _classify_with_llm(self, content: str) -> Optional[tuple[str, float]]: response_format={"type": "json_object"}, **extra_params, ) + if self._circuit is not None: + self._circuit.record_success() raw_content = response.choices[0].message.content if not raw_content: diff --git a/automem/config.py b/automem/config.py index 5241c62..555e9a9 100644 --- a/automem/config.py +++ b/automem/config.py @@ -178,6 +178,8 @@ } # Target length for summarized content MEMORY_SUMMARY_TARGET_LENGTH = int(os.getenv("MEMORY_SUMMARY_TARGET_LENGTH", "300")) +# Cooldown after a definitive LLM quota exhaustion response. +ENRICHMENT_CIRCUIT_COOLDOWN_SECONDS = float(os.getenv("ENRICHMENT_CIRCUIT_COOLDOWN_SECONDS", "300")) # Memory types for classification MEMORY_TYPES = {"Decision", "Pattern", "Preference", "Style", "Habit", "Insight", "Context"} diff --git a/automem/service_state.py b/automem/service_state.py index e583182..31eb63e 100644 --- a/automem/service_state.py +++ b/automem/service_state.py @@ -4,15 +4,69 @@ from queue import Queue from threading import Event, Lock, Thread from typing import Any, Dict, Optional, Set +import time from falkordb import FalkorDB from qdrant_client import QdrantClient -from automem.config import VECTOR_SIZE +from automem.config import ENRICHMENT_CIRCUIT_COOLDOWN_SECONDS, VECTOR_SIZE from automem.embedding.provider import EmbeddingProvider from automem.utils.time import utc_now +class EnrichmentCircuit: + """Fail-soft cooldown for definitive LLM quota failures.""" + + def __init__(self, cooldown_seconds: float = 300, clock: Any = time.monotonic) -> None: + self._cooldown_seconds = cooldown_seconds + self._clock = clock + self._opened_until = 0.0 + self._probe_pending = False + self._lock = Lock() + self.circuit_open_skips = 0 + self.recoveries = 0 + + def allow_request(self) -> bool: + with self._lock: + now = self._clock() + if now < self._opened_until: + self.circuit_open_skips += 1 + return False + if self._opened_until: + if self._probe_pending: + self.circuit_open_skips += 1 + return False + self._probe_pending = True + return True + + def record_failure(self, error: str) -> bool: + if "insufficient_quota" not in error.lower() and "quota" not in error.lower(): + with self._lock: + if self._probe_pending: + self._opened_until = 0.0 + self._probe_pending = False + return False + with self._lock: + self._opened_until = self._clock() + self._cooldown_seconds + self._probe_pending = False + return True + + def record_success(self) -> None: + with self._lock: + if self._opened_until: + self.recoveries += 1 + self._opened_until = 0.0 + self._probe_pending = False + + def to_dict(self) -> Dict[str, Any]: + with self._lock: + return { + "circuit_open_skips": self.circuit_open_skips, + "recoveries": self.recoveries, + "open": self._clock() < self._opened_until, + } + + @dataclass class EnrichmentStats: processed_total: int = 0 @@ -106,6 +160,9 @@ class ServiceState: enrichment_thread: Optional[Thread] = None enrichment_stats: EnrichmentStats = field(default_factory=EnrichmentStats) classification_stats: ClassificationStats = field(default_factory=ClassificationStats) + enrichment_circuit: EnrichmentCircuit = field( + default_factory=lambda: EnrichmentCircuit(ENRICHMENT_CIRCUIT_COOLDOWN_SECONDS) + ) enrichment_inflight: Set[str] = field(default_factory=set) enrichment_pending: Set[str] = field(default_factory=set) enrichment_lock: Lock = field(default_factory=Lock) diff --git a/automem/utils/text.py b/automem/utils/text.py index 2c4a4c6..ce15fdf 100644 --- a/automem/utils/text.py +++ b/automem/utils/text.py @@ -122,6 +122,7 @@ def summarize_content( openai_client: Any, model: str, target_length: int = 300, + circuit: Any = None, ) -> Optional[str]: """Summarize content using an LLM to fit within target length. @@ -137,9 +138,11 @@ def summarize_content( if openai_client is None: logger.warning("Cannot summarize: OpenAI client not available") return None - if not content or len(content) <= target_length: return content + if circuit is not None and not circuit.allow_request(): + logger.info("Skipping summarization while enrichment circuit is open") + return None try: system_prompt = SUMMARIZE_SYSTEM_PROMPT.format(target_length=target_length) @@ -163,6 +166,8 @@ def summarize_content( ], **extra_params, ) + if circuit is not None: + circuit.record_success() summary = response.choices[0].message.content.strip() @@ -183,8 +188,10 @@ def summarize_content( ) return None - except Exception: + except Exception as exc: logger.exception("Memory summarization failed") + if circuit is not None: + circuit.record_failure(str(exc)) return None diff --git a/tests/test_enrichment_circuit.py b/tests/test_enrichment_circuit.py new file mode 100644 index 0000000..04cfbbd --- /dev/null +++ b/tests/test_enrichment_circuit.py @@ -0,0 +1,89 @@ +from __future__ import annotations + +from types import SimpleNamespace +import logging + +from automem.classification.memory_classifier import MemoryClassifier +from automem.service_state import EnrichmentCircuit +from automem.utils.text import summarize_content + + +def test_quota_failure_opens_circuit_and_skips_requests(): + circuit = EnrichmentCircuit(cooldown_seconds=60, clock=lambda: 100.0) + + assert circuit.allow_request() is True + assert circuit.record_failure("429 insufficient_quota") is True + assert circuit.allow_request() is False + assert circuit.to_dict()["circuit_open_skips"] == 1 + + +def test_non_quota_failure_does_not_open_circuit(): + circuit = EnrichmentCircuit(cooldown_seconds=60, clock=lambda: 100.0) + + assert circuit.record_failure("connection reset") is False + assert circuit.allow_request() is True + + +def test_successful_probe_closes_circuit_and_records_recovery(): + now = [100.0] + circuit = EnrichmentCircuit(cooldown_seconds=60, clock=lambda: now[0]) + + circuit.record_failure("insufficient_quota") + now[0] += 60 + assert circuit.allow_request() is True + circuit.record_success() + + assert circuit.allow_request() is True + assert circuit.to_dict()["recoveries"] == 1 + + +def test_failed_probe_does_not_permanently_block_future_requests(): + now = [100.0] + circuit = EnrichmentCircuit(cooldown_seconds=60, clock=lambda: now[0]) + + circuit.record_failure("insufficient_quota") + now[0] += 60 + assert circuit.allow_request() is True + assert circuit.record_failure("connection reset") is False + assert circuit.allow_request() is True + + +def test_classifier_makes_no_second_request_while_circuit_is_open(): + calls = [] + + def create(*args, **kwargs): + calls.append(1) + raise RuntimeError("insufficient_quota") + + client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=create))) + circuit = EnrichmentCircuit(cooldown_seconds=60, clock=lambda: 100.0) + classifier = MemoryClassifier( + normalize_memory_type=lambda raw: (raw, False), + ensure_openai_client=lambda: None, + get_openai_client=lambda: client, + classification_model="gpt-4o-mini", + logger=logging.getLogger(__name__), + circuit=circuit, + ) + + classifier.classify("qwxz flibber jabberwock snorkelblatt") + classifier.classify("qwxz flibber jabberwock snorkelblatt") + + assert len(calls) == 1 + + +def test_summarizer_makes_no_second_request_while_circuit_is_open(): + calls = [] + + def create(*args, **kwargs): + calls.append(1) + raise RuntimeError("insufficient_quota") + + client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=create))) + circuit = EnrichmentCircuit(cooldown_seconds=60, clock=lambda: 100.0) + content = "x" * 600 + + summarize_content(content, client, "gpt-4o-mini", 300, circuit) + summarize_content(content, client, "gpt-4o-mini", 300, circuit) + + assert len(calls) == 1