Merge pull request #130 from ks6088ts-labs/feature/issue-126_support-cosmosdb

support Azure Cosmos DB No SQL

Author: ks6088ts

infra/modules/cosmosDb.bicep

@@ -72,7 +72,7 @@ var locations = [
   }
 ]
 
-resource account 'Microsoft.DocumentDB/databaseAccounts@2024-02-15-preview' = {
+resource account 'Microsoft.DocumentDB/databaseAccounts@2024-05-15' = {
   name: accountName
   location: location
   tags: tags
@@ -82,6 +82,11 @@ resource account 'Microsoft.DocumentDB/databaseAccounts@2024-02-15-preview' = {
     locations: locations
     databaseAccountOfferType: 'Standard'
     enableAutomaticFailover: systemManagedFailover
+    capabilities: [
+      {
+        name: 'EnableNoSQLVectorSearch'
+      }
+    ]
   }
 }
 

poetry.lock

@@ -67,8 +67,10 @@ pymupdf = "^1.24.7"
 faiss-cpu = "^1.8.0.post1"
 langchain = "^0.2.10"
 langchain-openai = "^0.1.17"
+langchain-community = "^0.2.9"
 pandas = "^2.2.2"
 azure-search-documents = "^11.5.0"
+azure-cosmos = "^4.7.0"
 
 [build-system]
 requires = ["poetry-core"]

pyproject.toml

@@ -7,6 +7,11 @@ make install-deps-dev
 # Help
 poetry run python scripts/manage_vector_store.py --help
 
+VECTOR_STORE_TYPE=faiss
+
 # Create a new vector store in local
-poetry run python scripts/manage_vector_store.py create-vector-store
+poetry run python scripts/manage_vector_store.py create --vector-store-type $VECTOR_STORE_TYPE
+
+# Search for similar vectors
+poetry run python scripts/manage_vector_store.py search --vector-store-type $VECTOR_STORE_TYPE
 ```

scripts/README.md

@@ -8,3 +8,7 @@ AZURE_OPENAI_MODEL_CHAT = "gpt-4o"
 # Azure AI Search
 AZURE_AI_SEARCH_ENDPOINT = "https://<your-aisearch-name>.search.windows.net"
 AZURE_AI_SEARCH_API_KEY = "<api-key>"
+
+# Azure Cosmos DB
+AZURE_COSMOSDB_ENDPOINT = "https://<your-cosmosdb-name>.documents.azure.com:443/"
+AZURE_COSMOSDB_KEY = "<api-key>"

scripts/manage_vector_store.env.sample

@@ -5,9 +5,11 @@
 from typing import Annotated
 
 import typer
+from azure.cosmos import CosmosClient, PartitionKey
 from dotenv import load_dotenv
 from langchain_community.document_loaders.csv_loader import CSVLoader
 from langchain_community.vectorstores import FAISS, VectorStore
+from langchain_community.vectorstores.azure_cosmos_db_no_sql import AzureCosmosDBNoSqlVectorSearch
 from langchain_community.vectorstores.azuresearch import AzureSearch
 from langchain_core.documents import Document
 from langchain_openai import AzureOpenAIEmbeddings
@@ -16,8 +18,9 @@
 
 
 class VectorStoreType(str, Enum):
-    AzureAISearch = "azureaisearch"
     Faiss = "faiss"
+    AzureAISearch = "azureaisearch"
+    AzureCosmosDbNoSql = "azurecosmosdbnosql"
 
 
 def get_log_level(debug: bool) -> int:
@@ -41,10 +44,39 @@ def get_embeddings():
     )
 
 
+def get_cosmos_client():
+    return CosmosClient(
+        url=getenv("AZURE_COSMOSDB_ENDPOINT"),
+        credential=getenv("AZURE_COSMOSDB_KEY"),
+    )
+
+
 def get_local_vector_store_path(identifier: str):
     return f"./artifacts/vectorstore/{identifier}"
 
 
+def get_vector_embedding_policy():
+    return {
+        "vectorEmbeddings": [
+            {
+                "path": "/embedding",
+                "dataType": "float32",
+                "distanceFunction": "cosine",
+                "dimensions": 1536,
+            }
+        ]
+    }
+
+
+def get_indexing_policy():
+    return {
+        "indexingMode": "consistent",
+        "includedPaths": [{"path": "/*"}],
+        "excludedPaths": [{"path": '/"_etag"/?'}],
+        "vectorIndexes": [{"path": "/embedding", "type": "quantizedFlat"}],
+    }
+
+
 def create_azure_search(index_name: str) -> AzureSearch:
     return AzureSearch(
         azure_search_endpoint=getenv("AZURE_AI_SEARCH_ENDPOINT"),
@@ -59,40 +91,69 @@ def get_vector_store(
     vector_store_type: VectorStoreType,
     identifier: str,
 ) -> VectorStore:
-    if vector_store_type == VectorStoreType.AzureAISearch:
-        logging.info("Creating Azure AI Search vector store")
-        return create_azure_search(identifier)
-    elif vector_store_type == VectorStoreType.Faiss:
+    if vector_store_type == VectorStoreType.Faiss:
         logging.info("Creating Faiss vector store")
         return FAISS.load_local(
             folder_path=get_local_vector_store_path(identifier),
             embeddings=get_embeddings(),
             allow_dangerous_deserialization=True,
         )
+    elif vector_store_type == VectorStoreType.AzureAISearch:
+        logging.info("Creating Azure AI Search vector store")
+        return create_azure_search(identifier)
+    elif vector_store_type == VectorStoreType.AzureCosmosDbNoSql:
+        logging.info("Creating Azure Cosmos DB NoSQL vector store")
+        cosmos_database_name = "langchain_python_db"
+        return AzureCosmosDBNoSqlVectorSearch(
+            cosmos_client=get_cosmos_client(),
+            embedding=get_embeddings(),
+            vector_embedding_policy=get_vector_embedding_policy(),
+            indexing_policy=get_indexing_policy(),
+            cosmos_container_properties={"partition_key": PartitionKey(path="/id")},
+            cosmos_database_properties={"id": cosmos_database_name},
+            database_name=cosmos_database_name,
+            container_name="langchain_python_container",
+        )
 
 
-def _create_vector_store(
+def create_vector_store(
     vector_store_type: VectorStoreType,
     identifier: str,
     documents: list[Document],
 ) -> VectorStore:
-    if vector_store_type == VectorStoreType.AzureAISearch:
+    if vector_store_type == VectorStoreType.Faiss:
+        logging.info("Creating Faiss vector store")
+        vector_store: FAISS = FAISS.from_documents(
+            documents=documents,
+            embedding=get_embeddings(),
+        )
+        vector_store.save_local(folder_path=get_local_vector_store_path(identifier))
+        return
+    elif vector_store_type == VectorStoreType.AzureAISearch:
         logging.info("Creating Azure AI Search vector store")
         vector_store = create_azure_search(identifier)
         vector_store.add_documents(documents=documents)
         return
-    elif vector_store_type == VectorStoreType.Faiss:
-        logging.info("Creating Faiss vector store")
-        vector_store: FAISS = FAISS.from_documents(
+    elif vector_store_type == VectorStoreType.AzureCosmosDbNoSql:
+        logging.info("Creating Azure Cosmos DB NoSQL vector store")
+        cosmos_database_name = "langchain_python_db"
+
+        AzureCosmosDBNoSqlVectorSearch.from_documents(
             documents=documents,
             embedding=get_embeddings(),
+            cosmos_client=get_cosmos_client(),
+            database_name=cosmos_database_name,
+            container_name="langchain_python_container",
+            vector_embedding_policy=get_vector_embedding_policy(),
+            indexing_policy=get_indexing_policy(),
+            cosmos_database_properties={"id": cosmos_database_name},
+            cosmos_container_properties={"partition_key": PartitionKey(path="/id")},
         )
-        vector_store.save_local(folder_path=get_local_vector_store_path(identifier))
         return
 
 
 @app.command()
-def create_vector_store(
+def create(
     input_csv_file_path: Annotated[str, typer.Option(help="Path to the input CSV file")] = "./data/contoso_rules.csv",
     identifier="contoso_rules",
     vector_store_type: Annotated[VectorStoreType, typer.Option(case_sensitive=False)] = VectorStoreType.Faiss,
@@ -108,7 +169,7 @@ def create_vector_store(
         return
 
     # Create vector store
-    _create_vector_store(
+    create_vector_store(
         vector_store_type=vector_store_type,
         identifier=identifier,
         documents=documents,
@@ -120,6 +181,7 @@ def search(
     identifier="contoso_rules",
     vector_store_type: Annotated[VectorStoreType, typer.Option(case_sensitive=False)] = VectorStoreType.Faiss,
     query: Annotated[str, typer.Option(help="Query to search")] = "社内の機密情報は外部に漏らさないでください",
+    k: Annotated[int, typer.Option(help="Number of documents to retrieve")] = 5,
     debug: Annotated[bool, typer.Option(help="Enable debug mode")] = False,
 ):
     setup_logging(debug)
@@ -131,12 +193,12 @@ def search(
     )
 
     # Search
-    result = vector_store.similarity_search_with_relevance_scores(
+    got_documents = vector_store.similarity_search(
         query=query,
-        k=5,
-        score_threshold=0.5,
+        k=k,
     )
-    pprint(result)
+    for document in got_documents:
+        pprint(document.page_content)
 
 
 if __name__ == "__main__":