Use chunk index in document_key in ModelSource

Author: tomusher

src/django_ai_core/contrib/index/source.py

@@ -175,8 +175,8 @@ def provides_object(self, obj: object) -> bool:
     def provides_document(self, document: Document) -> bool:
         return document.document_key.split(":")[0] == self.source_id
 
-    def get_document_key(self, obj) -> str:
-        return f"{self.source_id}:{obj.pk}"
+    def get_document_key(self, obj, chunk) -> str:
+        return f"{self.source_id}:{obj.pk}:{chunk}"
 
     def _object_to_documents(self, obj: models.Model) -> Iterable[Document]:
         if not self.provides_object(obj):
@@ -185,9 +185,9 @@ def _object_to_documents(self, obj: models.Model) -> Iterable[Document]:
         metadata = self.get_metadata(obj)
         content = self.get_content(obj)
 
-        for document in self.chunk_transformer.transform(content):
+        for chunk, document in enumerate(self.chunk_transformer.transform(content)):
             yield Document(
-                document_key=self.get_document_key(obj),
+                document_key=self.get_document_key(obj, chunk),
                 content=document,
                 metadata=metadata,
             )

tests/unit/contrib/index/test_model_source.py

@@ -0,0 +1,66 @@
+import random
+
+import pytest
+from testapp.models import Book
+
+from django_ai_core.contrib.index.source import ModelSource
+
+
+def generate_long_string(length=2000):
+    words = [
+        "the",
+        "quick",
+        "brown",
+        "fox",
+        "jumps",
+        "over",
+        "lazy",
+        "dog",
+        "artificial",
+        "intelligence",
+        "machine",
+        "learning",
+        "data",
+        "science",
+        "python",
+        "django",
+        "wagtail",
+        "testing",
+        "development",
+        "software",
+        "application",
+        "framework",
+        "database",
+        "model",
+        "query",
+        "index",
+        "search",
+        "semantic",
+        "vector",
+        "embedding",
+        "algorithm",
+        "neural",
+        "network",
+        "training",
+        "optimization",
+        "performance",
+        "scalability",
+    ]
+    text = []
+    current_length = 0
+    while current_length < length:
+        word = random.choice(words)
+        text.append(word)
+        current_length += len(word) + 1
+    return " ".join(text).capitalize() + "."
+
+
+@pytest.mark.django_db
+def test_model_source_returns_unique_keys():
+    for _ in range(5):
+        Book.objects.create(title="Book Title", description=generate_long_string())
+
+    model_source = ModelSource(model=Book)
+    documents = model_source.get_documents()
+    document_keys = [doc.document_key for doc in documents]
+    assert len(document_keys) == len(set(document_keys))