diff --git a/comps/dataprep/src/integrations/redis.py b/comps/dataprep/src/integrations/redis.py index fa733ad118..0528cd6210 100644 --- a/comps/dataprep/src/integrations/redis.py +++ b/comps/dataprep/src/integrations/redis.py @@ -52,6 +52,10 @@ TIMEOUT_SECONDS = int(os.getenv("TIMEOUT_SECONDS", 600)) SEARCH_BATCH_SIZE = int(os.getenv("SEARCH_BATCH_SIZE", 10)) +# Vector Schema Configuration +DEFAULT_VECTOR_SCHEMA = {"algorithm": "HNSW", "m": 16, "ef_construction": 200} +VECTOR_SCHEMA = os.getenv("VECTOR_SCHEMA", json.dumps(DEFAULT_VECTOR_SCHEMA)) + # Redis Connection Information REDIS_HOST = os.getenv("REDIS_HOST", "localhost") REDIS_PORT = int(os.getenv("REDIS_PORT", 6379)) @@ -200,6 +204,12 @@ async def ingest_chunks_to_redis(file_name: str, chunks: List, embedder, index_n # if data will be saved to a different index name than the default one ingest_index_name = index_name if index_name else INDEX_NAME + # Parse vector schema + try: + vector_schema = json.loads(VECTOR_SCHEMA) + except json.JSONDecodeError as e: + logger.error(f"Invalid VECTOR_SCHEMA format: {e}") + vector_schema = DEFAULT_VECTOR_SCHEMA file_ids = [] for i in range(0, num_chunks, batch_size): @@ -214,6 +224,7 @@ async def ingest_chunks_to_redis(file_name: str, chunks: List, embedder, index_n embedding=embedder, index_name=ingest_index_name, redis_url=REDIS_URL, + vector_schema=vector_schema, ) if logflag: logger.info(f"[ redis ingest chunks ] keys: {keys}") diff --git a/comps/retrievers/src/integrations/config.py b/comps/retrievers/src/integrations/config.py index 070a9f541d..d791d9fcf3 100644 --- a/comps/retrievers/src/integrations/config.py +++ b/comps/retrievers/src/integrations/config.py @@ -46,7 +46,9 @@ def get_boolean_env_var(var_name, default_value=False): LOCAL_EMBEDDING_MODEL = os.getenv("LOCAL_EMBEDDING_MODEL", "maidalun1020/bce-embedding-base_v1") TEI_EMBEDDING_ENDPOINT = os.getenv("TEI_EMBEDDING_ENDPOINT", "") BRIDGE_TOWER_EMBEDDING = os.getenv("BRIDGE_TOWER_EMBEDDING", False) + HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN", "") +ENABLE_SCHEMA = get_boolean_env_var("ENABLE_SCHEMA", False) # OpenAI OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") diff --git a/comps/retrievers/src/integrations/redis.py b/comps/retrievers/src/integrations/redis.py index 38dd4f4f3e..81cadc7173 100644 --- a/comps/retrievers/src/integrations/redis.py +++ b/comps/retrievers/src/integrations/redis.py @@ -26,6 +26,7 @@ from .config import ( BRIDGE_TOWER_EMBEDDING, EMBED_MODEL, + ENABLE_SCHEMA, HF_TOKEN, INDEX_NAME, INDEX_SCHEMA, @@ -100,7 +101,13 @@ async def _initialize_client(self, index_name=INDEX_NAME) -> Redis: client = Redis( embedding=self.embeddings, index_name=index_name, index_schema=INDEX_SCHEMA, redis_url=REDIS_URL ) + elif ENABLE_SCHEMA: + logger.info(f"generate redis instance with index_schema:{INDEX_SCHEMA}") + client = Redis( + embedding=self.embeddings, index_name=index_name, index_schema=INDEX_SCHEMA, redis_url=REDIS_URL + ) else: + logger.info(f"generate redis instance with index_name:{INDEX_NAME}") client = Redis(embedding=self.embeddings, index_name=index_name, redis_url=REDIS_URL) return client except Exception as e: diff --git a/comps/retrievers/src/integrations/redis_schema_hnsw.yml b/comps/retrievers/src/integrations/redis_schema_hnsw.yml new file mode 100644 index 0000000000..92f0fae3a7 --- /dev/null +++ b/comps/retrievers/src/integrations/redis_schema_hnsw.yml @@ -0,0 +1,9 @@ +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +vector: + - name: content_vector + algorithm: HNSW + datatype: FLOAT32 + dims: 768 + distance_metric: COSINE