Bring YouRM .forward signature in-line with Retrieve child class .forward methods (add K kwarg)

dspy/retrieve/retrieve.py

@@ -1,3 +1,5 @@
+from typing import List, Union, Optional
+
 import dsp
 import random
 
@@ -13,31 +15,29 @@ class Retrieve(Parameter):
     def __init__(self, k=3):
         self.stage = random.randbytes(8).hex()
         self.k = k
-    
+
     def reset(self):
         pass
-    
+
     def dump_state(self):
         state_keys = ["k"]
         return {k: getattr(self, k) for k in state_keys}
 
     def load_state(self, state):
         for name, value in state.items():
             setattr(self, name, value)
-    
+
     def __call__(self, *args, **kwargs):
         return self.forward(*args, **kwargs)
-    
-    def forward(self, query_or_queries):
+
+    def forward(self, query_or_queries: Union[str, List[str]], k: Optional[int] = None) -> Prediction:
         queries = [query_or_queries] if isinstance(query_or_queries, str) else query_or_queries
         queries = [query.strip().split('\n')[0].strip() for query in queries]
 
-
         # print(queries)
         # TODO: Consider removing any quote-like markers that surround the query too.
-
-        passages = dsp.retrieveEnsemble(queries, k=self.k)
+        k = k if k is not None else self.k
+        passages = dsp.retrieveEnsemble(queries, k=k)
         return Prediction(passages=passages)
-
 
-# TODO: Consider doing Prediction.from_completions with the individual sets of passages (per query) too.
+# TODO: Consider doing Prediction.from_completions with the individual sets of passages (per query) too.

dspy/retrieve/you_rm.py

@@ -2,7 +2,7 @@
 import os
 import requests
 
-from typing import Union, List
+from typing import Union, List, Optional
 
 
 class YouRM(dspy.Retrieve):
@@ -15,15 +15,19 @@ def __init__(self, ydc_api_key=None, k=3):
         else:
             self.ydc_api_key = os.environ["YDC_API_KEY"]
 
-    def forward(self, query_or_queries: Union[str, List[str]]) -> dspy.Prediction:
+    def forward(self, query_or_queries: Union[str, List[str]], k: Optional[int] = None) -> dspy.Prediction:
         """Search with You.com for self.k top passages for query or queries
 
         Args:
             query_or_queries (Union[str, List[str]]): The query or queries to search for.
+            k (Optional[int]): The number of context strings to return, if not already specified in self.k
 
         Returns:
             dspy.Prediction: An object containing the retrieved passages.
         """
+
+        k = k if k is not None else self.k
+
         queries = (
             [query_or_queries]
             if isinstance(query_or_queries, str)
@@ -36,7 +40,7 @@ def forward(self, query_or_queries: Union[str, List[str]]) -> dspy.Prediction:
                 f"https://api.ydc-index.io/search?query={query}",
                 headers=headers,
             ).json()
-            for hit in results["hits"][:self.k]:
+            for hit in results["hits"][:k]:
                 for snippet in hit["snippets"]:
                     docs.append(snippet)
         return dspy.Prediction(passages=docs)