chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,252 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Annotated
|
||||
from uuid import uuid4
|
||||
|
||||
from azure.identity import AzureCliCredential
|
||||
|
||||
from samples.concepts.memory.utils import print_record
|
||||
from samples.concepts.resources.utils import Colors, print_with_color
|
||||
from semantic_kernel import Kernel
|
||||
from semantic_kernel.connectors.ai.open_ai import AzureTextEmbedding, OpenAITextEmbedding
|
||||
from semantic_kernel.connectors.memory import (
|
||||
AzureAISearchCollection,
|
||||
ChromaCollection,
|
||||
CosmosMongoCollection,
|
||||
CosmosNoSqlCollection,
|
||||
FaissCollection,
|
||||
InMemoryCollection,
|
||||
PineconeCollection,
|
||||
PostgresCollection,
|
||||
QdrantCollection,
|
||||
RedisHashsetCollection,
|
||||
RedisJsonCollection,
|
||||
SqlServerCollection,
|
||||
WeaviateCollection,
|
||||
)
|
||||
from semantic_kernel.data.vector import (
|
||||
SearchType,
|
||||
VectorSearchProtocol,
|
||||
VectorStoreCollection,
|
||||
VectorStoreField,
|
||||
vectorstoremodel,
|
||||
)
|
||||
|
||||
# This is a rather complex sample, showing how to use the vector store
|
||||
# with a number of different collections.
|
||||
# It also shows how to use the vector store with a number of different data models.
|
||||
# It also uses all the types of search available in the vector store.
|
||||
# For a simpler example, see "simple_memory.py"
|
||||
|
||||
|
||||
# Depending on the vector database, the index kind and distance function may need to be adjusted
|
||||
# since not all combinations are supported by all databases.
|
||||
# The values below might need to be changed for your collection to work.
|
||||
@vectorstoremodel(collection_name="test")
|
||||
@dataclass
|
||||
class DataModel:
|
||||
title: Annotated[str, VectorStoreField("data", is_full_text_indexed=True)]
|
||||
content: Annotated[str, VectorStoreField("data", is_full_text_indexed=True)]
|
||||
embedding: Annotated[
|
||||
list[float] | str | None,
|
||||
VectorStoreField("vector", dimensions=1536, type="float"),
|
||||
] = None
|
||||
id: Annotated[
|
||||
str,
|
||||
VectorStoreField(
|
||||
"key",
|
||||
),
|
||||
] = field(default_factory=lambda: str(uuid4()))
|
||||
tag: Annotated[str | None, VectorStoreField("data", type="str", is_indexed=True)] = None
|
||||
|
||||
def __post_init__(self, **kwargs):
|
||||
if self.embedding is None:
|
||||
self.embedding = f"{self.title} {self.content}"
|
||||
if self.tag is None:
|
||||
self.tag = "general"
|
||||
|
||||
|
||||
# A list of VectorStoreRecordCollection that can be used.
|
||||
# Available collections are:
|
||||
# - ai_search: Azure AI Search
|
||||
# - postgres: PostgreSQL
|
||||
# - redis_json: Redis JSON
|
||||
# - redis_hashset: Redis Hashset
|
||||
# - qdrant: Qdrant
|
||||
# - in_memory: In-memory store
|
||||
# - weaviate: Weaviate
|
||||
# Please either configure the weaviate settings via environment variables or provide them through the constructor.
|
||||
# Note that embed mode is not supported on Windows: https://github.com/weaviate/weaviate/issues/3315
|
||||
# - azure_cosmos_nosql: Azure Cosmos NoSQL
|
||||
# https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/how-to-create-account?tabs=azure-portal
|
||||
# Please see the link above to learn how to set up an Azure Cosmos NoSQL account.
|
||||
# https://learn.microsoft.com/en-us/azure/cosmos-db/how-to-develop-emulator?tabs=windows%2Cpython&pivots=api-nosql
|
||||
# Please see the link above to learn how to set up the Azure Cosmos NoSQL emulator on your machine.
|
||||
# For this sample to work with Azure Cosmos NoSQL, please adjust the index_kind of the data model to QUANTIZED_FLAT.
|
||||
# - azure_cosmos_mongodb: Azure Cosmos MongoDB
|
||||
# https://learn.microsoft.com/en-us/azure/cosmos-db/mongodb/introduction
|
||||
# - chroma: Chroma
|
||||
# The chroma collection is currently only available for in-memory versions
|
||||
# Client-Server mode and Chroma Cloud are not yet supported.
|
||||
# More info on Chroma here: https://docs.trychroma.com/docs/overview/introduction
|
||||
# - faiss: Faiss - in-memory with optimized indexes.
|
||||
# - pinecone: Pinecone
|
||||
# - sql_server: SQL Server, can connect to any SQL Server compatible database, like Azure SQL.
|
||||
# This is represented as a mapping from the collection name to a
|
||||
# function which returns the collection.
|
||||
# Using a function allows for lazy initialization of the collection,
|
||||
# so that settings for unused collections do not cause validation errors.
|
||||
collections: dict[str, Callable[[], VectorStoreCollection]] = {
|
||||
"ai_search": lambda: AzureAISearchCollection[str, DataModel](record_type=DataModel),
|
||||
"postgres": lambda: PostgresCollection[str, DataModel](record_type=DataModel),
|
||||
"redis_json": lambda: RedisJsonCollection[str, DataModel](
|
||||
record_type=DataModel,
|
||||
prefix_collection_name_to_key_names=True,
|
||||
),
|
||||
"redis_hash": lambda: RedisHashsetCollection[str, DataModel](
|
||||
record_type=DataModel,
|
||||
prefix_collection_name_to_key_names=True,
|
||||
),
|
||||
"qdrant": lambda: QdrantCollection[str, DataModel](
|
||||
record_type=DataModel,
|
||||
prefer_grpc=True,
|
||||
named_vectors=False,
|
||||
),
|
||||
"in_memory": lambda: InMemoryCollection[str, DataModel](record_type=DataModel),
|
||||
"weaviate": lambda: WeaviateCollection[str, DataModel](record_type=DataModel),
|
||||
"azure_cosmos_nosql": lambda: CosmosNoSqlCollection[str, DataModel](
|
||||
record_type=DataModel,
|
||||
create_database=True,
|
||||
),
|
||||
"azure_cosmos_mongodb": lambda: CosmosMongoCollection[str, DataModel](record_type=DataModel),
|
||||
"faiss": lambda: FaissCollection[str, DataModel](record_type=DataModel),
|
||||
"chroma": lambda: ChromaCollection[str, DataModel](record_type=DataModel),
|
||||
"pinecone": lambda: PineconeCollection[str, DataModel](record_type=DataModel),
|
||||
"sql_server": lambda: SqlServerCollection[str, DataModel](record_type=DataModel),
|
||||
}
|
||||
|
||||
|
||||
async def cleanup(record_collection):
|
||||
print("-" * 30)
|
||||
delete = input("Do you want to delete the collection? (y/n): ")
|
||||
if delete.lower() != "y":
|
||||
print("Skipping deletion.")
|
||||
return
|
||||
print_with_color("Deleting collection!", Colors.CBLUE)
|
||||
await record_collection.ensure_collection_deleted()
|
||||
print_with_color("Done!", Colors.CGREY)
|
||||
|
||||
|
||||
async def main(collection: str, use_azure_openai: bool):
|
||||
print("-" * 30)
|
||||
kernel = Kernel()
|
||||
embedder = (
|
||||
AzureTextEmbedding(service_id="embedding", credential=AzureCliCredential())
|
||||
if use_azure_openai
|
||||
else OpenAITextEmbedding(service_id="embedding")
|
||||
)
|
||||
kernel.add_service(embedder)
|
||||
async with collections[collection]() as record_collection:
|
||||
assert isinstance(record_collection, VectorSearchProtocol) # nosec
|
||||
record_collection.embedding_generator = embedder
|
||||
print_with_color(f"Creating {collection} collection!", Colors.CGREY)
|
||||
# cleanup any existing collection
|
||||
await record_collection.ensure_collection_deleted()
|
||||
# create a new collection
|
||||
await record_collection.ensure_collection_exists()
|
||||
|
||||
record1 = DataModel(
|
||||
content="Semantic Kernel is awesome",
|
||||
id="e6103c03-487f-4d7d-9c23-4723651c17f4",
|
||||
title="Overview",
|
||||
)
|
||||
record2 = DataModel(
|
||||
content="Semantic Kernel is available in dotnet, python and Java.",
|
||||
id="09caec77-f7e1-466a-bcec-f1d51c5b15be",
|
||||
title="Semantic Kernel Languages",
|
||||
)
|
||||
record3 = DataModel(
|
||||
content="```python\nfrom semantic_kernel import Kernel\nkernel = Kernel()\n```",
|
||||
id="d5c9913a-e015-4944-b960-5d4a84bca002",
|
||||
title="Code sample",
|
||||
tag="code",
|
||||
)
|
||||
|
||||
print_with_color("Adding records!", Colors.CBLUE)
|
||||
records = [record1, record2, record3]
|
||||
keys = await record_collection.upsert(records)
|
||||
print(f" Upserted {keys=}")
|
||||
print_with_color("Getting records!", Colors.CBLUE)
|
||||
results = await record_collection.get(top=10, order_by="content")
|
||||
if results:
|
||||
[print_record(record=result) for result in results]
|
||||
else:
|
||||
print("Nothing found...")
|
||||
options = {
|
||||
"vector_property_name": "embedding",
|
||||
"additional_property_name": "content",
|
||||
"filter": lambda x: x.tag == "general",
|
||||
}
|
||||
|
||||
print("-" * 30)
|
||||
print_with_color("Now we can start searching.", Colors.CBLUE)
|
||||
print_with_color(" For each type of search, enter a search term, for instance `python`.", Colors.CBLUE)
|
||||
print_with_color(" Enter exit to exit, and skip or nothing to skip this search.", Colors.CBLUE)
|
||||
print("-" * 30)
|
||||
print_with_color(
|
||||
"This collection supports the following search types: "
|
||||
f"{', '.join(search.value for search in record_collection.supported_search_types)}",
|
||||
Colors.CBLUE,
|
||||
)
|
||||
if SearchType.KEYWORD_HYBRID in record_collection.supported_search_types:
|
||||
search_text = input("Enter search text for hybrid text search: ")
|
||||
if search_text.lower() == "exit":
|
||||
await cleanup(record_collection)
|
||||
return
|
||||
if search_text and search_text.lower() != "skip":
|
||||
print("-" * 30)
|
||||
print_with_color(
|
||||
"Using hybrid text search: ",
|
||||
Colors.CBLUE,
|
||||
)
|
||||
search_results = await record_collection.hybrid_search(values=search_text, **options)
|
||||
if search_results.total_count == 0:
|
||||
print("\nNothing found...\n")
|
||||
else:
|
||||
[print_record(result) async for result in search_results.results]
|
||||
print("-" * 30)
|
||||
|
||||
if SearchType.VECTOR in record_collection.supported_search_types:
|
||||
search_text = input("Enter search text for vector search: ")
|
||||
if search_text.lower() == "exit":
|
||||
await cleanup(record_collection)
|
||||
return
|
||||
if search_text and search_text.lower() != "skip":
|
||||
print("-" * 30)
|
||||
print_with_color(
|
||||
"Using vector search: ",
|
||||
Colors.CBLUE,
|
||||
)
|
||||
search_results = await record_collection.search(search_text, **options)
|
||||
if search_results.total_count == 0:
|
||||
print("\nNothing found...\n")
|
||||
else:
|
||||
[print_record(result) async for result in search_results.results]
|
||||
print("-" * 30)
|
||||
await cleanup(record_collection)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
argparse.ArgumentParser()
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--collection", default="in_memory", choices=collections.keys(), help="What collection to use.")
|
||||
# Option of whether to use OpenAI or Azure OpenAI.
|
||||
parser.add_argument("--use-azure-openai", action="store_true", help="Use Azure OpenAI instead of OpenAI.")
|
||||
args = parser.parse_args()
|
||||
args.collection = "ai_search"
|
||||
asyncio.run(main(collection=args.collection, use_azure_openai=args.use_azure_openai))
|
||||
Reference in New Issue
Block a user