feat(client): add timeout parameter to MemFuse clients for improved request handling

savourylie · savourylie · commit 5aab9ef63e7e · 2025-08-19T19:37:39.000+08:00
- Introduced a `timeout` parameter in both `AsyncMemFuse` and `MemFuse` classes to allow customization of request timeouts, defaulting to 10 seconds.
- Updated relevant methods to utilize the new timeout setting for HTTP requests, enhancing reliability during long-running operations.
- Adjusted tests to reflect the new timeout parameter, ensuring consistent behavior across different scenarios.
diff --git a/benchmarks/utils.py b/benchmarks/utils.py
@@ -382,15 +382,36 @@ def convert_messages_for_memfuse(messages: List[Dict[str, Any]], dataset_type: s
         # LoCoMo messages already have speaker names embedded in content (e.g., "[CAROLINE]: text")
         memfuse_messages = []
         for msg in messages:
+            content = msg.get("content", "")
+            role = msg.get("role", "user")
+            
+            # Skip messages with empty content or role
+            if not content or not content.strip() or not role or not role.strip():
+                continue
+                
             memfuse_msg = {
-                "role": msg.get("role", "user"),
-                "content": msg.get("content", "")
+                "role": role,
+                "content": content
             }
             memfuse_messages.append(memfuse_msg)
         return memfuse_messages
     else:
-        # MSC and LME use standard format already
-        return messages
+        # MSC and LME need to filter out extra fields and empty messages to match MemFuse format
+        memfuse_messages = []
+        for msg in messages:
+            content = msg.get("content", "")
+            role = msg.get("role", "user")
+            
+            # Skip messages with empty content or role
+            if not content or not content.strip() or not role or not role.strip():
+                continue
+                
+            memfuse_msg = {
+                "role": role,
+                "content": content
+            }
+            memfuse_messages.append(memfuse_msg)
+        return memfuse_messages
 
 
 async def load_dataset_to_memfuse(
diff --git a/src/memfuse/client.py b/src/memfuse/client.py
@@ -23,15 +23,17 @@ class AsyncMemFuse:
     # Class variable to track all instances
     _instances = set()
 
-    def __init__(self, base_url: str = "http://localhost:8000", api_key: Optional[str] = None):
+    def __init__(self, base_url: str = "http://localhost:8000", api_key: Optional[str] = None, timeout: int = 10):
         """Initialize the MemFuse client.
 
         Args:
             base_url: URL of the MemFuse server API
             api_key: API key for authentication (optional for local usage)
+            timeout: Request timeout in seconds (default: 10)
         """
         self.base_url = base_url.rstrip("/")
         self.api_key = api_key or os.environ.get("MEMFUSE_API_KEY")
+        self.timeout = timeout
         self.session = None
 
         # Initialize ASYNC API clients using the classes from .api
@@ -69,7 +71,7 @@ async def _check_server_health(self) -> bool:
         await self._ensure_session()
         try:
             url = f"{self.base_url}/api/v1/health"
-            async with self.session.get(url, timeout=10) as response:
+            async with self.session.get(url, timeout=self.timeout) as response:
                 if response.status == 200:
                     return True
                 return False
@@ -108,7 +110,7 @@ async def _request(
 
             url = f"{self.base_url}{endpoint}"
 
-            async with getattr(self.session, method.lower())(url, json=data) as response:
+            async with getattr(self.session, method.lower())(url, json=data, timeout=self.timeout) as response:
                 response_data = await response.json()
                 if response.status >= 400:
                     error_message = response_data.get("message", "Unknown error")
@@ -257,15 +259,17 @@ async def _thread_safe_coro_runner(self, coro):
 class MemFuse:
     """Synchronous MemFuse client for communicating with the MemFuse server."""
 
-    def __init__(self, base_url: str = "http://localhost:8000", api_key: Optional[str] = None):
+    def __init__(self, base_url: str = "http://localhost:8000", api_key: Optional[str] = None, timeout: int = 10):
         """Initialize the synchronous MemFuse client.
 
         Args:
             base_url: URL of the MemFuse server API
             api_key: API key for authentication (optional for local usage)
+            timeout: Request timeout in seconds (default: 10)
         """
         self.base_url = base_url.rstrip("/")
         self.api_key = api_key or os.environ.get("MEMFUSE_API_KEY")
+        self.timeout = timeout
         self.sync_session = None  # requests session for sync requests
 
         # Initialize API clients using the classes from .api
@@ -300,7 +304,7 @@ def _check_server_health_sync(self) -> bool:
         self._ensure_sync_session()
         try:
             url = f"{self.base_url}/api/v1/health"
-            response = self.sync_session.get(url, timeout=10)
+            response = self.sync_session.get(url, timeout=self.timeout)
             return response.status_code == 200
         except Exception:
             return False
@@ -336,7 +340,7 @@ def _request_sync(
 
             url = f"{self.base_url}{endpoint}"
 
-            response = getattr(self.sync_session, method.lower())(url, json=data, timeout=10)
+            response = getattr(self.sync_session, method.lower())(url, json=data, timeout=self.timeout)
             response_data = response.json()
             if response.status_code >= 400:
                 error_message = response_data.get("message", "Unknown error")
diff --git a/tests/benchmarks/test_msc_accuracy.py b/tests/benchmarks/test_msc_accuracy.py
@@ -54,8 +54,11 @@
 
 # Suppress Pydantic serialization warnings from LiteLLM
 warnings.filterwarnings("ignore", message=".*Pydantic serializer warnings.*")
+warnings.filterwarnings("ignore", message=".*Expected.*fields but got.*")
+warnings.filterwarnings("ignore", message=".*PydanticSerializationUnexpectedValue.*")
 warnings.filterwarnings("ignore", category=UserWarning, module=".*pydantic.*")
 warnings.filterwarnings("ignore", category=UserWarning, module=".*litellm.*")
+warnings.filterwarnings("ignore", category=UserWarning, module=".*litellm_core_utils.*")
 
 # Test configuration
 DEFAULT_NUM_QUESTIONS = 20
@@ -103,7 +106,6 @@ async def load_msc_to_memfuse(dataset, logger):
                 else:
                     logger.error(f"Failed to add messages to session '{session_id}' for Q{question_number}: {add_result}")
                     continue
-                await memory_instance.close()
             successfully_loaded_count += 1
             logger.info(f"--- Successfully loaded Question {question_number}/{len(dataset)} ---")
         except ConnectionError as e:
@@ -120,7 +122,6 @@ async def run_msc_benchmark_with_results(dataset, logger):
     """Query MemFuse memory with MSC questions and return structured results."""
     results = MSCBenchmarkResults()
     memfuse_client = AsyncMemFuse()
-    all_created_mem_instances = []
     total_start_time = time.perf_counter()
     for i, data_sample in enumerate(dataset):
         question_number = i + 1
@@ -143,7 +144,6 @@ async def run_msc_benchmark_with_results(dataset, logger):
                 user=user_name_for_test,
                 agent=agent_name_for_test,
             )
-            all_created_mem_instances.append(query_memory_instance)
             start_time = time.perf_counter()
             memory_response = await query_memory_instance.query(query=question_text, top_k=3)
             end_time = time.perf_counter()
@@ -207,12 +207,7 @@ async def run_msc_benchmark_with_results(dataset, logger):
     except Exception as e:
         logger.error(f"A critical unexpected error occurred: {e}", exc_info=True)
     finally:
-        logger.info("Closing AsyncMemFuse client and query session instances...")
-        for mem_instance in all_created_mem_instances:
-            try:
-                await mem_instance.close()
-            except Exception as e_close:
-                logger.error(f"Error closing memory instance for session {mem_instance.session if mem_instance else 'N/A'}: {e_close}")
+        logger.info("Closing AsyncMemFuse client...")
         if memfuse_client:
             await memfuse_client.close()
         total_end_time = time.perf_counter()
diff --git a/tests/e2e/test_e2e_async_memory_followup.py b/tests/e2e/test_e2e_async_memory_followup.py
@@ -56,7 +56,7 @@ async def test_async_memory_followup_includes_mars_reference():
     # Arrange – create AsyncMemFuse session & AsyncOpenAI client with memory
     # ---------------------------------------------------------------------
     try:
-        memfuse = AsyncMemFuse(base_url=memfuse_base_url)
+        memfuse = AsyncMemFuse(base_url=memfuse_base_url, timeout=30)
     except Exception as exc:
         pytest.skip(f"Cannot connect to MemFuse server at {memfuse_base_url}: {exc}")
 
@@ -174,7 +174,7 @@ async def test_async_memory_with_context_manager():
     # Test using async context manager (recommended approach)
     # ---------------------------------------------------------------------
     try:
-        async with AsyncMemFuse(base_url=memfuse_base_url) as memfuse:
+        async with AsyncMemFuse(base_url=memfuse_base_url, timeout=30) as memfuse:
             memory = await memfuse.init(user="e2e_async_context_test_user", session="context_test")
             
             client = AsyncOpenAI(
diff --git a/tests/e2e/test_e2e_memory_followup.py b/tests/e2e/test_e2e_memory_followup.py
@@ -54,7 +54,7 @@ def test_memory_followup_includes_mars_reference():
     # Arrange – create MemFuse session & OpenAI client with memory attached
     # ---------------------------------------------------------------------
     try:
-        memfuse = MemFuse(base_url=memfuse_base_url)
+        memfuse = MemFuse(base_url=memfuse_base_url, timeout=30)
     except Exception as exc:
         pytest.skip(f"Cannot connect to MemFuse server at {memfuse_base_url}: {exc}")
 
diff --git a/tests/e2e/test_e2e_multi_turn_memory.py b/tests/e2e/test_e2e_multi_turn_memory.py
@@ -58,8 +58,8 @@ def test_multi_turn_conversation_memory():
     openai_api_key = os.getenv("OPENAI_API_KEY")
     
     try:
-        # Initialize MemFuse client
-        memfuse = MemFuse(base_url=memfuse_base_url)
+        # Initialize MemFuse client with longer timeout for multi-turn operations
+        memfuse = MemFuse(base_url=memfuse_base_url, timeout=30)
         
         # Create a unique session for this test
         test_session = f"test_multi_turn_{int(time.time())}"
