@@ -0,0 +1,44 @@ +import redis + +from django.core.management.base import BaseCommand + +from django.conf import settings +from redis.commands.search.field import TagField, TextField, VectorField +from redis.commands.search.indexDefinition import IndexDefinition, IndexType + + +class Command(BaseCommand): + def handle(self, *args, **options) -> None: + """ + A command for creating indexes for storing a chunk's data (content, vectors, etc.) + """ + redis_client = redis.Redis(host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=0) + index_configs = ( + ('ml_model-index', 3072), + ('ml_model-index-1536', 1536), + ) + + for index_name, dim in index_configs: + try: + redis_client.ft(index_name).info() + self.stdout.write(f'Index {index_name} already exists') + except Exception: + message_uid = TagField('message_uid') + chunk_id = TextField('chunk_id') + section_text = TextField('section_text') + section_embeddings = VectorField( + 'section_embeddings', + 'FLAT', + { + 'TYPE': 'FLOAT32', + 'DIM': dim, + 'DISTANCE_METRIC': 'COSINE', + 'INITIAL_CAP': 10_000, + }, + ) + fields = [message_uid, chunk_id, section_text, section_embeddings] + redis_client.ft(index_name).create_index( + fields=fields, + definition=IndexDefinition(prefix=['ml_model:messages:'], index_type=IndexType.HASH), + ) + self.stdout.write(f'Index {index_name} has been successfully created') @@ -10,9 +10,6 @@ class MLModelConfig(AppConfig): def ready(self): from .signals import create_settings - from .utils import create_redis_search_index setting_changed.connect(create_settings) - create_redis_search_index() - return super().ready() @@ -61,36 +61,3 @@ def count_openrouter_tokens(model_name: str, messages: List[Dict[str, Any]], out output_tokens = len(encoding.encode(output)) return (input_tokens, output_tokens) - -def create_redis_search_index() -> None: - """ - A method for creating an index for storing a chunk's data (content, vectors, etc.) - """ - redis_client = redis.Redis(host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=0) - index_configs = ( - ('ml_model-index', 3072), - ('ml_model-index-1536', 1536), - ) - - for index_name, dim in index_configs: - try: - redis_client.ft(index_name).info() - except Exception: - message_uid = TagField('message_uid') - chunk_id = TextField('chunk_id') - section_text = TextField('section_text') - section_embeddings = VectorField( - 'section_embeddings', - 'FLAT', - { - 'TYPE': 'FLOAT32', - 'DIM': dim, - 'DISTANCE_METRIC': 'COSINE', - 'INITIAL_CAP': 10_000, - }, - ) - fields = [message_uid, chunk_id, section_text, section_embeddings] - redis_client.ft(index_name).create_index( - fields=fields, - definition=IndexDefinition(prefix=['ml_model:messages:'], index_type=IndexType.HASH), - ) @@ -32,6 +32,7 @@ services: - /bin/sh - -c - | + python manage.py create_indexes python manage.py compilemessages python -m uvicorn backend.asgi:application --host 0.0.0.0 --ws wsproto --http httptools --lifespan off --log-level info networks: