add parameter count filters to model search and trending endpoints

main.py

@@ -366,29 +366,64 @@ async def find_similar_datasets(
 @cache(ttl=CACHE_TTL)
 async def search_models(
     query: str,
-    k: int = Query(default=5, ge=1, le=100),
+    k: int = Query(default=5, ge=1, le=100, description="Number of results to return"),
     sort_by: str = Query(
-        default="similarity", enum=["similarity", "likes", "downloads", "trending"]
+        default="similarity",
+        enum=["similarity", "likes", "downloads", "trending"],
+        description="Sort method for results",
+    ),
+    min_likes: int = Query(default=0, ge=0, description="Minimum likes filter"),
+    min_downloads: int = Query(default=0, ge=0, description="Minimum downloads filter"),
+    min_param_count: int = Query(
+        default=0,
+        ge=0,
+        description="Minimum parameter count (models with param_count=0 will be excluded if any param filter is used)",
+    ),
+    max_param_count: Optional[int] = Query(
+        default=None,
+        ge=0,
+        description="Maximum parameter count (None means no upper limit)",
     ),
-    min_likes: int = Query(default=0, ge=0),
-    min_downloads: int = Query(default=0, ge=0),
 ):
+    """
+    Search for models based on a text query with optional filtering.
+
+    - When min_param_count > 0 or max_param_count is specified, models with param_count=0 are excluded
+    - param_count=0 indicates missing/unknown parameter count in the dataset
+    """
     try:
         collection = client.get_collection(
             name="model_cards", embedding_function=get_embedding_function()
         )
 
+        where_conditions = []
+        if min_likes > 0:
+            where_conditions.append({"likes": {"$gte": min_likes}})
+        if min_downloads > 0:
+            where_conditions.append({"downloads": {"$gte": min_downloads}})
+
+        # Add parameter count filters
+        using_param_filters = min_param_count > 0 or max_param_count is not None
+        if using_param_filters:
+            # Always exclude zero param count when using any parameter filters
+            where_conditions.append({"param_count": {"$gt": 0}})
+
+        if min_param_count > 0:
+            where_conditions.append({"param_count": {"$gte": min_param_count}})
+        if max_param_count is not None:
+            where_conditions.append({"param_count": {"$lte": max_param_count}})
+
+        # Handle where clause creation based on number of conditions
+        where_clause = None
+        if len(where_conditions) > 1:
+            where_clause = {"$and": where_conditions}
+        elif len(where_conditions) == 1:
+            where_clause = where_conditions[0]  # Single condition without $and
+
         results = collection.query(
             query_texts=[f"search_query: {query}"],
             n_results=k * 4 if sort_by != "similarity" else k,
-            where={
-                "$and": [
-                    {"likes": {"$gte": min_likes}},
-                    {"downloads": {"$gte": min_downloads}},
-                ]
-            }
-            if min_likes > 0 or min_downloads > 0
-            else None,
+            where=where_clause,
         )
 
         query_results = await process_search_results(results, "model", k, sort_by)
@@ -404,13 +439,31 @@ async def search_models(
 @cache(ttl=CACHE_TTL)
 async def find_similar_models(
     model_id: str,
-    k: int = Query(default=5, ge=1, le=100),
+    k: int = Query(default=5, ge=1, le=100, description="Number of results to return"),
     sort_by: str = Query(
-        default="similarity", enum=["similarity", "likes", "downloads", "trending"]
+        default="similarity",
+        enum=["similarity", "likes", "downloads", "trending"],
+        description="Sort method for results",
+    ),
+    min_likes: int = Query(default=0, ge=0, description="Minimum likes filter"),
+    min_downloads: int = Query(default=0, ge=0, description="Minimum downloads filter"),
+    min_param_count: int = Query(
+        default=0,
+        ge=0,
+        description="Minimum parameter count (models with param_count=0 will be excluded if any param filter is used)",
+    ),
+    max_param_count: Optional[int] = Query(
+        default=None,
+        ge=0,
+        description="Maximum parameter count (None means no upper limit)",
     ),
-    min_likes: int = Query(default=0, ge=0),
-    min_downloads: int = Query(default=0, ge=0),
 ):
+    """
+    Find similar models to a specified model with optional filtering.
+
+    - When min_param_count > 0 or max_param_count is specified, models with param_count=0 are excluded
+    - param_count=0 indicates missing/unknown parameter count in the dataset
+    """
     try:
         collection = client.get_collection("model_cards")
 
@@ -421,17 +474,34 @@ async def find_similar_models(
                 status_code=404, detail=f"Model ID '{model_id}' not found"
             )
 
+        where_conditions = []
+        if min_likes > 0:
+            where_conditions.append({"likes": {"$gte": min_likes}})
+        if min_downloads > 0:
+            where_conditions.append({"downloads": {"$gte": min_downloads}})
+
+        # Add parameter count filters
+        using_param_filters = min_param_count > 0 or max_param_count is not None
+        if using_param_filters:
+            # Always exclude zero param count when using any parameter filters
+            where_conditions.append({"param_count": {"$gt": 0}})
+
+        if min_param_count > 0:
+            where_conditions.append({"param_count": {"$gte": min_param_count}})
+        if max_param_count is not None:
+            where_conditions.append({"param_count": {"$lte": max_param_count}})
+
+        # Handle where clause creation based on number of conditions
+        where_clause = None
+        if len(where_conditions) > 1:
+            where_clause = {"$and": where_conditions}
+        elif len(where_conditions) == 1:
+            where_clause = where_conditions[0]  # Single condition without $and
+
         results = collection.query(
             query_embeddings=[results["embeddings"][0]],
             n_results=k * 4 if sort_by != "similarity" else k + 1,
-            where={
-                "$and": [
-                    {"likes": {"$gte": min_likes}},
-                    {"downloads": {"$gte": min_downloads}},
-                ]
-            }
-            if min_likes > 0 or min_downloads > 0
-            else None,
+            where=where_clause,
         )
 
         query_results = await process_search_results(
@@ -538,6 +608,8 @@ async def get_trending_models_with_summaries(
     limit: int = 10,
     min_likes: int = 0,
     min_downloads: int = 0,
+    min_param_count: int = 0,
+    max_param_count: Optional[int] = None,
 ) -> List[ModelQueryResult]:
     """Fetch trending models and combine with summaries from database"""
     try:
@@ -573,13 +645,30 @@ async def get_trending_models_with_summaries(
         for model in trending_models:
             if model["modelId"] in id_to_summary:
                 metadata = id_to_metadata.get(model["modelId"], {})
+                param_count = metadata.get("param_count", 0)
+
+                # Apply parameter count filters
+                using_param_filters = min_param_count > 0 or max_param_count is not None
+
+                # Skip if param_count is 0 and we're using param filters
+                if using_param_filters and param_count == 0:
+                    continue
+
+                # Skip if param_count is less than min_param_count
+                if min_param_count > 0 and param_count < min_param_count:
+                    continue
+
+                # Skip if param_count is greater than max_param_count
+                if max_param_count is not None and param_count > max_param_count:
+                    continue
+
                 result = ModelQueryResult(
                     model_id=model["modelId"],
                     similarity=1.0,  # Not applicable for trending
                     summary=id_to_summary[model["modelId"]],
                     likes=model.get("likes", 0),
                     downloads=model.get("downloads", 0),
-                    param_count=metadata.get("param_count"),
+                    param_count=param_count,
                 )
                 results.append(result)
 
@@ -592,13 +681,34 @@ async def get_trending_models_with_summaries(
 
 @app.get("/trending/models", response_model=ModelQueryResponse)
 async def get_trending_models(
-    limit: int = Query(default=10, ge=1, le=100),
-    min_likes: int = Query(default=0, ge=0),
-    min_downloads: int = Query(default=0, ge=0),
+    limit: int = Query(
+        default=10, ge=1, le=100, description="Number of results to return"
+    ),
+    min_likes: int = Query(default=0, ge=0, description="Minimum likes filter"),
+    min_downloads: int = Query(default=0, ge=0, description="Minimum downloads filter"),
+    min_param_count: int = Query(
+        default=0,
+        ge=0,
+        description="Minimum parameter count (models with param_count=0 will be excluded if any param filter is used)",
+    ),
+    max_param_count: Optional[int] = Query(
+        default=None,
+        ge=0,
+        description="Maximum parameter count (None means no upper limit)",
+    ),
 ):
-    """Get trending models with their summaries"""
+    """
+    Get trending models with their summaries and optional filtering.
+
+    - When min_param_count > 0 or max_param_count is specified, models with param_count=0 are excluded
+    - param_count=0 indicates missing/unknown parameter count in the dataset
+    """
     results = await get_trending_models_with_summaries(
-        limit=limit, min_likes=min_likes, min_downloads=min_downloads
+        limit=limit,
+        min_likes=min_likes,
+        min_downloads=min_downloads,
+        min_param_count=min_param_count,
+        max_param_count=max_param_count,
     )
     return ModelQueryResponse(results=results)