diff --git a/docs/modules/indexes/vectorstores/examples/azuresearch.ipynb b/docs/modules/indexes/vectorstores/examples/azuresearch.ipynb new file mode 100644 index 00000000..36a6abb6 --- /dev/null +++ b/docs/modules/indexes/vectorstores/examples/azuresearch.ipynb @@ -0,0 +1,245 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Azure Cognitive Search" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Install Azure Cognitive Search SDK" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!pip install --index-url=https://pkgs.dev.azure.com/azure-sdk/public/_packaging/azure-sdk-for-python/pypi/simple/ azure-search-documents==11.4.0a20230509004\n", + "!pip install azure-identity" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Import required libraries" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "import os, json\n", + "import openai\n", + "from dotenv import load_dotenv\n", + "from langchain.embeddings.openai import OpenAIEmbeddings\n", + "from langchain.schema import BaseRetriever\n", + "from langchain.vectorstores.azuresearch import AzureSearch" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Configure OpenAI settings\n", + "Configure the OpenAI settings to use Azure OpenAI or OpenAI" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [], + "source": [ + "# Load environment variables from a .env file using load_dotenv():\n", + "load_dotenv()\n", + "\n", + "openai.api_type = \"azure\"\n", + "openai.api_base = \"YOUR_OPENAI_ENDPOINT\"\n", + "openai.api_version = \"2023-05-15\"\n", + "openai.api_key = \"YOUR_OPENAI_API_KEY\"\n", + "model: str = \"text-embedding-ada-002\"" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Configure vector store settings\n", + " \n", + "Set up the vector store settings using environment variables:" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [], + "source": [ + "vector_store_address: str = 'YOUR_AZURE_SEARCH_ENDPOINT'\n", + "vector_store_password: str = 'YOUR_AZURE_SEARCH_ADMIN_KEY'\n", + "index_name: str = \"langchain-vector-demo\"" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Create embeddings and vector store instances\n", + " \n", + "Create instances of the OpenAIEmbeddings and AzureSearch classes:" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [], + "source": [ + "embeddings: OpenAIEmbeddings = OpenAIEmbeddings(model=model, chunk_size=1) \n", + "vector_store: AzureSearch = AzureSearch(azure_search_endpoint=vector_store_address, \n", + " azure_search_key=vector_store_password, \n", + " index_name=index_name, \n", + " embedding_function=embeddings.embed_query) \n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Insert text and embeddings into vector store\n", + " \n", + "Add texts and metadata from the JSON data to the vector store:" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [], + "source": [ + "from langchain.document_loaders import TextLoader\n", + "from langchain.text_splitter import CharacterTextSplitter\n", + "loader = TextLoader('../../../state_of_the_union.txt', encoding='utf-8')\n", + "\n", + "documents = loader.load()\n", + "text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)\n", + "docs = text_splitter.split_documents(documents)\n", + "\n", + "vector_store.add_documents(documents=docs)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Perform a vector similarity search\n", + " \n", + "Execute a pure vector similarity search using the similarity_search() method:" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while you’re at it, pass the Disclose Act so Americans can know who is funding our elections. \n", + "\n", + "Tonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. \n", + "\n", + "One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n", + "\n", + "And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence.\n" + ] + } + ], + "source": [ + "# Perform a similarity search\n", + "docs = vector_store.similarity_search(query=\"What did the president say about Ketanji Brown Jackson\", k=3, search_type='similarity')\n", + "print(docs[0].page_content)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Perform a Hybrid Search\n", + "\n", + "Execute hybrid search using the hybrid_search() method:" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while you’re at it, pass the Disclose Act so Americans can know who is funding our elections. \n", + "\n", + "Tonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. \n", + "\n", + "One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n", + "\n", + "And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence.\n" + ] + } + ], + "source": [ + "# Perform a hybrid search \n", + "docs = vector_store.similarity_search(query=\"What did the president say about Ketanji Brown Jackson\", k=3)\n", + "print(docs[0].page_content)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3.9.13 ('.venv': venv)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.3" + }, + "orig_nbformat": 4, + "vscode": { + "interpreter": { + "hash": "645053d6307d413a1a75681b5ebb6449bb2babba4bcb0bf65a1ddc3dbefb108a" + } + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/langchain/vectorstores/__init__.py b/langchain/vectorstores/__init__.py index c59b06fd..1bb3f901 100644 --- a/langchain/vectorstores/__init__.py +++ b/langchain/vectorstores/__init__.py @@ -3,6 +3,7 @@ from langchain.vectorstores.analyticdb import AnalyticDB from langchain.vectorstores.annoy import Annoy from langchain.vectorstores.atlas import AtlasDB from langchain.vectorstores.awadb import AwaDB +from langchain.vectorstores.azuresearch import AzureSearch from langchain.vectorstores.base import VectorStore from langchain.vectorstores.chroma import Chroma from langchain.vectorstores.clickhouse import Clickhouse, ClickhouseSettings @@ -31,6 +32,7 @@ from langchain.vectorstores.weaviate import Weaviate from langchain.vectorstores.zilliz import Zilliz __all__ = [ + "AzureSearch", "Redis", "ElasticVectorSearch", "FAISS", diff --git a/langchain/vectorstores/azuresearch.py b/langchain/vectorstores/azuresearch.py new file mode 100644 index 00000000..94ac1172 --- /dev/null +++ b/langchain/vectorstores/azuresearch.py @@ -0,0 +1,507 @@ +"""Wrapper around Azure Cognitive Search.""" +from __future__ import annotations + +import base64 +import json +import logging +import uuid +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + Iterable, + List, + Optional, + Tuple, + Type, +) + +import numpy as np +from pydantic import BaseModel, root_validator + +from langchain.docstore.document import Document +from langchain.embeddings.base import Embeddings +from langchain.schema import BaseRetriever +from langchain.utils import get_from_env +from langchain.vectorstores.base import VectorStore + +logger = logging.getLogger() + +if TYPE_CHECKING: + from azure.search.documents import SearchClient + + +# Allow overriding field names for Azure Search +FIELDS_ID = get_from_env( + key="AZURESEARCH_FIELDS_ID", env_key="AZURESEARCH_FIELDS_ID", default="id" +) +FIELDS_CONTENT = get_from_env( + key="AZURESEARCH_FIELDS_CONTENT", + env_key="AZURESEARCH_FIELDS_CONTENT", + default="content", +) +FIELDS_CONTENT_VECTOR = get_from_env( + key="AZURESEARCH_FIELDS_CONTENT_VECTOR", + env_key="AZURESEARCH_FIELDS_CONTENT_VECTOR", + default="content_vector", +) +FIELDS_METADATA = get_from_env( + key="AZURESEARCH_FIELDS_TAG", env_key="AZURESEARCH_FIELDS_TAG", default="metadata" +) + +MAX_UPLOAD_BATCH_SIZE = 1000 + + +def _get_search_client( + endpoint: str, + key: str, + index_name: str, + embedding_function: Callable, + semantic_configuration_name: Optional[str] = None, +) -> SearchClient: + from azure.core.credentials import AzureKeyCredential + from azure.core.exceptions import ResourceNotFoundError + from azure.identity import DefaultAzureCredential + from azure.search.documents import SearchClient + from azure.search.documents.indexes import SearchIndexClient + from azure.search.documents.indexes.models import ( + PrioritizedFields, + SearchableField, + SearchField, + SearchFieldDataType, + SearchIndex, + SemanticConfiguration, + SemanticField, + SemanticSettings, + SimpleField, + VectorSearch, + VectorSearchAlgorithmConfiguration, + ) + + if key is None: + credential = DefaultAzureCredential() + else: + credential = AzureKeyCredential(key) + index_client: SearchIndexClient = SearchIndexClient( + endpoint=endpoint, credential=credential + ) + try: + index_client.get_index(name=index_name) + except ResourceNotFoundError: + # Fields configuration + fields = [ + SimpleField( + name=FIELDS_ID, + type=SearchFieldDataType.String, + key=True, + filterable=True, + ), + SearchableField( + name=FIELDS_CONTENT, + type=SearchFieldDataType.String, + searchable=True, + retrievable=True, + ), + SearchField( + name=FIELDS_CONTENT_VECTOR, + type=SearchFieldDataType.Collection(SearchFieldDataType.Single), + searchable=True, + dimensions=len(embedding_function("Text")), + vector_search_configuration="default", + ), + SearchableField( + name=FIELDS_METADATA, + type=SearchFieldDataType.String, + searchable=True, + retrievable=True, + ), + ] + # Vector search configuration + vector_search = VectorSearch( + algorithm_configurations=[ + VectorSearchAlgorithmConfiguration( + name="default", + kind="hnsw", + hnsw_parameters={ + "m": 4, + "efConstruction": 400, + "efSearch": 500, + "metric": "cosine", + }, + ) + ] + ) + # Create the semantic settings with the configuration + semantic_settings = ( + None + if semantic_configuration_name is None + else SemanticSettings( + configurations=[ + SemanticConfiguration( + name=semantic_configuration_name, + prioritized_fields=PrioritizedFields( + prioritized_content_fields=[ + SemanticField(field_name=FIELDS_CONTENT) + ], + ), + ) + ] + ) + ) + # Create the search index with the semantic settings and vector search + index = SearchIndex( + name=index_name, + fields=fields, + vector_search=vector_search, + semantic_settings=semantic_settings, + ) + index_client.create_index(index) + # Create the search client + return SearchClient(endpoint=endpoint, index_name=index_name, credential=credential) + + +class AzureSearch(VectorStore): + def __init__( + self, + azure_search_endpoint: str, + azure_search_key: str, + index_name: str, + embedding_function: Callable, + search_type: str = "hybrid", + semantic_configuration_name: Optional[str] = None, + semantic_query_language: str = "en-us", + **kwargs: Any, + ): + """Initialize with necessary components.""" + # Initialize base class + self.embedding_function = embedding_function + self.client = _get_search_client( + azure_search_endpoint, + azure_search_key, + index_name, + embedding_function, + semantic_configuration_name, + ) + self.search_type = search_type + self.semantic_configuration_name = semantic_configuration_name + self.semantic_query_language = semantic_query_language + + def add_texts( + self, + texts: Iterable[str], + metadatas: Optional[List[dict]] = None, + **kwargs: Any, + ) -> List[str]: + """Add texts data to an existing index.""" + keys = kwargs.get("keys") + ids = [] + # Write data to index + data = [] + for i, text in enumerate(texts): + # Use provided key otherwise use default key + key = keys[i] if keys else str(uuid.uuid4()) + # Encoding key for Azure Search valid characters + key = base64.urlsafe_b64encode(bytes(key, "utf-8")).decode("ascii") + metadata = metadatas[i] if metadatas else {} + # Add data to index + data.append( + { + "@search.action": "upload", + FIELDS_ID: key, + FIELDS_CONTENT: text, + FIELDS_CONTENT_VECTOR: np.array( + self.embedding_function(text), dtype=np.float32 + ).tolist(), + FIELDS_METADATA: json.dumps(metadata), + } + ) + ids.append(key) + # Upload data in batches + if len(data) == MAX_UPLOAD_BATCH_SIZE: + response = self.client.upload_documents(documents=data) + # Check if all documents were successfully uploaded + if not all([r.succeeded for r in response]): + raise Exception(response) + # Reset data + data = [] + + # Considering case where data is an exact multiple of batch-size entries + if len(data) == 0: + return ids + + # Upload data to index + response = self.client.upload_documents(documents=data) + # Check if all documents were successfully uploaded + if all([r.succeeded for r in response]): + return ids + else: + raise Exception(response) + + def similarity_search( + self, query: str, k: int = 4, **kwargs: Any + ) -> List[Document]: + search_type = kwargs.get("search_type", self.search_type) + if search_type == "similarity": + docs = self.vector_search(query, k=k) + elif search_type == "hybrid": + docs = self.hybrid_search(query, k=k) + elif search_type == "semantic_hybrid": + docs = self.semantic_hybrid_search(query, k=k) + else: + raise ValueError(f"search_type of {search_type} not allowed.") + return docs + + def vector_search(self, query: str, k: int = 4, **kwargs: Any) -> List[Document]: + """ + Returns the most similar indexed documents to the query text. + + Args: + query (str): The query text for which to find similar documents. + k (int): The number of documents to return. Default is 4. + + Returns: + List[Document]: A list of documents that are most similar to the query text. + """ + docs_and_scores = self.vector_search_with_score( + query, k=k, filters=kwargs.get("filters", None) + ) + return [doc for doc, _ in docs_and_scores] + + def vector_search_with_score( + self, query: str, k: int = 4, filters: Optional[str] = None + ) -> List[Tuple[Document, float]]: + """Return docs most similar to query. + + Args: + query: Text to look up documents similar to. + k: Number of Documents to return. Defaults to 4. + + Returns: + List of Documents most similar to the query and score for each + """ + from azure.search.documents.models import Vector + + results = self.client.search( + search_text="", + vector=Vector( + value=np.array( + self.embedding_function(query), dtype=np.float32 + ).tolist(), + k=k, + fields=FIELDS_CONTENT_VECTOR, + ), + select=[f"{FIELDS_ID},{FIELDS_CONTENT},{FIELDS_METADATA}"], + filter=filters, + ) + # Convert results to Document objects + docs = [ + ( + Document( + page_content=result[FIELDS_CONTENT], + metadata=json.loads(result[FIELDS_METADATA]), + ), + float(result["@search.score"]), + ) + for result in results + ] + return docs + + def hybrid_search(self, query: str, k: int = 4, **kwargs: Any) -> List[Document]: + """ + Returns the most similar indexed documents to the query text. + + Args: + query (str): The query text for which to find similar documents. + k (int): The number of documents to return. Default is 4. + + Returns: + List[Document]: A list of documents that are most similar to the query text. + """ + docs_and_scores = self.hybrid_search_with_score( + query, k=k, filters=kwargs.get("filters", None) + ) + return [doc for doc, _ in docs_and_scores] + + def hybrid_search_with_score( + self, query: str, k: int = 4, filters: Optional[str] = None + ) -> List[Tuple[Document, float]]: + """Return docs most similar to query with an hybrid query. + + Args: + query: Text to look up documents similar to. + k: Number of Documents to return. Defaults to 4. + + Returns: + List of Documents most similar to the query and score for each + """ + from azure.search.documents.models import Vector + + results = self.client.search( + search_text=query, + vector=Vector( + value=np.array( + self.embedding_function(query), dtype=np.float32 + ).tolist(), + k=k, + fields=FIELDS_CONTENT_VECTOR, + ), + select=[f"{FIELDS_ID},{FIELDS_CONTENT},{FIELDS_METADATA}"], + filter=filters, + top=k, + ) + # Convert results to Document objects + docs = [ + ( + Document( + page_content=result[FIELDS_CONTENT], + metadata=json.loads(result[FIELDS_METADATA]), + ), + float(result["@search.score"]), + ) + for result in results + ] + return docs + + def semantic_hybrid_search( + self, query: str, k: int = 4, **kwargs: Any + ) -> List[Document]: + """ + Returns the most similar indexed documents to the query text. + + Args: + query (str): The query text for which to find similar documents. + k (int): The number of documents to return. Default is 4. + + Returns: + List[Document]: A list of documents that are most similar to the query text. + """ + docs_and_scores = self.semantic_hybrid_search_with_score( + query, k=k, filters=kwargs.get("filters", None) + ) + return [doc for doc, _ in docs_and_scores] + + def semantic_hybrid_search_with_score( + self, query: str, k: int = 4, filters: Optional[str] = None + ) -> List[Tuple[Document, float]]: + """Return docs most similar to query with an hybrid query. + + Args: + query: Text to look up documents similar to. + k: Number of Documents to return. Defaults to 4. + + Returns: + List of Documents most similar to the query and score for each + """ + from azure.search.documents.models import Vector + + results = self.client.search( + search_text=query, + vector=Vector( + value=np.array( + self.embedding_function(query), dtype=np.float32 + ).tolist(), + k=50, # Hardcoded value to maximize L2 retrieval + fields=FIELDS_CONTENT_VECTOR, + ), + select=[f"{FIELDS_ID},{FIELDS_CONTENT},{FIELDS_METADATA}"], + filter=filters, + query_type="semantic", + query_language=self.semantic_query_language, + semantic_configuration_name=self.semantic_configuration_name, + query_caption="extractive", + query_answer="extractive", + top=k, + ) + # Get Semantic Answers + semantic_answers = results.get_answers() + semantic_answers_dict = {} + for semantic_answer in semantic_answers: + semantic_answers_dict[semantic_answer.key] = { + "text": semantic_answer.text, + "highlights": semantic_answer.highlights, + } + # Convert results to Document objects + docs = [ + ( + Document( + page_content=result["content"], + metadata={ + **json.loads(result["metadata"]), + **{ + "captions": { + "text": result.get("@search.captions", [{}])[0].text, + "highlights": result.get("@search.captions", [{}])[ + 0 + ].highlights, + } + if result.get("@search.captions") + else {}, + "answers": semantic_answers_dict.get( + json.loads(result["metadata"]).get("key"), "" + ), + }, + }, + ), + float(result["@search.score"]), + ) + for result in results + ] + return docs + + @classmethod + def from_texts( + cls: Type[AzureSearch], + texts: List[str], + embedding: Embeddings, + metadatas: Optional[List[dict]] = None, + azure_search_endpoint: str = "", + azure_search_key: str = "", + index_name: str = "langchain-index", + **kwargs: Any, + ) -> AzureSearch: + # Creating a new Azure Search instance + azure_search = cls( + azure_search_endpoint, + azure_search_key, + index_name, + embedding.embed_query, + ) + azure_search.add_texts(texts, metadatas, **kwargs) + return azure_search + + +class AzureSearchVectorStoreRetriever(BaseRetriever, BaseModel): + vectorstore: AzureSearch + search_type: str = "hybrid" + k: int = 4 + + class Config: + """Configuration for this pydantic object.""" + + arbitrary_types_allowed = True + + @root_validator() + def validate_search_type(cls, values: Dict) -> Dict: + """Validate search type.""" + if "search_type" in values: + search_type = values["search_type"] + if search_type not in ("similarity", "hybrid", "semantic_hybrid"): + raise ValueError(f"search_type of {search_type} not allowed.") + return values + + def get_relevant_documents(self, query: str) -> List[Document]: + if self.search_type == "similarity": + docs = self.vectorstore.vector_search(query, k=self.k) + elif self.search_type == "hybrid": + docs = self.vectorstore.hybrid_search(query, k=self.k) + elif self.search_type == "semantic_hybrid": + docs = self.vectorstore.semantic_hybrid_search(query, k=self.k) + else: + raise ValueError(f"search_type of {self.search_type} not allowed.") + return docs + + async def aget_relevant_documents(self, query: str) -> List[Document]: + raise NotImplementedError( + "AzureSearchVectorStoreRetriever does not support async" + ) diff --git a/poetry.lock b/poetry.lock index 9ed0db75..6d5bb1e0 100644 --- a/poetry.lock +++ b/poetry.lock @@ -702,6 +702,28 @@ msal = ">=1.20.0,<2.0.0" msal-extensions = ">=0.3.0,<2.0.0" six = ">=1.12.0" +[[package]] +name = "azure-search-documents" +version = "11.4.0a20230509004" +description = "Microsoft Azure Cognitive Search Client Library for Python" +category = "main" +optional = true +python-versions = ">=3.7" +files = [ + {file = "azure-search-documents-11.4.0a20230509004.zip", hash = "sha256:6cca144573161a10aa0fcd13927264453e79c63be6a53cf2ec241c9c8c22f6b5"}, + {file = "azure_search_documents-11.4.0a20230509004-py3-none-any.whl", hash = "sha256:6215e9a4f9e935ff3eac1b7d5519c6c0789b4497eb11242d376911aaefbb0359"}, +] + +[package.dependencies] +azure-common = ">=1.1,<2.0" +azure-core = ">=1.24.0,<2.0.0" +isodate = ">=0.6.0" + +[package.source] +type = "legacy" +url = "https://pkgs.dev.azure.com/azure-sdk/public/_packaging/azure-sdk-for-python/pypi/simple" +reference = "azure-sdk-dev" + [[package]] name = "babel" version = "2.12.1" @@ -11451,7 +11473,7 @@ cffi = ["cffi (>=1.11)"] [extras] all = ["anthropic", "cohere", "openai", "nlpcloud", "huggingface_hub", "jina", "manifest-ml", "elasticsearch", "opensearch-py", "google-search-results", "faiss-cpu", "sentence-transformers", "transformers", "spacy", "nltk", "wikipedia", "beautifulsoup4", "tiktoken", "torch", "jinja2", "pinecone-client", "pinecone-text", "pymongo", "weaviate-client", "redis", "google-api-python-client", "google-auth", "wolframalpha", "qdrant-client", "tensorflow-text", "pypdf", "networkx", "nomic", "aleph-alpha-client", "deeplake", "pgvector", "psycopg2-binary", "pyowm", "pytesseract", "html2text", "atlassian-python-api", "gptcache", "duckduckgo-search", "arxiv", "azure-identity", "clickhouse-connect", "azure-cosmos", "lancedb", "langkit", "lark", "pexpect", "pyvespa", "O365", "jq", "docarray", "steamship", "pdfminer-six", "lxml", "requests-toolbelt", "neo4j", "openlm", "azure-ai-formrecognizer", "azure-ai-vision", "azure-cognitiveservices-speech", "momento", "singlestoredb", "tigrisdb", "nebula3-python", "awadb"] -azure = ["azure-identity", "azure-cosmos", "openai", "azure-core", "azure-ai-formrecognizer", "azure-ai-vision", "azure-cognitiveservices-speech"] +azure = ["azure-identity", "azure-cosmos", "openai", "azure-core", "azure-ai-formrecognizer", "azure-ai-vision", "azure-cognitiveservices-speech", "azure-search-documents"] cohere = ["cohere"] docarray = ["docarray"] embeddings = ["sentence-transformers"] @@ -11464,4 +11486,4 @@ text-helpers = ["chardet"] [metadata] lock-version = "2.0" python-versions = ">=3.8.1,<4.0" -content-hash = "7a39130af070d4a4fe6b0af5d6b70615c868ab0b1867e404060ff00eacd10f5f" +content-hash = "17e9c7a2ae2d0ef7cf45bc232ebeb7fd3eee2760bb2a19b34a63dcddafd3e4ad" diff --git a/pyproject.toml b/pyproject.toml index fc6edb5b..0278d4ab 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -107,7 +107,7 @@ tigrisdb = {version = "^1.0.0b6", optional = true} nebula3-python = {version = "^3.4.0", optional = true} langchainplus-sdk = ">=0.0.7" awadb = {version = "^0.3.2", optional = true} - +azure-search-documents = {version = "11.4.0a20230509004", source = "azure-sdk-dev", optional = true} [tool.poetry.group.docs.dependencies] autodoc_pydantic = "^1.8.0" @@ -218,7 +218,16 @@ text_helpers = ["chardet"] cohere = ["cohere"] docarray = ["docarray"] embeddings = ["sentence-transformers"] -azure = ["azure-identity", "azure-cosmos", "openai", "azure-core", "azure-ai-formrecognizer", "azure-ai-vision", "azure-cognitiveservices-speech"] +azure = [ + "azure-identity", + "azure-cosmos", + "openai", + "azure-core", + "azure-ai-formrecognizer", + "azure-ai-vision", + "azure-cognitiveservices-speech", + "azure-search-documents", +] all = [ "anthropic", "cohere", @@ -320,6 +329,11 @@ extended_testing = [ "openai" ] +[[tool.poetry.source]] +name = "azure-sdk-dev" +url = "https://pkgs.dev.azure.com/azure-sdk/public/_packaging/azure-sdk-for-python/pypi/simple/" +secondary = true + [tool.ruff] select = [ "E", # pycodestyle diff --git a/tests/integration_tests/vectorstores/test_azuresearch.py b/tests/integration_tests/vectorstores/test_azuresearch.py new file mode 100644 index 00000000..b5c3b720 --- /dev/null +++ b/tests/integration_tests/vectorstores/test_azuresearch.py @@ -0,0 +1,93 @@ +import os +import time + +import openai +import pytest +from dotenv import load_dotenv + +from langchain.embeddings.openai import OpenAIEmbeddings +from langchain.vectorstores.azuresearch import AzureSearch + +load_dotenv() + +# Azure OpenAI settings +openai.api_type = "azure" +openai.api_base = os.getenv("OPENAI_API_BASE", "") +openai.api_version = "2023-05-15" +openai.api_key = os.getenv("OPENAI_API_KEY", "") +model: str = os.getenv("OPENAI_EMBEDDINGS_ENGINE_DOC", "text-embedding-ada-002") + +# Vector store settings +vector_store_address: str = os.getenv("AZURE_SEARCH_ENDPOINT", "") +vector_store_password: str = os.getenv("AZURE_SEARCH_ADMIN_KEY", "") +index_name: str = "embeddings-vector-store-test" + + +@pytest.fixture +def similarity_search_test() -> None: + """Test end to end construction and search.""" + # Create Embeddings + embeddings: OpenAIEmbeddings = OpenAIEmbeddings(model=model, chunk_size=1) + # Create Vector store + vector_store: AzureSearch = AzureSearch( + azure_search_endpoint=vector_store_address, + azure_search_key=vector_store_password, + index_name=index_name, + embedding_function=embeddings.embed_query, + ) + # Add texts to vector store and perform a similarity search + vector_store.add_texts( + ["Test 1", "Test 2", "Test 3"], + [ + {"title": "Title 1", "any_metadata": "Metadata 1"}, + {"title": "Title 2", "any_metadata": "Metadata 2"}, + {"title": "Title 3", "any_metadata": "Metadata 3"}, + ], + ) + time.sleep(1) + res = vector_store.similarity_search(query="Test 1", k=3) + assert len(res) == 3 + + +def from_text_similarity_search_test() -> None: + """Test end to end construction and search.""" + # Create Embeddings + embeddings: OpenAIEmbeddings = OpenAIEmbeddings(model=model, chunk_size=1) + # Create Vector store + vector_store: AzureSearch = AzureSearch.from_texts( + azure_search_endpoint=vector_store_address, + azure_search_key=vector_store_password, + index_name=index_name, + texts=["Test 1", "Test 2", "Test 3"], + embedding=embeddings, + ) + time.sleep(1) + # Perform a similarity search + res = vector_store.similarity_search(query="Test 1", k=3) + assert len(res) == 3 + + +def test_semantic_hybrid_search() -> None: + """Test end to end construction and search.""" + # Create Embeddings + embeddings: OpenAIEmbeddings = OpenAIEmbeddings(model=model, chunk_size=1) + # Create Vector store + vector_store: AzureSearch = AzureSearch( + azure_search_endpoint=vector_store_address, + azure_search_key=vector_store_password, + index_name=index_name, + embedding_function=embeddings.embed_query, + semantic_configuration_name="default", + ) + # Add texts to vector store and perform a semantic hybrid search + vector_store.add_texts( + ["Test 1", "Test 2", "Test 3"], + [ + {"title": "Title 1", "any_metadata": "Metadata 1"}, + {"title": "Title 2", "any_metadata": "Metadata 2"}, + {"title": "Title 3", "any_metadata": "Metadata 3"}, + ], + ) + time.sleep(1) + res = vector_store.semantic_hybrid_search(query="What's Azure Search?", k=3) + assert len(res) == 3