Phase 9.1 Step 4: Merge count_documents into get_collection_info

- Removed standalone count_documents tool
- Added include_count parameter to get_collection_info
- Updated E2E tests to use get_collection_info with include_count=True
- Updated tool list validation tests

Tool count: 22 → 21

Breaking change: count_documents tool removed
Migration: Use get_collection_info(database="db", include_count=True)

Signed-off-by: Nigel Jones <jonesn@uk.ibm.com>

Author: planetf1

src/maestro_mcp/server.py

@@ -1422,6 +1422,10 @@ async def get_collection_info(
             default=None,
             description="Name of the collection to get info for. If not provided, uses the default collection.",
         ),
+        include_count: bool = Field(
+            default=False,
+            description="Include document count in the response",
+        ),
     ) -> str:
         """Get information about a collection in a vector database."""
         db = get_database_by_name(database)
@@ -1443,6 +1447,15 @@ async def get_collection_info(
             return str(info_any)
         info: dict[str, Any] = cast("dict[str, Any]", info_any)
 
+        # Add document count if requested
+        if include_count:
+            ok_count, count_any = await run_with_timeout(
+                db.count_documents(), "count_documents", get_timeout("read")
+            )
+            if ok_count:
+                count = int(count_any) if count_any is not None else 0
+                info["document_count"] = count
+
         return (
             f"Collection information for '{info.get('name')}' in vector database "
             f"'{database}':\n{json.dumps(info, indent=2)}"

tests/e2e/test_functions.py

@@ -159,10 +159,12 @@ async def run_document_operations_tests(client: Client, backend_name: str) -> No
             docs_list = []
     assert docs_list is not None, f"list_documents failed: {res}"
 
-    # Test count_documents
-    res = await client.call_tool("count_documents", {"database": db_name})
+    # Test get_collection_info with count
+    res = await client.call_tool(
+        "get_collection_info", {"database": db_name, "include_count": True}
+    )
     if not hasattr(res, "data") and isinstance(res, str):
-        assert res, f"count_documents failed: {res}"
+        assert res, f"get_collection_info with count failed: {res}"
 
     # Test delete_document (get a document ID first from list_documents)
     first_doc_id = None
@@ -702,10 +704,14 @@ async def run_full_flow_test(client: Client, backend_name: str) -> None:
         if not hasattr(res, "data"):
             pytest.fail(f"list_documents failed for {backend_name}: {res}")
 
-        # Count documents
-        res = await client.call_tool("count_documents", {"database": db_name})
+        # Count documents via get_collection_info
+        res = await client.call_tool(
+            "get_collection_info", {"database": db_name, "include_count": True}
+        )
         if not hasattr(res, "data"):
-            pytest.fail(f"count_documents failed for {backend_name}: {res}")
+            pytest.fail(
+                f"get_collection_info with count failed for {backend_name}: {res}"
+            )
 
         # Get collection info
         res = await client.call_tool("get_collection_info", {"database": db_name})

tests/test_mcp_server.py

@@ -59,7 +59,6 @@ async def test_tool_definitions() -> None:
                 "write_documents",
                 "write_document",
                 "list_documents",
-                "count_documents",
                 "delete_documents",
                 "delete_document",
                 "delete_collection",

tests/test_phase1_schema_validation.py

@@ -127,7 +127,6 @@ async def test_all_tools_accessible() -> None:
         "write_document_to_collection",
         "list_documents",
         "list_documents_in_collection",
-        "count_documents",
         "delete_documents",
         "delete_document",
         "delete_document_from_collection",