Nebius Retriever

The NebiusRetriever enables efficient similarity search using embeddings from Nebius AI Studio. It leverages high-quality embedding models to enable semantic search over documents.

This retriever is optimized for scenarios where you need to perform similarity search over a collection of documents, but don't need to persist the vectors to a vector database. It performs vector similarity search in-memory using matrix operations, making it efficient for medium-sized document collections.

Installation

%pip install --upgrade langchain-nebius

Requirement already satisfied: langchain-nebius in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (0.1.0)
Requirement already satisfied: langchain-core==0.3.35 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from langchain-nebius) (0.3.35)
Requirement already satisfied: langchain_openai==0.3.5 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from langchain-nebius) (0.3.5)
Requirement already satisfied: openai<2.0.0,>=1.0.0 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from langchain-nebius) (1.79.0)
Requirement already satisfied: langsmith<0.4,>=0.1.125 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from langchain-core==0.3.35->langchain-nebius) (0.3.42)
Requirement already satisfied: tenacity!=8.4.0,<10.0.0,>=8.1.0 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from langchain-core==0.3.35->langchain-nebius) (9.1.2)
Requirement already satisfied: jsonpatch<2.0,>=1.33 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from langchain-core==0.3.35->langchain-nebius) (1.33)
Requirement already satisfied: PyYAML>=5.3 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from langchain-core==0.3.35->langchain-nebius) (6.0.2)
Requirement already satisfied: packaging<25,>=23.2 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from langchain-core==0.3.35->langchain-nebius) (24.2)
Requirement already satisfied: typing-extensions>=4.7 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from langchain-core==0.3.35->langchain-nebius) (4.13.2)
Requirement already satisfied: pydantic<3.0.0,>=2.7.4 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from langchain-core==0.3.35->langchain-nebius) (2.11.4)
Requirement already satisfied: tiktoken<1,>=0.7 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from langchain_openai==0.3.5->langchain-nebius) (0.9.0)
Requirement already satisfied: anyio<5,>=3.5.0 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from openai<2.0.0,>=1.0.0->langchain-nebius) (4.9.0)
Requirement already satisfied: distro<2,>=1.7.0 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from openai<2.0.0,>=1.0.0->langchain-nebius) (1.9.0)
Requirement already satisfied: httpx<1,>=0.23.0 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from openai<2.0.0,>=1.0.0->langchain-nebius) (0.28.1)
Requirement already satisfied: jiter<1,>=0.4.0 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from openai<2.0.0,>=1.0.0->langchain-nebius) (0.10.0)
Requirement already satisfied: sniffio in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from openai<2.0.0,>=1.0.0->langchain-nebius) (1.3.1)
Requirement already satisfied: tqdm>4 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from openai<2.0.0,>=1.0.0->langchain-nebius) (4.67.1)
Requirement already satisfied: idna>=2.8 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from anyio<5,>=3.5.0->openai<2.0.0,>=1.0.0->langchain-nebius) (3.10)
Requirement already satisfied: certifi in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from httpx<1,>=0.23.0->openai<2.0.0,>=1.0.0->langchain-nebius) (2025.4.26)
Requirement already satisfied: httpcore==1.* in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from httpx<1,>=0.23.0->openai<2.0.0,>=1.0.0->langchain-nebius) (1.0.9)
Requirement already satisfied: h11>=0.16 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from httpcore==1.*->httpx<1,>=0.23.0->openai<2.0.0,>=1.0.0->langchain-nebius) (0.16.0)
Requirement already satisfied: jsonpointer>=1.9 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from jsonpatch<2.0,>=1.33->langchain-core==0.3.35->langchain-nebius) (3.0.0)
Requirement already satisfied: orjson<4.0.0,>=3.9.14 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from langsmith<0.4,>=0.1.125->langchain-core==0.3.35->langchain-nebius) (3.10.18)
Requirement already satisfied: requests<3,>=2 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from langsmith<0.4,>=0.1.125->langchain-core==0.3.35->langchain-nebius) (2.32.3)
Requirement already satisfied: requests-toolbelt<2.0.0,>=1.0.0 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from langsmith<0.4,>=0.1.125->langchain-core==0.3.35->langchain-nebius) (1.0.0)
Requirement already satisfied: zstandard<0.24.0,>=0.23.0 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from langsmith<0.4,>=0.1.125->langchain-core==0.3.35->langchain-nebius) (0.23.0)
Requirement already satisfied: annotated-types>=0.6.0 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from pydantic<3.0.0,>=2.7.4->langchain-core==0.3.35->langchain-nebius) (0.7.0)
Requirement already satisfied: pydantic-core==2.33.2 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from pydantic<3.0.0,>=2.7.4->langchain-core==0.3.35->langchain-nebius) (2.33.2)
Requirement already satisfied: typing-inspection>=0.4.0 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from pydantic<3.0.0,>=2.7.4->langchain-core==0.3.35->langchain-nebius) (0.4.0)
Requirement already satisfied: regex>=2022.1.18 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from tiktoken<1,>=0.7->langchain_openai==0.3.5->langchain-nebius) (2024.11.6)
Requirement already satisfied: charset-normalizer<4,>=2 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from requests<3,>=2->langsmith<0.4,>=0.1.125->langchain-core==0.3.35->langchain-nebius) (3.4.2)
Requirement already satisfied: urllib3<3,>=1.21.1 in /Users/akimtsvigun/Anaconda/nebius/repos/merge/langchain/.venv/lib/python3.12/site-packages (from requests<3,>=2->langsmith<0.4,>=0.1.125->langchain-core==0.3.35->langchain-nebius) (2.4.0)

