Ecosystem Cookbook Agent Blueprints Builder devsite

Cookbook repo ↗

← CookbookView source on GitHub ↗

RAG With llama-index + Milvus + Qwen - Part 2

References

• https://tokenfactory.nebius.com/
• https://docs.llamaindex.ai/en/stable/examples/vector_stores/MilvusIndexDemo/
• https://docs.llamaindex.ai/en/stable/api_reference/storage/vector_store/milvus/?h=milvusvectorstore#llama_index.vector_stores.milvus.MilvusVectorStore

Step-1: Configuration

import os
from dotenv import load_dotenv
load_dotenv()

if os.getenv('NEBIUS_API_KEY'):
    print ("✅ Found NEBIUS_API_KEY in environment, using it")
else:
    raise ValueError("❌ NEBIUS_API_KEY not found in environment. Please set it in .env file before running this script.")

✅ Found NEBIUS_API_KEY in environment, using it

Step-2: Setup Embedding Model

We have a choice of local embedding model (fast) or running it on the cloud

If running locally:

• choose smaller models
• less accuracy but faster

If running on the cloud

• We can run large models (billions of params)

import os
from llama_index.core import Settings

# Option 1: Running embedding models on Nebius cloud
from llama_index.embeddings.nebius import NebiusEmbedding
EMBEDDING_MODEL = 'Qwen/Qwen3-Embedding-8B'  # 8B params
EMBEDDING_LENGTH = 4096  # Length of the embedding vector
Settings.embed_model = NebiusEmbedding(
                        model_name=EMBEDDING_MODEL,
                        embed_batch_size=50,  # Batch size for embedding (default is 10)
                        api_key=os.getenv("NEBIUS_API_KEY") # if not specfified here, it will get taken from env variable
                       )

## Option 2: Running embedding models locally
# from llama_index.embeddings.huggingface import HuggingFaceEmbedding
# os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
# Settings.embed_model = HuggingFaceEmbedding(
#     # model_name = 'sentence-transformers/all-MiniLM-L6-v2' # 23 M params
#     model_name = 'BAAI/bge-small-en-v1.5'  # 33M params
#     # model_name = 'Qwen/Qwen3-Embedding-0.6B'  # 600M params
#     # model_name = 'BAAI/bge-en-icl'  # 7B params
#     #model_name = 'intfloat/multilingual-e5-large-instruct'  # 560M params
# )

Step-3: Connect to Milvus

from pymilvus import MilvusClient

DB_URI = './rag.db'  # For embedded instance
COLLECTION_NAME = 'rag'

milvus_client = MilvusClient(DB_URI)
print ("✅ Connected to Milvus instance: ", DB_URI)

/home/sujee/my-stuff/projects/nebius/token-factory-cookbook-1/rag/rag-milvus-1/.venv/lib/python3.11/site-packages/google/protobuf/runtime_version.py:98: UserWarning: Protobuf gencode version 5.27.2 is exactly one major version older than the runtime version 6.31.1 at schema.proto. Please update the gencode to avoid compatibility violations in the next runtime release.
  warnings.warn(
/home/sujee/my-stuff/projects/nebius/token-factory-cookbook-1/rag/rag-milvus-1/.venv/lib/python3.11/site-packages/google/protobuf/runtime_version.py:98: UserWarning: Protobuf gencode version 5.27.2 is exactly one major version older than the runtime version 6.31.1 at common.proto. Please update the gencode to avoid compatibility violations in the next runtime release.
  warnings.warn(
/home/sujee/my-stuff/projects/nebius/token-factory-cookbook-1/rag/rag-milvus-1/.venv/lib/python3.11/site-packages/google/protobuf/runtime_version.py:98: UserWarning: Protobuf gencode version 5.27.2 is exactly one major version older than the runtime version 6.31.1 at milvus.proto. Please update the gencode to avoid compatibility violations in the next runtime release.
  warnings.warn(
/home/sujee/my-stuff/projects/nebius/token-factory-cookbook-1/rag/rag-milvus-1/.venv/lib/python3.11/site-packages/google/protobuf/runtime_version.py:98: UserWarning: Protobuf gencode version 5.27.2 is exactly one major version older than the runtime version 6.31.1 at rg.proto. Please update the gencode to avoid compatibility violations in the next runtime release.
  warnings.warn(
/home/sujee/my-stuff/projects/nebius/token-factory-cookbook-1/rag/rag-milvus-1/.venv/lib/python3.11/site-packages/google/protobuf/runtime_version.py:98: UserWarning: Protobuf gencode version 5.27.2 is exactly one major version older than the runtime version 6.31.1 at feder.proto. Please update the gencode to avoid compatibility violations in the next runtime release.
  warnings.warn(
/home/sujee/my-stuff/projects/nebius/token-factory-cookbook-1/rag/rag-milvus-1/.venv/lib/python3.11/site-packages/google/protobuf/runtime_version.py:98: UserWarning: Protobuf gencode version 5.27.2 is exactly one major version older than the runtime version 6.31.1 at msg.proto. Please update the gencode to avoid compatibility violations in the next runtime release.
  warnings.warn(
/home/sujee/my-stuff/projects/nebius/token-factory-cookbook-1/rag/rag-milvus-1/.venv/lib/python3.11/site-packages/milvus_lite/__init__.py:15: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
  from pkg_resources import DistributionNotFound, get_distribution

✅ Connected to Milvus instance:  ./rag.db

# connect to vector db
from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.vector_stores.milvus import MilvusVectorStore

vector_store = MilvusVectorStore(
    uri = DB_URI ,
    dim = EMBEDDING_LENGTH ,
    collection_name = COLLECTION_NAME,
    overwrite=False
)
storage_context = StorageContext.from_defaults(vector_store=vector_store)

print ("✅ Connected Llama-index to Milvus instance: ", DB_URI )

✅ Connected Llama-index to Milvus instance:  ./rag.db

Step-4: Load Document Index from DB

%%time

from llama_index.core import VectorStoreIndex

index = VectorStoreIndex.from_vector_store(
    vector_store=vector_store, storage_context=storage_context)

print ("✅ Loaded index from vector db:", DB_URI )

✅ Loaded index from vector db: ./rag.db
CPU times: user 94 ms, sys: 20.7 ms, total: 115 ms
Wall time: 113 ms

Step-5: Setup LLM

from llama_index.llms.nebius import NebiusLLM
from llama_index.core import Settings

Settings.llm = NebiusLLM(
                model='openai/gpt-oss-120b',
                # model='openai/gpt-oss-20b',
                # model='deepseek-ai/DeepSeek-R1-0528',
                api_key=os.getenv("NEBIUS_API_KEY") # if not specfified, it will get taken from env variable
    )

/home/sujee/my-stuff/projects/nebius/token-factory-cookbook-1/rag/rag-milvus-1/.venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm
None of PyTorch, TensorFlow >= 2.0, or Flax have been found. Models won't be available and only tokenizers, configuration and file/data utilities can be used.

Step-6: Query

query_engine = index.as_query_engine()
res = query_engine.query("What was Uber's revenue for 2020?")
print(res)

Uber’s revenue for 2020 was $11.139 billion.

Making sure the model uses context

Let's ask a generic factual question "When was the moon landing".

Now the model should know this generic factual answer.

But since we are querying documents, we want to the model to find answers from within the documents.

It should come back with something like "provided context does not have information about moon landing"

query_engine = index.as_query_engine()
res = query_engine.query("When was the moon landing?")
print(res)

The provided context does not contain information about the date of the moon landing.

Recipes © Nebius · nebius/token-factory-cookbook (MIT). Rendered for an ecosystem demo. Ecosystem & Partners