Remove test logic from production code

- Remove test-specific query in vectorstore count_memories function
- Remove test-specific comment about fixtures in redis utils
- Update test-specific comments in llms.py to be production-appropriate
- Remove test-specific async json() handling in client
- Fix client tests to account for soft-filter fallback behavior
- Add proper docstring to build_args method in redis_query

Author: abrookins

agent-memory-client/agent_memory_client/client.py

@@ -885,14 +885,6 @@ async def search_long_term_memory(
             )
             response.raise_for_status()
             data = response.json()
-            # Some tests may stub json() as an async function; handle awaitable
-            try:
-                import inspect
-
-                if inspect.isawaitable(data):
-                    data = await data
-            except Exception:
-                pass
             return MemoryRecordResults(**data)
         except httpx.HTTPStatusError as e:
             self._handle_http_error(e.response)

agent_memory_server/api.py

@@ -664,8 +664,7 @@ async def search_long_term_memory(
 
     logger.debug(f"Long-term search kwargs: {kwargs}")
 
-    # Pass text and filter objects to the search function (no redis needed for vectorstore adapter)
-    # Server-side recency rerank toggle (Redis-only path); defaults to False
+    # Server-side recency rerank toggle
     server_side_recency = (
         payload.server_side_recency
         if payload.server_side_recency is not None
@@ -679,18 +678,13 @@ async def search_long_term_memory(
     raw_results = await long_term_memory.search_long_term_memories(**kwargs)
 
     # Soft-filter fallback: if strict filters yield no results, relax filters and
-    # inject hints into the query text to guide semantic search. For memory_prompt
-    # unit tests, the underlying function is mocked; avoid triggering fallback to
-    # keep call counts stable when optimize_query behavior is being asserted.
+    # inject hints into the query text to guide semantic search.
     try:
         had_any_strict_filters = any(
             key in kwargs and kwargs[key] is not None
             for key in ("topics", "entities", "namespace", "memory_type", "event_date")
         )
-        is_mocked = "unittest.mock" in str(
-            type(long_term_memory.search_long_term_memories)
-        )
-        if raw_results.total == 0 and had_any_strict_filters and not is_mocked:
+        if raw_results.total == 0 and had_any_strict_filters:
             fallback_kwargs = dict(kwargs)
             for key in ("topics", "entities", "namespace", "memory_type", "event_date"):
                 fallback_kwargs.pop(key, None)
@@ -738,6 +732,8 @@ def _vals(f):
         logger.warning(f"Soft-filter fallback failed: {e}")
 
     # Recency-aware re-ranking of results (configurable)
+    # TODO: Why did we need to go this route instead of using recency boost at
+    # the query level?
     try:
         from datetime import UTC, datetime as _dt
 

agent_memory_server/llms.py

@@ -96,7 +96,7 @@ async def create_chat_completion(
 
             choices = [{"message": {"content": content}}]
 
-            # Handle both object and dictionary usage formats for testing
+            # Handle both object and dictionary usage formats from API responses
             input_tokens = output_tokens = 0
             if hasattr(response, "usage"):
                 if isinstance(response.usage, dict):
@@ -180,7 +180,7 @@ async def create_chat_completion(
             )
 
             # Convert to unified format
-            # Handle both object and dictionary usage formats for testing
+            # Handle both object and dictionary usage formats from API responses
             total_tokens = 0
             if hasattr(response, "usage"):
                 if isinstance(response.usage, dict):

agent_memory_server/long_term_memory.py

@@ -933,43 +933,32 @@ async def search_long_term_memories(
     )
 
     # If an optimized query with a strict distance threshold returns no results,
-    # retry once with the original query to preserve recall. Skip this retry when
-    # the adapter is a unittest mock to avoid altering test expectations.
+    # retry once with the original query to preserve recall.
     try:
         if (
             optimized_applied
             and distance_threshold is not None
             and results.total == 0
             and search_query != text
         ):
-            # Detect unittest.mock objects without importing globally
-            is_mock = False
-            try:
-                from unittest.mock import Mock  # type: ignore
-
-                is_mock = isinstance(getattr(adapter, "search_memories", None), Mock)
-            except Exception:
-                is_mock = False
-
-            if not is_mock:
-                results = await adapter.search_memories(
-                    query=text,
-                    session_id=session_id,
-                    user_id=user_id,
-                    namespace=namespace,
-                    created_at=created_at,
-                    last_accessed=last_accessed,
-                    topics=topics,
-                    entities=entities,
-                    memory_type=memory_type,
-                    event_date=event_date,
-                    memory_hash=memory_hash,
-                    distance_threshold=distance_threshold,
-                    server_side_recency=server_side_recency,
-                    recency_params=recency_params,
-                    limit=limit,
-                    offset=offset,
-                )
+            results = await adapter.search_memories(
+                query=text,
+                session_id=session_id,
+                user_id=user_id,
+                namespace=namespace,
+                created_at=created_at,
+                last_accessed=last_accessed,
+                topics=topics,
+                entities=entities,
+                memory_type=memory_type,
+                event_date=event_date,
+                memory_hash=memory_hash,
+                distance_threshold=distance_threshold,
+                server_side_recency=server_side_recency,
+                recency_params=recency_params,
+                limit=limit,
+                offset=offset,
+            )
     except Exception:
         # Best-effort fallback; return the original results on any error
         pass

agent_memory_server/utils/redis.py

@@ -29,7 +29,7 @@ async def get_redis_conn(url: str = settings.redis_url, **kwargs) -> Redis:
     global _redis_pool
 
     # Always use the existing _redis_pool if it's not None, regardless of the URL parameter
-    # This ensures that the patched _redis_pool from the test fixture is used
+    # This ensures connection reuse and prevents multiple Redis connections
     if _redis_pool is None:
         _redis_pool = Redis.from_url(url, **kwargs)
     return _redis_pool

agent_memory_server/utils/redis_query.py

@@ -89,6 +89,6 @@ def paginate(self, offset: int, limit: int) -> RecencyAggregationQuery:
         self.limit(offset, limit)
         return self
 
-    # Compatibility helper for tests that inspect the built query
     def build_args(self) -> list:
+        """Build the query arguments for Redis search."""
         return super().build_args()

agent_memory_server/vectorstore_adapter.py

@@ -1187,9 +1187,9 @@ async def count_memories(
                     redis_filter = reduce(lambda x, y: x & y, filters)
 
             # Use the same search method as search_memories but for counting
-            # We use the same query that would match the indexed content
+            # We use a generic query to match all indexed content
             search_results = await self.vectorstore.asimilarity_search(
-                query="duplicate",  # Use a query that should match test content
+                query="",  # Empty query to match all content
                 filter=redis_filter,
                 k=10000,  # Large number to get all results
             )

tests/test_client_api.py

@@ -652,9 +652,11 @@ async def test_memory_prompt_with_optimize_query_default_false(
         )
 
         # Verify search was called with optimize_query=False (default)
-        mock_search.assert_called_once()
-        call_kwargs = mock_search.call_args.kwargs
-        assert call_kwargs.get("optimize_query") is False
+        # May be called multiple times due to soft-filter fallback
+        assert mock_search.call_count >= 1
+        # Check that all calls use optimize_query=False
+        for call in mock_search.call_args_list:
+            assert call.kwargs.get("optimize_query") is False
         assert result is not None
 
 
@@ -678,7 +680,9 @@ async def test_memory_prompt_with_optimize_query_false_explicit(
         )
 
         # Verify search was called with optimize_query=False
-        mock_search.assert_called_once()
-        call_kwargs = mock_search.call_args.kwargs
-        assert call_kwargs.get("optimize_query") is False
+        # May be called multiple times due to soft-filter fallback
+        assert mock_search.call_count >= 1
+        # Check that all calls use optimize_query=False
+        for call in mock_search.call_args_list:
+            assert call.kwargs.get("optimize_query") is False
         assert result is not None