chore: Initialize collection on startup and refresh data

main.py

@@ -8,7 +8,7 @@ from fastapi import FastAPI, HTTPException, Query
 from pydantic import BaseModel
 from starlette.responses import RedirectResponse
 
-from load_data import get_save_path, refresh_data
+from load_data import get_embedding_function, get_save_path, refresh_data
 
 # Set up logging
 logging.basicConfig(
@@ -22,7 +22,7 @@ cache.setup("mem://?check_interval=10&size=10000")
 # Initialize Chroma client
 SAVE_PATH = get_save_path()
 client = chromadb.PersistentClient(path=SAVE_PATH)
-collection = client.get_collection("dataset_cards")
+collection = None
 
 
 class QueryResult(BaseModel):
@@ -36,13 +36,23 @@ class QueryResponse(BaseModel):
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    # Startup: refresh data
+    global collection
+    # Startup: refresh data and initialize collection
     logger.info("Starting up the application")
     try:
+        # Create or get the collection
+        embedding_function = get_embedding_function()
+        collection = client.get_or_create_collection(
+            name="dataset_cards", embedding_function=embedding_function
+        )
+        logger.info("Collection initialized successfully")
+
+        # Refresh data
         refresh_data()
         logger.info("Data refresh completed successfully")
     except Exception as e:
-        logger.error(f"Error during data refresh: {str(e)}")
+        logger.error(f"Error during startup: {str(e)}")
+        raise
 
     yield  # Here the app is running and handling requests