[DI-133] feat(document-index): integrate chunk overlap configuration

- you can now specify a chunk overlap when configuring an index
- add validation to ensure 0 <= overlap < chunk_size
- add validation to ensure chunk_size is specified and valid

CHANGELOG.md

@@ -7,6 +7,7 @@
 
 ### Features
 - Add `StudioClient` as connector to PhariaStudio for submitting traces.
+- You can now specify a `chunk_overlap` when creating an index in the Document Index.
 
 ### Fixes
 ...

src/intelligence_layer/connectors/document_index/document_index.py

@@ -6,7 +6,7 @@
 from urllib.parse import quote
 
 import requests
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, model_validator
 from requests import HTTPError
 
 from intelligence_layer.connectors.base.json_serializable import JsonSerializable
@@ -29,11 +29,20 @@ class IndexConfiguration(BaseModel):
 
     Args:
         embedding_type: "symmetric" or "asymmetric" embedding type.
+        chunk_overlap: The maximum number of tokens of overlap between consecutive chunks. Must be
+            less than `chunk_size`.
         chunk_size: The maximum size of the chunks in tokens to be used for the index.
     """
 
     embedding_type: Literal["symmetric", "asymmetric"]
-    chunk_size: int
+    chunk_overlap: int = Field(default=0, ge=0)
+    chunk_size: int = Field(..., gt=0, le=2046)
+
+    @model_validator(mode="after")
+    def check_chunk_values(cls, values):
+        if not values.chunk_overlap < values.chunk_size:
+            raise ValueError('chunk_overlap must be less than chunk_size')
+        return values
 
 
 class DocumentContents(BaseModel):
@@ -427,6 +436,7 @@ def index_configuration(self, index_path: IndexPath) -> IndexConfiguration:
         response_json: Mapping[str, Any] = response.json()
         return IndexConfiguration(
             embedding_type=response_json["embedding_type"],
+            chunk_overlap=response_json["chunk_overlap"],
             chunk_size=response_json["chunk_size"],
         )
 

tests/connectors/document_index/test_document_index.py

@@ -1,6 +1,7 @@
 from http import HTTPStatus
 
 import pytest
+from pydantic import ValidationError
 from pytest import fixture, raises
 
 from intelligence_layer.connectors.document_index.document_index import (
@@ -9,6 +10,7 @@
     DocumentFilterQueryParams,
     DocumentIndexClient,
     DocumentPath,
+    IndexConfiguration,
     IndexPath,
     ResourceNotFound,
     SearchQuery,
@@ -232,6 +234,17 @@ def test_document_path_is_immutable() -> None:
 
     assert dictionary[path] == 1
 
+def test_index_configuration_rejects_invalid_chunk_overlap() -> None:
+    try:
+        IndexConfiguration(
+                chunk_size=128,
+                chunk_overlap=128,
+                embedding_type="asymmetric"
+        )
+    except ValidationError as e:
+        assert 'chunk_overlap must be less than chunk_size' in str(e)
+    else:
+        raise AssertionError("ValidationError was not raised")
 
 def test_document_indexes_are_returned(
     document_index: DocumentIndexClient, collection_path: CollectionPath
@@ -243,4 +256,5 @@ def test_document_indexes_are_returned(
     )
 
     assert index_configuration.embedding_type == "asymmetric"
+    assert index_configuration.chunk_overlap == 0
     assert index_configuration.chunk_size == 512