evoagentx/rag/rag.py

import os
import json
import asyncio
from uuid import uuid4
from datetime import datetime
from typing import List, Union, Optional, Sequence, Dict, Any, Tuple
from llama_index.core.schema import NodeWithScore, TextNode, ImageNode, RelatedNodeInfo

from .rag_config import RAGConfig
from .readers import LLamaIndexReader, MultimodalReader
from .indexings import IndexFactory, BaseIndexWrapper
from .chunkers import ChunkFactory
from .embeddings import EmbeddingFactory, EmbeddingProvider
from .retrievers import RetrieverFactory, BaseRetrieverWrapper
from .postprocessors import PostprocessorFactory
from .indexings.base import IndexType
from .retrievers.base import RetrieverType
from .schema import Corpus, ChunkMetadata, IndexMetadata, Query, RagResult, ImageChunk, TextChunk
from evoagentx.storages.base import StorageHandler
from evoagentx.storages.schema import IndexStore
from evoagentx.models.base_model import BaseLLM
from evoagentx.core.logging import logger


class RAGEngine:
    def __init__(self, config: RAGConfig, storage_handler: StorageHandler, llm: Optional[BaseLLM] = None):
        self.config = config
        self.storage_handler = storage_handler  # Maybe reinit the vector_store by the load funcion.
        self.embedding_factory = EmbeddingFactory()
        self.index_factory = IndexFactory()
        self.chunk_factory = ChunkFactory()
        self.retriever_factory = RetrieverFactory()
        self.postprocessor_factory = PostprocessorFactory()
        self.llm = llm  # LLM for entity extractor
        
        # Set chunk class based on modality config
        logger.info(f"RAGEngine modality config: {self.config.modality}")
        
        if self.config.modality == "multimodal":
            self.chunk_class = ImageChunk
        else:
            self.chunk_class = TextChunk

        # Initialize reader based on modality
        if self.config.modality == "multimodal":
            self.reader = MultimodalReader(
                recursive=self.config.reader.recursive,
                exclude_hidden=self.config.reader.exclude_hidden,
                num_files_limits=self.config.reader.num_files_limit,
                errors=self.config.reader.errors
            )
        else:
            self.reader = LLamaIndexReader(
                recursive=self.config.reader.recursive,
                exclude_hidden=self.config.reader.exclude_hidden,
                num_workers=self.config.num_workers,
                num_files_limits=self.config.reader.num_files_limit,
                custom_metadata_function=self.config.reader.custom_metadata_function,
                extern_file_extractor=self.config.reader.extern_file_extractor,
                errors=self.config.reader.errors,
                encoding=self.config.reader.encoding
            )

        # Initialize embedding model. 
        self.embed_model = self.embedding_factory.create(
            provider=self.config.embedding.provider,
            model_config=self.config.embedding.model_dump(exclude_unset=True),
        )

        # Dynamic Check the dimensions in StorageHandler
        if (self.storage_handler.vector_store is not None) and (self.embed_model.dimensions is not None):
            if self.storage_handler.storageConfig.vectorConfig.dimensions != self.embed_model.dimensions:
                logger.warning("The dimensions in vector_store is not equal with embed_model. Reiniliaze vector_store.")
                self.storage_handler.storageConfig.vectorConfig.dimensions = self.embed_model.dimensions
                self.storage_handler._init_vector_store()

        # Initialize chunker (skip for multimodal)
        if self.config.modality == "multimodal":
            self.chunker = None  # No chunking for images
        else:
            self.chunker = self.chunk_factory.create(
                strategy=self.config.chunker.strategy,
                embed_model=self.embed_model.get_embedding_model(),
                chunker_config={
                    "chunk_size": self.config.chunker.chunk_size,
                    "chunk_overlap": self.config.chunker.chunk_overlap,
                    "max_chunks": self.config.chunker.max_chunks
                }
            )

        # Initialize indices and retrievers
        self.indices: Dict[str, Dict[str, BaseIndexWrapper]] = {}  # Nested: {corpus_id: {index_type: index}}
        self.retrievers: Dict[str, Dict[str, BaseRetrieverWrapper]] = {}  # Nested: {corpus_id: {index_type: retriever}}

    def read(self, file_paths: Union[Sequence[str], str], 
             exclude_files: Optional[Union[str, List, Tuple, Sequence]] = None,
             filter_file_by_suffix: Optional[Union[str, List, Tuple, Sequence]] = None,
             merge_by_file: bool = False,
             show_progress: bool = False,
             corpus_id: str = None) -> Corpus:
        """Load and chunk documents from files.

        Reads files from specified paths, processes them into documents, and chunks them into a Corpus.

        Args:
            file_paths (Union[Sequence[str], str]): Path(s) to files or directories.
            exclude_files (Optional[Union[str, List, Tuple, Sequence]]): Files to exclude.
            filter_file_by_suffix (Optional[Union[str, List, Tuple, Sequence]]): Filter files by suffix (e.g., '.pdf').
            merge_by_file (bool): Merge documents by file.
            show_progress (bool): Show loading progress.
            corpus_id (Optional[str]): Identifier for the corpus. Defaults to a UUID if None.

        Returns:
            Corpus: The chunked corpus containing processed document chunks.

        Raises:
            Exception: If document reading or chunking fails.
        """
        try:
            corpus_id = corpus_id or str(uuid4())
            documents = self.reader.load(
                file_paths=file_paths,
                exclude_files=exclude_files,
                filter_file_by_suffix=filter_file_by_suffix,
                merge_by_file=merge_by_file,
                show_progress=show_progress
            )
            if self.config.modality == "multimodal":
                # No chunking - convert ImageDocuments directly to ImageChunks
                image_chunks = []
                for doc in documents:
                    # Get image path from document - try multiple possible attributes
                    image_path = getattr(doc, 'image_path', None) or doc.metadata.get('file_path')
                    image_mimetype = getattr(doc, 'image_mimetype', None)
                    
                    image_chunk = self.chunk_class(
                        image_path=image_path,
                        image_mimetype=image_mimetype,
                        chunk_id=doc.metadata.get('file_name', f'img_{len(image_chunks)}'),
                        metadata=ChunkMetadata(
                            doc_id=doc.metadata.get('file_name', f'doc_{len(image_chunks)}'),
                            corpus_id=corpus_id,
                            **doc.metadata
                        )
                    )
                    image_chunks.append(image_chunk)
                corpus = Corpus(chunks=image_chunks, corpus_id=corpus_id)
                logger.info(f"Read {len(documents)} multimodal documents (no chunking) for corpus {corpus_id}")
            else:
                corpus = self.chunker.chunk(documents)
                corpus.corpus_id = corpus_id
                logger.info(f"Read {len(documents)} documents and created {len(corpus.chunks)} chunks for corpus {corpus_id}")
            return corpus
        except Exception as e:
            logger.error(f"Failed to read documents for corpus {corpus_id}: {str(e)}")
            raise

    def add(self, index_type: str, nodes: Union[Corpus, List[NodeWithScore], List[TextNode], List[ImageNode]], 
            corpus_id: str = None) -> None:
        """Add nodes to an index for a specific corpus.

        Initializes an index if it doesn't exist and inserts nodes, updating metadata with corpus_id and index_type.

        Args:
            index_type (str): Type of index (e.g., VECTOR, GRAPH).
            nodes (Union[Corpus, List[NodeWithScore], List[TextNode]]): Nodes or Corpus to add.
            corpus_id (str, optional): Identifier for the corpus. Defaults to a UUID if None.

        Return:
            return a sequence with id of each added node.
            
        Raises:
            Exception: If index creation or node insertion fails.
        """
        try:
            corpus_id = corpus_id or str(uuid4())
            if corpus_id not in self.indices:
                self.indices[corpus_id] = {}
                self.retrievers[corpus_id] = {}

            if index_type not in self.indices[corpus_id]:
                index = self.index_factory.create(
                    index_type=index_type,
                    embed_model=self.embed_model.get_embedding_model(),
                    storage_handler=self.storage_handler,
                    index_config=self.config.index.model_dump(exclude_unset=True) if self.config.index else {},
                    llm=self.llm,
                )
                self.indices[corpus_id][index_type] = index
                self.retrievers[corpus_id][index_type] = self.retriever_factory.create(
                    retriever_type=self.config.retrieval.retrivel_type,
                    llm=self.llm,
                    index=index.get_index(),
                    graph_store=index.get_index().storage_context.graph_store,
                    embed_model=self.embed_model.get_embedding_model(),
                    query=Query(query_str="", top_k=self.config.retrieval.top_k if self.config.retrieval else 5),
                    storage_handler=self.storage_handler,
                    chunk_class=self.chunk_class  
                )

            nodes_to_insert = nodes.to_llama_nodes() if isinstance(nodes, Corpus) else nodes
            for node in nodes_to_insert:
                node.metadata.update({"corpus_id": corpus_id, "index_type": index_type})
            nodes_ids = self.indices[corpus_id][index_type].insert_nodes(nodes_to_insert)
            logger.info(f"Added {len(nodes_to_insert)} nodes to {index_type} index for corpus {corpus_id}")
            return nodes_ids
        except Exception as e:
            logger.error(f"Failed to add nodes to {index_type} index for corpus {corpus_id}: {str(e)}")
            return []

    def delete(self, corpus_id: str, index_type: Optional[str] = None, 
               node_ids: Optional[Union[str, List[str]]] = None, 
               metadata_filters: Optional[Dict[str, Any]] = None) -> None:
        """Delete nodes or an entire index from a corpus.

        Removes specific nodes by ID or metadata filters, or deletes the entire index if no filters are provided.

        Args:
            corpus_id (str): Identifier for the corpus.
            index_type (Optional[IndexType]): Specific index type to delete from. If None, affects all indices.
            node_ids (Union[str, Optional[List[str]]]): List of node IDs to delete.
            metadata_filters (Optional[Dict[str, Any]]): Metadata filters to select nodes for deletion.

        Raises:
            Exception: If deletion fails.
        """
        try:
            if corpus_id not in self.indices:
                logger.warning(f"No indices found for corpus {corpus_id}")
                return

            target_indices = [index_type] if index_type else self.indices[corpus_id].keys()
            for idx_type in list(target_indices):  # Use list to avoid runtime modification issues
                if idx_type not in self.indices[corpus_id]:
                    logger.warning(f"Index type {idx_type} not found for corpus {corpus_id}")
                    continue

                index = self.indices[corpus_id][idx_type]
                if node_ids or metadata_filters:
                    # Convert single node_id to list for consistency
                    node_ids_list = [node_ids] if isinstance(node_ids, str) else node_ids
                    index.delete_nodes(node_ids=node_ids_list, metadata_filters=metadata_filters)
                    logger.info(f"Deleted nodes from {idx_type} index for corpus {corpus_id}")
                else:
                    # Delete entire index
                    index.clear()
                    del self.indices[corpus_id][idx_type]
                    del self.retrievers[corpus_id][idx_type]
                    logger.info(f"Deleted entire {idx_type} index for corpus {corpus_id}")

            # Clean up corpus if no indices remain
            if not self.indices[corpus_id]:
                del self.indices[corpus_id]
                del self.retrievers[corpus_id]
                logger.info(f"Removed empty corpus {corpus_id}")

        except Exception as e:
            logger.error(f"Failed to delete from corpus {corpus_id}, index {index_type}: {str(e)}")
            raise

    def clear(self, corpus_id: Optional[str] = None) -> None:
        """Clear all indices for a specific corpus or all corpora.

        Args:
            corpus_id (Optional[str]): Specific corpus to clear. If None, clears all corpora.

        Raises:
            Exception: If clearing fails.
        """
        try:
            target_corpora = [corpus_id] if corpus_id else list(self.indices.keys())
            for cid in target_corpora:
                if cid not in self.indices:
                    logger.warning(f"No indices found for corpus {cid}")
                    continue

                for idx_type in list(self.indices[cid].keys()):
                    index = self.indices[cid][idx_type]
                    index.clear()
                    del self.indices[cid][idx_type]
                    del self.retrievers[cid][idx_type]
                    logger.info(f"Cleared {idx_type} index for corpus {cid}")

                # Clean up corpus if no indices remain
                del self.indices[cid]
                del self.retrievers[cid]
                logger.info(f"Cleared corpus {cid}")

        except Exception as e:
            logger.error(f"Failed to clear indices for corpus {corpus_id or 'all'}: {str(e)}")
            raise

    def save(self, output_path: Optional[str] = None, corpus_id: Optional[str] = None, 
                index_type: Optional[str] = None, table: Optional[str] = None, 
                graph_exported: bool = False) -> None:
        """Save indices to files or database.

        Serializes corpus chunks to JSONL files and metadata to JSON files if output_path is provided,
        or saves to the SQLite database via StorageHandler if output_path is None.

        Args:
            output_path (Optional[str]): Directory to save JSONL and JSON files. If None, saves to database.
            corpus_id (Optional[str]): Specific corpus to save. If None, saves all corpora.
            index_type (Optional[str]): Specific index type to save. If None, saves all indices.
            table (Optional[str]): Database table name for index data. Defaults to 'indexing' if None.
            graph_exported (bool): If True, export graph nodes and relations for graph indices. Defaults to False.

        Raises:
            Exception: If saving fails or file operations encounter errors.
        """
        try:
            target_corpora = [corpus_id] if corpus_id else list(self.indices.keys())
            table = table or "indexing"

            for cid in target_corpora:
                if cid not in self.indices:
                    logger.warning(f"No indices found for corpus {cid}")
                    continue

                target_indices = [index_type] if index_type and index_type in self.indices[cid] else self.indices[cid].keys()
                for idx_type in target_indices:
                    index = self.indices[cid][idx_type]

                    # Skip saving for graph indices unless graph_exported is True
                    if idx_type == IndexType.GRAPH and not graph_exported:
                        logger.warning(f"Skipping save for graph index {idx_type} in corpus {cid} as graph_exported is False")
                        continue

                    # For graph indices, include kg_nodes and kg_rels if graph_exported is True
                    if idx_type == IndexType.GRAPH and graph_exported:
                        index.build_kv_store()  # FIXME: Computer's memory may increase dramatically

                    # Convert index nodes to Corpus
                    chunks = [
                        self.chunk_class.from_llama_node(node_data)
                        for node_id, node_data in index.id_to_node.items()
                    ]
                    corpus = Corpus(chunks=chunks, corpus_id=cid)

                    # Prepare metadata
                    vector_config = self.storage_handler.storageConfig.vectorConfig.model_dump() if self.storage_handler.storageConfig.vectorConfig else {}
                    graph_config = self.storage_handler.storageConfig.graphConfig.model_dump() if self.storage_handler.storageConfig.graphConfig else {}
                    metadata = IndexMetadata(
                        corpus_id=cid,
                        index_type=idx_type,
                        collection_name=vector_config.get("qdrant_collection_name", "default_collection"),
                        dimension=self.embed_model.dimensions,
                        vector_db_type=vector_config.get("vector_name", None),
                        graph_db_type=graph_config.get("graph_name", None),
                        embedding_model_name=self.config.embedding.model_name,
                        date=str(datetime.now()),
                    )

                    if output_path:
                        # File-based saving
                        os.makedirs(output_path, exist_ok=True)
                        safe_cid = "".join(c if c.isalnum() or c in ["-", "_"] else "_" for c in cid)
                        safe_idx_type = "".join(c if c.isalnum() or c in ["-", "_"] else "_" for c in idx_type)
                        nodes_file = os.path.join(output_path, f"{safe_cid}_{safe_idx_type}_nodes.jsonl")
                        metadata_file = os.path.join(output_path, f"{safe_cid}_{safe_idx_type}_metadata.json")

                        # Save corpus as JSONL
                        corpus.to_jsonl(nodes_file, indent=0)
                        logger.info(f"Saved {len(corpus.chunks)} chunks to {nodes_file}")

                        # Save metadata as JSON
                        with open(metadata_file, "w", encoding="utf-8") as f:
                            json.dump(metadata.model_dump(), f, indent=2, ensure_ascii=False)
                        logger.info(f"Saved metadata to {metadata_file}")
                    else:
                        # Database saving
                        index_data = {
                            "corpus_id": cid,
                            "content": corpus.model_dump(),
                            "date": str(datetime.now()),
                            "metadata": metadata.model_dump()
                        }
                        self.storage_handler.save_index(index_data, table=table)
                        logger.info(f"Saved {idx_type} index with {len(corpus.chunks)} chunks for corpus {cid} to database table {table}")

        except Exception as e:
            logger.error(f"Failed to save indices for corpus {corpus_id or 'all'}: {str(e)}")
            raise

    def load(self, source: Optional[str] = None, corpus_id: Optional[str] = None, 
             index_type: Optional[str] = None, table: Optional[str] = None) -> None:
        """Load indices from files or database.

        Reconstructs indices and retrievers from JSONL/JSON files or SQLite database records.
        Validates the embedding model name and dimension before reinitializing the embedding model.

        Args:
            source (Optional[str]): Directory containing JSONL/JSON files. If None, loads from database.
            corpus_id (Optional[str]): Specific corpus to load. If None, loads all corpora.
            index_type (Optional[str]): Specific index type to load. If None, loads all indices.
            table (Optional[str]): Database table name for index data. Defaults to 'indexing' if None.

        Returns:
            The Sequence with id of loaded chunk.
        
        Raises:
            Exception: If loading fails due to file or database errors, invalid data, or unsupported embedding model/dimension.
        
        Warning:
            Try to call this function may cause some Bugs, when you load the nodes from file or database storage systems at twice. 
            Because All the indexing share the same storage backend from storageHandler.
            For example:
            The vector database (.e.g Faiss) can insert again, even thougt there is a same node.
        """
        try:
            table = table or "indexing"
            config_dimension = self.storage_handler.storageConfig.vectorConfig.dimensions
            loaded_chunk_ids: List[str] = []

            if source:
                # File-based loading
                if not os.path.exists(source):
                    logger.error(f"Source directory {source} does not exist")
                    raise FileNotFoundError(f"Source directory {source} does not exist")

                for file_name in os.listdir(source):
                    if not file_name.endswith("_metadata.json"):
                        continue
                    parts = file_name.split("_")
                    if len(parts) < 3:
                        logger.warning(f"Skipping invalid metadata file: {file_name}")
                        continue
                    cid = "_".join(parts[:-2])
                    idx_type = parts[-2]
                    
                    if (corpus_id and corpus_id != cid) or (index_type and index_type != idx_type):
                        continue

                    metadata_file = os.path.join(source, file_name)
                    nodes_file = os.path.join(source, f"{cid}_{idx_type}_nodes.jsonl")

                    # Load metadata
                    with open(metadata_file, "r", encoding="utf-8") as f:
                        metadata = IndexMetadata.model_validate(json.load(f))

                    # Validate embedding model
                    if not self.embed_model.validate_model(self.config.embedding.provider, metadata.embedding_model_name):
                        raise ValueError(
                            f"Embedding model '{metadata.embedding_model_name}' is not supported by provider '{self.config.embedding.provider}'. "
                            f"Supported models: {EmbeddingProvider.SUPPORTED_MODELS.get(self.config.embedding.provider, [])}"
                        )

                    # Validate dimension
                    if metadata.dimension != config_dimension:
                        raise ValueError(
                            f"Embedding dimension {metadata.dimension} in metadata does not match configured dimension {config_dimension}."
                        )

                    # Load corpus
                    if not os.path.exists(nodes_file):
                        logger.warning(f"Nodes file {nodes_file} not found for metadata {metadata_file}")
                        continue
                    corpus = Corpus.from_jsonl(nodes_file, corpus_id=cid)

                    # Reinitialize embedding model if needed
                    if metadata.embedding_model_name != self.config.embedding.model_name:
                        logger.info(f"Reinitializing embedding model to {metadata.embedding_model_name}")
                        self.embed_model = self.embedding_factory.create(
                            provider=self.config.embedding.provider,
                            model_config=self.config.embedding.model_dump(exclude_unset=True)
                        )

                    # Load index
                    chunk_ids = self._load_index(corpus, cid, idx_type)
                    loaded_chunk_ids.extend(chunk_ids)
                    logger.info(f"Loaded {idx_type} index with {len(corpus.chunks)} chunks for corpus {cid} from {nodes_file}")
            else:
                # Database loading
                records = self.storage_handler.load(tables=[table]).get(table, [])

                if not records:
                    logger.warning(f"No records found in table {table}")
                    return

                for record in records:
                    parsed = self.storage_handler.parse_result(record, IndexStore)
                    cid = parsed["corpus_id"]
                    idx_type = parsed["metadata"]["index_type"]
                    if (corpus_id and corpus_id != cid) or (index_type and index_type != idx_type):
                        continue

                    # Reconstruct corpus
                    chunks = []
                    for chunk_data in parsed["content"]["chunks"]:
                        metadata = ChunkMetadata.model_validate(chunk_data["metadata"])
                        
                        if self.config.modality == "multimodal":
                            # Create ImageChunk for multimodal mode
                            chunk = ImageChunk(
                                chunk_id=chunk_data["chunk_id"],
                                image_path=chunk_data["image_path"],
                                image_mimetype=chunk_data.get("image_mimetype"),
                                metadata=metadata,
                                embedding=chunk_data["embedding"],
                                excluded_embed_metadata_keys=chunk_data["excluded_embed_metadata_keys"],
                                excluded_llm_metadata_keys=chunk_data["excluded_llm_metadata_keys"],
                                relationships={k: RelatedNodeInfo(**v) for k, v in chunk_data["relationships"].items()}
                            )
                        else:
                            # Create TextChunk for text mode
                            chunk = TextChunk(
                                chunk_id=chunk_data["chunk_id"],
                                text=chunk_data["text"],
                                metadata=metadata,
                                embedding=chunk_data["embedding"],
                                start_char_idx=chunk_data["start_char_idx"],
                                end_char_idx=chunk_data["end_char_idx"],
                                excluded_embed_metadata_keys=chunk_data["excluded_embed_metadata_keys"],
                                excluded_llm_metadata_keys=chunk_data["excluded_llm_metadata_keys"],
                                relationships={k: RelatedNodeInfo(**v) for k, v in chunk_data["relationships"].items()}
                            )
                        chunks.append(chunk)
                    
                    corpus = Corpus(
                        chunks=chunks,
                        corpus_id=cid,
                        metadata=IndexMetadata.model_validate(parsed["metadata"])
                    )

                    # Validate embedding model
                    metadata = IndexMetadata.model_validate(parsed["metadata"])
                    if not self.embed_model.validate_model(self.config.embedding.provider, metadata.embedding_model_name):
                        raise ValueError(
                            f"Embedding model '{metadata.embedding_model_name}' is not supported by provider '{self.config.embedding.provider}'. "
                            f"Supported models: {EmbeddingProvider.SUPPORTED_MODELS.get(self.config.embedding.provider, [])}"
                        )

                    # Validate dimension
                    if metadata.dimension != config_dimension:
                        raise ValueError(
                            f"Embedding dimension {metadata.dimension} in metadata does not match configured dimension {config_dimension}."
                        )

                    # Reinitialize embedding model if needed
                    if metadata.embedding_model_name != self.config.embedding.model_name:
                        logger.info(f"Reinitializing embedding model to {metadata.embedding_model_name}")
                        self.embed_model = self.embedding_factory.create(
                            provider=self.config.embedding.provider,
                            model_config=self.config.embedding.model_dump(exclude_unset=True)
                        )

                    # Load index
                    chunk_ids = self._load_index(corpus, cid, idx_type)
                    loaded_chunk_ids.extend(chunk_ids)
                    logger.info(f"Loaded {idx_type} index with {len(corpus.chunks)} chunks for corpus {cid} from database table {table}")
            
            return loaded_chunk_ids
        except Exception as e:
            logger.error(f"Failed to load indices: {str(e)}")
            raise

    def _load_index(self, corpus: Corpus, corpus_id: str, index_type: str) -> Sequence[str]:
        """Helper method to load an index and its retriever."""
        try:
            if corpus_id not in self.indices:
                self.indices[corpus_id] = {}
                self.retrievers[corpus_id] = {}

            if index_type not in self.indices[corpus_id]:
                index = self.index_factory.create(
                    index_type=index_type,
                    embed_model=self.embed_model.get_embedding_model(),
                    storage_handler=self.storage_handler,
                    index_config=self.config.index.model_dump(exclude_unset=True) if self.config.index else {},
                    llm=self.llm
                )
                self.indices[corpus_id][index_type] = index
                retriever_type = RetrieverType.GRAPH if index_type == IndexType.GRAPH else RetrieverType.VECTOR
                self.retrievers[corpus_id][index_type] = self.retriever_factory.create(
                    retriever_type=retriever_type,
                    llm=self.llm,
                    index=index.get_index(),
                    graph_store=index.get_index().storage_context.graph_store,
                    embed_model=self.embed_model.get_embedding_model(),
                    query=Query(query_str="", top_k=self.config.retrieval.top_k if self.config.retrieval else 5),
                    storage_handler=self.storage_handler
                )

            nodes = corpus.to_llama_nodes()
           
            for node in nodes:
                node.metadata.update({"corpus_id": corpus_id, "index_type": index_type})
           
            chunk_ids = self.indices[corpus_id][index_type].load(nodes)
            logger.info(f"Inserted {len(nodes)} nodes into {index_type} index for corpus {corpus_id}")
            return chunk_ids
        except Exception as e:
            logger.error(f"Failed to load index for corpus {corpus_id}, index_type {index_type}: {str(e)}")
            raise

    async def aget(self, corpus_id: str, index_type: str, node_ids: List[str]) -> List[Union[TextChunk, ImageChunk]]:
        """Retrieve chunks by node_ids from the index."""
        try:
            chunks = await self.indices[corpus_id][index_type].get(node_ids=node_ids)
            logger.info(f"Retrieved {len(chunks)} chunks for node_ids: {node_ids}")
            return chunks
        except Exception as e:
            logger.error(f"Failed to get chunks: {str(e)}")
            return []

    async def query_async(self, query: Union[str, Query], corpus_id: Optional[str] = None,
                        query_transforms: Optional[List] = None) -> RagResult:
        """Execute a query across indices and return processed results asynchronously.

        Performs query preprocessing, asynchronous retrieval, and post-processing.

        Args:
            query (Union[str, Query]): Query string or Query object.
            corpus_id (Optional[str]): Specific corpus to query. If None, queries all corpora.
            query_transforms (Optional[List]): Query Transforms is used to augment query in pre-processing.

        Returns:
            RagResult: Retrieved chunks with scores and metadata.

        Raises:
            Exception: If query processing fails.
        """
        try:
            if isinstance(query, str):
                query = Query(query_str=query, top_k=self.config.retrieval.top_k)
            
            if not self.indices or (corpus_id and corpus_id not in self.indices):
                logger.warning(f"No indices found for corpus {corpus_id or 'any'}")
                return RagResult(corpus=Corpus(chunks=[]), scores=[], metadata={"query": query.query_str})

            # Pre-Processing
            if query_transforms and query_transforms is not None:
                for t in query_transforms:
                    query = t(query)

            results = []
            target_corpora = [corpus_id] if corpus_id else self.indices.keys()

            # Create all retrieval tasks
            tasks = []
            for cid in target_corpora:
                for idx_type, retriever in self.retrievers[cid].items():
                    if query.metadata_filters and query.metadata_filters.get("index_type") and \
                    query.metadata_filters["index_type"] != idx_type:
                        continue
                    
                    task = retriever.aretrieve(
                        Query(
                            query_str=query.query_str,
                            top_k=query.top_k or self.config.retrieval.top_k,
                            similarity_cutoff=query.similarity_cutoff,
                            keyword_filters=query.keyword_filters,
                            metadata_filters=query.metadata_filters
                        )
                    )
                    tasks.append((task, cid, idx_type))

            # Run all retrievals concurrently
            retrieval_tasks = [task for task, _, _ in tasks]
            retrieval_results = await asyncio.gather(*retrieval_tasks, return_exceptions=True)

            # Process results
            for (_, cid, idx_type), result in zip(tasks, retrieval_results):
                if isinstance(result, Exception):
                    logger.error(f"Retrieval failed for {idx_type} in corpus {cid}: {str(result)}")
                else:
                    results.append(result)
                    logger.info(f"Retrieved {len(result.corpus.chunks)} chunks from {idx_type} retriever for corpus {cid}")

            if not results:
                return RagResult(corpus=Corpus(chunks=[]), scores=[], metadata={"query": query.query_str})

            # Check the 'similarity_cutoff' and 'keyword_filters' in query
            query.similarity_cutoff = self.config.retrieval.similarity_cutoff if query.similarity_cutoff is None else query.similarity_cutoff
            query.keyword_filters = self.config.retrieval.keyword_filters if query.keyword_filters is None else query.keyword_filters

            postprocessor = self.postprocessor_factory.create(
                self.config.retrieval.postprocessor_type,
                query=query
            )
            final_result = postprocessor.postprocess(query, results)

            if query.metadata_filters:
                final_result.corpus.chunks = [
                    chunk for chunk in final_result.corpus.chunks
                    if all(chunk.metadata.model_dump().get(k) == v for k, v in query.metadata_filters.items())
                ]
                final_result.scores = [chunk.metadata.similarity_score for chunk in final_result.corpus.chunks]
                logger.info(f"Applied metadata filters, retained {len(final_result.corpus.chunks)} chunks")

            logger.info(f"Query returned {len(final_result.corpus.chunks)} chunks after post-processing")
            return final_result
        except Exception as e:
            logger.error(f"Query failed: {str(e)}")
            raise

    def query(self, query: Union[str, Query], corpus_id: Optional[str] = None,
            query_transforms: Optional[List] = None) -> RagResult:
        """Synchronous wrapper for the async query method."""
        return asyncio.run(self.query_async(query, corpus_id, query_transforms))