restored 'k' parameter when creating strategies (#159)

* restored 'k' parameter when creating strategies

* lint

Author: epinzur

packages/graph-retriever/src/graph_retriever/strategies/base.py

@@ -10,6 +10,8 @@
 from graph_retriever.content import Content
 from graph_retriever.types import Node
 
+DEFAULT_SELECT_K = 5
+
 
 class NodeTracker:
     """Helper class for tracking traversal progress."""
@@ -95,16 +97,27 @@ class Strategy(abc.ABC):
         If `None`, there is no limit.
     max_depth :
         Maximum traversal depth. If `None`, there is no limit.
+    k:
+        Deprecated: Use `select_k` instead.
+        Maximum number of nodes to select and return during traversal.
     """
 
-    select_k: int = 5
+    select_k: int = dataclasses.field(default=DEFAULT_SELECT_K)
     start_k: int = 4
     adjacent_k: int = 10
     max_traverse: int | None = None
     max_depth: int | None = None
+    k: int = dataclasses.field(default=DEFAULT_SELECT_K, repr=False)
 
     _query_embedding: list[float] = dataclasses.field(default_factory=list)
 
+    def __post_init__(self):
+        """Allow passing the deprecated 'k' value instead of 'select_k'."""
+        if self.select_k == DEFAULT_SELECT_K and self.k != DEFAULT_SELECT_K:
+            self.select_k = self.k
+        else:
+            self.k = self.select_k
+
     @abc.abstractmethod
     def iteration(self, *, nodes: Iterable[Node], tracker: NodeTracker) -> None:
         """

packages/graph-retriever/src/graph_retriever/strategies/eager.py

@@ -30,6 +30,9 @@ class Eager(Strategy):
         Number of documents to fetch for each outgoing edge.
     max_depth :
         Maximum traversal depth. If `None`, there is no limit.
+    k:
+        Deprecated: Use `select_k` instead.
+        Maximum number of nodes to select and return during traversal.
     """
 
     @override

packages/graph-retriever/src/graph_retriever/strategies/mmr.py

@@ -67,6 +67,9 @@ class Mmr(Strategy):
     min_mmr_score :
         Only nodes with a score greater than or equal to this value will be
         selected.
+    k:
+        Deprecated: Use `select_k` instead.
+        Maximum number of nodes to select and return during traversal.
     """
 
     lambda_mult: float = 0.5

packages/graph-retriever/src/graph_retriever/strategies/scored.py

@@ -19,7 +19,33 @@ def __lt__(self, other: "_ScoredNode") -> bool:
 
 @dataclasses.dataclass
 class Scored(Strategy):
-    """Strategy selecting nodes using a scoring function."""
+    """
+    Scored traversal strategy.
+
+    This strategy uses a scoring function to select nodes using a local maximum
+    approach. In each iteration, it chooses the top scoring nodes available and
+    then traverses the connected nodes.
+
+    Parameters
+    ----------
+    scorer:
+        A callable function that returns the score of a node.
+    select_k :
+        Maximum number of nodes to retrieve during traversal.
+    start_k :
+        Number of documents to fetch via similarity for starting the traversal.
+        Added to any initial roots provided to the traversal.
+    adjacent_k :
+        Number of documents to fetch for each outgoing edge.
+    max_depth :
+        Maximum traversal depth. If `None`, there is no limit.
+    per_iteration_limit:
+        Maximum number of nodes to select and traverse during a single
+        iteration.
+    k:
+        Deprecated: Use `select_k` instead.
+        Maximum number of nodes to select and return during traversal.
+    """
 
     scorer: Callable[[Node], float]
     _nodes: list[_ScoredNode] = dataclasses.field(default_factory=list)

packages/graph-retriever/tests/strategies/test_base.py

@@ -102,3 +102,8 @@ def test_build_strategy_base_override_mmr():
         Strategy.build(
             base_strategy=override_strategy, strategy=base_strategy, lambda_mult=0.2
         )
+
+
+def test_setting_k_sets_select_k():
+    assert Eager(select_k=4) == Eager(k=4)
+    assert Mmr(select_k=3) == Mmr(k=3)