| from collections import defaultdict | |
| from typing import List, Union | |
| import dspy | |
| try: | |
| import marqo | |
| except ImportError: | |
| raise ImportError( | |
| "The 'marqo' extra is required to use MarqoRM. Install it with `pip install dspy-ai[marqo]`" | |
| ) | |
| class MarqoRM(dspy.Retrieve): | |
| """ | |
| A retrieval module that uses Marqo to return the top passages for a given query. | |
| Assumes that a Marqo index has been created and populated with the following payload: | |
| - document: The text of the passage | |
| Args: | |
| marqo_index_name (str): The name of the marqo index. | |
| marqo_client (marqo.client.Client): A marqo client instance. | |
| k (int, optional): The number of top passages to retrieve. Defaults to 3. | |
| Examples: | |
| Below is a code snippet that shows how to use Marqo as the default retriver: | |
| ```python | |
| import marqo | |
| marqo_client = marqo.Client(url="http://0.0.0.0:8882") | |
| llm = dspy.OpenAI(model="gpt-3.5-turbo") | |
| retriever_model = MarqoRM("my_index_name", marqo_client=marqo_client) | |
| dspy.settings.configure(lm=llm, rm=retriever_model) | |
| ``` | |
| Below is a code snippet that shows how to use Marqo in the forward() function of a module | |
| ```python | |
| self.retrieve = MarqoRM("my_index_name", marqo_client=marqo_client, k=num_passages) | |
| ``` | |
| """ | |
| def __init__( | |
| self, | |
| marqo_index_name: str, | |
| marqo_client: marqo.client.Client, | |
| k: int = 3, | |
| ): | |
| self._marqo_index_name = marqo_index_name | |
| self._marqo_client = marqo_client | |
| super().__init__(k=k) | |
| def forward(self, query_or_queries: Union[str, List[str]]) -> dspy.Prediction: | |
| """Search with Marqo for self.k top passages for query | |
| Args: | |
| query_or_queries (Union[str, List[str]]): The query or queries to search for. | |
| Returns: | |
| dspy.Prediction: An object containing the retrieved passages. | |
| """ | |
| queries = ( | |
| [query_or_queries] | |
| if isinstance(query_or_queries, str) | |
| else query_or_queries | |
| ) | |
| queries = [q for q in queries if q] | |
| all_query_results = [] | |
| for query in queries: | |
| _result = self._marqo_client.index(self._marqo_index_name).search( | |
| q=query, | |
| limit=self.k | |
| ) | |
| all_query_results.append(_result) | |
| passages = defaultdict(float) | |
| for result_dict in all_query_results: | |
| for result in result_dict['hits']: | |
| passages[result['document']] += result['_score'] | |
| sorted_passages = sorted( | |
| passages.items(), key=lambda x: x[1], reverse=True)[:self.k] | |
| return dspy.Prediction(passages=[passage for passage, _ in sorted_passages]) | |