refactor: Fix typo and rename attribute in MultipleChunkRetrieverQa (…

…#738)
Aleph-Alpha · Apr 15, 2024 · 3bf8667 · 3bf8667
1 parent 4889aae
commit 3bf8667
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/src/intelligence_layer/use_cases/__init__.py b/src/intelligence_layer/use_cases/__init__.py
@@ -50,10 +50,10 @@
 from .qa.multiple_chunk_qa import MultipleChunkQaOutput as MultipleChunkQaOutput
 from .qa.multiple_chunk_qa import Subanswer as Subanswer
 from .qa.multiple_chunk_retriever_qa import (
-    MulMultipleChunkRetrieverQaOutput as MulMultipleChunkRetrieverQaOutput,
+    MultipleChunkRetrieverQa as MultipleChunkRetrieverQa,
 )
 from .qa.multiple_chunk_retriever_qa import (
-    MultipleChunkRetrieverQa as MultipleChunkRetrieverQa,
+    MultipleChunkRetrieverQaOutput as MultipleChunkRetrieverQaOutput,
 )
 from .qa.retriever_based_qa import EnrichedSubanswer as EnrichedSubanswer
 from .qa.retriever_based_qa import RetrieverBasedQa as RetrieverBasedQa

diff --git a/src/intelligence_layer/use_cases/qa/multiple_chunk_retriever_qa.py b/src/intelligence_layer/use_cases/qa/multiple_chunk_retriever_qa.py
@@ -18,17 +18,17 @@
 
 
 class AnswerSource(BaseModel, Generic[ID]):
-    chunk: SearchResult[ID]
+    search_result: SearchResult[ID]
     highlights: Sequence[ScoredTextHighlight]
 
 
-class MulMultipleChunkRetrieverQaOutput(BaseModel, Generic[ID]):
+class MultipleChunkRetrieverQaOutput(BaseModel, Generic[ID]):
     answer: Optional[str]
     sources: Sequence[AnswerSource[ID]]
 
 
 class MultipleChunkRetrieverQa(
-    Task[RetrieverBasedQaInput, MulMultipleChunkRetrieverQaOutput[ID]], Generic[ID]
+    Task[RetrieverBasedQaInput, MultipleChunkRetrieverQaOutput[ID]], Generic[ID]
 ):
     """Answer a question based on documents found by a retriever.
 
@@ -113,7 +113,7 @@ def _get_highlights_per_chunk(
 
     def do_run(
         self, input: RetrieverBasedQaInput, task_span: TaskSpan
-    ) -> MulMultipleChunkRetrieverQaOutput[ID]:
+    ) -> MultipleChunkRetrieverQaOutput[ID]:
         search_output = self._search.run(
             SearchInput(query=input.question), task_span
         ).results
@@ -140,11 +140,11 @@ def do_run(
             chunk_start_indices, single_chunk_qa_output.highlights
         )
 
-        return MulMultipleChunkRetrieverQaOutput(
+        return MultipleChunkRetrieverQaOutput(
             answer=single_chunk_qa_output.answer,
             sources=[
                 AnswerSource(
-                    chunk=chunk,
+                    search_result=chunk,
                     highlights=highlights,
                 )
                 for chunk, highlights in zip(sorted_search_output, highlights_per_chunk)