Fix: Resolve OOM DoS via payload truncation/memory bounds and optimize FAISS cache

Author: Adar5

chatbot-core/api/models/embedding_model.py

@@ -6,4 +6,5 @@
 
 logger = LoggerFactory.instance().get_logger("api")
 
-EMBEDDING_MODEL = load_embedding_model(CONFIG["retrieval"]["embedding_model_name"], logger)
+EMBEDDING_MODEL = load_embedding_model(
+    CONFIG["retrieval"]["embedding_model_name"])

chatbot-core/api/routes/chatbot.py

@@ -17,6 +17,7 @@
 # Third-party imports
 # =========================
 from typing import List, Optional
+from urllib import request
 from fastapi import (
     APIRouter,
     HTTPException,
@@ -106,13 +107,19 @@ async def chatbot_stream(websocket: WebSocket, session_id: str):
             message_data = json.loads(data)
             user_message = message_data.get("message", "")
 
+            if len(user_message) > 2000:
+                logger.warning(
+                    f"Truncated massive WebSocket payload from session {session_id}")
+                user_message = user_message[:2000]
+
             if not user_message:
                 continue
 
             async for token in get_chatbot_reply_stream(
                 session_id,
                 user_message,
             ):
+
                 await websocket.send_text(
                     json.dumps({"token": token})
                 )
@@ -166,6 +173,7 @@ def start_chat(response: Response):
     )
     return SessionResponse(session_id=session_id)
 
