From ea148e5e15e0b37ac98cf49a71470bc15b54f721 Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Sun, 23 Aug 2026 12:14:29 +0800 Subject: [PATCH 1/4] fix(enrichment): circuit-break exhausted quota calls --- app.py | 1 + automem/api/enrichment.py | 1 + automem/api/memory.py | 2 + automem/classification/memory_classifier.py | 9 ++++ automem/config.py | 2 + automem/service_state.py | 55 ++++++++++++++++++++- automem/utils/text.py | 11 ++++- tests/test_enrichment_circuit.py | 34 +++++++++++++ 8 files changed, 112 insertions(+), 3 deletions(-) create mode 100644 tests/test_enrichment_circuit.py diff --git a/app.py b/app.py index b24e1b4e..35e6a96d 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 ba5b2ae9..c8841ed4 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 947df7bb..4f9fd330 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 17f6509e..4869f966 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 5241c621..555e9a9c 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 e5831825..af532bb7 100644 --- a/automem/service_state.py +++ b/automem/service_state.py @@ -4,15 +4,65 @@ 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(): + 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 +156,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 2c4a4c64..ce15fdfa 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 00000000..618cc1cf --- /dev/null +++ b/tests/test_enrichment_circuit.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from datetime import timedelta + +from automem.service_state import EnrichmentCircuit + + +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 From 6566675564e4a65f4d9453ecd9f0f42bbdc96930 Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Sun, 23 Aug 2026 12:30:29 +0800 Subject: [PATCH 2/4] test(enrichment): cover probe recovery and request suppression --- automem/service_state.py | 4 +++ tests/test_enrichment_circuit.py | 57 +++++++++++++++++++++++++++++++- 2 files changed, 60 insertions(+), 1 deletion(-) diff --git a/automem/service_state.py b/automem/service_state.py index af532bb7..31eb63e6 100644 --- a/automem/service_state.py +++ b/automem/service_state.py @@ -41,6 +41,10 @@ def allow_request(self) -> bool: 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 diff --git a/tests/test_enrichment_circuit.py b/tests/test_enrichment_circuit.py index 618cc1cf..fd6690d0 100644 --- a/tests/test_enrichment_circuit.py +++ b/tests/test_enrichment_circuit.py @@ -1,8 +1,11 @@ from __future__ import annotations -from datetime import timedelta +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(): @@ -32,3 +35,55 @@ def test_successful_probe_closes_circuit_and_records_recovery(): 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 + + +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 + assert circuit.record_failure("connection reset") is False + assert circuit.allow_request() is True From 97d67fa1d1ad6a5a992410410850a21b6dc83b1b Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Sun, 23 Aug 2026 19:01:17 +0800 Subject: [PATCH 3/4] test(enrichment): fix probe recovery assertions --- tests/test_enrichment_circuit.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_enrichment_circuit.py b/tests/test_enrichment_circuit.py index fd6690d0..bf299934 100644 --- a/tests/test_enrichment_circuit.py +++ b/tests/test_enrichment_circuit.py @@ -31,6 +31,8 @@ def test_successful_probe_closes_circuit_and_records_recovery(): 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 circuit.record_success() assert circuit.allow_request() is True @@ -85,5 +87,3 @@ def create(*args, **kwargs): summarize_content(content, client, "gpt-4o-mini", 300, circuit) assert len(calls) == 1 - assert circuit.record_failure("connection reset") is False - assert circuit.allow_request() is True From 83df3525aac5df53121a95940786942a06424aa1 Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Sun, 23 Aug 2026 20:29:12 +0800 Subject: [PATCH 4/4] test(enrichment): separate probe success and failure cases --- tests/test_enrichment_circuit.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_enrichment_circuit.py b/tests/test_enrichment_circuit.py index bf299934..04cfbbd9 100644 --- a/tests/test_enrichment_circuit.py +++ b/tests/test_enrichment_circuit.py @@ -31,8 +31,6 @@ def test_successful_probe_closes_circuit_and_records_recovery(): 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 circuit.record_success() assert circuit.allow_request() is True @@ -46,6 +44,8 @@ def test_failed_probe_does_not_permanently_block_future_requests(): 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():