使用ElasticsearchRetriever进行高效数据检索

最新推荐文章于 2025-04-25 09:49:51 发布

eahba

最新推荐文章于 2025-04-25 09:49:51 发布

阅读量277

点赞数 4

文章标签： jenkins 运维 python

本文链接：https://blog.youkuaiyun.com/eahba/article/details/146466004

版权

使用ElasticsearchRetriever进行高效数据检索

Elasticsearch 是一个分布式、RESTful 搜索和分析引擎，适用于关键词搜索、向量搜索、混合搜索和复杂过滤。ElasticsearchRetriever 是一个灵活的封装器，允许通过 Query DSL 使用所有 Elasticsearch 特性。本文将详细介绍如何使用 ElasticsearchRetriever 进行数据检索。

技术背景介绍

Elasticsearch 提供了一种强大的搜索引擎技术，它不仅支持传统的关键词搜索，还能够处理现代的向量搜索和混合搜索。ElasticsearchRetriever 通过封装 Elasticsearch 的 API，使得开发者可以轻松地实现复杂的数据检索任务。

核心原理解析

ElasticsearchRetriever 提供了一种便捷的方式来访问 Elasticsearch 的各种检索功能，包括关键词搜索（BM25）、向量搜索、混合搜索和复杂过滤。这些功能结合使用可以处理大量数据并提供高效的检索结果。

代码实现演示

设置 Elasticsearch 实例

我们可以选择使用本地实例或 Elastic Cloud 来设置 Elasticsearch。在本示例中，我们使用本地实例。

from elasticsearch import Elasticsearch

# 连接到本地的 Elasticsearch 实例
es_url = "http://localhost:9200"
es_client = Elasticsearch(hosts=[es_url])
es_client.info()

创建索引和数据

首先，我们定义一些示例数据并将其索引到 Elasticsearch。

from elasticsearch.helpers import bulk
from langchain_community.embeddings import DeterministicFakeEmbedding
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings

# 使用假嵌入进行演示
embeddings = DeterministicFakeEmbedding(size=3)

# 定义示例数据
index_name = "test-langchain-retriever"
text_field = "text"
dense_vector_field = "fake_embedding"
num_characters_field = "num_characters"
texts = [
    "foo",
    "bar",
    "world",
    "hello world",
    "hello",
    "foo bar",
    "bla bla foo",
]

# 创建索引和索引数据
def create_index(es_client, index_name, text_field, dense_vector_field, num_characters_field):
    es_client.indices.create(
        index=index_name,
        mappings={
            "properties": {
                text_field: {"type": "text"},
                dense_vector_field: {"type": "dense_vector"},
                num_characters_field: {"type": "integer"},
            }
        },
    )

def index_data(es_client, index_name, text_field, dense_vector_field, embeddings, texts, refresh=True):
    create_index(es_client, index_name, text_field, dense_vector_field, num_characters_field)
    vectors = embeddings.embed_documents(list(texts))
    requests = [
        {
            "_op_type": "index",
            "_index": index_name,
            "_id": i,
            text_field: text,
            dense_vector_field: vector,
            num_characters_field: len(text),
        }
        for i, (text, vector) in enumerate(zip(texts, vectors))
    ]
    bulk(es_client, requests)
    if refresh:
        es_client.indices.refresh(index=index_name)

# 索引数据
index_data(es_client, index_name, text_field, dense_vector_field, embeddings, texts)

使用 ElasticsearchRetriever 进行检索

我们展示了如何通过不同的检索方式从 Elasticsearch 中获取文档。

向量搜索

from typing import Dict
from langchain_elasticsearch import ElasticsearchRetriever

def vector_query(search_query: str) -> Dict:
    vector = embeddings.embed_query(search_query)
    return {
        "knn": {
            "field": dense_vector_field,
            "query_vector": vector,
            "k": 5,
            "num_candidates": 10,
        }
    }

vector_retriever = ElasticsearchRetriever.from_es_params(
    index_name=index_name,
    body_func=vector_query,
    content_field=text_field,
    url=es_url,
)

# 执行向量搜索
results = vector_retriever.invoke("foo")
print(results)

BM25 关键词搜索

def bm25_query(search_query: str) -> Dict:
    return {
        "query": {
            "match": {
                text_field: search_query,
            },
        },
    }

bm25_retriever = ElasticsearchRetriever.from_es_params(
    index_name=index_name,
    body_func=bm25_query,
    content_field=text_field,
    url=es_url,
)

# 执行关键词搜索
results = bm25_retriever.invoke("foo")
print(results)

混合搜索

def hybrid_query(search_query: str) -> Dict:
    vector = embeddings.embed_query(search_query)
    return {
        "query": {
            "match": {
                text_field: search_query,
            },
        },
        "knn": {
            "field": dense_vector_field,
            "query_vector": vector,
            "k": 5,
            "num_candidates": 10,
        },
        "rank": {"rrf": {}},
    }

hybrid_retriever = ElasticsearchRetriever.from_es_params(
    index_name=index_name,
    body_func=hybrid_query,
    content_field=text_field,
    url=es_url,
)

# 执行混合搜索
results = hybrid_retriever.invoke("foo")
print(results)

复杂过滤

def filter_query_func(search_query: str) -> Dict:
    return {
        "query": {
            "bool": {
                "must": [
                    {"range": {num_characters_field: {"gte": 5}}},
                ],
                "must_not": [
                    {"prefix": {text_field: "bla"}},
                ],
                "should": [
                    {"match": {text_field: search_query}},
                ],
            }
        }
    }

filtering_retriever = ElasticsearchRetriever.from_es_params(
    index_name=index_name,
    body_func=filter_query_func,
    content_field=text_field,
    url=es_url,
)

# 执行复杂过滤
results = filtering_retriever.invoke("foo")
print(results)