From 46d295c571732bcb358ce547a767380610d88a21 Mon Sep 17 00:00:00 2001
From: Rivaaj Jumna <50959956+rivms@users.noreply.github.com>
Date: Tue, 5 Sep 2023 01:23:11 +0800
Subject: [PATCH] Added Azure Cognitive Search vector store (#7469)
---
.vscode/settings.json | 5 +-
CHANGELOG.md | 5 +
docs/community/integrations/vector_stores.md | 39 +-
.../data_modules/storage/customization.md | 11 +-
.../data_modules/storage/vector_stores.md | 5 +-
.../CognitiveSearchIndexDemo.ipynb | 543 ++++++++++++++++
llama_index/vector_stores/__init__.py | 2 +
llama_index/vector_stores/cogsearch.py | 607 ++++++++++++++++++
tests/vector_stores/test_cogsearch.py | 133 ++++
9 files changed, 1340 insertions(+), 10 deletions(-)
create mode 100644 docs/examples/vector_stores/CognitiveSearchIndexDemo.ipynb
create mode 100644 llama_index/vector_stores/cogsearch.py
create mode 100644 tests/vector_stores/test_cogsearch.py
diff --git a/.vscode/settings.json b/.vscode/settings.json
index 781eafd2ec..0cee07a33d 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -1,7 +1,10 @@
{
- "python.formatting.provider": "black",
+ "python.formatting.provider": "none",
"editor.formatOnSave": true,
"editor.codeActionsOnSave": {
"source.organizeImports": true,
},
+ "[python]": {
+ "editor.defaultFormatter": "ms-python.black-formatter"
+ },
}
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9657808640..3e556cf4e7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,10 @@
# ChangeLog
+## Unreleased
+
+### New Features
+- Added support for Azure Cognitive Search vector store (#7469)
+
## [0.8.20] - 2023-09-04
### New Features
diff --git a/docs/community/integrations/vector_stores.md b/docs/community/integrations/vector_stores.md
index ca1fe81ab8..73f511df54 100644
--- a/docs/community/integrations/vector_stores.md
+++ b/docs/community/integrations/vector_stores.md
@@ -12,6 +12,7 @@ LlamaIndex offers multiple integration points with vector stores / vector databa
LlamaIndex also supports different vector stores
as the storage backend for `VectorStoreIndex`.
+- Azure Cognitive Search (`CognitiveSearchVectorStore`). [Quickstart](https://learn.microsoft.com/en-us/azure/search/search-get-started-vector)
- [Apache Cassandra®](https://cassandra.apache.org/) and compatible databases such as [Astra DB](https://www.datastax.com/press-release/datastax-adds-vector-search-to-astra-db-on-google-cloud-for-building-real-time-generative-ai-applications) (`CassandraVectorStore`)
- Chroma (`ChromaVectorStore`) [Installation](https://docs.trychroma.com/getting-started)
- DeepLake (`DeepLakeVectorStore`) [Installation](https://docs.deeplake.ai/en/latest/Installation.html)
@@ -162,8 +163,8 @@ Zep stores texts, metadata, and embeddings. All are returned in search results.
from llama_index.vector_stores import ZepVectorStore
vector_store = ZepVectorStore(
- api_url="<api_url>",
- api_key="<api_key>",
+ api_url="<api_url>",
+ api_key="<api_key>",
collection_name="<unique_collection_name>", # Can either be an existing collection or a new one
embedding_dimensions=1536 # Optional, required if creating a new collection
)
@@ -178,7 +179,6 @@ retriever = index.as_retriever(filters=filters)
result = retriever.retrieve("What is inception about?")
```
-
**Pinecone**
```python
@@ -346,7 +346,7 @@ vector_store = MyScaleVectorStore(
```python
from llama_index.vector_stores import (
- DocArrayHnswVectorStore,
+ DocArrayHnswVectorStore,
DocArrayInMemoryVectorStore,
)
@@ -358,6 +358,7 @@ vector_store = DocArrayInMemoryVectorStore()
```
**MongoDBAtlas**
+
```python
# Provide URI to constructor, or use environment variable
import pymongo
@@ -379,6 +380,34 @@ uber_docs = SimpleDirectoryReader(input_files=["../data/10k/uber_2021.pdf"]).loa
index = VectorStoreIndex.from_documents(uber_docs, storage_context=storage_context)
```
+**Azure Cognitive Search**
+
+```python
+from azure.search.documents import SearchClient
+from llama_index.vector_stores import ChromaVectorStore
+from azure.core.credentials import AzureKeyCredential
+
+service_endpoint = f"https://{search_service_name}.search.windows.net"
+index_name = "quickstart"
+cognitive_search_credential = AzureKeyCredential("<API key>")
+
+search_client = SearchClient(
+ endpoint=service_endpoint,
+ index_name=index_name,
+ credential=cognitive_search_credential,
+)
+
+# construct vector store
+vector_store = CognitiveSearchVectorStore(
+ search_client,
+ id_field_key="id",
+ chunk_field_key="content",
+ embedding_field_key="embedding",
+ metadata_field_key="li_jsonMetadata",
+ doc_id_field_key="li_doc_id",
+)
+```
+
[Example notebooks can be found here](https://github.com/jerryjliu/llama_index/tree/main/docs/examples/vector_stores).
## Loading Data from Vector Stores using Data Connector
@@ -483,7 +512,6 @@ documents = reader.load_data(
[Example notebooks can be found here](https://github.com/jerryjliu/llama_index/tree/main/docs/examples/data_connectors).
-
```{toctree}
---
caption: Examples
@@ -514,4 +542,5 @@ maxdepth: 1
../../examples/vector_stores/MongoDBAtlasVectorSearch.ipynb
../../examples/vector_stores/postgres.ipynb
../../examples/vector_stores/AwadbDemo.ipynb
+../../examples/vector_stores/CognitiveSearchIndexDemo.ipynb
```
diff --git a/docs/core_modules/data_modules/storage/customization.md b/docs/core_modules/data_modules/storage/customization.md
index 7b2fec1376..3c1c73192d 100644
--- a/docs/core_modules/data_modules/storage/customization.md
+++ b/docs/core_modules/data_modules/storage/customization.md
@@ -1,6 +1,7 @@
# Customizing Storage
By default, LlamaIndex hides away the complexities and let you query your data in under 5 lines of code:
+
```python
from llama_index import VectorStoreIndex, SimpleDirectoryReader
@@ -12,22 +13,25 @@ response = query_engine.query("Summarize the documents.")
Under the hood, LlamaIndex also supports a swappable **storage layer** that allows you to customize where ingested documents (i.e., `Node` objects), embedding vectors, and index metadata are stored.
-
### Low-Level API
+
To do this, instead of the high-level API,
+
```python
index = VectorStoreIndex.from_documents(documents)
```
+
we use a lower-level API that gives more granular control:
+
```python
from llama_index.storage.docstore import SimpleDocumentStore
from llama_index.storage.index_store import SimpleIndexStore
from llama_index.vector_stores import SimpleVectorStore
from llama_index.node_parser import SimpleNodeParser
-# create parser and parse document into nodes
+# create parser and parse document into nodes
parser = SimpleNodeParser.from_defaults()
nodes = parser.get_nodes_from_documents(documents)
@@ -79,6 +83,7 @@ Most of our vector store integrations store the entire index (vectors + text) in
The vector stores that support this practice are:
+- CognitiveSearchVectorStore
- ChatGPTRetrievalPluginClient
- CassandraVectorStore
- ChromaVectorStore
@@ -125,7 +130,7 @@ documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
```
-If you have an existing vector store with data already loaded in,
+If you have an existing vector store with data already loaded in,
you can connect to it and directly create a `VectorStoreIndex` as follows:
```python
diff --git a/docs/core_modules/data_modules/storage/vector_stores.md b/docs/core_modules/data_modules/storage/vector_stores.md
index 3280794feb..96b96ffa7c 100644
--- a/docs/core_modules/data_modules/storage/vector_stores.md
+++ b/docs/core_modules/data_modules/storage/vector_stores.md
@@ -14,7 +14,7 @@ LlamaIndex supports over 20 different vector store options.
We are actively adding more integrations and improving feature coverage for each.
| Vector Store | Type | Metadata Filtering | Hybrid Search | Delete | Store Documents | Async |
-|--------------------------|---------------------|--------------------|---------------|--------|-----------------|-------|
+| ------------------------ | ------------------- | ------------------ | ------------- | ------ | --------------- | ----- |
| Pinecone | cloud | ✓ | ✓ | ✓ | ✓ | |
| Weaviate | self-hosted / cloud | ✓ | ✓ | ✓ | ✓ | |
| Postgres | self-hosted / cloud | ✓ | ✓ | ✓ | ✓ | ✓ |
@@ -28,6 +28,7 @@ We are actively adding more integrations and improving feature coverage for each
| Redis | self-hosted / cloud | ✓ | | ✓ | ✓ | |
| Deeplake | self-hosted / cloud | ✓ | | ✓ | ✓ | |
| OpenSearch | self-hosted / cloud | ✓ | | ✓ | ✓ | |
+| Azure Cognitive Search | cloud | | ✓ | ✓ | ✓ | |
| DynamoDB | cloud | | | ✓ | | |
| LanceDB | cloud | ✓ | | ✓ | ✓ | |
| Metal | cloud | ✓ | | ✓ | ✓ | |
@@ -37,6 +38,7 @@ We are actively adding more integrations and improving feature coverage for each
| FAISS | in-memory | | | | | |
| ChatGPT Retrieval Plugin | aggregator | | | ✓ | ✓ | |
| DocArray | aggregator | ✓ | | ✓ | ✓ | |
+| Azure Cognitive Search | cloud | ✓ | ✓ | ✓ | ✓ | |
For more details, see [Vector Store Integrations](/community/integrations/vector_stores.md).
@@ -69,4 +71,5 @@ maxdepth: 1
/examples/vector_stores/DocArrayInMemoryIndexDemo.ipynb
/examples/vector_stores/MongoDBAtlasVectorSearch.ipynb
/examples/vector_stores/CassandraIndexDemo.ipynb
+/examples/vector_stores/CognitiveSearchIndexDemo.ipynb
```
diff --git a/docs/examples/vector_stores/CognitiveSearchIndexDemo.ipynb b/docs/examples/vector_stores/CognitiveSearchIndexDemo.ipynb
new file mode 100644
index 0000000000..fc8a76cf71
--- /dev/null
+++ b/docs/examples/vector_stores/CognitiveSearchIndexDemo.ipynb
@@ -0,0 +1,543 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Azure Cognitive Search"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Basic Example\n",
+ "\n",
+ "In this basic example, we take a Paul Graham essay, split it into chunks, embed it using an OpenAI embedding model, load it into an Azure Cognitive Search index, and then query it."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import logging\n",
+ "import sys\n",
+ "from IPython.display import Markdown, display\n",
+ "\n",
+ "# logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
+ "# logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
+ "\n",
+ "# logger = logging.getLogger(__name__)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "#!{sys.executable} -m pip install llama-index\n",
+ "#!{sys.executable} -m pip install azure-search-documents==11.4.0b8\n",
+ "#!{sys.executable} -m pip install azure-identity"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# set up OpenAI\n",
+ "import os\n",
+ "import getpass\n",
+ "\n",
+ "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"OpenAI API Key:\")\n",
+ "import openai\n",
+ "\n",
+ "openai.api_key = os.environ[\"OPENAI_API_KEY\"]"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# set up Azure Cognitive Search\n",
+ "from azure.search.documents.indexes import SearchIndexClient\n",
+ "from azure.search.documents import SearchClient\n",
+ "from azure.core.credentials import AzureKeyCredential\n",
+ "\n",
+ "search_service_name = getpass.getpass(\"Azure Cognitive Search Service Name\")\n",
+ "\n",
+ "key = getpass.getpass(\"Azure Cognitive Search Key\")\n",
+ "\n",
+ "cognitive_search_credential = AzureKeyCredential(key)\n",
+ "\n",
+ "service_endpoint = f\"https://{search_service_name}.search.windows.net\"\n",
+ "\n",
+ "# Index name to use\n",
+ "index_name = \"quickstart\"\n",
+ "\n",
+ "# Use index client to demonstrate creating an index\n",
+ "index_client = SearchIndexClient(\n",
+ " endpoint=service_endpoint,\n",
+ " credential=cognitive_search_credential,\n",
+ ")\n",
+ "\n",
+ "# Use search client to demonstration using existing index\n",
+ "search_client = SearchClient(\n",
+ " endpoint=service_endpoint,\n",
+ " index_name=index_name,\n",
+ " credential=cognitive_search_credential,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Create Index (if it does not exist)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Demonstrates creating a vector index named quickstart01 if one doesn't exist. The index has the following fields:\n",
+ "- id (Edm.String)\n",
+ "- content (Edm.String)\n",
+ "- embedding (Edm.SingleCollection)\n",
+ "- li_jsonMetadata (Edm.String)\n",
+ "- li_doc_id (Edm.String)\n",
+ "- author (Edm.String)\n",
+ "- theme (Edm.String)\n",
+ "- director (Edm.String)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 5,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from azure.search.documents import SearchClient\n",
+ "from llama_index.vector_stores import CognitiveSearchVectorStore\n",
+ "from llama_index.vector_stores.cogsearch import (\n",
+ " IndexManagement,\n",
+ " MetadataIndexFieldType,\n",
+ " CognitiveSearchVectorStore,\n",
+ ")\n",
+ "\n",
+ "# Example of a complex mapping, metadata field 'theme' is mapped to a differently name index field 'topic' with its type explicitly set\n",
+ "metadata_fields = {\n",
+ " \"author\": \"author\",\n",
+ " \"theme\": (\"topic\", MetadataIndexFieldType.STRING),\n",
+ " \"director\": \"director\",\n",
+ "}\n",
+ "\n",
+ "# A simplified metadata specification is available if all metadata and index fields are similarly named\n",
+ "# metadata_fields = {\"author\", \"theme\", \"director\"}\n",
+ "\n",
+ "\n",
+ "vector_store = CognitiveSearchVectorStore(\n",
+ " search_or_index_client=index_client,\n",
+ " index_name=index_name,\n",
+ " filterable_metadata_field_keys=metadata_fields,\n",
+ " index_management=IndexManagement.CREATE_IF_NOT_EXISTS,\n",
+ " id_field_key=\"id\",\n",
+ " chunk_field_key=\"content\",\n",
+ " embedding_field_key=\"embedding\",\n",
+ " metadata_string_field_key=\"li_jsonMetadata\",\n",
+ " doc_id_field_key=\"li_doc_id\",\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 6,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# define embedding function\n",
+ "from llama_index.embeddings import OpenAIEmbedding\n",
+ "from llama_index import (\n",
+ " SimpleDirectoryReader,\n",
+ " StorageContext,\n",
+ " ServiceContext,\n",
+ " VectorStoreIndex,\n",
+ ")\n",
+ "\n",
+ "embed_model = OpenAIEmbedding()\n",
+ "\n",
+ "# load documents\n",
+ "documents = SimpleDirectoryReader(\n",
+ " \"../../../examples/paul_graham_essay/data\"\n",
+ ").load_data()\n",
+ "\n",
+ "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
+ "service_context = ServiceContext.from_defaults(embed_model=embed_model)\n",
+ "index = VectorStoreIndex.from_documents(\n",
+ " documents, storage_context=storage_context, service_context=service_context\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 7,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/markdown": [
+ "<b>The author wrote short stories and programmed on an IBM 1401 computer during their time in school. They later got their own microcomputer, a TRS-80, and started programming games and a word processor.</b>"
+ ],
+ "text/plain": [
+ "<IPython.core.display.Markdown object>"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "# Query Data\n",
+ "query_engine = index.as_query_engine(similarity_top_k=3)\n",
+ "response = query_engine.query(\"What did the author do growing up?\")\n",
+ "display(Markdown(f\"<b>{response}</b>\"))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 8,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/markdown": [
+ "<b>The author learned several things during their time at Interleaf. They learned that it's better for technology companies to be run by product people than sales people, that code edited by too many people leads to bugs, that cheap office space is not worth it if it's depressing, that planned meetings are inferior to corridor conversations, that big bureaucratic customers can be a dangerous source of money, and that there's not much overlap between conventional office hours and the optimal time for hacking. However, the most important thing the author learned is that the low end eats the high end, meaning that it's better to be the \"entry level\" option because if you're not, someone else will be and will surpass you.</b>"
+ ],
+ "text/plain": [
+ "<IPython.core.display.Markdown object>"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "response = query_engine.query(\n",
+ " \"What did the author learn?\",\n",
+ ")\n",
+ "display(Markdown(f\"<b>{response}</b>\"))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Use Existing Index"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 9,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.vector_stores import CognitiveSearchVectorStore\n",
+ "from llama_index.vector_stores.cogsearch import (\n",
+ " IndexManagement,\n",
+ " MetadataIndexFieldType,\n",
+ " CognitiveSearchVectorStore,\n",
+ ")\n",
+ "\n",
+ "\n",
+ "index_name = \"quickstart\"\n",
+ "\n",
+ "metadata_fields = {\n",
+ " \"author\": \"author\",\n",
+ " \"theme\": (\"topic\", MetadataIndexFieldType.STRING),\n",
+ " \"director\": \"director\",\n",
+ "}\n",
+ "vector_store = CognitiveSearchVectorStore(\n",
+ " search_or_index_client=search_client,\n",
+ " filterable_metadata_field_keys=metadata_fields,\n",
+ " index_management=IndexManagement.NO_VALIDATION,\n",
+ " id_field_key=\"id\",\n",
+ " chunk_field_key=\"content\",\n",
+ " embedding_field_key=\"embedding\",\n",
+ " metadata_string_field_key=\"li_jsonMetadata\",\n",
+ " doc_id_field_key=\"li_doc_id\",\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 10,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# define embedding function\n",
+ "from llama_index.embeddings import OpenAIEmbedding\n",
+ "from llama_index import (\n",
+ " SimpleDirectoryReader,\n",
+ " StorageContext,\n",
+ " ServiceContext,\n",
+ " VectorStoreIndex,\n",
+ ")\n",
+ "\n",
+ "embed_model = OpenAIEmbedding()\n",
+ "\n",
+ "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
+ "service_context = ServiceContext.from_defaults(embed_model=embed_model)\n",
+ "index = VectorStoreIndex.from_documents(\n",
+ " [], storage_context=storage_context, service_context=service_context\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 11,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/markdown": [
+ "<b>The author experienced a difficult moment when their mother had a stroke and was put in a nursing home. The stroke destroyed her balance, and the author and their sister were determined to help her get out of the nursing home and back to her house.</b>"
+ ],
+ "text/plain": [
+ "<IPython.core.display.Markdown object>"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "query_engine = index.as_query_engine()\n",
+ "response = query_engine.query(\"What was a hard moment for the author?\")\n",
+ "display(Markdown(f\"<b>{response}</b>\"))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 12,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/markdown": [
+ "<b>The author of the given context is Paul Graham.</b>"
+ ],
+ "text/plain": [
+ "<IPython.core.display.Markdown object>"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "response = query_engine.query(\"Who is the author?\")\n",
+ "display(Markdown(f\"<b>{response}</b>\"))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 13,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "At Interleaf, there was a group called Release Engineering that seemed to be as big as the group that actually wrote the software. The software at Interleaf had to be updated on the server, and there was a lot of emphasis on high production values to make the online store builders look legitimate.\n",
+ "\n",
+ "Streamed output at 20.953424485215063 tokens/s\n"
+ ]
+ }
+ ],
+ "source": [
+ "import time\n",
+ "\n",
+ "query_engine = index.as_query_engine(streaming=True)\n",
+ "response = query_engine.query(\"What happened at interleaf?\")\n",
+ "\n",
+ "start_time = time.time()\n",
+ "\n",
+ "token_count = 0\n",
+ "for token in response.response_gen:\n",
+ " print(token, end=\"\")\n",
+ " token_count += 1\n",
+ "\n",
+ "time_elapsed = time.time() - start_time\n",
+ "tokens_per_second = token_count / time_elapsed\n",
+ "\n",
+ "print(f\"\\n\\nStreamed output at {tokens_per_second} tokens/s\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Adding a document to existing index"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 14,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/markdown": [
+ "<b>The color of the sky can vary depending on various factors such as time of day, weather conditions, and location. It can range from shades of blue during the day to hues of orange, pink, and purple during sunrise or sunset.</b>"
+ ],
+ "text/plain": [
+ "<IPython.core.display.Markdown object>"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "response = query_engine.query(\"What colour is the sky?\")\n",
+ "display(Markdown(f\"<b>{response}</b>\"))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 15,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index import Document\n",
+ "\n",
+ "index.insert_nodes([Document(text=\"The sky is indigo today\")])"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 16,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/markdown": [
+ "<b>The colour of the sky is indigo.</b>"
+ ],
+ "text/plain": [
+ "<IPython.core.display.Markdown object>"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "response = query_engine.query(\"What colour is the sky?\")\n",
+ "display(Markdown(f\"<b>{response}</b>\"))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Filtering"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 17,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.schema import TextNode\n",
+ "\n",
+ "nodes = [\n",
+ " TextNode(\n",
+ " text=\"The Shawshank Redemption\",\n",
+ " metadata={\n",
+ " \"author\": \"Stephen King\",\n",
+ " \"theme\": \"Friendship\",\n",
+ " },\n",
+ " ),\n",
+ " TextNode(\n",
+ " text=\"The Godfather\",\n",
+ " metadata={\n",
+ " \"director\": \"Francis Ford Coppola\",\n",
+ " \"theme\": \"Mafia\",\n",
+ " },\n",
+ " ),\n",
+ " TextNode(\n",
+ " text=\"Inception\",\n",
+ " metadata={\n",
+ " \"director\": \"Christopher Nolan\",\n",
+ " },\n",
+ " ),\n",
+ "]"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 18,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "index.insert_nodes(nodes)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 19,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/plain": [
+ "[NodeWithScore(node=TextNode(id_='5a97da0c-8f04-4c63-b90b-8c474d8c273d', embedding=None, metadata={'director': 'Francis Ford Coppola', 'theme': 'Mafia'}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='81cf4b9e847ba42e83fc401e31af8e17d629f0d5cf9c0c320ec7ac69dd0257e1', text='The Godfather', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), score=0.81316805)]"
+ ]
+ },
+ "execution_count": 19,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "from llama_index.vector_stores.types import ExactMatchFilter, MetadataFilters\n",
+ "\n",
+ "\n",
+ "filters = MetadataFilters(filters=[ExactMatchFilter(key=\"theme\", value=\"Mafia\")])\n",
+ "\n",
+ "retriever = index.as_retriever(filters=filters)\n",
+ "retriever.retrieve(\"What is inception about?\")"
+ ]
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "llamaindextest01",
+ "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.4"
+ },
+ "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/llama_index/vector_stores/__init__.py b/llama_index/vector_stores/__init__.py
index 78b27cce9d..ea9428ed32 100644
--- a/llama_index/vector_stores/__init__.py
+++ b/llama_index/vector_stores/__init__.py
@@ -29,6 +29,7 @@ from llama_index.vector_stores.docarray import (
DocArrayHnswVectorStore,
DocArrayInMemoryVectorStore,
)
+from llama_index.vector_stores.cogsearch import CognitiveSearchVectorStore
from llama_index.vector_stores.awadb import AwaDBVectorStore
__all__ = [
@@ -57,4 +58,5 @@ __all__ = [
"ZepVectorStore",
"AwaDBVectorStore",
"BagelVectorStore",
+ "CognitiveSearchVectorStore",
]
diff --git a/llama_index/vector_stores/cogsearch.py b/llama_index/vector_stores/cogsearch.py
new file mode 100644
index 0000000000..d3522f253b
--- /dev/null
+++ b/llama_index/vector_stores/cogsearch.py
@@ -0,0 +1,607 @@
+"""Azure Cognitive Search vector store."""
+import logging
+from typing import Any, List, cast, Dict, Callable, Optional, Tuple, Union
+import enum
+from enum import auto
+
+from llama_index.schema import MetadataMode, TextNode
+from llama_index.vector_stores.types import (
+ MetadataFilters,
+ ExactMatchFilter,
+ NodeWithEmbedding,
+ VectorStore,
+ VectorStoreQuery,
+ VectorStoreQueryResult,
+ VectorStoreQueryMode,
+)
+
+from llama_index.vector_stores.utils import (
+ node_to_metadata_dict,
+ metadata_dict_to_node,
+ legacy_metadata_dict_to_node,
+)
+
+import json
+
+logger = logging.getLogger(__name__)
+
+
+class MetadataIndexFieldType(int, enum.Enum):
+ """
+ Enumeration representing the supported types for metadata fields in an
+ Azure Cognitive Search Index, corresponds with types supported in a flat
+ metadata dictionary
+ """
+
+ STRING = auto() # "Edm.String"
+ BOOLEAN = auto() # "Edm.Boolean"
+ INT32 = auto() # "Edm.Int32"
+ INT64 = auto() # "Edm.Int64"
+ DOUBLE = auto() # "Edm.Double"
+
+
+class IndexManagement(int, enum.Enum):
+ """Enumeration representing the supported index management operations"""
+
+ NO_VALIDATION = auto()
+ VALIDATE_INDEX = auto()
+ CREATE_IF_NOT_EXISTS = auto()
+
+
+class CognitiveSearchVectorStore(VectorStore):
+ stores_text: bool = True
+ flat_metadata: bool = True
+
+ def _normalise_metadata_to_index_fields(
+ self,
+ filterable_metadata_field_keys: Union[
+ List[str],
+ Dict[str, str],
+ Dict[str, Tuple[str, MetadataIndexFieldType]],
+ None,
+ ] = [],
+ ) -> Dict[str, Tuple[str, MetadataIndexFieldType]]:
+ index_field_spec: Dict[str, Tuple[str, MetadataIndexFieldType]] = {}
+
+ if isinstance(filterable_metadata_field_keys, List):
+ for field in filterable_metadata_field_keys:
+ # Index field name and the metadata field name are the same
+ # Use String as the default index field type
+ index_field_spec[field] = (field, MetadataIndexFieldType.STRING)
+
+ elif isinstance(filterable_metadata_field_keys, Dict):
+ for k, v in filterable_metadata_field_keys.items():
+ if isinstance(v, tuple):
+ # Index field name and metadata field name may differ
+ # The index field type used is as supplied
+ index_field_spec[k] = v
+ else:
+ # Index field name and metadata field name may differ
+ # Use String as the default index field type
+ index_field_spec[k] = (v, MetadataIndexFieldType.STRING)
+
+ return index_field_spec
+
+ def _create_index_if_not_exists(self, index_name: str) -> None:
+ if index_name not in self._index_client.list_index_names():
+ logger.info(f"Index {index_name} does not exist, creating index")
+ self._create_index(index_name)
+
+ def _create_metadata_index_fields(self) -> List[Any]:
+ """Create a list of index fields for storing metadata values"""
+ from azure.search.documents.indexes.models import SimpleField
+
+ index_fields = []
+
+ # create search fields
+ for k, v in self._metadata_to_index_field_map.items():
+ field_name, field_type = v
+
+ if field_type == MetadataIndexFieldType.STRING:
+ index_field_type = "Edm.String"
+ elif field_type == MetadataIndexFieldType.INT32:
+ index_field_type = "Edm.Int32"
+ elif field_type == MetadataIndexFieldType.INT64:
+ index_field_type = "Edm.Int64"
+ elif field_type == MetadataIndexFieldType.DOUBLE:
+ index_field_type = "Edm.Double"
+ elif field_type == MetadataIndexFieldType.BOOLEAN:
+ index_field_type = "Edm.Boolean"
+
+ field = SimpleField(name=field_name, type=index_field_type, filterable=True)
+ index_fields.append(field)
+
+ return index_fields
+
+ def _create_index(self, index_name: Optional[str]) -> None:
+ """
+ Creates a default index based on the supplied index name, key field names and
+ metadata filtering keys
+ """
+
+ from azure.search.documents.indexes.models import (
+ SimpleField,
+ SearchField,
+ SearchIndex,
+ SearchableField,
+ SearchFieldDataType,
+ HnswVectorSearchAlgorithmConfiguration,
+ HnswParameters,
+ VectorSearch,
+ SemanticSettings,
+ SemanticConfiguration,
+ PrioritizedFields,
+ SemanticField,
+ )
+
+ logger.info(f"Configuring {index_name} fields")
+ fields = [
+ SimpleField(name=self._field_mapping["id"], type="Edm.String", key=True),
+ SearchableField(
+ name=self._field_mapping["chunk"],
+ type="Edm.String",
+ analyzer_name="en.microsoft",
+ ),
+ SearchField(
+ name=self._field_mapping["embedding"],
+ type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
+ hidden=False,
+ searchable=True,
+ filterable=False,
+ sortable=False,
+ facetable=False,
+ vector_search_dimensions=1536,
+ vector_search_configuration="default",
+ ),
+ SimpleField(name=self._field_mapping["metadata"], type="Edm.String"),
+ SimpleField(
+ name=self._field_mapping["doc_id"], type="Edm.String", filterable=True
+ ),
+ ]
+
+ logger.info(f"Configuring {index_name} metadata fields")
+ # Add on the metadata fields
+
+ metadata_index_fields = self._create_metadata_index_fields()
+
+ fields.extend(metadata_index_fields)
+
+ logger.info(f"Configuring {index_name} vector search")
+
+ hnsw_param = HnswParameters(
+ m=4,
+ ef_construction=500,
+ ef_search=1000,
+ metric="cosine",
+ )
+
+ vector_search = VectorSearch(
+ algorithm_configurations=[
+ HnswVectorSearchAlgorithmConfiguration(
+ name="default",
+ kind="hnsw",
+ parameters=hnsw_param,
+ )
+ ]
+ )
+
+ logger.info(f"Configuring {index_name} semantic search")
+ semantic_settings = SemanticSettings(
+ configurations=[
+ SemanticConfiguration(
+ name="default",
+ prioritized_fields=PrioritizedFields(
+ title_field=None,
+ prioritized_content_fields=[
+ SemanticField(field_name=self._field_mapping["chunk"])
+ ],
+ ),
+ )
+ ]
+ )
+
+ index = SearchIndex(
+ name=index_name,
+ fields=fields,
+ semantic_settings=semantic_settings,
+ vector_search=vector_search,
+ )
+
+ logger.debug(f"Creating {index_name} search index")
+ self._index_client.create_index(index)
+
+ def _validate_index(self, index_name: Optional[str]) -> None:
+ if self._index_client and index_name:
+ if index_name not in self._index_client.list_index_names():
+ raise ValueError(
+ f"Validation failed, index {index_name} does not exist."
+ )
+
+ def __init__(
+ self,
+ search_or_index_client: Any,
+ id_field_key: str,
+ chunk_field_key: str,
+ embedding_field_key: str,
+ metadata_string_field_key: str,
+ doc_id_field_key: str,
+ filterable_metadata_field_keys: Optional[
+ Union[
+ List[str],
+ Dict[str, str],
+ Dict[str, Tuple[str, MetadataIndexFieldType]],
+ ]
+ ] = None,
+ index_name: Optional[str] = None,
+ index_mapping: Optional[
+ Callable[[Dict[str, str], Dict[str, Any]], Dict[str, str]]
+ ] = None,
+ index_management: IndexManagement = IndexManagement.NO_VALIDATION,
+ **kwargs: Any,
+ ) -> None:
+ # ruff: noqa: E501
+ """
+ Embeddings and documents are stored in an Azure Cognitive Search index,
+ a merge or upload approach is used when adding embeddings.
+ When adding multiple embeddings the index is updated by this vector store
+ in batches of 10 documents, very large nodes may result in failure due to
+ the batch byte size being exceeded.
+
+ Args:
+ search_client (azure.search.documents.SearchClient):
+ Client for index to populated / queried.
+ id_field_key (str): Index field storing the id
+ chunk_field_key (str): Index field storing the node text
+ embedding_field_key (str): Index field storing the embedding vector
+ metadata_string_field_key (str):
+ Index field storing node metadata as a json string.
+ Schema is arbitrary, to filter on metadata values they must be stored
+ as separate fields in the index, use filterable_metadata_field_keys
+ to specify the metadata values that should be stored in these filterable fields
+ doc_id_field_key (str): Index field storing doc_id
+ index_mapping:
+ Optional function with definition
+ (enriched_doc: Dict[str, str], metadata: Dict[str, Any]): Dict[str,str]
+ used to map document fields to the Cognitive search index fields
+ (return value of function).
+ If none is specified a default mapping is provided which uses
+ the field keys. The keys in the enriched_doc are
+ ["id", "chunk", "embedding", "metadata"]
+ The default mapping is:
+ - "id" to id_field_key
+ - "chunk" to chunk_field_key
+ - "embedding" to embedding_field_key
+ - "metadata" to metadata_field_key
+ *kwargs (Any): Additional keyword arguments.
+
+ Raises:
+ ImportError: Unable to import `azure-search-documents`
+ ValueError: If `search_or_index_client` is not provided
+ ValueError: If `index_name` is not provided and `search_or_index_client`
+ is of type azure.search.documents.SearchIndexClient
+ ValueError: If `index_name` is provided and `search_or_index_client`
+ is of type azure.search.documents.SearchClient
+ ValueError: If `create_index_if_not_exists` is true and
+ `search_or_index_client` is of type azure.search.documents.SearchClient
+ """
+
+ import_err_msg = (
+ "`azure-search-documents` package not found, please run "
+ "`pip install azure-search-documents==11.4.0b8`"
+ )
+
+ try:
+ import azure.search.documents # noqa: F401
+ from azure.search.documents import SearchClient # noqa: F401
+ from azure.search.documents.indexes import SearchIndexClient # noqa: F401
+ except ImportError:
+ raise ImportError(import_err_msg)
+
+ self._index_client: SearchIndexClient = cast(SearchIndexClient, None)
+ self._search_client: SearchClient = cast(SearchClient, None)
+
+ # Validate search_or_index_client
+ if search_or_index_client is not None:
+ if isinstance(search_or_index_client, SearchIndexClient):
+ # If SearchIndexClient is supplied so must index_name
+ self._index_client = cast(SearchIndexClient, search_or_index_client)
+
+ if not index_name:
+ raise ValueError(
+ "index_name must be supplied if search_or_index_client is of "
+ "type azure.search.documents.SearchIndexClient"
+ )
+
+ self._search_client = self._index_client.get_search_client(
+ index_name=index_name
+ )
+
+ elif isinstance(search_or_index_client, SearchClient):
+ self._search_client = cast(SearchClient, search_or_index_client)
+
+ # Validate index_name
+ if index_name:
+ raise ValueError(
+ "index_name cannot be supplied if search_or_index_client "
+ "is of type azure.search.documents.SearchClient"
+ )
+
+ if not self._index_client and not self._search_client:
+ raise ValueError(
+ "search_or_index_client must be of type "
+ "azure.search.documents.SearchClient or "
+ "azure.search.documents.SearchIndexClient"
+ )
+ else:
+ raise ValueError("search_or_index_client not specified")
+
+ if (
+ index_management == IndexManagement.CREATE_IF_NOT_EXISTS
+ and not self._index_client
+ ):
+ raise ValueError(
+ "index_management has value of IndexManagement.CREATE_IF_NOT_EXISTS "
+ "but search_or_index_client is not of type "
+ "azure.search.documents.SearchIndexClient"
+ )
+
+ self._index_management = index_management
+
+ # Default field mapping
+ field_mapping = {
+ "id": id_field_key,
+ "chunk": chunk_field_key,
+ "embedding": embedding_field_key,
+ "metadata": metadata_string_field_key,
+ "doc_id": doc_id_field_key,
+ }
+
+ self._field_mapping = field_mapping
+
+ self._index_mapping = (
+ self._default_index_mapping if index_mapping is None else index_mapping
+ )
+
+ # self._filterable_metadata_field_keys = filterable_metadata_field_keys
+ self._metadata_to_index_field_map = self._normalise_metadata_to_index_fields(
+ filterable_metadata_field_keys
+ )
+
+ if self._index_management == IndexManagement.CREATE_IF_NOT_EXISTS:
+ if index_name:
+ self._create_index_if_not_exists(index_name)
+
+ if self._index_management == IndexManagement.VALIDATE_INDEX:
+ self._validate_index(index_name)
+
+ @property
+ def client(self) -> Any:
+ """Get client."""
+ return self._search_client
+
+ def _default_index_mapping(
+ self, enriched_doc: Dict[str, str], metadata: Dict[str, Any]
+ ) -> Dict[str, str]:
+ index_doc: Dict[str, str] = {}
+
+ for field in self._field_mapping.keys():
+ index_doc[self._field_mapping[field]] = enriched_doc[field]
+
+ for metadata_field_name, (
+ index_field_name,
+ _,
+ ) in self._metadata_to_index_field_map.items():
+ metadata_value = metadata.get(metadata_field_name)
+ if metadata_value:
+ index_doc[index_field_name] = metadata_value
+
+ return index_doc
+
+ def add(
+ self,
+ embedding_results: List[NodeWithEmbedding],
+ ) -> List[str]:
+ """Add embedding results to index associated with the configured search client.
+
+ Args
+ embedding_results: List[NodeWithEmbedding]: list of embedding results
+
+ """
+
+ if not self._search_client:
+ raise ValueError("Search client not initialized")
+
+ documents = []
+ ids = []
+
+ for embedding in embedding_results:
+ logger.debug(f"Processing embedding: {embedding.id}")
+ ids.append(embedding.id)
+
+ index_document = self._create_index_document(embedding)
+
+ documents.append(index_document)
+
+ if len(documents) >= 10:
+ logger.info(
+ f"Uploading batch of size {len(documents)}, "
+ f"current progress {len(ids)} of {len(embedding_results)}"
+ )
+ self._search_client.merge_or_upload_documents(documents)
+ documents = []
+
+ # Upload remaining batch of less than 10 documents
+ if len(documents) > 0:
+ logger.info(
+ f"Uploading remaining batch of size {len(documents)}, "
+ f"current progress {len(ids)} of {len(embedding_results)}"
+ )
+ self._search_client.merge_or_upload_documents(documents)
+ documents = []
+
+ return ids
+
+ def _create_index_document(self, embedding: NodeWithEmbedding) -> Dict[str, Any]:
+ """Create Cognitive Search index document from embedding result"""
+ doc: Dict[str, Any] = {}
+ doc["id"] = embedding.id
+ doc["chunk"] = embedding.node.get_content(metadata_mode=MetadataMode.NONE) or ""
+ doc["embedding"] = embedding.embedding
+ doc["doc_id"] = embedding.ref_doc_id
+
+ node_metadata = node_to_metadata_dict(
+ embedding.node,
+ remove_text=True,
+ flat_metadata=self.flat_metadata,
+ )
+
+ doc["metadata"] = json.dumps(node_metadata)
+
+ index_document = self._index_mapping(doc, node_metadata)
+
+ return index_document
+
+ def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None:
+ """
+ Delete documents from the Cognitive Search Index
+ with doc_id_field_key field equal to ref_doc_id."""
+
+ # Locate documents to delete
+ filter = f'{self._field_mapping["doc_id"]} eq \'{ref_doc_id}\''
+ results = self._search_client.search(search_text="*", filter=filter)
+
+ logger.debug(f"Searching with filter {filter}")
+
+ docs_to_delete = []
+ for result in results:
+ doc = {}
+ doc["id"] = result[self._field_mapping["id"]]
+ logger.debug(f"Found document to delete: {doc}")
+ docs_to_delete.append(doc)
+
+ if len(docs_to_delete) > 0:
+ logger.debug(f"Deleting {len(docs_to_delete)} documents")
+ self._search_client.delete_documents(docs_to_delete)
+
+ def _create_odata_filter(self, metadata_filters: MetadataFilters) -> str:
+ """Generate an OData filter string using supplied metadata filters"""
+ odata_filter: List[str] = []
+ for f in metadata_filters.filters:
+ if not isinstance(f, ExactMatchFilter):
+ raise NotImplementedError(
+ "Only `ExactMatchFilter` filters are supported"
+ )
+
+ # Raise error if filtering on a metadata field that lacks a mapping to
+ # an index field
+ metadata_mapping = self._metadata_to_index_field_map.get(f.key)
+
+ if not metadata_mapping:
+ raise ValueError(
+ f"Metadata field '{f.key}' is missing a mapping to an index field, "
+ "provide entry in 'filterable_metadata_field_keys' for this "
+ "vector store"
+ )
+
+ index_field = metadata_mapping[0]
+
+ if len(odata_filter) > 0:
+ odata_filter.append(" and ")
+ if isinstance(f.value, str):
+ escaped_value = "".join([("''" if s == "'" else s) for s in f.value])
+ odata_filter.append(f"{index_field} eq '{escaped_value}'")
+ else:
+ odata_filter.append(f"{index_field} eq {f.value}")
+
+ odata_expr = "".join(odata_filter)
+
+ logger.info(f"Odata filter: {odata_expr}")
+
+ return odata_expr
+
+ def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResult:
+ """Query vector store."""
+ from azure.search.documents.models import Vector
+
+ select_fields = [
+ self._field_mapping["id"],
+ self._field_mapping["chunk"],
+ self._field_mapping["metadata"],
+ self._field_mapping["doc_id"],
+ ]
+
+ search_query = "*"
+ vectors = None
+
+ if query.mode in (VectorStoreQueryMode.SPARSE, VectorStoreQueryMode.HYBRID):
+ if query.query_str is None:
+ raise ValueError("Query missing query string")
+
+ search_query = query.query_str
+
+ logger.info(f"Hybrid search with search text: {search_query}")
+
+ if query.mode in (VectorStoreQueryMode.DEFAULT, VectorStoreQueryMode.HYBRID):
+ if not query.query_embedding:
+ raise ValueError("Query missing embedding")
+
+ vector = Vector(
+ value=query.query_embedding,
+ k=query.similarity_top_k,
+ fields=self._field_mapping["embedding"],
+ )
+ vectors = [vector]
+ logger.info("Vector search with supplied embedding")
+
+ odata_filter = None
+ if query.filters is not None:
+ odata_filter = self._create_odata_filter(query.filters)
+
+ results = self._search_client.search(
+ search_text=search_query,
+ vectors=vectors,
+ top=query.similarity_top_k,
+ select=select_fields,
+ filter=odata_filter,
+ )
+
+ id_result = []
+ node_result = []
+ score_result = []
+ for result in results:
+ node_id = result[self._field_mapping["id"]]
+ metadata = json.loads(result[self._field_mapping["metadata"]])
+ score = result["@search.score"]
+ chunk = result[self._field_mapping["chunk"]]
+
+ try:
+ node = metadata_dict_to_node(metadata)
+ node.set_content(chunk)
+ except Exception:
+ # NOTE: deprecated legacy logic for backward compatibility
+ metadata, node_info, relationships = legacy_metadata_dict_to_node(
+ metadata
+ )
+
+ node = TextNode(
+ text=chunk,
+ id_=node_id,
+ metadata=metadata,
+ start_char_idx=node_info.get("start", None),
+ end_char_idx=node_info.get("end", None),
+ relationships=relationships,
+ )
+
+ logger.debug(f"Retrieved node id {node_id} with node data of {node}")
+
+ id_result.append(node_id)
+ node_result.append(node)
+ score_result.append(score)
+
+ logger.debug(
+ f"Search query '{search_query}' returned {len(id_result)} results."
+ )
+
+ return VectorStoreQueryResult(
+ nodes=node_result, similarities=score_result, ids=id_result
+ )
diff --git a/tests/vector_stores/test_cogsearch.py b/tests/vector_stores/test_cogsearch.py
new file mode 100644
index 0000000000..9dfeb037c9
--- /dev/null
+++ b/tests/vector_stores/test_cogsearch.py
@@ -0,0 +1,133 @@
+from unittest.mock import MagicMock
+from typing import Any, List, Optional
+
+from llama_index.schema import NodeRelationship, RelatedNodeInfo, TextNode
+from llama_index.vector_stores.types import NodeWithEmbedding
+
+from llama_index.vector_stores import CognitiveSearchVectorStore
+from llama_index.vector_stores.cogsearch import IndexManagement
+import pytest
+
+try:
+ from azure.search.documents import SearchClient
+ from azure.search.documents.indexes import SearchIndexClient
+
+ cogsearch_installed = True
+except ImportError:
+ cogsearch_installed = False
+ search_client = None # type: ignore
+
+
+def create_mock_vector_store(
+ search_client: Any,
+ index_name: Optional[str] = None,
+ index_management: IndexManagement = IndexManagement.NO_VALIDATION,
+) -> CognitiveSearchVectorStore:
+ vector_store = CognitiveSearchVectorStore(
+ search_or_index_client=search_client,
+ id_field_key="id",
+ chunk_field_key="content",
+ embedding_field_key="embedding",
+ metadata_string_field_key="li_jsonMetadata",
+ doc_id_field_key="li_doc_id",
+ index_name=index_name,
+ index_management=index_management,
+ )
+ return vector_store
+
+
+def create_sample_documents(n: int) -> List[NodeWithEmbedding]:
+ nodes: List[NodeWithEmbedding] = []
+
+ for i in range(n):
+ nodes.append(
+ NodeWithEmbedding(
+ node=TextNode(
+ text=f"test node text {i}",
+ relationships={
+ NodeRelationship.SOURCE: RelatedNodeInfo(
+ node_id=f"test doc id {i}"
+ )
+ },
+ ),
+ embedding=[0.5, 0.5],
+ )
+ )
+
+ return nodes
+
+
+@pytest.mark.skipif(
+ not cogsearch_installed, reason="azure-search-documents package not installed"
+)
+def test_cogsearch_add_two_batches() -> None:
+ search_client = MagicMock(spec=SearchClient)
+ vector_store = create_mock_vector_store(search_client)
+
+ nodes = create_sample_documents(11)
+
+ ids = vector_store.add(nodes)
+
+ call_count = search_client.merge_or_upload_documents.call_count
+
+ assert ids is not None
+ assert len(ids) == 11
+ assert call_count == 2
+
+
+@pytest.mark.skipif(
+ not cogsearch_installed, reason="azure-search-documents package not installed"
+)
+def test_cogsearch_add_one_batch() -> None:
+ search_client = MagicMock(spec=SearchClient)
+ vector_store = create_mock_vector_store(search_client)
+
+ nodes = create_sample_documents(10)
+
+ ids = vector_store.add(nodes)
+
+ call_count = search_client.merge_or_upload_documents.call_count
+
+ assert ids is not None
+ assert len(ids) == 10
+ assert call_count == 1
+
+
+@pytest.mark.skipif(
+ not cogsearch_installed, reason="azure-search-documents package not installed"
+)
+def test_invalid_index_management_for_searchclient() -> None:
+ search_client = MagicMock(spec=SearchClient)
+
+ # No error
+ create_mock_vector_store(
+ search_client, index_management=IndexManagement.VALIDATE_INDEX
+ )
+
+ # Cannot supply index name
+ # ruff: noqa: E501
+ with pytest.raises(
+ ValueError,
+ match="index_name cannot be supplied if search_or_index_client is of type azure.search.documents.SearchClient",
+ ):
+ create_mock_vector_store(search_client, index_name="test01")
+
+ # SearchClient cannot create an index
+ with pytest.raises(ValueError):
+ create_mock_vector_store(
+ search_client,
+ index_management=IndexManagement.CREATE_IF_NOT_EXISTS,
+ )
+
+
+@pytest.mark.skipif(
+ not cogsearch_installed, reason="azure-search-documents package not installed"
+)
+def test_invalid_index_management_for_searchindexclient() -> None:
+ search_client = MagicMock(spec=SearchIndexClient)
+
+ # Index name must be supplied
+ with pytest.raises(ValueError):
+ create_mock_vector_store(
+ search_client, index_management=IndexManagement.VALIDATE_INDEX
+ )
--
GitLab