Improve memory efficiency and resource cleanup

app/routes/document_routes.py

@@ -123,19 +123,27 @@ async def load_file_content(
 
 
 def extract_text_from_documents(documents: List[Document], file_ext: str) -> str:
-    """Extract text content from loaded documents."""
-    text_content = ""
+    """
+    Extract text content from loaded documents.
+
+    Uses a list to collect parts and joins at the end to avoid O(n²)
+    string concatenation overhead for large documents.
+
+    :param documents: List of Document objects to extract text from
+    :param file_ext: File extension (e.g., 'pdf') to determine if cleaning is needed
+    :return: Combined text content from all documents
+    """
+    parts: List[str] = []
     if documents:
         for doc in documents:
             if hasattr(doc, "page_content"):
                 # Clean text if it's a PDF
                 if file_ext == "pdf":
-                    text_content += clean_text(doc.page_content) + "\n"
+                    parts.append(clean_text(doc.page_content))
                 else:
-                    text_content += doc.page_content + "\n"
+                    parts.append(doc.page_content)
 
-    # Remove trailing newline
-    return text_content.rstrip("\n")
+    return "\n".join(parts)
 
 
 async def cleanup_temp_file_async(file_path: str) -> None:
@@ -628,19 +636,21 @@ def _prepare_documents_sync(
         for doc in documents:
             doc.page_content = clean_text(doc.page_content)
 
-    # Preparing documents with page content and metadata for insertion.
-    return [
-        Document(
-            page_content=doc.page_content,
-            metadata={
+    # Mutate metadata in-place rather than creating new Document objects.
+    # This optimization reduces peak memory usage by avoiding duplicate allocations
+    # when processing large files with many chunks.
+    for doc in documents:
+        metadata = dict(doc.metadata or {})
+        metadata.update(
+            {
                 "file_id": file_id,
                 "user_id": user_id,
                 "digest": generate_digest(doc.page_content),
-                **(doc.metadata or {}),
-            },
+            }
         )
-        for doc in documents
-    ]
+        doc.metadata = metadata
+
+    return documents
 
 
 async def store_data_in_vector_db(
@@ -838,6 +848,10 @@ async def embed_file(
         )
     finally:
         await cleanup_temp_file_async(temp_file_path)
+        try:
+            await file.close()
+        except Exception:
+            pass
 
     return {
         "status": response_status,
@@ -949,6 +963,10 @@ async def embed_file_upload(
         )
     finally:
         await cleanup_temp_file_async(temp_file_path)
+        try:
+            await uploaded_file.close()
+        except Exception:
+            pass
 
     return {
         "status": True,
@@ -1065,3 +1083,7 @@ async def extract_text_from_file(
             )
     finally:
         await cleanup_temp_file_async(temp_file_path)
+        try:
+            await file.close()
+        except Exception:
+            pass

app/services/mongo_client.py

@@ -7,10 +7,17 @@
 logger = logging.getLogger(__name__)
 
 async def mongo_health_check() -> bool:
+    client = None
     try:
         client = MongoClient(ATLAS_MONGO_DB_URI)
         client.admin.command("ping")
         return True
     except PyMongoError as e:
         logger.error(f"MongoDB health check failed: {e}")
-        return False
+        return False
+    finally:
+        try:
+            if client is not None:
+                client.close()
+        except Exception:
+            pass

app/services/vector_store/atlas_mongo_vector.py

@@ -5,6 +5,21 @@
 from langchain_mongodb import MongoDBAtlasVectorSearch
 
 class AtlasMongoVector(MongoDBAtlasVectorSearch):
+    def close(self) -> None:
+        """
+        Close the underlying MongoDB client connection.
+
+        This method safely closes the MongoClient associated with this vector store
+        to release connection resources and prevent memory leaks.
+        """
+        collection = getattr(self, "_collection", None)
+        try:
+            client = collection.database.client if collection is not None else None
+            if client is not None:
+                client.close()
+        except Exception:
+            pass
+
     @property
     def embedding_function(self) -> Embeddings:
         return self.embeddings

app/services/vector_store/extended_pg_vector.py

@@ -165,3 +165,10 @@ def _delete_multiple(
                 stmt = stmt.where(self.EmbeddingStore.custom_id.in_(ids))
                 session.execute(stmt)
             session.commit()
+
+    def close(self) -> None:
+        """Release underlying SQLAlchemy engine resources."""
+        bind = getattr(self, "_bind", None)
+        dispose = getattr(bind, "dispose", None)
+        if callable(dispose):
+            dispose()

app/services/vector_store/factory.py

@@ -1,4 +1,5 @@
 from typing import Optional
+from functools import lru_cache
 from pymongo import MongoClient
 from langchain_core.embeddings import Embeddings
 
@@ -7,6 +8,20 @@
 from .extended_pg_vector import ExtendedPgVector
 
 
+@lru_cache(maxsize=4)
+def _get_mongo_client(connection_string: str) -> MongoClient:
+    """
+    Get or create a cached MongoClient for the given connection string.
+
+    Caches up to 4 clients to avoid creating new connections on every request,
+    which prevents socket/memory leaks from accumulating MongoClient instances.
+
+    :param connection_string: MongoDB connection URI
+    :return: Cached MongoClient instance
+    """
+    return MongoClient(connection_string)
+
+
 def get_vector_store(
     connection_string: str,
     embeddings: Embeddings,
@@ -27,7 +42,7 @@ def get_vector_store(
             collection_name=collection_name,
         )
     elif mode == "atlas-mongo":
-        mongo_db = MongoClient(connection_string).get_database()
+        mongo_db = _get_mongo_client(connection_string).get_database()
         mong_collection = mongo_db[collection_name]
         return AtlasMongoVector(
             collection=mong_collection, embedding=embeddings, index_name=search_index

app/utils/document_loader.py

@@ -93,8 +93,12 @@ def get_loader(filename: str, file_content_type: str, filepath: str):
                     with open(
                         filepath, "r", encoding=encoding, errors="replace"
                     ) as original_file:
-                        content = original_file.read()
-                        temp_file.write(content)
+                        # Stream conversion to avoid loading large CSVs into memory.
+                        while True:
+                            chunk = original_file.read(1024 * 1024)  # 1 MB
+                            if not chunk:
+                                break
+                            temp_file.write(chunk)
 
                     temp_filepath = temp_file.name
 
@@ -191,7 +195,16 @@ def remove_non_utf8(text: str) -> str:
 
 
 def process_documents(documents: List[Document]) -> str:
-    processed_text = ""
+    """
+    Process a list of documents into a single formatted string.
+
+    Uses a list to collect parts and joins at the end to avoid O(n²)
+    string concatenation overhead for large documents.
+
+    :param documents: List of Document objects to process
+    :return: Formatted string with page markers and content
+    """
+    parts: List[str] = []
     last_page: Optional[int] = None
     doc_basename = ""
 
@@ -200,21 +213,31 @@ def process_documents(documents: List[Document]) -> str:
             doc_basename = doc.metadata["source"].split("/")[-1]
             break
 
-    processed_text += f"{doc_basename}\n"
+    parts.append(f"{doc_basename}\n")
+
+    # Track only the tail to avoid repeatedly scanning a huge string.
+    tail = (parts[-1][-CHUNK_OVERLAP:] if CHUNK_OVERLAP > 0 else "")
 
     for doc in documents:
         current_page = doc.metadata.get("page")
         if current_page and current_page != last_page:
-            processed_text += f"\n# PAGE {doc.metadata['page']}\n\n"
+            header = f"\n# PAGE {doc.metadata['page']}\n\n"
+            parts.append(header)
+            if CHUNK_OVERLAP > 0:
+                tail = (tail + header)[-CHUNK_OVERLAP:]
             last_page = current_page
 
         new_content = doc.page_content
-        if processed_text.endswith(new_content[:CHUNK_OVERLAP]):
-            processed_text += new_content[CHUNK_OVERLAP:]
+        if CHUNK_OVERLAP > 0 and tail == new_content[:CHUNK_OVERLAP]:
+            to_add = new_content[CHUNK_OVERLAP:]
         else:
-            processed_text += new_content
+            to_add = new_content
+
+        parts.append(to_add)
+        if CHUNK_OVERLAP > 0:
+            tail = (tail + to_add)[-CHUNK_OVERLAP:]
 
-    return processed_text.strip()
+    return "".join(parts).strip()
 
 
 class SafePyPDFLoader:

main.py

@@ -18,6 +18,7 @@
     CHUNK_OVERLAP,
     PDF_EXTRACT_IMAGES,
     VECTOR_DB_TYPE,
+    vector_store,
     LogMiddleware,
     logger,
 )
@@ -51,6 +52,22 @@ async def lifespan(app: FastAPI):
     app.state.thread_pool.shutdown(wait=True)
     logger.info("Thread pool shutdown complete")
 
+    if VECTOR_DB_TYPE == VectorDBType.PGVECTOR:
+        try:
+            await PSQLDatabase.close_pool()
+            logger.info("PostgreSQL pool closed")
+        except Exception as e:
+            logger.warning("Failed to close PostgreSQL pool: %s", e)
+
+    # Close vector store resources (SQLAlchemy engine / Mongo client)
+    try:
+        close_fn = getattr(vector_store, "close", None)
+        if callable(close_fn):
+            close_fn()
+            logger.info("Vector store resources closed")
+    except Exception as e:
+        logger.warning("Failed to close vector store resources: %s", e)
+
 
 app = FastAPI(lifespan=lifespan, debug=debug_mode)
 
@@ -79,7 +96,12 @@ async def lifespan(app: FastAPI):
 
 @app.exception_handler(RequestValidationError)
 async def validation_exception_handler(request: Request, exc: RequestValidationError):
+    # Limit logged body size to prevent memory issues with large uploads.
+    # The full body is still read by FastAPI, but we truncate for logging only.
     body = await request.body()
+    max_log_bytes = 64 * 1024  # 64 KB
+    if len(body) > max_log_bytes:
+        body = body[:max_log_bytes] + b"... (truncated)"
     logger.debug(f"Validation error occurred")
     logger.debug(f"Raw request body: {body.decode()}")
     logger.debug(f"Validation errors: {exc.errors()}")

tests/services/test_vector_store.py

@@ -1,4 +1,5 @@
 from app.services.vector_store.extended_pg_vector import ExtendedPgVector
+from app.services.vector_store.atlas_mongo_vector import AtlasMongoVector
 
 # Create a dummy subclass that simulates DB responses.
 class DummyPgVector(ExtendedPgVector):
@@ -12,4 +13,44 @@ def get_all_ids(self) -> list[str]:
 def test_extended_pgvector_get_all_ids():
     dummy_vector = DummyPgVector()
     ids = dummy_vector.get_all_ids()
-    assert ids == ["id1", "id2"]
+    assert ids == ["id1", "id2"]
+
+
+def test_extended_pgvector_close_disposes_engine():
+    class DummyBind:
+        def __init__(self):
+            self.disposed = False
+
+        def dispose(self):
+            self.disposed = True
+
+    vec = DummyPgVector()
+    vec._bind = DummyBind()
+    vec.close()
+    assert vec._bind.disposed is True
+
+
+def test_atlas_mongo_vector_close_closes_client():
+    class DummyClient:
+        def __init__(self):
+            self.closed = False
+
+        def close(self):
+            self.closed = True
+
+    class DummyDatabase:
+        def __init__(self, client):
+            self.client = client
+
+    class DummyCollection:
+        def __init__(self, client):
+            self.database = DummyDatabase(client)
+
+    client = DummyClient()
+    collection = DummyCollection(client)
+
+    # Bypass base class init; we only need _collection for close()
+    vec = AtlasMongoVector.__new__(AtlasMongoVector)
+    vec._collection = collection
+    vec.close()
+    assert client.closed is True