WIP: start and end index for retriever SearchResult

Aleph-Alpha · Apr 5, 2024 · 8614af3 · 8614af3
1 parent c48e0be
commit 8614af3
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/src/intelligence_layer/connectors/retrievers/base_retriever.py b/src/intelligence_layer/connectors/retrievers/base_retriever.py
@@ -22,10 +22,14 @@ class DocumentChunk(BaseModel):
     Attributes:
         text: Chunk of the document that matched the search query.
         metadata: Any metadata added to the document.
+        start: Start index of the chunk within the document
+        end: End index of the chunk within the document
     """
 
     text: str
     metadata: Any = None
+    start: int
+    end: int
 
 
 ID = TypeVar("ID")

diff --git a/tests/use_cases/search/test_search.py b/tests/use_cases/search/test_search.py
@@ -35,3 +35,5 @@ def test_search(
     assert [to_document(r.document_chunk) for r in result.results] == [
         in_memory_retriever_documents[2]
     ]
+    assert result.results[0].document_chunk.start == 0
+    assert result.results[0].document_chunk.end == len(in_memory_retriever_documents[2].text) - 1