[Question]: How can I combine Vector DB and a new query engine?

[huangjia2019]

Question Validation

• [X] I have searched both the documentation and discord for an answer.

Question

For example, I want to use this query_engine: Joint QA Summary Query Engine https://docs.llamaindex.ai/en/stable/examples/query_engine/JointQASummary/

from llama_index.core.composability import QASummaryQueryEngineBuilder

query_engine_builder = QASummaryQueryEngineBuilder(
    llm=gpt4,
)
query_engine = query_engine_builder.build_from_documents(documents)

But I want to use a vertor DB, such as Milvus?

How to du that?

[dosubot[bot]]

To combine a Vector DB like Milvus with the Joint QA Summary Query Engine, you need to adapt the build_from_documents method in the QASummaryQueryEngineBuilder to utilize Milvus as your document store. Here's how you can do it:

1. Initialize the Milvus client.
2. Ingest document vectors into Milvus.
3. Build indices and query engines.

Here's an example of how you can modify the build_from_documents method:

from milvus import Milvus, DataType
from llama_index.core.indices.vector_store import VectorStoreIndex
from llama_index.core.indices.list.base import SummaryIndex
from llama_index.core.query_engine.router_query_engine import RouterQueryEngine
from llama_index.core.tools.query_engine import QueryEngineTool
from llama_index.core.ingestion import run_transformations
from llama_index.core.schema import Document
from llama_index.core.service_context import ServiceContext
from llama_index.core.storage.storage_context import StorageContext

class QASummaryQueryEngineBuilder:
    def __init__(
        self,
        storage_context: Optional[StorageContext] = None,
        service_context: Optional[ServiceContext] = None,
        summary_text: str = DEFAULT_SUMMARY_TEXT,
        qa_text: str = DEFAULT_QA_TEXT,
    ) -> None:
        self._storage_context = storage_context or StorageContext.from_defaults()
        self._service_context = service_context or ServiceContext.from_defaults()
        self._summary_text = summary_text
        self._qa_text = qa_text
        self._milvus_client = Milvus(host='localhost', port='19530')  # Initialize Milvus client

    def build_from_documents(
        self,
        documents: Sequence[Document],
    ) -> RouterQueryEngine:
        # Parse nodes
        nodes = run_transformations(
            documents, self._service_context.transformations  # type: ignore
        )

        # Ingest nodes into Milvus
        vectors = [node.vector for node in nodes]  # Assuming nodes have a 'vector' attribute
        ids = [node.id for node in nodes]  # Assuming nodes have an 'id' attribute

        # Create collection if not exists
        collection_name = 'document_collection'
        if not self._milvus_client.has_collection(collection_name):
            self._milvus_client.create_collection({
                'collection_name': collection_name,
                'fields': [
                    {'name': 'id', 'type': DataType.INT64, 'is_primary': True},
                    {'name': 'vector', 'type': DataType.FLOAT_VECTOR, 'params': {'dim': len(vectors[0])}}
                ]
            })

        # Insert vectors into Milvus
        self._milvus_client.insert(collection_name, records=vectors, ids=ids)

        # Build indices
        vector_index = VectorStoreIndex(
            nodes=nodes,
            transformations=self._service_context.transformations,
            embed_model=self._service_context.embed_model,
            service_context=self._service_context,
            storage_context=self._storage_context,
        )
        summary_index = SummaryIndex(
            nodes,
            service_context=self._service_context,
            storage_context=self._storage_context,
        )

        vector_query_engine = vector_index.as_query_engine(
            llm=self._service_context.llm, service_context=self._service_context
        )
        list_query_engine = summary_index.as_query_engine(
            llm=self._service_context.llm,
            service_context=self._service_context,
            response_mode="tree_summarize",
        )

        # Build query engine
        return RouterQueryEngine.from_defaults(
            llm=self._service_context.llm,
            query_engine_tools=[
                QueryEngineTool.from_defaults(
                    vector_query_engine, description=self._qa_text
                ),
                QueryEngineTool.from_defaults(
                    list_query_engine, description=self._summary_text
                ),
            ],
            service_context=self._service_context,
            select_multi=False,
        )

This example initializes the Milvus client, creates a collection if it doesn't exist, and inserts document vectors into Milvus. The rest of the method builds the necessary indices and query engines, integrating Milvus as the document store ^[1][2][3].

To continue talking to Dosu, mention @dosu.