feat: add drop index in redis and fix prefix generate logic (#1857)

# Description

Add `drop_index` for redis

RediSearch: [RediSearch quick
start](https://redis.io/docs/stack/search/quick_start/)

# How to use

```
from langchain.vectorstores.redis import Redis

Redis.drop_index(index_name="doc",delete_documents=False)
```
This commit is contained in:
Xin Qiu 2023-03-23 10:44:42 +08:00 committed by GitHub
parent 12f868b292
commit ea142f6a32
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -2,6 +2,7 @@
from __future__ import annotations from __future__ import annotations
import json import json
import logging
import uuid import uuid
from typing import Any, Callable, Iterable, List, Mapping, Optional from typing import Any, Callable, Iterable, List, Mapping, Optional
@ -13,6 +14,8 @@ from langchain.embeddings.base import Embeddings
from langchain.utils import get_from_dict_or_env from langchain.utils import get_from_dict_or_env
from langchain.vectorstores.base import VectorStore from langchain.vectorstores.base import VectorStore
logger = logging.getLogger()
def _check_redis_module_exist(client: RedisType, module: str) -> bool: def _check_redis_module_exist(client: RedisType, module: str) -> bool:
return module in [m["name"] for m in client.info().get("modules", {"name": ""})] return module in [m["name"] for m in client.info().get("modules", {"name": ""})]
@ -180,7 +183,7 @@ class Redis(VectorStore):
# name of the search index if not given # name of the search index if not given
if not index_name: if not index_name:
index_name = uuid.uuid4().hex index_name = uuid.uuid4().hex
prefix = "doc" # prefix for the document keys prefix = f"doc:{index_name}" # prefix for the document keys
distance_metric = ( distance_metric = (
"COSINE" # distance metric for the vectors (ex. COSINE, IP, L2) "COSINE" # distance metric for the vectors (ex. COSINE, IP, L2)
) )
@ -201,7 +204,7 @@ class Redis(VectorStore):
# Check if index exists # Check if index exists
try: try:
client.ft(index_name).info() client.ft(index_name).info()
print("Index already exists") logger.info("Index already exists")
except: # noqa except: # noqa
# Create Redis Index # Create Redis Index
client.ft(index_name).create_index( client.ft(index_name).create_index(
@ -211,7 +214,7 @@ class Redis(VectorStore):
pipeline = client.pipeline() pipeline = client.pipeline()
for i, text in enumerate(texts): for i, text in enumerate(texts):
key = f"{prefix}:{str(uuid.uuid4().hex)}" key = f"{prefix}:{i}"
metadata = metadatas[i] if metadatas else {} metadata = metadatas[i] if metadatas else {}
pipeline.hset( pipeline.hset(
key, key,
@ -225,3 +228,34 @@ class Redis(VectorStore):
) )
pipeline.execute() pipeline.execute()
return cls(redis_url, index_name, embedding.embed_query) return cls(redis_url, index_name, embedding.embed_query)
@staticmethod
def drop_index(
index_name: str,
delete_documents: bool,
**kwargs: Any,
) -> bool:
redis_url = get_from_dict_or_env(kwargs, "redis_url", "REDIS_URL")
try:
import redis
except ImportError:
raise ValueError(
"Could not import redis python package. "
"Please install it with `pip install redis`."
)
try:
# We need to first remove redis_url from kwargs,
# otherwise passing it to Redis will result in an error.
kwargs.pop("redis_url")
client = redis.from_url(url=redis_url, **kwargs)
except ValueError as e:
raise ValueError(f"Your redis connected error: {e}")
# Check if index exists
try:
client.ft(index_name).dropindex(delete_documents)
logger.info("Drop index")
return True
except: # noqa
# Index not exist
return False