demo_run.py

"""
RedisVL 完整演示：从 Schema 定义到向量搜索的全流程
"""
import numpy as np
from redisvl.index import SearchIndex
from redisvl.query import VectorQuery, FilterQuery
from redisvl.query.filter import Tag, Num, Text
from redisvl.schema import IndexSchema

REDIS_URL = "redis://localhost:6379"
INDEX_NAME = "demo-movie-idx"

# ──────────────────────────────────────────────
# Step 1: 定义 Schema
# ──────────────────────────────────────────────
print("=" * 60)
print("Step 1: 定义索引 Schema")
print("=" * 60)

schema = IndexSchema.from_dict({
    "index": {
        "name": INDEX_NAME,
        "prefix": "movie",
        "storage_type": "json",
    },
    "fields": [
        {"name": "title", "type": "text"},
        {"name": "genre", "type": "tag"},
        {"name": "year", "type": "numeric"},
        {"name": "rating", "type": "numeric"},
        {"name": "embedding", "type": "vector",
         "attrs": {
             "algorithm": "flat",
             "dims": 4,
             "distance_metric": "cosine",
             "datatype": "float32",
         }},
    ],
})
import yaml
print(yaml.dump(schema.to_dict(), default_flow_style=False, allow_unicode=True))

# ──────────────────────────────────────────────
# Step 2: 创建索引
# ──────────────────────────────────────────────
print("\n" + "=" * 60)
print("Step 2: 创建索引（连接 Redis）")
print("=" * 60)

# 用 context manager 自动管理连接
with SearchIndex.from_dict(schema.to_dict(), redis_url=REDIS_URL) as index:
    # 如果已存在则先删除
    if index.exists():
        index.delete(drop=True)
        print("  → 已删除旧索引")

    index.create()
    print(f"  ✅ 索引 '{index.name}' 创建成功")
    print(f"  key 前缀: {index.prefix}")
    print(f"  存储类型: {index.storage_type}")
    print(f"  Redis key 示例: {index.key('inception')}")

# ──────────────────────────────────────────────
# Step 3: 加载电影数据
# ──────────────────────────────────────────────
print("\n" + "=" * 60)
print("Step 3: 加载数据")
print("=" * 60)

movies = [
    {"title": "Inception", "genre": "sci-fi", "year": 2010, "rating": 8.8,
     "embedding": [0.1, 0.8, 0.1, 0.5]},
    {"title": "Interstellar", "genre": "sci-fi", "year": 2014, "rating": 8.7,
     "embedding": [0.2, 0.7, 0.2, 0.4]},
    {"title": "The Godfather", "genre": "crime", "year": 1972, "rating": 9.2,
     "embedding": [0.9, 0.1, 0.3, 0.1]},
    {"title": "The Dark Knight", "genre": "action", "year": 2008, "rating": 9.0,
     "embedding": [0.3, 0.6, 0.7, 0.2]},
    {"title": "Pulp Fiction", "genre": "crime", "year": 1994, "rating": 8.9,
     "embedding": [0.8, 0.2, 0.5, 0.1]},
    {"title": "The Matrix", "genre": "sci-fi", "year": 1999, "rating": 8.7,
     "embedding": [0.1, 0.9, 0.1, 0.3]},
    {"title": "Fight Club", "genre": "drama", "year": 1999, "rating": 8.8,
     "embedding": [0.5, 0.3, 0.8, 0.1]},
    {"title": "Goodfellas", "genre": "crime", "year": 1990, "rating": 8.7,
     "embedding": [0.9, 0.1, 0.1, 0.2]},
]

with SearchIndex.from_existing(INDEX_NAME, redis_url=REDIS_URL) as index:
    keys = index.load(movies, id_field="title")
    print(f"  ✅ 加载了 {len(keys)} 部电影")
    print(f"  示例 key: {keys[:2]}")

    # 验证：fetch 一部电影
    movie = index.fetch("Inception")
    print(f"  Fetch 'Inception': genre={movie['genre']}, year={movie['year']}, rating={movie['rating']}")

# ──────────────────────────────────────────────
# Step 4: 向量相似度搜索
# ──────────────────────────────────────────────
print("\n" + "=" * 60)
print("Step 4: 向量相似度搜索（模拟推荐）")
print("=" * 60)

