Repository: kaito-project/kaito

Issue:
RAGEngine Index Document API returns wrong doc_id's
**Describe the bug**
When calling `/index` with multiple documents you can hit a race condition where doc_id's [attached to documents on the response](https://github.com/kaito-project/kaito/blob/bb4635cfc63a78529bf35653e0e051f48b739aad/presets/ragengine/main.py#L99) are not the correct id's of the documents. Likely due to [this async gather](https://github.com/kaito-project/kaito/blob/main/presets/ragengine/vector_store/base.py#L77)

**Steps To Reproduce**
Debug the service and hit:

```
curl -X POST -H "Content-Type: application/json" -d '{"index_name":"test_index","documents":[{"text": "This is a test document"},{"text": "Another test document"}]}' localhost:5000/index
```
**Expected behavior**

Document Id's are for correct document

**Logs**

**Environment**

- Kubernetes version (use `kubectl version`): N/A
- OS (e.g: `cat /etc/os-release`): MAC OS 
- Install tools:
- Others:

**Additional context**
In this attached example, you an see the response from the server vs the docstore containing the documents show different ID's

![Image](https://github.com/user-attachments/assets/d4d91e1c-278a-4cab-8772-53bca6c3c8b5)

Discussion:
@kaito-pr-agent: ### :bulb: Auto-generated documentation-based answer:
Kaito is an operator that automates the AI/ML model inference or tuning workload in a Kubernetes cluster. It supports popular open-sourced large models and provides preset configurations, auto-provisioning of GPU nodes, and support for inference runtimes like vLLM and transformers.

#### Relevant Sources:

> - https://github.com/kaito-project/kaito/blob/main/README.md


Solution patch:
diff --git a/presets/ragengine/vector_store/base.py b/presets/ragengine/vector_store/base.py
index e493495..75468b0 100644
--- a/presets/ragengine/vector_store/base.py
+++ b/presets/ragengine/vector_store/base.py
@@ -58,9 +58,9 @@ class BaseVectorStore(ABC):
     async def _append_documents_to_index(self, index_name: str, documents: List[Document]) -> List[str]:
         """Common logic for appending documents to existing index."""
         logger.info(f"Index {index_name} already exists. Appending documents to existing index.")
-        indexed_doc_ids = set()
+        indexed_docs = [None] * len(documents)
 
-        async def handle_document(doc: Document):
+        async def handle_document(doc_index: int, doc: Document):
             doc_id = self.generate_doc_id(doc.text)
             if self.use_rwlock:
                 async with self.rwlock.reader_lock:
@@ -69,13 +69,13 @@ class BaseVectorStore(ABC):
                 retrieved_doc = await self.index_map[index_name].docstore.aget_ref_doc_info(doc_id)
             if not retrieved_doc:
                 await self.add_document_to_index(index_name, doc, doc_id)
-                indexed_doc_ids.add(doc_id)
             else:
                 logger.info(f"Document {doc_id} already exists in index {index_name}. Skipping.")
+            indexed_docs[doc_index] = doc_id
 
         # Gather all coroutines for processing documents
-        await asyncio.gather(*(handle_document(doc) for doc in documents))
-        return list(indexed_doc_ids)
+        await asyncio.gather(*(handle_document(idx, doc) for idx, doc in enumerate(documents)))
+        return indexed_docs
     
     @abstractmethod
     async def _create_new_index(self, index_name: str, documents: List[Document]) -> List[str]:



Test patch:
diff --git a/presets/ragengine/tests/vector_store/test_base_store.py b/presets/ragengine/tests/vector_store/test_base_store.py
index 49b17e6..8b761d6 100644
--- a/presets/ragengine/tests/vector_store/test_base_store.py
+++ b/presets/ragengine/tests/vector_store/test_base_store.py
@@ -108,6 +108,28 @@ class BaseVectorStoreTest(ABC):
         assert await vector_store_manager.document_exists("test_index", new_document[0],
                                                     BaseVectorStore.generate_doc_id("Fourth document"))
 
+    @pytest.mark.asyncio
+    async def test_add_document_on_existing_index(self, vector_store_manager):
+        # Create index with single doc
+        await vector_store_manager.index_documents("test_add_index", [Document(text=f"Initial Doc", metadata={"type": "text"})])
+
+        documents = [
+            Document(text=f"Document {i}", metadata={"type": "text"})
+            for i in range(10)
+        ]
+        # load documents into the same index
+        ids = await vector_store_manager.index_documents("test_add_index", documents)
+
+        # Check if the documents were indexed correctly
+        all_docs = await vector_store_manager.list_documents_in_index("test_add_index", limit=10, offset=1)
+        print(all_docs)
+
+        for idx, doc_id in enumerate(ids):
+            # list_documents_in_index returns documents with doc_id's that are not the same as the id's returned from index_documents
+            response_doc = [doc for doc in all_docs if doc['doc_id'] == vector_store_manager.index_map['test_add_index'].docstore.get_ref_doc_info(doc_id).node_ids[0]][0] or None
+            assert response_doc is not None
+            assert documents[idx].text == response_doc['text']
+
     @pytest.mark.asyncio
     async def test_persist_index(self, vector_store_manager):
         documents = [Document(text="Test document", metadata={"type": "text"})]