Add task/license/language/modified_after filters + /lookup endpoint; union_by_name load

app.py

@@ -97,19 +97,25 @@ def load_table(cfg: str, table: str, has_param: bool):
     plist = "[" + ",".join(f"'{p}'" for p in paths) + "]"
     param_sel = "COALESCE(param_count,0)::BIGINT AS param_count" if has_param else "0::BIGINT AS param_count"
     # slice to EMB_DIM, L2-renormalise, cast to fixed-size FLOAT[] array (nested, no correlated subquery)
+    # task/license/language are copied straight from the source cards; cast to
+    # VARCHAR because the models config's all-null `language` reads back as a
+    # NULL-typed column that would otherwise reject `=` filters at query time.
     con.execute(f"""
         CREATE OR REPLACE TABLE {table} AS
         SELECT id, summary,
                COALESCE(likes,0)::BIGINT AS likes,
                COALESCE(downloads,0)::BIGINT AS downloads,
                last_modified,
+               CAST(task AS VARCHAR) AS task,
+               CAST(license AS VARCHAR) AS license,
+               CAST(language AS VARCHAR) AS language,
                {param_sel},
                list_transform(s, x -> (x / nrm))::FLOAT[{EMB_DIM}] AS emb
         FROM (
           SELECT *, sqrt(list_dot_product(s, s)) AS nrm
           FROM (
             SELECT *, embedding[1:{EMB_DIM}] AS s
-            FROM read_parquet({plist})
+            FROM read_parquet({plist}, union_by_name=True)
             WHERE embedding IS NOT NULL
             QUALIFY row_number() OVER (PARTITION BY id ORDER BY last_modified DESC NULLS LAST) = 1
           )
@@ -184,8 +190,15 @@ class ModelQueryResponse(BaseModel):
 
 
 # --- core search -------------------------------------------------------------
-def _where(min_likes, min_downloads, min_param=0, max_param=None, use_param=False):
-    conds = []
+def _where(min_likes, min_downloads, min_param=0, max_param=None, use_param=False,
+           task=None, license=None, language=None, modified_after=None):
+    """Build a WHERE clause + its bind params.
+
+    Numeric bounds are inlined (ints are injection-safe); the free-text metadata
+    filters (task/license/language/modified_after) are bound as `?` params.
+    Returns ``(where_sql, params)`` for ``_knn``.
+    """
+    conds, params = [], []
     if min_likes > 0:
         conds.append(f"likes >= {int(min_likes)}")
     if min_downloads > 0:
@@ -196,17 +209,30 @@ def _where(min_likes, min_downloads, min_param=0, max_param=None, use_param=Fals
             conds.append(f"param_count >= {int(min_param)}")
         if max_param is not None:
             conds.append(f"param_count <= {int(max_param)}")
-    return ("WHERE " + " AND ".join(conds)) if conds else ""
+    if task:
+        conds.append("task = ?")
+        params.append(task)
+    if license:
+        conds.append("license = ?")
+        params.append(license)
+    if language:
+        conds.append("language = ?")
+        params.append(language)
+    if modified_after:
+        # last_modified is an ISO-8601 string; lexicographic >= is chronological.
+        conds.append("last_modified >= ?")
+        params.append(modified_after)
+    return (("WHERE " + " AND ".join(conds)) if conds else ""), params
 
 
 def _vec_literal(vec):
     return "[" + ",".join(repr(float(x)) for x in vec) + f"]::FLOAT[{EMB_DIM}]"
 
 
-def _knn(table, qvec, k, where_sql, cols):
+def _knn(table, qvec, k, where_sql, cols, params=None):
     q = f"""SELECT {cols}, array_cosine_similarity(emb, {_vec_literal(qvec)}) AS sim
             FROM {table} {where_sql} ORDER BY sim DESC LIMIT {int(k)}"""
-    return con.execute(q).fetchall()
+    return con.execute(q, params or []).fetchall()
 
 
 def _emb_of(table, repo_id):
@@ -273,19 +299,39 @@ async def health():
     }
 
 
+@app.get("/lookup/{repo_id:path}")
+async def lookup(repo_id: str):
+    """Resolve a repo id to its type in one call: {"id":..., "type": "dataset"|"model"}.
+
+    Saves clients the 404-fallthrough (try datasets, then models). 404 with a JSON
+    detail if the id is in neither index. `:path` so `org/name` ids keep their slash.
+    """
+    for table, typ in (("dataset_cards", "dataset"), ("model_cards", "model")):
+        if con.execute(f"SELECT 1 FROM {table} WHERE id = ? LIMIT 1", [repo_id]).fetchone():
+            return {"id": repo_id, "type": typ}
+    raise HTTPException(status_code=404, detail=f"'{repo_id}' not found as a dataset or model in the index")
+
+
 @app.get("/search/datasets", response_model=QueryResponse)
-@cache(ttl=CACHE_TTL, key="search:d:{query}:{k}:{sort_by}:{min_likes}:{min_downloads}")
+@cache(ttl=CACHE_TTL, key="search:d:{query}:{k}:{sort_by}:{min_likes}:{min_downloads}:{task}:{license}:{language}:{modified_after}")
 async def search_datasets(
     query: str,
     k: int = Query(default=5, ge=1, le=100),
     sort_by: str = Query(default="similarity", enum=["similarity", "likes", "downloads", "trending"]),
     min_likes: int = Query(default=0, ge=0),
     min_downloads: int = Query(default=0, ge=0),
+    task: Optional[str] = Query(default=None),
+    license: Optional[str] = Query(default=None),
+    language: Optional[str] = Query(default=None),
+    modified_after: Optional[str] = Query(default=None),
 ):
     try:
         qvec = embed_query(f"Instruct: {DS_TASK}\nQuery:{query}")
         n = k * 4 if sort_by != "similarity" else k
-        raw = _knn("dataset_cards", qvec, n, _where(min_likes, min_downloads), COLS)
+        where_sql, wparams = _where(min_likes, min_downloads,
+                                    task=task, license=license, language=language,
+                                    modified_after=modified_after)
+        raw = _knn("dataset_cards", qvec, n, where_sql, COLS, wparams)
         rows = await _sort_results(_rows_to_dataset(raw), "dataset", sort_by, k)
         return QueryResponse(results=[QueryResult(**{k2: v for k2, v in r.items() if k2 != "_id"}) for r in rows])
     except Exception as e:
@@ -294,20 +340,27 @@ async def search_datasets(
 
 
 @app.get("/similarity/datasets", response_model=QueryResponse)
-@cache(ttl=CACHE_TTL, key="sim:d:{dataset_id}:{k}:{sort_by}:{min_likes}:{min_downloads}")
+@cache(ttl=CACHE_TTL, key="sim:d:{dataset_id}:{k}:{sort_by}:{min_likes}:{min_downloads}:{task}:{license}:{language}:{modified_after}")
 async def find_similar_datasets(
     dataset_id: str,
     k: int = Query(default=5, ge=1, le=100),
     sort_by: str = Query(default="similarity", enum=["similarity", "likes", "downloads", "trending"]),
     min_likes: int = Query(default=0, ge=0),
     min_downloads: int = Query(default=0, ge=0),
+    task: Optional[str] = Query(default=None),
+    license: Optional[str] = Query(default=None),
+    language: Optional[str] = Query(default=None),
+    modified_after: Optional[str] = Query(default=None),
 ):
     emb = _emb_of("dataset_cards", dataset_id)
     if emb is None:
         raise HTTPException(status_code=404, detail=f"Dataset ID '{dataset_id}' not found")
     try:
         n = k * 4 if sort_by != "similarity" else k + 1
-        raw = _knn("dataset_cards", emb, n, _where(min_likes, min_downloads), COLS)
+        where_sql, wparams = _where(min_likes, min_downloads,
+                                    task=task, license=license, language=language,
+                                    modified_after=modified_after)
+        raw = _knn("dataset_cards", emb, n, where_sql, COLS, wparams)
         rows = [r for r in _rows_to_dataset(raw) if r["_id"] != dataset_id]
         rows = await _sort_results(rows, "dataset", sort_by, k)
         return QueryResponse(results=[QueryResult(**{k2: v for k2, v in r.items() if k2 != "_id"}) for r in rows])
@@ -319,7 +372,7 @@ async def find_similar_datasets(
 
 
 @app.get("/search/models", response_model=ModelQueryResponse)
-@cache(ttl=CACHE_TTL, key="search:m:{query}:{k}:{sort_by}:{min_likes}:{min_downloads}:{min_param_count}:{max_param_count}")
+@cache(ttl=CACHE_TTL, key="search:m:{query}:{k}:{sort_by}:{min_likes}:{min_downloads}:{min_param_count}:{max_param_count}:{task}:{license}:{language}:{modified_after}")
 async def search_models(
     query: str,
     k: int = Query(default=5, ge=1, le=100),
@@ -328,13 +381,19 @@ async def search_models(
     min_downloads: int = Query(default=0, ge=0),
     min_param_count: int = Query(default=0, ge=0),
     max_param_count: Optional[int] = Query(default=None, ge=0),
+    task: Optional[str] = Query(default=None),
+    license: Optional[str] = Query(default=None),
+    language: Optional[str] = Query(default=None),
+    modified_after: Optional[str] = Query(default=None),
 ):
     try:
         use_param = min_param_count > 0 or max_param_count is not None
         qvec = embed_query(f"search_query: {query}")
         n = k * 4 if sort_by != "similarity" else k
-        raw = _knn("model_cards", qvec, n,
-                   _where(min_likes, min_downloads, min_param_count, max_param_count, use_param), COLS)
+        where_sql, wparams = _where(min_likes, min_downloads, min_param_count, max_param_count, use_param,
+                                    task=task, license=license, language=language,
+                                    modified_after=modified_after)
+        raw = _knn("model_cards", qvec, n, where_sql, COLS, wparams)
         rows = await _sort_results(_rows_to_model(raw), "model", sort_by, k)
         return ModelQueryResponse(results=[ModelQueryResult(**{k2: v for k2, v in r.items() if k2 != "_id"}) for r in rows])
     except Exception as e:
@@ -343,7 +402,7 @@ async def search_models(
 
 
 @app.get("/similarity/models", response_model=ModelQueryResponse)
-@cache(ttl=CACHE_TTL, key="sim:m:{model_id}:{k}:{sort_by}:{min_likes}:{min_downloads}:{min_param_count}:{max_param_count}")
+@cache(ttl=CACHE_TTL, key="sim:m:{model_id}:{k}:{sort_by}:{min_likes}:{min_downloads}:{min_param_count}:{max_param_count}:{task}:{license}:{language}:{modified_after}")
 async def find_similar_models(
     model_id: str,
     k: int = Query(default=5, ge=1, le=100),
@@ -352,6 +411,10 @@ async def find_similar_models(
     min_downloads: int = Query(default=0, ge=0),
     min_param_count: int = Query(default=0, ge=0),
     max_param_count: Optional[int] = Query(default=None, ge=0),
+    task: Optional[str] = Query(default=None),
+    license: Optional[str] = Query(default=None),
+    language: Optional[str] = Query(default=None),
+    modified_after: Optional[str] = Query(default=None),
 ):
     emb = _emb_of("model_cards", model_id)
     if emb is None:
@@ -359,8 +422,10 @@ async def find_similar_models(
     try:
         use_param = min_param_count > 0 or max_param_count is not None
         n = k * 4 if sort_by != "similarity" else k + 1
-        raw = _knn("model_cards", emb, n,
-                   _where(min_likes, min_downloads, min_param_count, max_param_count, use_param), COLS)
+        where_sql, wparams = _where(min_likes, min_downloads, min_param_count, max_param_count, use_param,
+                                    task=task, license=license, language=language,
+                                    modified_after=modified_after)
+        raw = _knn("model_cards", emb, n, where_sql, COLS, wparams)
         rows = [r for r in _rows_to_model(raw) if r["_id"] != model_id]
         rows = await _sort_results(rows, "model", sort_by, k)
         return ModelQueryResponse(results=[ModelQueryResult(**{k2: v for k2, v in r.items() if k2 != "_id"}) for r in rows])