[1m[[0m[34;49mnotice[0m[1;39;49m][0m[39;49m A new release of pip is available: [0m[31;49m24.3.1[0m[39;49m -> [0m[32;49m25.1.1[0m
[1m[[0m[34;49mnotice[0m[1;39;49m][0m[39;49m To update, run: [0m[32;49mpip install --upgrade pip[0m
Note: you may need to restart the kernel to use updated packages.

Environment

Nebius requires an API key that can be passed as an initialization parameter api_key or set as the environment variable NEBIUS_API_KEY. You can obtain an API key by creating an account on Nebius AI Studio.

import os

# Make sure you've set your API key as an environment variable
os.environ["NEBIUS_API_KEY"] = "YOUR-NEBIUS-API-KEY"

Basic Usage

The NebiusRetriever requires a NebiusEmbeddings instance and a list of documents:

from langchain_core.documents import Document
from langchain_nebius import NebiusEmbeddings, NebiusRetriever

# Create sample documents
docs = [
    Document(
        page_content="Paris is the capital of France", metadata={"country": "France"}
    ),
    Document(
        page_content="Berlin is the capital of Germany", metadata={"country": "Germany"}
    ),
    Document(
        page_content="Rome is the capital of Italy", metadata={"country": "Italy"}
    ),
    Document(
        page_content="Madrid is the capital of Spain", metadata={"country": "Spain"}
    ),
    Document(
        page_content="London is the capital of the United Kingdom",
        metadata={"country": "UK"},
    ),
    Document(
        page_content="Moscow is the capital of Russia", metadata={"country": "Russia"}
    ),
    Document(
        page_content="Washington DC is the capital of the United States",
        metadata={"country": "USA"},
    ),
    Document(
        page_content="Tokyo is the capital of Japan", metadata={"country": "Japan"}
    ),
    Document(
        page_content="Beijing is the capital of China", metadata={"country": "China"}
    ),
    Document(
        page_content="Canberra is the capital of Australia",
        metadata={"country": "Australia"},
    ),
]

# Initialize embeddings
embeddings = NebiusEmbeddings()

# Create retriever
retriever = NebiusRetriever(
    embeddings=embeddings,
    docs=docs,
    k=3,  # Number of documents to return
)

API Reference:Document

Retrieve Relevant Documents

You can use the retriever to find documents related to a query:

# Query for European capitals
query = "What are some capitals in Europe?"
results = retriever.invoke(query)

print(f"Query: {query}")
print(f"Top {len(results)} results:")
for i, doc in enumerate(results):
    print(f"{i+1}. {doc.page_content} (Country: {doc.metadata['country']})")

Query: What are some capitals in Europe?
Top 3 results:
1. Paris is the capital of France (Country: France)
2. Berlin is the capital of Germany (Country: Germany)
3. Rome is the capital of Italy (Country: Italy)

Using get_relevant_documents

You can also use the get_relevant_documents method directly (though invoke is the preferred interface):

# Query for Asian countries
query = "What are the capitals in Asia?"
results = retriever.get_relevant_documents(query)

print(f"Query: {query}")
print(f"Top {len(results)} results:")
for i, doc in enumerate(results):
    print(f"{i+1}. {doc.page_content} (Country: {doc.metadata['country']})")

Query: What are the capitals in Asia?
Top 3 results:
1. Beijing is the capital of China (Country: China)
2. Tokyo is the capital of Japan (Country: Japan)
3. Canberra is the capital of Australia (Country: Australia)
``````output
/var/folders/5j/jybstkcx4nvgdnx887f6y_680000gn/T/ipykernel_72588/4132599787.py:3: LangChainDeprecationWarning: The method `BaseRetriever.get_relevant_documents` was deprecated in langchain-core 0.1.46 and will be removed in 1.0. Use :meth:`~invoke` instead.
  results = retriever.get_relevant_documents(query)

Customizing Number of Results

You can adjust the number of results at query time by passing k as a parameter:

# Query for a specific country, with custom k
query = "Where is France?"
results = retriever.invoke(query, k=1)  # Override default k

print(f"Query: {query}")
print(f"Top {len(results)} result:")
for i, doc in enumerate(results):
    print(f"{i+1}. {doc.page_content} (Country: {doc.metadata['country']})")

Query: Where is France?
Top 1 result:
1. Paris is the capital of France (Country: France)

Async Support

NebiusRetriever supports async operations:

import asyncio


async def retrieve_async():
    query = "What are some capital cities?"
    results = await retriever.ainvoke(query)

    print(f"Async query: {query}")
    print(f"Top {len(results)} results:")
    for i, doc in enumerate(results):
        print(f"{i+1}. {doc.page_content} (Country: {doc.metadata['country']})")


await retrieve_async()

Async query: What are some capital cities?
Top 3 results:
1. Washington DC is the capital of the United States (Country: USA)
2. Canberra is the capital of Australia (Country: Australia)
3. Paris is the capital of France (Country: France)

Handling Empty Documents

# Create a retriever with empty documents
empty_retriever = NebiusRetriever(
    embeddings=embeddings,
    docs=[],
    k=2,  # Empty document list
)

# Test the retriever with empty docs
results = empty_retriever.invoke("What are the capitals of European countries?")
print(f"Number of results: {len(results)}")

Number of results: 0

Using in a RAG Pipeline

NebiusRetriever works seamlessly in LangChain RAG pipelines:

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_nebius import ChatNebius

# Initialize LLM
llm = ChatNebius(model="meta-llama/Llama-3.3-70B-Instruct-fast")

# Create a prompt template
prompt = ChatPromptTemplate.from_template(
    """
Answer the question based only on the following context:

Context:
{context}

Question: {question}
"""
)


# Format documents function
def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)


# Create RAG chain
rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

# Run the chain
answer = rag_chain.invoke("What are three European capitals?")
print(answer)

API Reference:StrOutputParser | ChatPromptTemplate | RunnablePassthrough

Based on the context, three European capitals are:

Paris
Berlin
Rome

Creating a Search Tool

You can use the NebiusRetrievalTool to create a tool for agents:

from langchain_nebius import NebiusRetrievalTool

# Create a retrieval tool
tool = NebiusRetrievalTool(
    retriever=retriever,
    name="capital_search",
    description="Search for information about capital cities around the world",
)

# Use the tool
result = tool.invoke({"query": "capitals in Europe", "k": 3})
print("Tool results:")
print(result)

Tool results:
Document 1:
Paris is the capital of France

Document 2:
Berlin is the capital of Germany

Document 3:
Rome is the capital of Italy

How It Works

The NebiusRetriever works by:

During initialization:
- It stores the provided documents
- It uses the provided NebiusEmbeddings to compute embeddings for all documents
- These embeddings are stored in memory for quick retrieval
During retrieval (invoke or get_relevant_documents):
- It embeds the query using the same embedding model
- It computes similarity scores between the query embedding and all document embeddings
- It returns the top-k documents sorted by similarity

This approach is efficient for medium-sized document collections, as it avoids the need for a separate vector database while still providing high-quality semantic search.

API Reference

For more details about the Nebius AI Studio API, visit the Nebius AI Studio Documentation.

Retriever conceptual guide
Retriever how-to guides

Installation​

Environment​

Basic Usage​

Retrieve Relevant Documents​

Using get_relevant_documents​

Customizing Number of Results​

Async Support​

Handling Empty Documents​

Using in a RAG Pipeline​

Creating a Search Tool​

How It Works​

API Reference​

Related​

Was this page helpful?