diff --git a/src/uipath_langchain/agent/tools/context_tool.py b/src/uipath_langchain/agent/tools/context_tool.py
index ff5fd128b..a27b667dd 100644
--- a/src/uipath_langchain/agent/tools/context_tool.py
+++ b/src/uipath_langchain/agent/tools/context_tool.py
@@ -273,7 +273,8 @@ async def context_tool_fn(
             "documents": [
                 {"metadata": doc.metadata, "page_content": doc.page_content}
                 for doc in docs
-            ]
+            ],
+            "index_id": retriever.index_id,
         }
 
     if arg_props or has_arg_folder:
diff --git a/src/uipath_langchain/retrievers/context_grounding_retriever.py b/src/uipath_langchain/retrievers/context_grounding_retriever.py
index dc448c6ad..5a9b51a1e 100644
--- a/src/uipath_langchain/retrievers/context_grounding_retriever.py
+++ b/src/uipath_langchain/retrievers/context_grounding_retriever.py
@@ -17,6 +17,7 @@ class ContextGroundingRetriever(BaseRetriever):
     threshold: float = 0.0
     scope_folder: str | None = None
     scope_extension: str | None = None
+    index_id: str | None = None
 
     def _build_scope(self) -> UnifiedSearchScope | None:
         if self.scope_folder or self.scope_extension:
@@ -44,6 +45,7 @@ def _get_relevant_documents(
             folder_key=self.folder_key,
         )
 
+        self.index_id = result.index_id
         values = result.semantic_results.values if result.semantic_results else []
 
         return [
@@ -80,6 +82,7 @@ async def _aget_relevant_documents(
             folder_key=self.folder_key,
         )
 
+        self.index_id = result.index_id
         values = result.semantic_results.values if result.semantic_results else []
 
         return [