Azure AI Search (formerly known asAzure Cognitive Search) is a Microsoft cloud search service that gives developers infrastructure, APIs, and tools for information retrieval of vector, keyword, and hybrid queries at scale.

AzureAISearchRetriever is an integration module that returns documents from an unstructured query. It's based on the BaseRetriever class and it targets the 2023-11-01 stable REST API version of Azure AI Search, which means it supports vector indexing and queries.

This guide will help you get started with the Azure AI Searchretriever. For detailed documentation of allAzureAISearchRetriever features and configurations head to theAPI reference.

AzureAISearchRetriever replacesAzureCognitiveSearchRetriever, which will soon be deprecated. We recommend switching to the newer version that's based on the most recent stable version of the search APIs.

Integration details​

Setup​

To use this module, you need:

• An Azure AI Search service. You cancreate one for free if you sign up for the Azure trial. A free service has lower quotas, but it's sufficient for running the code in this notebook.

• An existing index with vector fields. There are several ways to create one, including using thevector store module. Or,try the Azure AI Search REST APIs.

• An API key or Azure AD Token.

  • API keys are generated when you create the search service. If you're just querying an index, you can use the query API key, otherwise use an admin API key. SeeFind your API keys for details.
  • Azure AD Token can be used with Azure Managed Identity. SeeConnect your app to Azure AI Search using identities for details.

We can then set the search service name, index name, and API key as environment variables (alternatively, you can pass them as arguments toAzureAISearchRetriever). The search index provides the searchable content.

With an API Key

import os

os.environ["AZURE_AI_SEARCH_SERVICE_NAME"]="<YOUR_SEARCH_SERVICE_NAME>"
os.environ["AZURE_AI_SEARCH_INDEX_NAME"]="<YOUR_SEARCH_INDEX_NAME>"
os.environ["AZURE_AI_SEARCH_API_KEY"]="<YOUR_API_KEY>"

With an Azure AD Token

import os

os.environ["AZURE_AI_SEARCH_SERVICE_NAME"]="<YOUR_SEARCH_SERVICE_NAME>"
os.environ["AZURE_AI_SEARCH_INDEX_NAME"]="<YOUR_SEARCH_INDEX_NAME>"
os.environ["AZURE_AI_SEARCH_AD_TOKEN"]="<YOUR_AZURE_AD_TOKEN>"

If you want to get automated tracing from individual queries, you can also set yourLangSmith API key by uncommenting below:

# os.environ["LANGSMITH_API_KEY"] = getpass.getpass("Enter your LangSmith API key: ")
# os.environ["LANGSMITH_TRACING"] = "true"

Installation​

This retriever lives in thelangchain-community package. We will need some additional dependencies as well:

%pip install--upgrade--quiet langchain-community
%pip install--upgrade--quiet langchain-openai
%pip install--upgrade--quiet  azure-search-documents>=11.4
%pip install--upgrade--quiet  azure-identity

Instantiation​

ForAzureAISearchRetriever, provide anindex_name,content_key, andtop_k set to the number of number of results you'd like to retrieve. Settingtop_k to zero (the default) returns all results.

from langchain_community.retrieversimport AzureAISearchRetriever

retriever= AzureAISearchRetriever(
    content_key="content", top_k=1, index_name="langchain-vector-demo"
)

Usage​

Now you can use it to retrieve documents from Azure AI Search.This is the method you would call to do so. It will return all documents relevant to the query.

retriever.invoke("here is my unstructured query string")

Example​

This section demonstrates using the retriever over built-in sample data. You can skip this step if you already have a vector index on your search service.

Start by providing the endpoints and keys. Since we're creating a vector index in this step, specify a text embedding model to get a vector representation of the text. This example assumes Azure OpenAI with a deployment of text-embedding-ada-002. Because this step creates an index, be sure to use an admin API key for your search service.

import os

from langchain_community.document_loadersimport DirectoryLoader, TextLoader
from langchain_community.retrieversimport AzureAISearchRetriever
from langchain_community.vectorstoresimport AzureSearch
from langchain_openaiimport AzureOpenAIEmbeddings, OpenAIEmbeddings
from langchain_text_splittersimport TokenTextSplitter

os.environ["AZURE_AI_SEARCH_SERVICE_NAME"]="<YOUR_SEARCH_SERVICE_NAME>"
os.environ["AZURE_AI_SEARCH_INDEX_NAME"]="langchain-vector-demo"
os.environ["AZURE_AI_SEARCH_API_KEY"]="<YOUR_SEARCH_SERVICE_ADMIN_API_KEY>"
azure_endpoint:str="<YOUR_AZURE_OPENAI_ENDPOINT>"
azure_openai_api_key:str="<YOUR_AZURE_OPENAI_API_KEY>"
azure_openai_api_version:str="2023-05-15"
azure_deployment:str="text-embedding-ada-002"

We'll use an embedding model from Azure OpenAI to turn our documents into embeddings stored in the Azure AI Search vector store. We'll also set the index name tolangchain-vector-demo. This will create a new vector store associated with that index name.

embeddings= AzureOpenAIEmbeddings(
    model=azure_deployment,
    azure_endpoint=azure_endpoint,
    openai_api_key=azure_openai_api_key,
)

vector_store: AzureSearch= AzureSearch(
    embedding_function=embeddings.embed_query,
    azure_search_endpoint=os.getenv("AZURE_AI_SEARCH_SERVICE_NAME"),
    azure_search_key=os.getenv("AZURE_AI_SEARCH_API_KEY"),
    index_name="langchain-vector-demo",
)

Next, we'll load data into our newly created vector store. For this example, we load thestate_of_the_union.txt file. We'll split the text in 400 token chunks with no overlap. Finally, the documents are added to our vector store as emeddings.

from langchain_community.document_loadersimport TextLoader
from langchain_text_splittersimport CharacterTextSplitter

loader= TextLoader("../../how_to/state_of_the_union.txt", encoding="utf-8")

documents= loader.load()
text_splitter= CharacterTextSplitter(chunk_size=400, chunk_overlap=0)
docs= text_splitter.split_documents(documents)

vector_store.add_documents(documents=docs)

Next, we'll create a retriever. The currentindex_name variable islangchain-vector-demo from the last step. If you skipped vector store creation, provide your index name in the parameter. In this query, the top result is returned.

retriever= AzureAISearchRetriever(
    content_key="content", top_k=1, index_name="langchain-vector-demo"
)

Now we can retrieve the data that is relevant to our query from the documents we uploaded.

retriever.invoke("does the president have a plan for covid-19?")

Use within a chain​

from langchain_core.output_parsersimport StrOutputParser
from langchain_core.promptsimport ChatPromptTemplate
from langchain_core.runnablesimport RunnablePassthrough
from langchain_openaiimport ChatOpenAI

prompt= ChatPromptTemplate.from_template(
"""Answer the question based only on the context provided.

Context: {context}

Question: {question}"""
)

llm= ChatOpenAI(model="gpt-4o-mini")


defformat_docs(docs):
return"\n\n".join(doc.page_contentfor docin docs)


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

chain.invoke("does the president have a plan for covid-19?")

API reference​

For detailed documentation of allAzureAISearchRetriever features and configurations head to theAPI reference.

Related​

• Retrieverconceptual guide
• Retrieverhow-to guides