# 模拟一个用户偏好向量（偏向 sci-fi 的）
query_vector = np.array([0.15, 0.75, 0.15, 0.4], dtype=np.float32).tobytes()

with SearchIndex.from_existing(INDEX_NAME, redis_url=REDIS_URL) as index:
    vq = VectorQuery(
        vector=query_vector,
        vector_field_name="embedding",
        return_fields=["title", "genre", "year", "rating", "vector_distance"],
        num_results=5,
    )
    results = index.query(vq)

    print(f"  查询向量偏好: sci-fi")
    for i, doc in enumerate(results):
        dist = doc.get("vector_distance", "N/A")
        print(f"  {i+1}. {doc['title']:20s} | {doc['genre']:8s} | {doc['year']} | ⭐{doc['rating']} | dist={float(dist):.4f}")

# ──────────────────────────────────────────────
# Step 5: 混合过滤查询
# ──────────────────────────────────────────────
print("\n" + "=" * 60)
print("Step 5: 过滤查询 — 找出 2000 年后的犯罪片")
print("=" * 60)

with SearchIndex.from_existing(INDEX_NAME, redis_url=REDIS_URL) as index:
    fq = FilterQuery(
        return_fields=["title", "genre", "year", "rating"],
        filter_expression=(Tag("genre") == "crime") & (Num("year") >= 2000),
        num_results=10,
    )
    results = index.query(fq)

    if results:
        for i, doc in enumerate(results):
            print(f"  {i+1}. {doc['title']:20s} | {doc['genre']:6s} | {doc['year']} | ⭐{doc['rating']}")
    else:
        print("  (无匹配结果 — 所有犯罪片都在 2000 年前)")

# ──────────────────────────────────────────────
# Step 6: 全文搜索
# ──────────────────────────────────────────────
print("\n" + "=" * 60)
print("Step 6: 全文搜索 — 标题含 'Dark' 的电影")
print("(💡 'The' 是 Redis 默认停用词, 所以搜 Dark)")
print("=" * 60)

with SearchIndex.from_existing(INDEX_NAME, redis_url=REDIS_URL) as index:
    fq = FilterQuery(
        return_fields=["title", "genre", "year", "rating"],
        filter_expression=Text("title") % "Dark*",
        num_results=10,
    )
    results = index.query(fq)

    for i, doc in enumerate(results):
        print(f"  {i+1}. {doc['title']:20s} | {doc['genre']:6s} | {doc['year']} | ⭐{doc['rating']}")

# ──────────────────────────────────────────────
# Step 7: 计数查询
# ──────────────────────────────────────────────
print("\n" + "=" * 60)
print("Step 7: 聚合统计")
print("=" * 60)

with SearchIndex.from_existing(INDEX_NAME, redis_url=REDIS_URL) as index:
    # 总体统计
    info = index.info()
    print(f"  索引信息:")
    print(f"    文档数量: {info.get('num_docs', 'N/A')}")
    print(f"    索引大小: {info.get('num_bytes', 0) / 1024:.2f} KB")

# ──────────────────────────────────────────────
# Step 8: 数据保留（不清理）
# ──────────────────────────────────────────────
print("\n" + "=" * 60)
print("Step 8: 数据保留在 Redis 中")
print("=" * 60)

with SearchIndex.from_existing(INDEX_NAME, redis_url=REDIS_URL) as index:
    info = index.info()
    print(f"  📌 索引名: {INDEX_NAME}")
    print(f"  📌 文档数: {info.get('num_docs', 'N/A')}")
    print(f"  📌 Redis key 前缀: {index.prefix}")
    print(f"  📌 连接地址: {REDIS_URL}")

print()
print("  🔍 现在可以用 redis-cli 探索数据：")
print()
print("     redis-cli")
print("     > FT._LIST")
print("     > FT.INFO demo-movie-idx")
print("     > FT.SEARCH demo-movie-idx '*' LIMIT 0 10")
print("     > JSON.GET movie:Inception $")
print("     > FT.SEARCH demo-movie-idx '@genre:{sci-fi}'")
print()

print("🎉 RedisVL 演示完成！数据已保留在 Redis 中。")