+
 @router.delete(
     "/sessions/{session_id}",
     response_model=DeleteResponse,
@@ -191,7 +199,6 @@ def delete_chat(session_id: str):
 # Chat Endpoint
 @router.post("/sessions/{session_id}/message", response_model=ChatResponse)
 def chatbot_reply(session_id: str, request: ChatRequest, _background_tasks: BackgroundTasks):
-
     """
     POST endpoint to handle chatbot replies.
 
@@ -210,11 +217,16 @@ def chatbot_reply(session_id: str, request: ChatRequest, _background_tasks: Back
             status_code=404,
             detail="Session not found.",
         )
-    reply =  get_chatbot_reply(session_id, request.message)
+
+    if len(request.message) > 2000:
+        logger.warning(f"Truncated massive payload from session {session_id}")
+        request.message = request.message[:2000]
+
+    reply = get_chatbot_reply(session_id, request.message)
     _background_tasks.add_task(
         persist_session,
         session_id,
-        )
+    )
 
     return reply
 
@@ -263,6 +275,10 @@ async def chatbot_reply_with_files(
             status_code=422,
             detail="Either message or files must be provided.",
         )
+    if has_message and len(message) > 2000:
+        logger.warning(
+            f"Truncated massive file upload message from session {session_id}")
+        message = message[:2000]
 
     # Process uploaded files
     processed_files: List[FileAttachment] = []

chatbot-core/api/services/chat_service.py

@@ -62,7 +62,8 @@ def get_chatbot_reply(
 
     memory = get_session(session_id)
     if memory is None:
-        raise RuntimeError(f"Session '{session_id}' not found in the memory store.")
+        raise RuntimeError(
+            f"Session '{session_id}' not found in the memory store.")
 
     context = retrieve_context(user_input)
     logger.info("Context retrieved: %s", context)
@@ -333,7 +334,8 @@ def _execute_search_tools(tool_calls) -> str:
         })
 
     return "\n\n".join(
-        f"[Result of the search tool {res['tool']}]:\n{res.get('output', '')}".strip()
+        f"[Result of the search tool {res['tool']}]:\n{res.get('output', '')}".strip(
+        )
         for res in retrieved_results
     )
 
@@ -381,7 +383,6 @@ def retrieve_context(user_input: str) -> str:
     data_retrieved, _ = get_relevant_documents(
         user_input,
         EMBEDDING_MODEL,
-        logger=logger,
         source_name="plugins",
         top_k=retrieval_config["top_k"]
     )
@@ -434,10 +435,12 @@ def generate_answer(prompt: str, max_tokens: Optional[int] = None) -> str:
         logger.error("LLM provider unavailable: %s", e)
         return "LLM is not available. Please install llama-cpp-python and configure a model."
     except (ValueError, RuntimeError) as exc:
-        logger.error("LLM generation failed for prompt: %r. Error: %r", prompt, exc)
+        logger.error(
+            "LLM generation failed for prompt: %r. Error: %r", prompt, exc)
         return "Sorry, I'm having trouble generating a response right now."
     except Exception:  # pylint: disable=broad-except
-        logger.exception("Unexpected error during LLM generation for prompt: %r", prompt)
+        logger.exception(
+            "Unexpected error during LLM generation for prompt: %r", prompt)
         return "Sorry, an unexpected error occurred. Please contact support."
 
 
@@ -532,6 +535,7 @@ def _extract_relevance_score(response: str) -> str:
 
     return relevance_score
 
+
 def _generate_search_query_from_logs(log_text: str) -> str:
     """
     Uses the LLM to extract a concise error signature from the logs

chatbot-core/api/services/memory.py

@@ -6,9 +6,10 @@
 import uuid
 from datetime import datetime, timedelta
 from threading import Lock
-from langchain.memory import ConversationBufferMemory
+from typing import Optional
+from langchain.memory import ConversationBufferWindowMemory
 from api.config.loader import CONFIG
-from api.services.sessionmanager import(
+from api.services.sessionmanager import (
     delete_session_file,
     load_session,
     session_exists_in_json,
@@ -31,13 +32,13 @@ def init_session() -> str:
     session_id = str(uuid.uuid4())
     with _lock:
         _sessions[session_id] = {
-            "memory": ConversationBufferMemory(return_messages=True),
+            "memory": ConversationBufferWindowMemory(k=10, return_messages=True),
             "last_accessed": datetime.now()
         }
     return session_id
 
 
-def get_session(session_id: str) -> ConversationBufferMemory | None:
+def get_session(session_id: str) -> Optional[ConversationBufferWindowMemory]:
     """
     Retrieve the chat session memory for the given session ID.
     Lazily restores from disk if missing in memory.
@@ -46,24 +47,24 @@ def get_session(session_id: str) -> ConversationBufferMemory | None:
         session_id (str): The session identifier.
 
     Returns:
-        ConversationBufferMemory | None: The memory object if found, else None.
+        Optional[ConversationBufferWindowMemory]: The memory object if found, else None.
     """
 
     with _lock:
 
         session_data = _sessions.get(session_id)
 
-        if session_data :
+        if session_data:
             session_data["last_accessed"] = datetime.now()
             return session_data["memory"]
 
         history = load_session(session_id)
         if not history:
             return None
 
-        memory = ConversationBufferMemory(return_messages=True)
+        memory = ConversationBufferWindowMemory(k=10, return_messages=True)
         for msg in history:
-            memory.chat_memory.add_message(# pylint: disable=no-member
+            memory.chat_memory.add_message(  # pylint: disable=no-member
                 {
                     "role": msg["role"],
                     "content": msg["content"],
@@ -77,14 +78,15 @@ def get_session(session_id: str) -> ConversationBufferMemory | None:
 
         return memory
 
-async def get_session_async(session_id: str) -> ConversationBufferMemory | None:
+
+async def get_session_async(session_id: str) -> Optional[ConversationBufferWindowMemory]:
     """
     Async wrapper for get_session to prevent event loop blocking.
     """
     return await asyncio.to_thread(get_session, session_id)
 
 
-def persist_session(session_id: str)-> None:
+def persist_session(session_id: str) -> None:
     """
     Persist the current session messages to disk.
 
@@ -97,7 +99,6 @@ def persist_session(session_id: str)-> None:
         append_message(session_id, messages)
 
 
-
 def delete_session(session_id: str) -> bool:
     """
     Delete a chat session and its persisted data.
@@ -138,7 +139,8 @@ def reset_sessions():
     with _lock:
         _sessions.clear()
 
-def get_last_accessed(session_id: str) -> datetime | None:
+
+def get_last_accessed(session_id: str) -> Optional[datetime]:
     """
     Get the last accessed timestamp for a given session.
 
@@ -157,9 +159,9 @@ def get_last_accessed(session_id: str) -> datetime | None:
         if not history:
             return None
 
-
     return history["last_accessed"]
 
+
 def set_last_accessed(session_id: str, timestamp: datetime) -> bool:
     """
     Set the last accessed timestamp for a given session (for testing purposes).
@@ -186,6 +188,7 @@ def set_last_accessed(session_id: str, timestamp: datetime) -> bool:
 
     return False
 
+
 def get_session_count() -> int:
     """
     Get the total number of active sessions (for testing purposes).
@@ -196,6 +199,7 @@ def get_session_count() -> int:
     with _lock:
         return len(_sessions)
 
+
 def cleanup_expired_sessions() -> int:
     """
     Remove sessions that have not been accessed within the configured timeout period.

chatbot-core/conftest.py

@@ -0,0 +1 @@
+pytest_plugins = ["tests.unit.mocks.test_env"]

chatbot-core/rag/embedding/embedding_utils.py

@@ -3,8 +3,12 @@
 """
 
 from sentence_transformers import SentenceTransformer
+import logging
 
-def load_embedding_model(model_name, logger):
+logger = logging.getLogger(__name__)
+
+
+def load_embedding_model(model_name):
     """
     Load the sentence transformer model for generating text embeddings.
 
@@ -14,7 +18,8 @@ def load_embedding_model(model_name, logger):
     logger.info(f"Loading embedding model: {model_name}")
     return SentenceTransformer(model_name)
 
-def embed_documents(texts, model, logger, batch_size=32):
+
+def embed_documents(texts, model, batch_size=32):
     """
     Embed a list of text documents into dense vector representations using the given model.
 

chatbot-core/rag/retriever/retrieve.py

@@ -5,8 +5,12 @@
 from rag.embedding.embedding_utils import embed_documents
 from rag.retriever.retriever_utils import load_vector_index, search_index
 from api.config.loader import CONFIG
+import logging
 
-def get_relevant_documents(query, model, logger, source_name, top_k=5):
+logger = logging.getLogger(__name__)
+
+
+def get_relevant_documents(query, model, source_name, top_k=5):
     """
     Retrieve the top-k most relevant chunks for a given natural language query.
 
@@ -24,13 +28,13 @@ def get_relevant_documents(query, model, logger, source_name, top_k=5):
         logger.warning("Empty query received.")
         return [], []
 
-    index, metadata = load_vector_index(logger, source_name)
+    index, metadata = load_vector_index(source_name)
 
     if not index or not metadata:
         return [], []
 
-    query_vector = embed_documents([query], model, logger)[0]
-    data, scores = search_index(query_vector, index, metadata, logger, top_k)
+    query_vector = embed_documents([query], model)[0]
+    data, scores = search_index(query_vector, index, metadata, top_k)
 
     filtered = [(d, s) for d, s in zip(data, scores)
                 if s <= CONFIG["retrieval"]["semantic_threshold"]]

chatbot-core/rag/retriever/retriever_utils.py

@@ -3,13 +3,21 @@
 to retrieve relevant document chunks based on a query vector.
 """
 
+
 import os
 import numpy as np
+from functools import lru_cache
+import logging
 from rag.vectorstore.vectorstore_utils import load_faiss_index, load_metadata
 
-VECTOR_STORE_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "data", "embeddings")
+VECTOR_STORE_DIR = os.path.join(os.path.dirname(
+    __file__), "..", "..", "data", "embeddings")
+
+logger = logging.getLogger(__name__)
 
-def load_vector_index(logger, source_name):
+
+@lru_cache(maxsize=1)
+def load_vector_index(source_name):
     """
     Load the FAISS index and associated metadata from disk.
 
@@ -24,14 +32,16 @@ def load_vector_index(logger, source_name):
         logger.warning("No source name provided. Returning empty results.")
         return [], []
     index_path = os.path.join(VECTOR_STORE_DIR, f"{source_name}_index.idx")
-    metadata_path = os.path.join(VECTOR_STORE_DIR, f"{source_name}_metadata.pkl")
+    metadata_path = os.path.join(
+        VECTOR_STORE_DIR, f"{source_name}_metadata.pkl")
 
-    index = load_faiss_index(index_path, logger)
-    metadata = load_metadata(metadata_path, logger)
+    index = load_faiss_index(index_path)
+    metadata = load_metadata(metadata_path)
 
     return index, metadata
 
-def search_index(query_vector, index, metadata, logger, top_k):
+
+def search_index(query_vector, index, metadata, top_k):
     """
     Search the FAISS index with a query vector and return the top-k closest metadata results.
 
@@ -54,7 +64,7 @@ def search_index(query_vector, index, metadata, logger, top_k):
 
     if index.ntotal != len(metadata):
         logger.warning(
-            "Index contains %d vectors but metadata has %d entries." \
+            "Index contains %d vectors but metadata has %d entries."
             " Some results may be missing or inconsistent.",
             index.ntotal,
             len(metadata)
@@ -73,9 +83,9 @@ def search_index(query_vector, index, metadata, logger, top_k):
             })
         else:
             logger.error("FAISS returned index %d out of range (metadata size: %d)",
-                idx,
-                len(metadata)
-            )
+                         idx,
+                         len(metadata)
+                         )
 
     data = []
     scores = []