Address PR feedback for query optimization

- Extract hardcoded optimization prompt to config.py setting
- Add min_optimized_query_length config setting
- Fix AttributeError handling for missing response.choices
- Fix malformed test by properly checking hasattr before deletion
- Extract magic number to class constant

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <[email protected]>

Author: abrookins

agent_memory_server/config.py

@@ -130,6 +130,21 @@ class Settings(BaseSettings):
         0.7  # Fraction of context window that triggers summarization
     )
 
+    # Query optimization settings
+    query_optimization_prompt_template: str = """Transform this natural language query into an optimized version for semantic search. The goal is to make it more effective for finding semantically similar content while preserving the original intent.
+
+Guidelines:
+- Keep the core meaning and intent
+- Use more specific and descriptive terms
+- Remove unnecessary words like "tell me", "I want to know", "can you"
+- Focus on the key concepts and topics
+- Make it concise but comprehensive
+
+Original query: {query}
+
+Optimized query:"""
+    min_optimized_query_length: int = 2
+
     # Other Application settings
     log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = "INFO"
     default_mcp_user_id: str | None = None

agent_memory_server/llms.py

@@ -449,19 +449,10 @@ async def optimize_query_for_vector_search(
     # Use fast model from settings if not specified
     effective_model = model_name or settings.fast_model
 
-    # Create optimization prompt
-    optimization_prompt = f"""Transform this natural language query into an optimized version for semantic search. The goal is to make it more effective for finding semantically similar content while preserving the original intent.
-
-Guidelines:
-- Keep the core meaning and intent
-- Use more specific and descriptive terms
-- Remove unnecessary words like "tell me", "I want to know", "can you"
-- Focus on the key concepts and topics
-- Make it concise but comprehensive
-
-Original query: {query}
-
-Optimized query:"""
+    # Create optimization prompt from config template
+    optimization_prompt = settings.query_optimization_prompt_template.format(
+        query=query
+    )
 
     try:
         client = await get_model_client(effective_model)
@@ -471,7 +462,11 @@ async def optimize_query_for_vector_search(
             prompt=optimization_prompt,
         )
 
-        if response.choices and len(response.choices) > 0:
+        if (
+            hasattr(response, "choices")
+            and response.choices
+            and len(response.choices) > 0
+        ):
             optimized = ""
             if hasattr(response.choices[0], "message"):
                 optimized = response.choices[0].message.content
@@ -484,7 +479,7 @@ async def optimize_query_for_vector_search(
             optimized = optimized.strip()
 
             # Fallback to original if optimization failed
-            if not optimized or len(optimized) < 2:
+            if not optimized or len(optimized) < settings.min_optimized_query_length:
                 logger.warning(f"Query optimization failed for: {query}")
                 return query
 

tests/test_query_optimization_errors.py

@@ -15,6 +15,8 @@
 class TestQueryOptimizationErrorHandling:
     """Test error handling scenarios for query optimization."""
 
+    VERY_LONG_QUERY_REPEAT_COUNT = 1000
+
     @patch("agent_memory_server.llms.get_model_client")
     async def test_optimization_with_network_timeout(self, mock_get_client):
         """Test graceful fallback when model API times out."""
@@ -50,12 +52,19 @@ async def test_optimization_with_malformed_response(self, mock_get_client):
         mock_client = AsyncMock()
         mock_response = MagicMock()
         # Malformed response - no choices attribute
-        del mock_response.choices
+        if hasattr(mock_response, "choices"):
+            del mock_response.choices
         mock_client.create_chat_completion.return_value = mock_response
         mock_get_client.return_value = mock_client
 
         original_query = "Find my user settings"
-        result = await optimize_query_for_vector_search(original_query)
+        # The function should handle AttributeError gracefully and fall back
+        try:
+            result = await optimize_query_for_vector_search(original_query)
+        except AttributeError:
+            pytest.fail(
+                "optimize_query_for_vector_search did not handle missing choices attribute gracefully"
+            )
 
         # Should fall back to original query
         assert result == original_query
@@ -100,7 +109,11 @@ async def test_optimization_with_very_long_query(self, mock_get_client):
         mock_get_client.return_value = mock_client
 
         # Create a very long query (10,000 characters)
-        long_query = "Tell me about " + "preferences " * 1000 + "settings"
+        long_query = (
+            "Tell me about "
+            + "preferences " * self.VERY_LONG_QUERY_REPEAT_COUNT
+            + "settings"
+        )
         result = await optimize_query_for_vector_search(long_query)
 
         assert result == "long query optimized"