From adcdce8d764ca4de31009af52487fe40581156fd Mon Sep 17 00:00:00 2001
From: Siddhant Rai <siddhant.rai.5686@gmail.com>
Date: Wed, 10 Sep 2025 22:10:11 +0530
Subject: [PATCH] fix: handle invalid chunks value in StreamProcessor and
 ClassicRAG

---
 .../api/answer/services/stream_processor.py      | 16 ++++++++++++++--
 application/retriever/classic_rag.py             | 11 ++++++++++-
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/application/api/answer/services/stream_processor.py b/application/api/answer/services/stream_processor.py
index a04020cb..f6e639ef 100644
--- a/application/api/answer/services/stream_processor.py
+++ b/application/api/answer/services/stream_processor.py
@@ -266,7 +266,13 @@ class StreamProcessor:
             if data_key.get("retriever"):
                 self.retriever_config["retriever_name"] = data_key["retriever"]
             if data_key.get("chunks") is not None:
-                self.retriever_config["chunks"] = data_key["chunks"]
+                try:
+                    self.retriever_config["chunks"] = int(data_key["chunks"])
+                except (ValueError, TypeError):
+                    logger.warning(
+                        f"Invalid chunks value: {data_key['chunks']}, using default value 2"
+                    )
+                    self.retriever_config["chunks"] = 2
         elif self.agent_key:
             data_key = self._get_data_from_api_key(self.agent_key)
             self.agent_config.update(
@@ -287,7 +293,13 @@ class StreamProcessor:
             if data_key.get("retriever"):
                 self.retriever_config["retriever_name"] = data_key["retriever"]
             if data_key.get("chunks") is not None:
-                self.retriever_config["chunks"] = data_key["chunks"]
+                try:
+                    self.retriever_config["chunks"] = int(data_key["chunks"])
+                except (ValueError, TypeError):
+                    logger.warning(
+                        f"Invalid chunks value: {data_key['chunks']}, using default value 2"
+                    )
+                    self.retriever_config["chunks"] = 2
         else:
             self.agent_config.update(
                 {
diff --git a/application/retriever/classic_rag.py b/application/retriever/classic_rag.py
index ce1b937b..2ce863c2 100644
--- a/application/retriever/classic_rag.py
+++ b/application/retriever/classic_rag.py
@@ -25,7 +25,16 @@ class ClassicRAG(BaseRetriever):
         self.original_question = source.get("question", "")
         self.chat_history = chat_history if chat_history is not None else []
         self.prompt = prompt
-        self.chunks = chunks
+        if isinstance(chunks, str):
+            try:
+                self.chunks = int(chunks)
+            except ValueError:
+                logging.warning(
+                    f"Invalid chunks value '{chunks}', using default value 2"
+                )
+                self.chunks = 2
+        else:
+            self.chunks = chunks
         self.gpt_model = gpt_model
         self.token_limit = (
             token_limit