功能描述

由于supabase作为docStore不受支持，我尝试使用两个VectorStoreIndex实现一个AutoMergingRetriever。所做的更改如下：
向SupabaseVectorStore添加了get_document,使用：

def get_document(self, doc_id: str, **kwargs: Any) -> Any:
        """Get document by doc id.

Args:
doc_id (str): document id
"""
        filters = {"doc_id": {"$eq": doc_id}}

        results = self._collection.query(
            data=None,
            filters=filters,
            include_value=True,
            include_metadata=True,
            **kwargs,
        )

        if len(results) > 0:
            id_, distance, metadata = results[0]
            text = metadata.pop("text", None)

            try:
                node = metadata_dict_to_node(metadata)
            except Exception:
                # NOTE: deprecated legacy logic for backward compatibility
                metadata, node_info, relationships = legacy_metadata_dict_to_node(
                    metadata
                )
                node = TextNode(
                    id_=id_,
                    text=text,
                    metadata=metadata,
                    start_char_idx=node_info.get("start", None),
                    end_char_idx=node_info.get("end", None),
                    relationships=relationships,
                )

            return node

        return None

更改了auto_merging_retriever.py中的两行：
在函数_get_parents_and_merge中：

parent_node = self._storage_context.vector_store.get_document(
                    parent_node_id
                )

以及：
在函数_fill_in_nodes中：

next_node = self._storage_context.vector_store.get_document(
                    cur_node.next_node.node_id
                )

定义了一个自定义检索器：

from typing import Dict, List, Optional, Tuple

from llama_index.core.indices.query.schema import QueryBundle
from llama_index.core.indices.utils import truncate_text
from llama_index.core.indices.vector_store.retrievers.retriever import (
    VectorIndexRetriever,
)
from llama_index.core.schema import BaseNode, IndexNode, NodeWithScore
from typing import Dict, List, Optional, Tuple, cast

from llama_index.core.callbacks.base import CallbackManager
from llama_index.core.indices.query.schema import QueryBundle
from llama_index.core.indices.utils import truncate_text
from llama_index.core.indices.vector_store.retrievers.retriever import (
    VectorIndexRetriever,
)
from llama_index.core.schema import BaseNode, IndexNode, NodeWithScore, QueryBundle
from llama_index.core.storage.storage_context import StorageContext

class MyRetriever(AutoMergingRetriever):
    def __init__(
        self,
        node_vector_retriever: VectorIndexRetriever,
        leaf_node_vector_retriever: VectorIndexRetriever,
        storage_context: StorageContext,
        simple_ratio_thresh: float = 0.5,
        verbose: bool = False,
        callback_manager: Optional[CallbackManager] = None,
        object_map: Optional[dict] = None,
        objects: Optional[List[IndexNode]] = None,
    ) -> None:
        super().__init__(
            vector_retriever=node_vector_retriever,
            storage_context=storage_context,
            simple_ratio_thresh=simple_ratio_thresh,
            verbose=verbose,
            callback_manager=callback_manager,
            object_map=object_map,
            objects=objects,
        )
        self._leaf_node_vector_retriever = leaf_node_vector_retriever

    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        initial_nodes = self._vector_retriever.retrieve(query_bundle)
        initial_leaf_nodes = self._leaf_node_vector_retriever.retrieve(query_bundle)

        # Merge the initial nodes and leaf nodes
        initial_nodes.extend(initial_leaf_nodes)

        cur_nodes, is_changed = self._try_merging(initial_nodes)
        while is_changed:
            cur_nodes, is_changed = self._try_merging(cur_nodes)

        # sort by similarity
        cur_nodes.sort(key=lambda x: x.get_score(), reverse=True)

        return cur_nodes

使用方法：

from llama_index.core.retrievers import AutoMergingRetriever
from llama_index.core import VectorStoreIndex 
from llama_index.core.node_parser import get_leaf_nodes, get_root_nodes
 # Some more imports here

parent_vector_store = SupabaseVectorStore(
    postgres_connection_string=DB_CONNECTION,
    collection_name='parent_nodes',
    dimension=768
)
doc_storage_context = StorageContext.from_defaults(vector_store=parent_vector_store)
doc_index = VectorStoreIndex(nodes=nodes, storage_context=doc_storage_context)

leaf_nodes = get_leaf_nodes(nodes)

leaf_vector_store = SupabaseVectorStore(
    postgres_connection_string=DB_CONNECTION,
    collection_name='leaf_nodes',
    dimension=768
)
leaf_storage_context = StorageContext.from_defaults(vector_store=leaf_vector_store)
leaf_index = VectorStoreIndex(nodes=leaf_nodes, storage_context=leaf_storage_context)
leaf_retriever = leaf_index.as_retriever(similarity_top_k=6)
base_retriever = doc_index.as_retriever(similarity_top_k=6)
retriever = AutoMergingRetriever(leaf_retriever, storage_context=doc_storage_context, verbose=True)

my_retriever = MyRetriever(node_vector_retriever=base_retriever, leaf_node_vector_retriever=leaf_retriever, storage_context=doc_storage_context, verbose=True)

query_str = "What did the author do growing up?"
my_retrived_nodes = my_retriever.retrieve(query_str)

原因

使用Supabase的AutoMergingRetriever管道

功能价值

完整的Supabase RAG管道，可作为任何准备好供消费的前端框架的后端使用。