Update app.py

app.py

@@ -9,10 +9,13 @@ from utils.auth import token_required
 from dotenv import load_dotenv
 import os
 
-load_dotenv()  # Load environment variables from .env file
+# Load environment variables from .env file
+load_dotenv()
 
+# Initialize FastAPI application
 app = FastAPI()
 
+# Set the cache directory for Hugging Face
 os.environ["HF_HOME"] = "/tmp/huggingface_cache"
 
 # Ensure the cache directory exists
@@ -22,23 +25,38 @@ if not os.path.exists(cache_dir):
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
+
 # Load Hugging Face token from environment variable
-# huggingface_token = os.getenv('HUGGINGFACE_HUB_TOKEN')
+huggingface_token = os.getenv('HUGGINGFACE_HUB_TOKEN')
+if huggingface_token:
+    try:
+        login(token=huggingface_token, add_to_git_credential=True)
+        logging.info("Successfully logged into Hugging Face Hub.")
+    except Exception as e:
+        logging.error(f"Failed to log into Hugging Face Hub: {e}")
+        raise HTTPException(status_code=500, detail="Failed to log into Hugging Face Hub.")
+else:
+    raise ValueError("Hugging Face token is not set. Please set the HUGGINGFACE_HUB_TOKEN environment variable.")
 
+# Initialize the Qdrant searcher
+qdrant_url = os.getenv('QDRANT_URL')
+access_token = os.getenv('QDRANT_ACCESS_TOKEN')
 
-# if huggingface_token:
-#     login(token=huggingface_token, add_to_git_credential=True)
-# else:
-#     raise ValueError("Hugging Face token is not set. Please set the HUGGINGFACE_HUB_TOKEN environment variable.")
+if not qdrant_url or not access_token:
+    raise ValueError("Qdrant URL or Access Token is not set. Please set the QDRANT_URL and QDRANT_ACCESS_TOKEN environment variables.")
 
+# Initialize the SentenceTransformer model
+try:
+    encoder = SentenceTransformer('nomic-ai/nomic-embed-text-v1.5')
+    logging.info("Successfully loaded the SentenceTransformer model.")
+except Exception as e:
+    logging.error(f"Failed to load the SentenceTransformer model: {e}")
+    raise HTTPException(status_code=500, detail="Failed to load the SentenceTransformer model.")
 
 # Initialize the Qdrant searcher
-qdrant_url = os.getenv('QDRANT_URL')
-access_token = os.getenv('QDRANT_ACCESS_TOKEN')
-encoder = SentenceTransformer('nomic-ai/nomic-embed-text-v1.5')  
 searcher = QdrantSearcher(encoder, qdrant_url, access_token)
 
-# Request body models
+# Define the request body models
 class SearchDocumentsRequest(BaseModel):
     query: str
     limit: int = 3
@@ -46,6 +64,7 @@ class SearchDocumentsRequest(BaseModel):
 class GenerateRAGRequest(BaseModel):
     search_query: str
 
+# Define the search documents endpoint
 @app.post("/api/search-documents")
 async def search_documents(
     body: SearchDocumentsRequest,
@@ -53,15 +72,13 @@ async def search_documents(
 ):
     customer_id, user_id = credentials
 
-    # Check if customer_id or user_id is missing
     if not customer_id or not user_id:
         logging.error("Failed to extract customer_id or user_id from the JWT token.")
         raise HTTPException(status_code=401, detail="Invalid token: missing customer_id or user_id")
 
     logging.info("Received request to search documents")
     try:
-        collection_name = "my_embeddings"
-        hits, error = searcher.search_documents(collection_name, body.query, user_id, body.limit)
+        hits, error = searcher.search_documents("documents", body.query, user_id, body.limit)
         
         if error:
             logging.error(f"Search documents error: {error}")
@@ -72,6 +89,7 @@ async def search_documents(
         logging.error(f"Unexpected error: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
+# Define the generate RAG response endpoint
 @app.post("/api/generate-rag-response")
 async def generate_rag_response_api(
     body: GenerateRAGRequest,
@@ -79,15 +97,13 @@ async def generate_rag_response_api(
 ):
     customer_id, user_id = credentials
 
-    # Check if customer_id or user_id is missing
     if not customer_id or not user_id:
         logging.error("Failed to extract customer_id or user_id from the JWT token.")
         raise HTTPException(status_code=401, detail="Invalid token: missing customer_id or user_id")
 
     logging.info("Received request to generate RAG response")
     try:
-        collection_name = "my_embeddings"
-        hits, error = searcher.search_documents(collection_name, body.search_query, user_id)
+        hits, error = searcher.search_documents("documents", body.search_query, user_id)
         
         if error:
             logging.error(f"Search documents error: {error}")
@@ -106,4 +122,4 @@ async def generate_rag_response_api(
 
 if __name__ == '__main__':
     import uvicorn
-    uvicorn.run(app, host='0.0.0.0', port=8000)
+    uvicorn.run(app, host='0.0.0.0', port=8000)