add endpoint for viewer search

main.py

@@ -4,18 +4,20 @@ from typing import List, Optional
 
 import chromadb
 from cashews import cache
+from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
 from fastapi import FastAPI, HTTPException, Query
 from httpx import AsyncClient
 from huggingface_hub import DatasetCard
 from pydantic import BaseModel
 from starlette.responses import RedirectResponse
 from starlette.status import (
+    HTTP_403_FORBIDDEN,
     HTTP_404_NOT_FOUND,
     HTTP_500_INTERNAL_SERVER_ERROR,
-    HTTP_403_FORBIDDEN,
 )
 
 from load_card_data import get_embedding_function, get_save_path, refresh_card_data
+from load_viewer_data import refresh_viewer_data
 
 # Set up logging
 logging.basicConfig(
@@ -43,20 +45,29 @@ async def lifespan(app: FastAPI):
     logger.info("Starting up the application")
     try:
         # Create or get the collection
+        logger.info("Initializing embedding function")
         embedding_function = get_embedding_function()
+        logger.info("Creating or getting collection")
         collection = client.get_or_create_collection(
             name="dataset_cards", embedding_function=embedding_function
         )
         logger.info("Collection initialized successfully")
 
         # Refresh data
+        logger.info("Starting refresh of card data")
         refresh_card_data()
+        logger.info("Card data refresh completed")
+
+        logger.info("Starting refresh of viewer data")
+        await refresh_viewer_data()
+        logger.info("Viewer data refresh completed")
+
         logger.info("Data refresh completed successfully")
     except Exception as e:
         logger.error(f"Error during startup: {str(e)}")
-        raise
+        logger.warning("Application starting with potential data issues")
 
-    yield  # Here the app is running and handling requests
+    yield
 
     # Shutdown: perform any cleanup
     logger.info("Shutting down the application")
@@ -171,7 +182,7 @@ 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)
+@app.post("/similar-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:
@@ -209,6 +220,48 @@ async def api_query_by_text(query: str, n: int = Query(default=10, ge=1, le=100)
         ) from e
 
 
+@app.post("/search-viewer", response_model=QueryResponse)
+@cache(ttl="1h")
+async def api_search_viewer(query: str, n: int = Query(default=10, ge=1, le=100)):
+    try:
+        embedding_function = SentenceTransformerEmbeddingFunction(
+            model_name="davanstrien/dataset-viewer-descriptions-processed-st",
+            trust_remote_code=True,
+        )
+        collection = client.get_collection(
+            name="dataset-viewer-descriptions",
+            embedding_function=embedding_function,
+        )
+        query = f"USER_QUERY: {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