S3VectorStore 集成

这是一个为LlamaIndex设计的向量存储集成，使用S3Vectors。

本笔记本假设您已经创建了一个S3向量存储桶（可能还包括索引）。

安装

%pip install llama-index-vector-stores-s3 llama-index-embeddings-openai

使用方法

创建向量存储对象

您可以在现有的S3存储桶中创建新的向量索引。

如果您没有在环境中配置S3凭据，可以提供一个带有凭据的boto3会话。

from llama_index.vector_stores.s3 import S3VectorStore
import boto3

vector_store = S3VectorStore.create_index_from_bucket(
    # S3 bucket name or ARN
    bucket_name_or_arn="test-bucket",
    # Name for the new index
    index_name="my-index",
    # Vector dimension (e.g., 1536 for OpenAI embeddings)
    dimension=1536,
    # Distance metric: "cosine", "euclidean", etc.
    distance_metric="cosine",
    # Data type for vectors
    data_type="float32",
    # Batch size for inserting vectors (max 500)
    insert_batch_size=500,
    # Metadata keys that won't be filterable
    non_filterable_metadata_keys=["custom_field"],
    # Optional: provide a boto3 session for custom AWS configuration
    # sync_session=boto3.Session(...)
)

或者，您可以在现有的S3存储桶中使用现有的向量索引。

from llama_index.vector_stores.s3 import S3VectorStore
import boto3

vector_store = S3VectorStore(
    # Index name or ARN
    index_name_or_arn="my-index",
    # S3 bucket name or ARN
    bucket_name_or_arn="test-bucket",
    # Data type for vectors (must match index)
    data_type="float32",
    # Distance metric (must match index)
    distance_metric="cosine",
    # Batch size for inserting vectors (max 500)
    insert_batch_size=500,
    # Optional: specify metadata field containing text if you already have a populated index
    # text_field="content",
    # Optional: provide a boto3 session for custom AWS configuration
    # sync_session=boto3.Session(...),
)

使用带索引的向量存储

一旦你有了向量存储，就可以将其与索引一起使用：

from llama_index.core import (
    VectorStoreIndex,
    StorageContext,
    SimpleDirectoryReader,
    Document,
)
from llama_index.embeddings.openai import OpenAIEmbedding

# Load some data
# documents = SimpleDirectoryReader("data").load_data()

# Or create new documents
documents = [
    Document(text="Hello, world!", metadata={"key": "1"}),
    Document(text="Hello, world! 2", metadata={"key": "2"}),
]

# Create a new index
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=StorageContext.from_defaults(vector_store=vector_store),
    # optional: set the embed model
    embed_model=OpenAIEmbedding(model="text-embedding-3-small", api_key="..."),
)

# Or reload from an existing index
# index = VectorStoreIndex.from_vector_store(
#     vector_store=vector_store,
#     # optional: set the embed model
#     # embed_model=embed_model,
# )

nodes = index.as_retriever(similarity_top_k=1).retrieve("Hello, world!")
print(nodes[0].text)

Hello, world!

您也可以使用过滤器来帮助您检索正确的节点！

from llama_index.core.vector_stores.types import (
    MetadataFilters,
    MetadataFilter,
    FilterOperator,
    FilterCondition,
)

nodes = index.as_retriever(
    similarity_top_k=2,
    filters=MetadataFilters(
        filters=[
            MetadataFilter(
                key="key",
                value="2",
                operator=FilterOperator.EQ,
            ),
        ],
        condition=FilterCondition.AND,
    ),
).retrieve("Hello, world!")

print(nodes[0].text)

Hello, world! 2

直接使用向量存储

你也可以直接使用向量存储：

from llama_index.core.schema import TextNode
from llama_index.core.vector_stores.types import VectorStoreQuery
from llama_index.embeddings.openai import OpenAIEmbedding

# embed nodes
nodes = [
    TextNode(text="Hello, world!"),
    TextNode(text="Hello, world! 2"),
]

embed_model = OpenAIEmbedding(model="text-embedding-3-small", api_key="...")
embeddings = embed_model.get_text_embedding_batch([n.text for n in nodes])
for node, embedding in zip(nodes, embeddings):
    node.embedding = embedding

# add nodes to the vector store
vector_store.add(nodes)

# query the vector store
query = VectorStoreQuery(
    query_embedding=embed_model.get_query_embedding("Hello, world!"),
    similarity_top_k=2,
)
results = vector_store.query(query)
print(results.nodes[0].text)

Hello, world!