From f52f1d6b6878415078cd499ebc1bcd2775163ff7 Mon Sep 17 00:00:00 2001
From: Niklas Finken <71568758+NickyHavoc@users.noreply.github.com>
Date: Mon, 22 Apr 2024 10:38:06 +0200
Subject: [PATCH] fix: `DocumentPath` & `CollectionPath` are immutable (#768)

* fix: `DocumentPath` & `CollectionPath` are immutable

* adjust changelog
---
 CHANGELOG.md                                           |  2 +-
 .../connectors/document_index/document_index.py        |  4 ++--
 tests/connectors/document_index/test_document_index.py | 10 ++++++++++
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b6c3e3e0f..25a10c3e4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,11 +6,11 @@
 - breaking change: `ExpandChunksOutput` now returns `ChunkWithStartEndIndices` instead of `TextChunk`
 - breaking change: `MultipleChunkRetrieverQa`'s `AnswerSource` now contains `EnrichedChunk` instead of just the `TextChunk`
 
-
 ### New Features
 
 ### Fixes
 - fix: `ChunkWithIndices` now additionally returns end_index
+- fix: `DocumentPath` and `CollectionPath` are now immutable
 
 
 ## 0.9.1
diff --git a/src/intelligence_layer/connectors/document_index/document_index.py b/src/intelligence_layer/connectors/document_index/document_index.py
index e4d85aff7..57db24fe1 100644
--- a/src/intelligence_layer/connectors/document_index/document_index.py
+++ b/src/intelligence_layer/connectors/document_index/document_index.py
@@ -51,7 +51,7 @@ def _to_modalities_json(self) -> Mapping[str, Any]:
         }
 
 
-class CollectionPath(BaseModel):
+class CollectionPath(BaseModel, frozen=True):
     """Path to a collection.
 
     Args:
@@ -64,7 +64,7 @@ class CollectionPath(BaseModel):
     collection: str
 
 
-class DocumentPath(BaseModel):
+class DocumentPath(BaseModel, frozen=True):
     """Path to a document.
 
     Args:
diff --git a/tests/connectors/document_index/test_document_index.py b/tests/connectors/document_index/test_document_index.py
index 36209ccbb..80a27b636 100644
--- a/tests/connectors/document_index/test_document_index.py
+++ b/tests/connectors/document_index/test_document_index.py
@@ -178,3 +178,13 @@ def test_document_list_documents_with_matching_prefix(
 
     assert len(filter_result) == 1
     assert filter_result[0].document_path.document_name.startswith(prefix)
+
+
+def test_document_path_is_immutable() -> None:
+    path = DocumentPath(
+        collection_path=CollectionPath(namespace="1", collection="2"), document_name="3"
+    )
+    dictionary = {}
+    dictionary[path] = 1
+
+    assert dictionary[path] == 1