refactor: update database connection handling and dependency injection for PostgreSQL and Qdrant

api/chat.py

@@ -1,13 +1,15 @@
 import json
+
+from fastapi.params import Depends
 from services.prompt import FALLBACK_RESPONSE
 from services.generation import craft_prompt, is_grounded_response, request_llm_response, parse_llm_json, relevant_chunks_to_json
 from services.resource_manager import get_llm_client
 from fastapi import APIRouter, Form, Request
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
-from db.qdrant import get_relevant_chunks
+from db.qdrant import get_qdrant, get_relevant_chunks
 from services.embedding import embed_question
-from db.postgres import update_document_activity
+from db.postgres import get_pg_pool, update_document_activity
 from services.rate_limiter import limiter
 
 class ChatRequest(BaseModel):
@@ -18,11 +20,11 @@ route = APIRouter(prefix="/chat", tags=["Chat"])
 
 @route.post("/ask")
 @limiter.limit("5/minute")
-async def ask_question(request: Request, body: ChatRequest = Form(...)):
+async def ask_question(request: Request, body: ChatRequest = Form(...), pg_pool = Depends(get_pg_pool), qdrant = Depends(get_qdrant)):
 	question = embed_question(body.question)
 	try:
-		await update_document_activity(body.uuid)
-		chunks, distances = await get_relevant_chunks(question, body.uuid)
+		await update_document_activity(body.uuid, pg_pool)
+		chunks, distances = await get_relevant_chunks(question, body.uuid, qdrant)
 		if not chunks:
 			return JSONResponse(content={"ok": False, "error": "No chunks found.", "data": None, "relevant_chunks": []}, status_code=404)
 		if distances and max(distances) < 0.52:  # Threshold for relevance, based on empirical testing

api/documents.py

@@ -1,5 +1,6 @@
 import asyncio
 from fastapi import APIRouter, Request, UploadFile, File, BackgroundTasks
+from fastapi.params import Depends
 from fastapi.responses import JSONResponse
 
 # Import service functions
@@ -11,14 +12,14 @@ from services.embedding import process_embeddings_background
 from services.rate_limiter import limiter
 
 #database imports
-from db.qdrant import get_chunk_context_by_index, get_document_chunks_by_uuid
-from db.postgres import get_document_by_hash, get_document_text, insert_document, get_document_by_uuid, document_exists, update_document_activity
+from db.qdrant import get_chunk_context_by_index, get_document_chunks_by_uuid, get_qdrant
+from db.postgres import get_document_by_hash, get_document_text, get_pg_pool, insert_document, get_document_by_uuid, document_exists, update_document_activity
 
 route = APIRouter(prefix="/documents", tags=["Documents"])
 
 @route.post("/upload")
 @limiter.limit("3/minute")
-async def upload(request: Request, file: UploadFile = File(...), background_tasks: BackgroundTasks = None):
+async def upload(request: Request, file: UploadFile = File(...), background_tasks: BackgroundTasks = None, pg_pool = Depends(get_pg_pool), qdrant = Depends(get_qdrant)):
 
 	file.filename = sanitize_for_display(file.filename)
 	#Check file type
