From c79cf43ff800c3ed9d26bc6bfc60d9821f9584a8 Mon Sep 17 00:00:00 2001 From: ARIA Date: Sat, 15 Aug 2026 19:52:25 +0200 Subject: [PATCH] feat: same-turn memory retrieval in provider mode prefetch() now reuses the in-flight background search when it matches the current query, or falls back to a synchronous search (capped at 3s, split across both requests in case-insensitive mode) so memories are injected on the same turn instead of one turn behind. - Track _prefetch_query to scope cached results to their query - Reset _prefetch_result when starting a new background search to prevent stale cross-turn injection after a timed-out join - Guard background thread writes against superseded queries - Add optional per-request timeout override to LocalMem0Client - Update AGENTS.md and README to describe the same-turn flow --- AGENTS.md | 4 ++-- README.md | 11 +++++------ __init__.py | 57 ++++++++++++++++++++++++++++++++++++++++++----------- client.py | 17 ++++++++++------ 4 files changed, 64 insertions(+), 25 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index dcd69a7..8cea255 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -43,8 +43,8 @@ curl http://localhost:8000/health **Memory injection flow:** 1. `queue_prefetch()` — Spawns background thread on user message -2. `prefetch()` — Returns cached result next turn -3. `pre_llm_call` hook — Sync prefetch for immediate injection (fallback context) +2. `prefetch()` — Sync retrieval for the current query (same-turn injection); falls back to cached background result when no query is provided +3. `pre_llm_call` hook — Sync prefetch for immediate injection (plugin mode) 4. `sync_turn()` — Non-blocking server-side fact extraction **Circuit breaker:** 5 consecutive failures → 120s cooldown. diff --git a/README.md b/README.md index 1df98a0..5947c89 100644 --- a/README.md +++ b/README.md @@ -5,7 +5,7 @@ Self-hosted Mem0 memory provider for Hermes-Agent. Provides semantic memory sear ## Features - **Local Mem0 server** — No cloud dependency, full data privacy -- **Async prefetch** — Memory retrieval happens in background (~40ms) +- **Same-turn prefetch** — Memory retrieved synchronously before the LLM call (~40ms) - **Context injection** — Relevant memories injected directly into LLM prompt - **Automatic fact extraction** — Server-side LLM extracts facts from conversations - **Semantic search** — Find memories by meaning, not keywords @@ -138,16 +138,15 @@ hermes gateway restart ### How It Works -1. **User message received** → `queue_prefetch()` spawns background thread -2. **Mem0 search** → Semantic search for relevant memories (~40ms) -3. **Context injection** → Results injected via `pre_llm_call` hook -4. **LLM receives** → User message + memory context (no tool call needed!) +1. **User message received** → `prefetch()` searches Mem0 synchronously (~40ms) +2. **Context injection** → Results injected before the LLM call +3. **LLM receives** → User message + memory context (no tool call needed!) **Example**: ``` User: "Hey, is a new episode out from my favorite anime?" -↓ [Background: mem0.prefetch() searches for "favorite anime"] +↓ [mem0.prefetch() searches for "favorite anime"] LLM receives: """ diff --git a/__init__.py b/__init__.py index af0fdb9..9c6f690 100644 --- a/__init__.py +++ b/__init__.py @@ -227,6 +227,7 @@ class Mem0LocalMemoryProvider(MemoryProvider): ], } self._prefetch_result = "" + self._prefetch_query = "" self._prefetch_lock = threading.Lock() self._prefetch_thread = None self._sync_thread = None @@ -520,20 +521,39 @@ class Mem0LocalMemoryProvider(MemoryProvider): ) def prefetch(self, query: str = "", *, session_id: str = "") -> str: - """Return cached prefetch result from previous turn. + """Return memory context for the current turn. + + Reuses the in-flight background prefetch when it was started for the + same query, otherwise performs a synchronous search (bounded by the + shorter of the configured client timeout and 3s) so memories are + injected on the same turn. Falls back to the cached background + result when no query is available. Args: - query: Deprecated, kept for API compatibility. + query: Current user message, used for synchronous retrieval. session_id: Session identifier. """ - if self._prefetch_thread and self._prefetch_thread.is_alive(): - self._prefetch_thread.join(timeout=3.0) - with self._prefetch_lock: - result = self._prefetch_result - self._prefetch_result = "" + if query: + with self._prefetch_lock: + query_matches = self._prefetch_query == query + if query_matches and self._prefetch_thread and self._prefetch_thread.is_alive(): + # Background search for this exact query is in flight — wait for it + self._prefetch_thread.join(timeout=3.0) + if query_matches: + with self._prefetch_lock: + result = self._prefetch_result + self._prefetch_result = "" + else: + # No cached result for this query — search synchronously + result = self.queue_prefetch_and_get(query) + else: + if self._prefetch_thread and self._prefetch_thread.is_alive(): + self._prefetch_thread.join(timeout=3.0) + with self._prefetch_lock: + result = self._prefetch_result + self._prefetch_result = "" if not result: return "" - # Check if it's an error message if result.startswith("ERROR:"): return f"\n{result[6:]}\n" return f"\n{result}\n" @@ -570,11 +590,17 @@ class Mem0LocalMemoryProvider(MemoryProvider): return "" try: client = self._get_client() + # Cap total blocking on the LLM hot path at ~3s + prefetch_timeout = min(client.timeout, 3.0) + if self._case_insensitive: + # case-insensitive search runs two sequential requests + prefetch_timeout /= 2 results = client.search( query=query, user_id=self._user_id, limit=self._prefetch_limit, case_insensitive=self._case_insensitive, + timeout=prefetch_timeout, ) # Filter by score threshold threshold = self._prefetch_score_threshold / 100.0 @@ -602,14 +628,20 @@ class Mem0LocalMemoryProvider(MemoryProvider): """ if self._is_breaker_open(): with self._prefetch_lock: + self._prefetch_query = query self._prefetch_result = "ERROR:Memory service temporarily unavailable. Please try again later." return if self._is_trivial_prompt(query): with self._prefetch_lock: + self._prefetch_query = query self._prefetch_result = "" return + with self._prefetch_lock: + self._prefetch_query = query + self._prefetch_result = "" + def _run(): try: client = self._get_client() @@ -625,16 +657,19 @@ class Mem0LocalMemoryProvider(MemoryProvider): if filtered: formatted = self._format_search_results(filtered, categorize=self._categorize_enabled) with self._prefetch_lock: - self._prefetch_result = formatted + if self._prefetch_query == query: + self._prefetch_result = formatted else: with self._prefetch_lock: - self._prefetch_result = "" + if self._prefetch_query == query: + self._prefetch_result = "" self._record_success() except Exception as e: self._record_failure() logger.debug("Mem0 prefetch failed: %s", e) with self._prefetch_lock: - self._prefetch_result = "ERROR:Memory service temporarily unavailable. Please try again later." + if self._prefetch_query == query: + self._prefetch_result = "ERROR:Memory service temporarily unavailable. Please try again later." self._prefetch_thread = threading.Thread( target=_run, daemon=True, name="mem0-local-prefetch" diff --git a/client.py b/client.py index 38fdf59..11d0726 100644 --- a/client.py +++ b/client.py @@ -39,17 +39,19 @@ class LocalMem0Client: endpoint: str, json: Optional[Dict] = None, params: Optional[Dict] = None, + timeout: Optional[float] = None, ) -> Dict: """Make HTTP request with error handling.""" url = f"{self.base_url}{endpoint}" + effective_timeout = timeout if timeout is not None else self.timeout try: resp = self.session.request( - method, url, json=json, params=params, timeout=self.timeout + method, url, json=json, params=params, timeout=effective_timeout ) resp.raise_for_status() return resp.json() except requests.exceptions.Timeout: - logger.error("Mem0 request timed out after %ss", self.timeout) + logger.error("Mem0 request timed out after %ss", effective_timeout) raise except requests.exceptions.ConnectionError as e: logger.error("Failed to connect to Mem0 server at %s: %s", self.base_url, e) @@ -66,6 +68,7 @@ class LocalMem0Client: user_id: Optional[str] = None, limit: int = 5, case_insensitive: bool = False, + timeout: Optional[float] = None, ) -> List[Dict]: """Search memories by semantic similarity. @@ -78,18 +81,19 @@ class LocalMem0Client: user_id: User identifier limit: Max results case_insensitive: If True, search with both original and lowercase query + timeout: Optional per-request timeout override (seconds) """ if not case_insensitive: payload = {"query": query, "limit": limit} if user_id: payload["user_id"] = user_id - result = self._request("POST", "/search", json=payload) + result = self._request("POST", "/search", json=payload, timeout=timeout) return result.get("results", []) # Case-insensitive mode: search with both original and lowercase # Fetch 2x limit to ensure we get top N after merging - results_original = self._search_with_query(query, user_id, limit * 2) - results_lower = self._search_with_query(query.lower(), user_id, limit * 2) + results_original = self._search_with_query(query, user_id, limit * 2, timeout) + results_lower = self._search_with_query(query.lower(), user_id, limit * 2, timeout) # Merge and deduplicate, keeping highest score merged = {} @@ -109,12 +113,13 @@ class LocalMem0Client: query: str, user_id: Optional[str] = None, limit: int = 5, + timeout: Optional[float] = None, ) -> List[Dict]: """Internal search helper for case-insensitive mode.""" payload = {"query": query, "limit": limit} if user_id: payload["user_id"] = user_id - result = self._request("POST", "/search", json=payload) + result = self._request("POST", "/search", json=payload, timeout=timeout) return result.get("results", []) def get_all(self, user_id: Optional[str] = None) -> List[Dict]: -- 2.54.0