feat(controller): use better abstract methods

Author: ThePyProgrammer

src/modular_search/controller/codebase.py

@@ -1,7 +1,8 @@
-from typing import List
+from typing import List, Dict
 from modular_search.controller.core import SearchController
 from modular_search.blocks.codebase import CodebaseSearchBlock, CodebaseSearchResult
 
+
 class CodebaseSearchController(SearchController[CodebaseSearchResult]):
     """
     Codebase Search Controller Class
@@ -12,7 +13,15 @@ def __init__(self, search_block: CodebaseSearchBlock):
             "CodebaseSearchBlock": search_block
         })
 
-    def search(self, query: str) -> List[CodebaseSearchResult]:
-        # since it's just one block, it makes our life easier
-        all_results = super().internal_search(query, ["CodebaseSearchBlock"])
-        return all_results["CodebaseSearchBlock"]
+    def select_blocks(self, query: str) -> List[str]:
+        """
+        Selects the unit search blocks to be used for the given query.
+        For CodebaseSearchController, we always use the CodebaseSearchBlock.
+        """
+        return ["CodebaseSearchBlock"]
+    
+    def aggregate(self, search_results: Dict[str, List[CodebaseSearchResult]]) -> List[CodebaseSearchResult]:
+        """ Aggregates the search results from the CodebaseSearchBlock.
+        Since we only have one block, we can return the results directly.
+        """
+        return search_results["CodebaseSearchBlock"]

src/modular_search/controller/core.py

@@ -76,16 +76,17 @@ def __init__(self, unit_blocks: Union[Dict[str, UnitSearchBlock[O]], List[UnitSe
             unit_blocks = block_dict
 
         self.unit_blocks = unit_blocks
-    
+        
     @abstractmethod
-    def search(self, query: str) -> List[O]:
+    def select_blocks(self, query: str) -> List[str]:
+        """
+        Selects the unit search blocks to be used for the given query.
+        This method should be implemented by subclasses to define
+        how blocks are selected based on the query.
+        """
         pass
 
-    def internal_search(self, query: str, active_blocks: List[str]) -> Dict[str, List[O]]:
-        missing_blocks = set(active_blocks) - set(self.unit_blocks.keys())
-        if missing_blocks:
-            raise ValueError(f"Active blocks {missing_blocks} are not registered in the controller.")
-
+    def block_search(self, query: str, active_blocks: List[str]) -> Dict[str, List[O]]:
         # Dispatch to active Unit Search Blocks
         all_results = {}
         for block_name in active_blocks:
@@ -94,4 +95,21 @@ def internal_search(self, query: str, active_blocks: List[str]) -> Dict[str, Lis
             all_results[block_name] = results
 
         return all_results
+    
+    @abstractmethod
+    def aggregate(self, search_results: Dict[str, List[O]]) -> List[O]:
+        pass
+    
+    def search(self, query: str) -> List[O]:
+        active_blocks = self.select_blocks(query)
+        
+        missing_blocks = set(active_blocks) - set(self.unit_blocks.keys())
+        if missing_blocks:
+            raise ValueError(f"Active blocks {missing_blocks} are not registered in the controller.")
+        
+        search_results = self.block_search(query, active_blocks)
+        
+        aggregated_results = self.aggregate(search_results)
+        
+        return aggregated_results