152 lines
5.9 KiB
Python
152 lines
5.9 KiB
Python
"""
|
|
Vector storage configuration for the RAG solution using LlamaIndex and Qdrant.
|
|
|
|
This module provides initialization and configuration for:
|
|
- Qdrant vector storage connection
|
|
- Embedding model based on configured strategy
|
|
- Automatic collection creation
|
|
"""
|
|
|
|
import os
|
|
from typing import Optional
|
|
|
|
from llama_index.core import VectorStoreIndex
|
|
from llama_index.vector_stores.qdrant import QdrantVectorStore
|
|
from loguru import logger
|
|
from qdrant_client import QdrantClient
|
|
|
|
# Import the new configuration module
|
|
from config import get_embedding_model
|
|
|
|
|
|
def initialize_vector_storage(
|
|
collection_name: str = "documents_llamaindex",
|
|
host: str = "localhost",
|
|
port: int = 6333,
|
|
grpc_port: int = 6334,
|
|
) -> tuple[QdrantVectorStore, VectorStoreIndex]:
|
|
"""
|
|
Initialize Qdrant vector storage with embedding model based on configured strategy.
|
|
|
|
Args:
|
|
collection_name: Name of the Qdrant collection
|
|
host: Qdrant host address
|
|
port: Qdrant REST API port
|
|
grpc_port: Qdrant gRPC API port
|
|
|
|
Returns:
|
|
Tuple of (QdrantVectorStore, VectorStoreIndex)
|
|
"""
|
|
logger.info(f"Initializing vector storage with collection: {collection_name}")
|
|
|
|
try:
|
|
# Initialize Qdrant client
|
|
client = QdrantClient(host=host, port=port)
|
|
|
|
# Get the embedding model based on the configured strategy
|
|
embed_model = get_embedding_model()
|
|
|
|
# Get a test embedding to determine the correct dimensions
|
|
test_embedding = embed_model.get_text_embedding("test")
|
|
embedding_dimension = len(test_embedding)
|
|
logger.info(f"Detected embedding dimension: {embedding_dimension}")
|
|
|
|
# Check if collection exists, create if it doesn't
|
|
collections = client.get_collections().collections
|
|
collection_names = [coll.name for coll in collections]
|
|
|
|
if collection_name not in collection_names:
|
|
logger.info(f"Collection '{collection_name}' does not exist, creating...")
|
|
client.create_collection(
|
|
collection_name=collection_name,
|
|
vectors_config={
|
|
"size": embedding_dimension, # Use the actual embedding size
|
|
"distance": "Cosine", # Cosine distance is commonly used
|
|
},
|
|
)
|
|
logger.info(
|
|
f"Collection '{collection_name}' created successfully with dimension {embedding_dimension}"
|
|
)
|
|
else:
|
|
logger.info(f"Collection '{collection_name}' already exists")
|
|
# Get the actual collection config to determine the vector size
|
|
collection_info = client.get_collection(collection_name)
|
|
# Access the vector configuration properly - handle different possible structures
|
|
if (
|
|
hasattr(collection_info.config.params, "vectors")
|
|
and collection_info.config.params.vectors is not None
|
|
):
|
|
existing_dimension = collection_info.config.params.vectors.size
|
|
if existing_dimension != embedding_dimension:
|
|
logger.warning(
|
|
f"Existing collection dimension ({existing_dimension}) doesn't match embedding dimension ({embedding_dimension}), recreating..."
|
|
)
|
|
# Delete and recreate the collection with the correct dimensions
|
|
client.delete_collection(collection_name)
|
|
client.create_collection(
|
|
collection_name=collection_name,
|
|
vectors_config={
|
|
"size": embedding_dimension, # Use the detected size
|
|
"distance": "Cosine",
|
|
},
|
|
)
|
|
logger.info(
|
|
f"Collection '{collection_name}' recreated with dimension {embedding_dimension}"
|
|
)
|
|
else:
|
|
logger.info(
|
|
f"Using existing collection with matching dimension: {embedding_dimension}"
|
|
)
|
|
else:
|
|
# Last resort: recreate the collection with the correct dimensions
|
|
logger.warning(
|
|
f"Could not determine vector dimension for existing collection, recreating..."
|
|
)
|
|
# Delete and recreate the collection with the correct dimensions
|
|
client.delete_collection(collection_name)
|
|
client.create_collection(
|
|
collection_name=collection_name,
|
|
vectors_config={
|
|
"size": embedding_dimension, # Use the detected size
|
|
"distance": "Cosine",
|
|
},
|
|
)
|
|
logger.info(
|
|
f"Collection '{collection_name}' recreated with dimension {embedding_dimension}"
|
|
)
|
|
|
|
# Initialize the Qdrant vector store
|
|
vector_store = QdrantVectorStore(client=client, collection_name=collection_name)
|
|
|
|
# Create index from vector store with the embedding model we already created
|
|
index = VectorStoreIndex.from_vector_store(
|
|
vector_store=vector_store, embed_model=embed_model
|
|
)
|
|
|
|
logger.info("Vector storage initialized successfully")
|
|
return vector_store, index
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize vector storage: {str(e)}")
|
|
raise
|
|
|
|
|
|
def get_vector_store_and_index() -> tuple[QdrantVectorStore, VectorStoreIndex]:
|
|
"""
|
|
Convenience function to get the initialized vector store and index.
|
|
|
|
Returns:
|
|
Tuple of (QdrantVectorStore, VectorStoreIndex)
|
|
"""
|
|
return initialize_vector_storage()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Example usage
|
|
logger.info("Testing vector storage initialization...")
|
|
try:
|
|
vector_store, index = get_vector_store_and_index()
|
|
logger.info("Vector storage test successful!")
|
|
except Exception as e:
|
|
logger.error(f"Vector storage test failed: {e}")
|