add fallback for vector space ID in usage

Author: MFA-X-AI

tests/test_management.py

@@ -164,16 +164,13 @@ def test_delete_vector_space():
 
 
 @pytest.mark.metrics
-class TestMetrics:
-
-    # Test getting monthly usage from Vecto
-    def test_usage(self): 
-        from datetime import datetime
-        today = datetime.now()
-        usage_response = user_vecto.usage(user_vecto.vector_space_id, today.year, today.month)
-        logger.info("Checking that usage returns a valid response.")
-        assert usage_response is not None
-        logger.info("Checking that usage for today is not empty.")
-        assert usage_response.usage.lookups.dailyMetrics[today.day-1].count > 0
-        assert usage_response.usage.indexing.dailyMetrics[today.day-1].count > 0
+def test_usage(): 
+    from datetime import datetime
+    today = datetime.now()
+    usage_response = user_vecto.usage(today.year, today.month)
+    logger.info("Checking that usage returns a valid response.")
+    assert usage_response is not None
+    logger.info("Checking that usage for today is not empty.")
+    assert usage_response.usage.lookups.dailyMetrics[today.day-1].count > 0
+    assert usage_response.usage.indexing.dailyMetrics[today.day-1].count > 0
 

vecto/vecto_requests.py

@@ -941,15 +941,25 @@ def _parse_vecto_usage_metrics(self, u: Dict[str, Any]) -> VectoUsageMetrics:
             indexing=self._parse_usage_metric(u['indexing'])
         )
 
-    def usage(self, vector_space_id: int, year: int, month: int, **kwargs) -> MonthlyUsageResponse:
+    def usage(self, year: int, month: int, vector_space_id: int = None, **kwargs) -> MonthlyUsageResponse:
         '''Return the usage metrics for the selected month
 
         Args:
-            vector_space_id (int): The ID of the vector space.
             year (int): The year for the usage data.
             month (int): The month for the usage data.
+            vector_space_id (int, optional): The ID of the vector space. Defaults to None.
             **kwargs: Other keyword arguments for clients other than `requests`
+        Returns:
+            MonthlyUsageResponse: Named tuple that contains the usage metrics for a specified vector space and month.
         '''
+        
+        # Use provided vector_space_id or fallback to self.vector_space_id
+        vector_space_id = vector_space_id or getattr(self, 'vector_space_id', None)
+        
+        # Raise an error if vector_space_id is still not available
+        if vector_space_id is None:
+            raise ValueError("A vector space ID must be provided either as a parameter or set in the instance.")
+
         url = f"/api/v0/space/{vector_space_id}/usage/{year}/{month}"
         response = self._client.get(url, **kwargs)
         response_data = response.json()
@@ -962,4 +972,4 @@ def usage(self, vector_space_id: int, year: int, month: int, **kwargs) -> Monthl
             usage=usage_metrics
         )
 
-        return monthly_usage_response
+        return monthly_usage_response