diff --git a/docs/modules/indexes/retrievers/examples/elastic_search_bm25.ipynb b/docs/modules/indexes/retrievers/examples/elastic_search_bm25.ipynb new file mode 100644 index 00000000..424de3dc --- /dev/null +++ b/docs/modules/indexes/retrievers/examples/elastic_search_bm25.ipynb @@ -0,0 +1,164 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "ab66dd43", + "metadata": {}, + "source": [ + "# ElasticSearch BM25\n", + "\n", + "This notebook goes over how to use a retriever that under the hood uses ElasticSearcha and BM25.\n", + "\n", + "For more information on the details of BM25 see [this blog post](https://www.elastic.co/blog/practical-bm25-part-2-the-bm25-algorithm-and-its-variables)." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "393ac030", + "metadata": {}, + "outputs": [], + "source": [ + "from langchain.retrievers import ElasticSearchBM25Retriever" + ] + }, + { + "cell_type": "markdown", + "id": "aaf80e7f", + "metadata": {}, + "source": [ + "## Create New Retriever" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "bcb3c8c2", + "metadata": {}, + "outputs": [], + "source": [ + "elasticsearch_url=\"http://localhost:9200\"\n", + "retriever = ElasticSearchBM25Retriever.create(elasticsearch_url, \"langchain-index-3\")" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "b605284d", + "metadata": {}, + "outputs": [], + "source": [ + "# Alternatively, you can load an existing index\n", + "# import elasticsearch\n", + "# elasticsearch_url=\"http://localhost:9200\"\n", + "# retriever = ElasticSearchBM25Retriever(elasticsearch.Elasticsearch(elasticsearch_url), \"langchain-index\")" + ] + }, + { + "cell_type": "markdown", + "id": "1c518c42", + "metadata": {}, + "source": [ + "## Add texts (if necessary)\n", + "\n", + "We can optionally add texts to the retriever (if they aren't already in there)" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "98b1c017", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "['386c76c9-4355-4c12-aaeb-7b80054caf93',\n", + " 'fffd279c-a0c9-4158-a904-6e242c517c99',\n", + " '7f5528a3-18d0-43b0-894d-f6770a002219',\n", + " 'e2ef5e32-d5bd-44e2-b045-cfc5a8e0a0a1',\n", + " 'cce8ba48-e473-4235-bca2-2c8d65e73ccf']" + ] + }, + "execution_count": 14, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "retriever.add_texts([\"foo\", \"bar\", \"world\", \"hello\", \"foo bar\"])" + ] + }, + { + "cell_type": "markdown", + "id": "08437fa2", + "metadata": {}, + "source": [ + "## Use Retriever\n", + "\n", + "We can now use the retriever!" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "c0455218", + "metadata": {}, + "outputs": [], + "source": [ + "result = retriever.get_relevant_documents(\"foo\")" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "7dfa5c29", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[Document(page_content='foo', metadata={}),\n", + " Document(page_content='foo bar', metadata={})]" + ] + }, + "execution_count": 16, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "result" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "74bd9256", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.1" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/langchain/retrievers/__init__.py b/langchain/retrievers/__init__.py index c7b7e60c..18224c62 100644 --- a/langchain/retrievers/__init__.py +++ b/langchain/retrievers/__init__.py @@ -1,4 +1,5 @@ from langchain.retrievers.chatgpt_plugin_retriever import ChatGPTPluginRetriever +from langchain.retrievers.elastic_search_bm25 import ElasticSearchBM25Retriever from langchain.retrievers.metal import MetalRetriever from langchain.retrievers.pinecone_hybrid_search import PineconeHybridSearchRetriever from langchain.retrievers.remote_retriever import RemoteLangChainRetriever @@ -8,4 +9,5 @@ __all__ = [ "RemoteLangChainRetriever", "PineconeHybridSearchRetriever", "MetalRetriever", + "ElasticSearchBM25Retriever", ] diff --git a/langchain/retrievers/elastic_search_bm25.py b/langchain/retrievers/elastic_search_bm25.py new file mode 100644 index 00000000..055e7c6e --- /dev/null +++ b/langchain/retrievers/elastic_search_bm25.py @@ -0,0 +1,126 @@ +"""Wrapper around Elasticsearch vector database.""" +from __future__ import annotations + +import uuid +from typing import Any, Iterable, List + +from langchain.docstore.document import Document +from langchain.schema import BaseRetriever + + +class ElasticSearchBM25Retriever(BaseRetriever): + """Wrapper around Elasticsearch using BM25 as a retrieval method. + + + To connect to an Elasticsearch instance that requires login credentials, + including Elastic Cloud, use the Elasticsearch URL format + https://username:password@es_host:9243. For example, to connect to Elastic + Cloud, create the Elasticsearch URL with the required authentication details and + pass it to the ElasticVectorSearch constructor as the named parameter + elasticsearch_url. + + You can obtain your Elastic Cloud URL and login credentials by logging in to the + Elastic Cloud console at https://cloud.elastic.co, selecting your deployment, and + navigating to the "Deployments" page. + + To obtain your Elastic Cloud password for the default "elastic" user: + + 1. Log in to the Elastic Cloud console at https://cloud.elastic.co + 2. Go to "Security" > "Users" + 3. Locate the "elastic" user and click "Edit" + 4. Click "Reset password" + 5. Follow the prompts to reset the password + + The format for Elastic Cloud URLs is + https://username:password@cluster_id.region_id.gcp.cloud.es.io:9243. + """ + + def __init__(self, client: Any, index_name: str): + self.client = client + self.index_name = index_name + + @classmethod + def create( + cls, elasticsearch_url: str, index_name: str, k1: float = 2.0, b: float = 0.75 + ) -> ElasticSearchBM25Retriever: + from elasticsearch import Elasticsearch + + # Create an Elasticsearch client instance + es = Elasticsearch(elasticsearch_url) + + # Define the index settings and mappings + index_settings = { + "settings": { + "analysis": {"analyzer": {"default": {"type": "standard"}}}, + "similarity": { + "custom_bm25": { + "type": "BM25", + "k1": k1, + "b": b, + } + }, + }, + "mappings": { + "properties": { + "content": { + "type": "text", + "similarity": "custom_bm25", # Use the custom BM25 similarity + } + } + }, + } + + # Create the index with the specified settings and mappings + es.indices.create(index=index_name, body=index_settings) + return cls(es, index_name) + + def add_texts( + self, + texts: Iterable[str], + refresh_indices: bool = True, + ) -> List[str]: + """Run more texts through the embeddings and add to the retriver. + + Args: + texts: Iterable of strings to add to the retriever. + refresh_indices: bool to refresh ElasticSearch indices + + Returns: + List of ids from adding the texts into the retriever. + """ + try: + from elasticsearch.helpers import bulk + except ImportError: + raise ValueError( + "Could not import elasticsearch python package. " + "Please install it with `pip install elasticsearch`." + ) + requests = [] + ids = [] + for i, text in enumerate(texts): + _id = str(uuid.uuid4()) + request = { + "_op_type": "index", + "_index": self.index_name, + "content": text, + "_id": _id, + } + ids.append(_id) + requests.append(request) + bulk(self.client, requests) + + if refresh_indices: + self.client.indices.refresh(index=self.index_name) + return ids + + def get_relevant_documents(self, query: str) -> List[Document]: + query_dict = {"query": {"match": {"content": query}}} + res = self.client.search(index=self.index_name, body=query_dict) + + docs = [] + for r in res["hits"]["hits"]: + docs.append(Document(page_content=r["_source"]["content"])) + return docs + + async def aget_relevant_documents(self, query: str) -> List[Document]: + raise NotImplementedError