refactor

main.py

@@ -16,8 +16,9 @@ from starlette.status import (
     HTTP_500_INTERNAL_SERVER_ERROR,
 )
 
-from load_card_data import get_embedding_function, get_save_path, refresh_card_data
+from load_card_data import card_embedding_function, refresh_card_data
 from load_viewer_data import refresh_viewer_data
+from utils import get_save_path, get_collection
 
 # Set up logging
 logging.basicConfig(
@@ -31,7 +32,7 @@ cache.setup("mem://?check_interval=10&size=1000")
 # Initialize Chroma client
 SAVE_PATH = get_save_path()
 client = chromadb.PersistentClient(path=SAVE_PATH)
-collection = None
+
 
 async_client = AsyncClient(
     follow_redirects=True,
@@ -40,33 +41,20 @@ async_client = AsyncClient(
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    global collection
     # Startup: refresh data and initialize collection
     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)}")
         logger.warning("Application starting with potential data issues")
-
     yield
 
     # Shutdown: perform any cleanup
@@ -123,6 +111,8 @@ class DatasetNotForAllAudiencesError(HTTPException):
 @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)):
+    embedding_function = card_embedding_function()
+    collection = get_collection(client, embedding_function, "dataset_cards")
     try:
         logger.info(f"Querying dataset: {dataset_id}")
         # Get the embedding for the given dataset_id
@@ -130,7 +120,6 @@ async def api_query_dataset(dataset_id: str, n: int = Query(default=10, ge=1, le
         if not result.get("embeddings"):
             logger.info(f"Dataset not found: {dataset_id}")
             try:
-                embedding_function = get_embedding_function()
                 card = await try_get_card(dataset_id)
                 if card is None:
                     raise DatasetCardNotFoundError(dataset_id)
@@ -182,13 +171,13 @@ async def api_query_dataset(dataset_id: str, n: int = Query(default=10, ge=1, le
         ) from e
 
 
-@app.post("/similar-text", response_model=QueryResponse)
+@app.get("/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:
         logger.info(f"Querying datasets by text: {query}")
         collection = client.get_collection(
-            name="dataset_cards", embedding_function=get_embedding_function()
+            name="dataset_cards", embedding_function=card_embedding_function()
         )
         print(query)
         query_result = collection.query(
@@ -220,7 +209,7 @@ 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)
+@app.get("/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: