rename

main.py

@@ -9,9 +9,13 @@ from httpx import AsyncClient
 from huggingface_hub import DatasetCard
 from pydantic import BaseModel
 from starlette.responses import RedirectResponse
-from starlette.status import HTTP_404_NOT_FOUND, HTTP_500_INTERNAL_SERVER_ERROR
+from starlette.status import (
+    HTTP_404_NOT_FOUND,
+    HTTP_500_INTERNAL_SERVER_ERROR,
+    HTTP_403_FORBIDDEN,
+)
 
-from load_data import get_embedding_function, get_save_path, refresh_data
+from load_card_data import get_embedding_function, get_save_path, refresh_data
 
 # Set up logging
 logging.basicConfig(
@@ -97,6 +101,14 @@ class DatasetCardNotFoundError(HTTPException):
         )
 
 
+class DatasetNotForAllAudiencesError(HTTPException):
+    def __init__(self, dataset_id: str):
+        super().__init__(
+            status_code=HTTP_403_FORBIDDEN,
+            detail=f"Dataset {dataset_id} is not for all audiences and not supported in this service.",
+        )
+
+
 @app.get("/similar", response_model=QueryResponse)
 @cache(ttl="1h")
 async def api_query_dataset(dataset_id: str, n: int = Query(default=10, ge=1, le=100)):
@@ -115,7 +127,9 @@ async def api_query_dataset(dataset_id: str, n: int = Query(default=10, ge=1, le
                 collection.upsert(ids=[dataset_id], embeddings=embeddings[0])
                 logger.info(f"Dataset {dataset_id} added to collection")
                 result = collection.get(ids=[dataset_id], include=["embeddings"])
-            except DatasetCardNotFoundError:
+                if result.get("not-for-all-audiences"):
+                    raise DatasetNotForAllAudiencesError(dataset_id)
+            except (DatasetCardNotFoundError, DatasetNotForAllAudiencesError):
                 raise
             except Exception as e:
                 logger.error(
@@ -157,6 +171,44 @@ async def api_query_dataset(dataset_id: str, n: int = Query(default=10, ge=1, le
         ) from e
 
 
+@app.post("/similar_by_text", response_model=QueryResponse)
+@cache(ttl="1h")
+async def api_query_by_text(query: str, n: int = Query(default=10, ge=1, le=100)):
+    try:
+        logger.info(f"Querying datasets by text: {query}")
+        collection = client.get_collection(
+            name="dataset_cards", embedding_function=get_embedding_function()
+        )
+        print(query)
+        query_result = collection.query(
+            query_texts=query, n_results=n, include=["distances"]
+        )
+        print(query_result)
+
+        if not query_result["ids"]:
+            logger.info(f"No similar datasets found for query: {query}")
+            raise HTTPException(
+                status_code=HTTP_404_NOT_FOUND, detail="No similar datasets found."
+            )
+
+        # Prepare the response
+        results = [
+            QueryResult(dataset_id=str(id), similarity=float(1 - distance))
+            for id, distance in zip(
+                query_result["ids"][0], query_result["distances"][0]
+            )
+        ]
+        logger.info(f"Found {len(results)} similar datasets for query: {query}")
+        return QueryResponse(results=results)
+
+    except Exception as e:
+        logger.error(f"Error querying datasets by text {query}: {str(e)}")
+        raise HTTPException(
+            status_code=HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="An unexpected error occurred.",
+        ) from e
+
+
 if __name__ == "__main__":
     import uvicorn