Allow Embeddings to return indices

Author: okhat

dspy/propose/utils.py

@@ -146,6 +146,7 @@ def get_dspy_source_code(module):
     base_code = ""
 
     # Don't get source code for Predict or ChainOfThought modules (NOTE we will need to extend this list as more DSPy.modules are added)
+    # TODO: if type(module).__name__ not in ["Predict", "ChainOfThought", "ReAct"]:
     if not type(module).__name__ == "Predict" and not type(module).__name__ == "ChainOfThought":
         try:
             base_code = inspect.getsource(type(module))

dspy/retrievers/embeddings.py

@@ -34,7 +34,9 @@ def __call__(self, query: str):
 
     def forward(self, query: str):
         import dspy
-        return dspy.Prediction(passages=self.search_fn(query))
+
+        passages, indices = self.search_fn(query)
+        return dspy.Prediction(passages=passages, indices=indices)
 
     def _batch_forward(self, queries: List[str]):
         q_embeds = self.embedder(queries)
@@ -76,7 +78,7 @@ def _rerank_and_predict(self, q_embeds: np.ndarray, candidate_indices: np.ndarra
         top_k_indices = np.argsort(-scores, axis=1)[:, :self.k]
         top_indices = candidate_indices[np.arange(len(q_embeds))[:, None], top_k_indices]
 
-        return [[self.corpus[idx] for idx in indices] for indices in top_indices]
+        return [([self.corpus[idx] for idx in indices], [idx for idx in indices]) for indices in top_indices]
 
     def _normalize(self, embeddings: np.ndarray):
         norms = np.linalg.norm(embeddings, axis=1, keepdims=True)