@@ -33,8 +34,9 @@ async def upload(request: Request, file: UploadFile = File(...), background_task
 	#Calculate file hash and check for duplicates
 	file_hash = await calculate_file_hash(content)
 	try :
-		existing_doc = await get_document_by_hash(file_hash)
-		if existing_doc:
+		existing_doc = await get_document_by_hash(file_hash, pg_pool)
+		chunks = await get_document_chunks_by_uuid(existing_doc.get("id"), qdrant, limit=1)
+		if existing_doc and len(chunks) > 0:
 			return JSONResponse(content={"success": True, "metadata": existing_doc}, status_code=200)
 	except Exception as e:
 		print(f"Database error: {e}")
@@ -51,13 +53,14 @@ async def upload(request: Request, file: UploadFile = File(...), background_task
 
 	#store document metadata and chunks in the database
 	try:
-		await insert_document(document_id, file.filename, validated_text, len(chunks), pages[-1]['page'], file_hash)
+		await insert_document(document_id, file.filename, validated_text, len(chunks), pages[-1]['page'], file_hash, pg_pool)
 	except Exception as e:
 		print(f"Database error: {e}")
 		return JSONResponse(content={"success": False, "error": "Database error, Please try again."}, status_code=500)
 
 	#embed chunks and store embeddings in the vector database
-	background_tasks.add_task(process_embeddings_background, chunks, document_id)
+	print(f"Scheduling background embedding for {document_id}...", flush=True)
+	background_tasks.add_task(process_embeddings_background, chunks, document_id, qdrant)
 
 	response = {
 		"success": True,
@@ -75,22 +78,23 @@ async def upload(request: Request, file: UploadFile = File(...), background_task
 
 @route.get("/{uuid}/status")
 @limiter.limit("20/minute")
-async def check_processing_status(request: Request, uuid: str):
+async def check_processing_status(request: Request, uuid: str, qdrant = Depends(get_qdrant)):
     """Checks if Qdrant has finished saving the chunks"""
     try:
-        chunks = await get_document_chunks_by_uuid(uuid)
+        chunks = await get_document_chunks_by_uuid(uuid, qdrant, limit=1)  # Fetch just one chunk to check if embedding is done
         is_ready = len(chunks) > 0
         return JSONResponse(content={"success": True, "ready": is_ready})
-    except Exception:
+    except Exception as e:
+        print(str(e))
         return JSONResponse(content={"success": True, "ready": False})
 
 @route.get("/{uuid}")
 @limiter.limit("30/minute")
-async def get_document(request: Request, uuid: str):
+async def get_document(request: Request, uuid: str, pg_pool = Depends(get_pg_pool)):
 	''' Placeholder for fetching document metadata and chunks from the database '''
 	try:
-		await update_document_activity(uuid)
-		document = await get_document_by_uuid(uuid)
+		await update_document_activity(uuid, pg_pool)
+		document = await get_document_by_uuid(uuid, pg_pool)
 	except Exception as e:
 		print(f"Database error: {e}")
 		return JSONResponse(content={"success": False, "message": "Database error"}, status_code=500)
@@ -102,11 +106,11 @@ async def get_document(request: Request, uuid: str):
 
 @route.get("/{uuid}/info")
 @limiter.limit("30/minute")
-async def get_document_info(request: Request, uuid: str):
+async def get_document_info(request: Request, uuid: str, pg_pool = Depends(get_pg_pool)):
 	''' Placeholder for fetching document metadata and chunks from the database '''
 	try:
-		await update_document_activity(uuid)
-		document = await get_document_by_uuid(uuid)  # Implement this function to fetch document metadata
+		await update_document_activity(uuid, pg_pool)
+		document = await get_document_by_uuid(uuid, pg_pool)  # Implement this function to fetch document metadata
 	except Exception as e:
 		print(f"Database error: {e}")
 		return JSONResponse(content={"success": False, "error": "Database error"}, status_code=500)
@@ -118,14 +122,14 @@ async def get_document_info(request: Request, uuid: str):
 
 @route.get("/{uuid}/preview")
 @limiter.limit("30/minute")
-async def get_document_chunks(request: Request, uuid: str):
+async def get_document_chunks(request: Request, uuid: str, pg_pool = Depends(get_pg_pool), qdrant = Depends(get_qdrant)):
 	''' Placeholder for fetching document chunks from the database '''
 	try:
-		if not await document_exists(uuid):
+		if not await document_exists(uuid, pg_pool):
 			return JSONResponse(content={"success": False, "error": "Document not found. Please try again."}, status_code=404)
-		await update_document_activity(uuid)
-		chunks = await get_document_chunks_by_uuid(uuid)
-		text_preview = await get_document_text(uuid)
+		await update_document_activity(uuid, pg_pool)
+		chunks = await get_document_chunks_by_uuid(uuid, qdrant)
+		text_preview = await get_document_text(uuid, pg_pool)
 	except Exception as e:
 		print(f"Database error: {e}")
 		return JSONResponse(content={"success": False, "error": "Database error. Please try again."}, status_code=500)
@@ -134,16 +138,16 @@ async def get_document_chunks(request: Request, uuid: str):
 
 @route.get("/{uuid}/{chunk_index}/context")
 @limiter.limit("30/minute")
-async def get_chunk_context(request: Request, uuid: str, chunk_index: int):
+async def get_chunk_context(request: Request, uuid: str, chunk_index: int, pg_pool = Depends(get_pg_pool), qdrant = Depends(get_qdrant)):
 	''' Placeholder for fetching document chunks from the database '''
 	try:
-		if not await document_exists(uuid):
+		if not await document_exists(uuid, pg_pool):
 			return JSONResponse(content={"success": False, "error": "Document not found. Please try again."}, status_code=404)
-		await update_document_activity(uuid)
-		index_pool = await get_chunk_context_by_index(uuid, chunk_index)
+		await update_document_activity(uuid, pg_pool)
+		index_pool = await get_chunk_context_by_index(uuid, chunk_index, qdrant)
 		if not index_pool:
 			return JSONResponse(content={"success": False, "error": "Chunk context not found. Please try again."}, status_code=404)
-		context = await get_document_text(uuid, min(index_pool), max(index_pool) - min(index_pool))
+		context = await get_document_text(uuid, min(index_pool), max(index_pool) - min(index_pool), pg_pool)
 		if not context:
 			return JSONResponse(content={"success": False, "error": "Chunk context not found. Please try again."}, status_code=404)
 	except Exception as e:

db/postgres.py

@@ -1,54 +1,37 @@
 from datetime import timedelta
 
-import asyncpg
-from click import UUID
+from fastapi import Request
 from services.storage import human_readable_size
 from config import PREVIEW_LENGTH
 
-# The global connection pool
-pool = None
+#CREATE TABLE IF NOT EXISTS documents (
+#     id TEXT PRIMARY KEY,
+#     original_filename TEXT NOT NULL,
+#     full_text TEXT NOT NULL,
+#     text_length INTEGER NOT NULL,
+#     file_hash TEXT UNIQUE NOT NULL,
+#     chunks_count INTEGER NOT NULL,
+#     pages_count INTEGER NOT NULL,
+#     created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
+#     last_activity_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
+# );
 
-INIT_QUERY = '''
-CREATE TABLE IF NOT EXISTS documents (
-    id TEXT PRIMARY KEY,
-    original_filename TEXT NOT NULL,
-    full_text TEXT NOT NULL,
-    text_length INTEGER NOT NULL,
-    file_hash TEXT UNIQUE NOT NULL,
-    chunks_count INTEGER NOT NULL,
-    pages_count INTEGER NOT NULL,
-    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
-    last_activity_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
-);
-'''
+async def get_pg_pool(request: Request):
+    return request.app.state.pg_pool
 
-async def initialize_database(db_url: str):
-	"""Initialize the database connection pool and create tables if they don't exist."""
-	global pool
-	# Create a connection pool
-	pool = await asyncpg.create_pool(db_url, min_size=2, max_size=10)
-	async with pool.acquire() as conn:
-		await conn.execute(INIT_QUERY)
-
-async def close_database():
-    """Close the database connection pool."""
-    global pool
-    if pool:
-        await pool.close()
-
-async def insert_document(document_id, original_filename, full_text, chunks_count, pages_count, file_hash):
+async def insert_document(document_id, original_filename, full_text, chunks_count, pages_count, file_hash, pg_pool):
     ''' Insert a new document into the database and return its ID '''
     query = '''
         INSERT INTO documents (id, original_filename, full_text, text_length, chunks_count, pages_count, file_hash)
         VALUES ($1, $2, $3, $4, $5, $6, $7)
     '''
-    async with pool.acquire() as conn:
+    async with pg_pool.acquire() as conn:
         await conn.execute(query, document_id, original_filename, full_text, len(full_text), chunks_count, pages_count, file_hash)
         return document_id
 
-async def get_document_by_uuid(document_id):
+async def get_document_by_uuid(document_id, pg_pool):
     ''' Retrieve document metadata by its UUID '''
-    async with pool.acquire() as conn:
+    async with pg_pool.acquire() as conn:
         query = 'SELECT id, original_filename, chunks_count, pages_count, text_length, created_at FROM documents WHERE id = $1'
         row = await conn.fetchrow(query, document_id)
         if row:
@@ -63,27 +46,27 @@ async def get_document_by_uuid(document_id):
             }
         return None
 
-async def document_exists(document_id):
+async def document_exists(document_id, pg_pool):
     ''' Check if a document with the given ID exists in the database '''
-    async with pool.acquire() as conn:
+    async with pg_pool.acquire() as conn:
         row = await conn.fetchval('SELECT 1 FROM documents WHERE id = $1', document_id)
         return row is not None
 
-async def get_document_text(document_id, start_char=0, length=PREVIEW_LENGTH):
+async def get_document_text(document_id, pg_pool, start_char=0, length=PREVIEW_LENGTH):
     ''' Retrieve a substring of the document's full text based on character offsets '''
-    async with pool.acquire() as conn:
+    async with pg_pool.acquire() as conn:
         query = 'SELECT SUBSTR(full_text, $1, $2) FROM documents WHERE id = $3'
         text = await conn.fetchval(query, start_char + 1, length, document_id)
         return text
 
-async def update_document_activity(document_id):
+async def update_document_activity(document_id, pg_pool):
     ''' Update the last activity timestamp for a document '''
-    async with pool.acquire() as conn:
+    async with pg_pool.acquire() as conn:
         await conn.execute('UPDATE documents SET last_activity_at = CURRENT_TIMESTAMP WHERE id = $1', document_id)
 
-async def delete_inactive_documents(inactivity_threshold_hours=48):
+async def delete_inactive_documents(pg_pool, inactivity_threshold_hours=48):
     ''' Delete documents that haven't been accessed within the specified inactivity threshold '''
-    async with pool.acquire() as conn:
+    async with pg_pool.acquire() as conn:
         # Pass a standard Python timedelta, and asyncpg translates it to a Postgres INTERVAL
         threshold = timedelta(seconds=inactivity_threshold_hours)
         query = '''
@@ -95,9 +78,9 @@ async def delete_inactive_documents(inactivity_threshold_hours=48):
 
         return [record['id'] for record in records]
     
-async def get_document_by_hash(file_hash):
+async def get_document_by_hash(file_hash, pg_pool):
     ''' Retrieve document metadata by its file hash '''
-    async with pool.acquire() as conn:
+    async with pg_pool.acquire() as conn:
         query = 'SELECT id, original_filename, chunks_count, pages_count, text_length, created_at FROM documents WHERE file_hash = $1'
         row = await conn.fetchrow(query, file_hash)
         if row:

db/qdrant.py

@@ -1,23 +1,17 @@
 import os
 import uuid
-from qdrant_client import AsyncQdrantClient
-from qdrant_client.models import Distance, Range, VectorParams, PointStruct, Filter, FieldCondition, MatchValue
+from fastapi import Request
+from qdrant_client.models import Range, PointStruct, Filter, FieldCondition, MatchValue
 from config import PREVIEW_LENGTH
 
-# Connect to your local Qdrant Docker container
-qdrant = AsyncQdrantClient(url=os.getenv("QDRANT_URL"))
 COLLECTION_NAME = os.getenv("QDRANT_COLLECTION")
 
-async def init_qdrant():
-	"""Call this in your FastAPI lifespan alongside your Postgres init"""
-	if not await qdrant.collection_exists(COLLECTION_NAME):
-		await qdrant.create_collection(
-			collection_name=COLLECTION_NAME,
-			vectors_config=VectorParams(size=384, distance=Distance.COSINE),
-		)
+async def get_qdrant(request: Request):
+    return request.app.state.qdrant
 
-async def save_embeddings(embeddings, chunks, document_id):
+async def save_embeddings(embeddings, chunks, document_id, qdrant):
 	"""Saves both the vector AND the text payload into Qdrant"""
+
 	points = []
 	
 	#loop through the text chunks and their matching vectors
@@ -42,7 +36,7 @@ async def save_embeddings(embeddings, chunks, document_id):
 		points=points
 	)
 
-async def get_relevant_chunks(question_embedding, document_id, top_k=3):
+async def get_relevant_chunks(question_embedding, document_id, qdrant, top_k=3):
 	"""Searches Qdrant and filters ONLY for the active document"""
 	
 	search_results = await qdrant.query_points(
@@ -65,7 +59,7 @@ async def get_relevant_chunks(question_embedding, document_id, top_k=3):
 	
 	return retrieved_chunks, scores
 
-async def get_chunk_context_by_index(document_id: str, target_index: int):
+async def get_chunk_context_by_index(document_id, target_index, qdrant):
 	"""Fetches ONLY the start and end characters for a chunk and its neighbors"""
 	
 	records, _ = await qdrant.scroll(
@@ -98,7 +92,7 @@ async def get_chunk_context_by_index(document_id: str, target_index: int):
 		for index in (record.payload['start_char'], record.payload['end_char'])
 	]
 
-async def get_document_chunks_by_uuid(document_id: str):
+async def get_document_chunks_by_uuid(document_id, qdrant, limit=30):
 	"""Fetches clean, contiguous preview chunks with exactly the keys needed"""
 	
 	records, _ = await qdrant.scroll(
@@ -115,7 +109,7 @@ async def get_document_chunks_by_uuid(document_id: str):
 				)
 			]
 		),
-		limit=30,
+		limit=limit,
 		with_payload=["chunk_index", "start_char", "end_char"],
 		with_vectors=False
 	)
@@ -131,7 +125,7 @@ async def get_document_chunks_by_uuid(document_id: str):
 		for record in records
 	]
 
-async def delete_document_chunks(document_id: str):
+async def delete_document_chunks(document_id, qdrant):
 	"""Deletes all chunks and vectors belonging to a specific document"""
 
 	await qdrant.delete(

main.py

@@ -1,11 +1,12 @@
 import asyncio
 import os
+import asyncpg
 from dotenv import load_dotenv
 from fastapi import FastAPI
 from contextlib import asynccontextmanager
+
+from qdrant_client import AsyncQdrantClient
 from api import chat, documents
-from db.qdrant import init_qdrant
-from db.postgres import close_database, initialize_database
 from services.resource_manager import clean_inactive_documents
 from fastapi.middleware.cors import CORSMiddleware
 from services.rate_limiter import limiter, custom_rate_limit_handler
@@ -19,23 +20,24 @@ origins = [
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    # STARTUP
-    try:
-        await initialize_database(os.getenv("POSTGRES_URI"))
-        await init_qdrant()
-    except Exception as e:
-        print(f"Error initializing database: {e}")
-        raise
-
-    asyncio.create_task(clean_inactive_documents())
+    print("Connecting to databases...")
+    
+    # Initialize PostgreSQL Connection Pool
+    app.state.pg_pool = await asyncpg.create_pool(os.getenv("POSTGRES_URI"))
+    
+    # Initialize Qdrant Client
+    app.state.qdrant = AsyncQdrantClient(
+        url=os.getenv("QDRANT_URL"),
+        api_key=os.getenv("QDRANT_API_KEY")
+    )
+
+    asyncio.create_task(clean_inactive_documents(app.state.pg_pool, app.state.qdrant))
 
     yield
 
-    # SHUTDOWN
-    try:
-        await close_database()
-    except Exception as e:
-        print(f"Error closing database: {e}")
+    print("Closing database connections...")
+    await app.state.pg_pool.close()
+    await app.state.qdrant.close()
 
 app = FastAPI(lifespan=lifespan)
 # Attach the limiter to the app and set up the custom rate limit handler

services/embedding.py

@@ -18,10 +18,11 @@ def embed_question(question):
 	embedding = model.encode([bge_query])
 	return embedding
 
-async def process_embeddings_background(chunks, document_id):
+async def process_embeddings_background(chunks, document_id, qdrant):
+    print(f"Starting background embedding for {document_id}...", flush=True)
     try:
         embeddings = await asyncio.to_thread(generate_embeddings, chunks)
-        await save_embeddings(embeddings, chunks, document_id)
-        print(f"Background embedding complete for {document_id}")
+        await save_embeddings(embeddings, chunks, document_id, qdrant)
+        print(f"Background embedding complete for {document_id}", flush=True)
     except Exception as e:
-        print(f"Background embedding failed: {e}")
+        print(f"Background embedding failed: {e}", flush=True)

services/resource_manager.py

@@ -18,15 +18,15 @@ def get_model():
 	print("Loading model...")
 	return SentenceTransformer('BAAI/bge-small-en-v1.5')
 
-async def clean_inactive_documents():
+async def clean_inactive_documents(pg_pool, qdrant):
 	''' clean up documents that haven't been accessed in a while '''
 	while True:
 		await asyncio.sleep(3600)  # Run cleanup every hour
 		print("Running cleanup of inactive documents...")
 		try:
-			ids = await delete_inactive_documents()
+			ids = await delete_inactive_documents(pg_pool)
 			for doc_id in ids:
-				await delete_document_chunks(doc_id)
+				await delete_document_chunks(doc_id, qdrant)
 			print(f"Deleted {len(ids)} inactive documents and their chunks.")
 		except Exception as e:
 			print(f"Error during cleanup: {e}")