Merge pull request stanfordnlp#1068 from sky-2002/epsilla-retriever

feat(dspy): Add epsilla retriever

Author: arnavsinghvi11

dspy/retrieve/epsilla_rm.py

@@ -0,0 +1,45 @@
+from collections import defaultdict  # noqa: F401
+from typing import Dict, List, Union  # noqa: UP035
+
+import dspy
+from dsp.utils import dotdict
+
+try:
+    from pyepsilla import vectordb
+except ImportError:
+    raise ImportError(  # noqa: B904
+        "The 'pyepsilla' extra is required to use EpsillaRM. Install it with `pip install dspy-ai[epsilla]`",
+    )
+
+
+class EpsillaRM(dspy.Retrieve):
+    def __init__(
+        self,
+        epsilla_client: vectordb.Client,
+        db_name: str,
+        db_path: str,
+        table_name: str,
+        k: int = 3,
+        page_content: str = "document",
+    ):
+        self._epsilla_client = epsilla_client
+        self._epsilla_client.load_db(db_name=db_name, db_path=db_path)
+        self._epsilla_client.use_db(db_name=db_name)
+        self.page_content = page_content
+        self.table_name = table_name
+
+        super().__init__(k=k)
+
+    def forward(self, query_or_queries: Union[str, List[str]], k: Union[int, None] = None, **kwargs) -> dspy.Prediction:  # noqa: ARG002
+        queries = [query_or_queries] if isinstance(query_or_queries, str) else query_or_queries
+        queries = [q for q in queries if q]
+        limit = k if k else self.k
+        all_query_results: list = []
+
+        passages: Dict = defaultdict(float)
+
+        for result_dict in all_query_results:
+            for result in result_dict:
+                passages[result[self.page_content]] += result["@distance"]
+        sorted_passages = sorted(passages.items(), key=lambda x: x[1], reverse=False)[:limit]
+        return dspy.Prediction(passages=[dotdict({"long_text": passage}) for passage, _ in sorted_passages])

poetry.lock

@@ -45,6 +45,7 @@ anthropic = ["anthropic~=0.18.0"]
 chromadb = ["chromadb~=0.4.14"]
 qdrant = ["qdrant-client>=1.6.2", "fastembed>=0.2.0"]
 marqo = ["marqo"]
+epsilla = ["pyepsilla~=0.3.7"]
 pinecone = ["pinecone-client~=2.2.4"]
 weaviate = ["weaviate-client~=4.5.4"]
 milvus = ["pymilvus~=2.3.7"]
@@ -99,6 +100,7 @@ anthropic = { version = "^0.18.0", optional = true }
 chromadb = { version = "^0.4.14", optional = true }
 fastembed = { version = ">=0.2.0", optional = true }
 marqo = { version = "*", optional = true }
+pyepsilla = {version = "^0.3.7", optional = true}
 qdrant-client = { version = "^1.6.2", optional = true }
 pinecone-client = { version = "^2.2.4", optional = true }
 weaviate-client = { version = "^4.5.4", optional = true }
@@ -140,6 +142,7 @@ ipykernel = "^6.29.4"
 chromadb = ["chromadb"]
 qdrant = ["qdrant-client", "fastembed"]
 marqo = ["marqo"]
+epsilla = ["pyepsilla"]
 pinecone = ["pinecone-client"]
 weaviate = ["weaviate-client"]
 milvus = ["pymilvus"]